More refactoring

This commit is contained in:
Travis Burtrum 2018-08-16 23:33:37 -04:00
parent d99445ab11
commit 2b62a7e1da
1 changed files with 58 additions and 40 deletions

View File

@ -32,7 +32,7 @@ use mpu::MultipartUpload;
const HOST: &'static str = "http://localhost:8000"; const HOST: &'static str = "http://localhost:8000";
const UPLOAD_MAX_SIZE: u64 = 8 * 1024 * 1024; type Result<T> = ::std::result::Result<T, Error>;
#[derive(Debug)] #[derive(Debug)]
pub enum Error { pub enum Error {
@ -67,7 +67,7 @@ impl From<toml::de::Error> for Error {
} }
} }
#[derive(Deserialize, Serialize, Default, Debug)] #[derive(Deserialize, Serialize, Debug)]
struct PasteInfo<'a> { struct PasteInfo<'a> {
// these never change // these never change
key: Option<PasteID<'a>>, // key to update/delete paste with key: Option<PasteID<'a>>, // key to update/delete paste with
@ -79,8 +79,21 @@ struct PasteInfo<'a> {
num_views: u32, // Only incremented if delete_after_num_views is Some, otherwise 0 num_views: u32, // Only incremented if delete_after_num_views is Some, otherwise 0
} }
impl<'a> Default for PasteInfo<'a> {
fn default() -> Self {
PasteInfo {
key: None,
delete_after: Some(SystemTime::now() + Duration::from_secs(2592000)), // default to 30 days
delete_after_num_views:None,
delete_if_not_viewed_in_last_seconds: None,
last_viewed: None,
num_views: 0,
}
}
}
impl<'a> PasteInfo<'a> { impl<'a> PasteInfo<'a> {
fn read<P: AsRef<Path>>(path: P) -> ::std::result::Result<PasteInfo<'static>,Error> { fn read<P: AsRef<Path>>(path: P) -> Result<PasteInfo<'static>> {
let mut f = File::open(path)?; let mut f = File::open(path)?;
let mut input = String::new(); let mut input = String::new();
f.read_to_string(&mut input)?; f.read_to_string(&mut input)?;
@ -88,7 +101,7 @@ impl<'a> PasteInfo<'a> {
Ok(paste_info) Ok(paste_info)
} }
fn write<P: AsRef<Path>>(&self, path: P) -> ::std::result::Result<(),Error> { fn write<P: AsRef<Path>>(&self, path: P) -> Result<()> {
let toml = toml::to_string(&self)?; let toml = toml::to_string(&self)?;
fs::write(path, toml)?; fs::write(path, toml)?;
Ok(()) Ok(())
@ -101,7 +114,7 @@ impl<'a> PasteInfo<'a> {
|| self.delete_if_not_viewed_in_last_seconds.map(|n| (SystemTime::now() - n) > self.last_viewed.unwrap()).unwrap_or(false) || self.delete_if_not_viewed_in_last_seconds.map(|n| (SystemTime::now() - n) > self.last_viewed.unwrap()).unwrap_or(false)
} }
fn mark_viewed_and_write<P: AsRef<Path>>(&mut self, path: P) -> ::std::result::Result<(),Error> { fn mark_viewed_and_write<P: AsRef<Path>>(&mut self, path: P) -> Result<()> {
let mut must_write = false; let mut must_write = false;
if self.delete_after_num_views.is_some() { if self.delete_after_num_views.is_some() {
must_write = true; must_write = true;
@ -178,27 +191,48 @@ impl Backend {
(format!("upload/{id}/file", id = id), format!("upload/{id}/info", id = id), format!("upload/{id}", id = id)) (format!("upload/{id}/file", id = id), format!("upload/{id}/info", id = id), format!("upload/{id}", id = id))
} }
fn upload(&self, paste: Data, _key: Option<PasteID>) -> io::Result<String> { fn upload(&self, paste: Data, _key: Option<PasteID>) -> Result<String> {
let (filename, info_filename, url) = self.new_paste(); let (filename, info_filename, url) = self.new_paste();
PasteInfo::default().write(info_filename)?;
paste.stream_to_file(Path::new(&filename))?; paste.stream_to_file(Path::new(&filename))?;
Ok(url) Ok(url)
} }
fn upload_multipart(&self, paste: MultipartUpload) -> io::Result<String> { fn upload_multipart(&self, paste: MultipartUpload) -> Result<String> {
let (filename, info_filename, url) = self.new_paste(); let (filename, info_filename, url) = self.new_paste();
PasteInfo::default().write(info_filename)?;
paste.stream_to_file(Path::new(&filename))?; paste.stream_to_file(Path::new(&filename))?;
Ok(url) Ok(url)
} }
fn upload_string(&self, paste: &PasteForm) -> ::std::result::Result<String,Error> { fn upload_string(&self, paste: &PasteForm) -> Result<String> {
let (filename, info_filename, url) = self.new_paste(); let (filename, info_filename, url) = self.new_paste();
PasteInfo::from(paste).write(info_filename)?; PasteInfo::from(paste).write(info_filename)?;
fs::write(filename, &paste.content)?; fs::write(filename, &paste.content)?;
Ok(url) Ok(url)
} }
fn get(&self, id: PasteID) -> ::std::result::Result<content::Plain<File>,Error> { fn upload_tcp_stream(&self, mut stream: TcpStream) -> Result<()>
{
let (filename, info_filename, url) = self.new_paste();
PasteInfo::default().write(info_filename)?;
let mut paste_file = File::create(&filename)?;
let timeout = Some(Duration::new(5, 0)); // todo: make this config store in struct
stream.set_read_timeout(timeout)?;
stream.set_write_timeout(timeout)?;
stream.write(&url.into_bytes())?;
stream.flush()?;
let upload_max_size: u64 = 8 * 1024 * 1024; // todo: make this config store in struct
copy(&mut stream, &mut paste_file, upload_max_size)?;
Ok(())
}
fn get(&self, id: PasteID) -> Result<content::Plain<File>> {
let (filename, info_filename, foldername) = self.file_paths(id); let (filename, info_filename, foldername) = self.file_paths(id);
let mut paste_info = PasteInfo::read(&info_filename)?; let mut paste_info = PasteInfo::read(&info_filename)?;
// first check if we should delete this // first check if we should delete this
@ -247,43 +281,43 @@ impl<'a> From<&'a PasteForm> for PasteInfo<'a> {
// todo: change /w to /, shouldn't conflict because of format, but it does currently // todo: change /w to /, shouldn't conflict because of format, but it does currently
#[post("/w", format = "application/x-www-form-urlencoded", data = "<paste>")] #[post("/w", format = "application/x-www-form-urlencoded", data = "<paste>")]
fn web_post(backend: &Backend, paste: LenientForm<PasteForm>) -> ::std::result::Result<String,Error> { fn web_post(backend: &Backend, paste: LenientForm<PasteForm>) -> Result<String> {
backend.upload_string(paste.get()) backend.upload_string(paste.get())
} }
// todo: change /w to /, shouldn't conflict because of format, but it does currently // todo: change /w to /, shouldn't conflict because of format, but it does currently
#[post("/m", format = "multipart/form-data", data = "<paste>")] #[post("/m", format = "multipart/form-data", data = "<paste>")]
fn mpu_post(backend: &Backend, paste: MultipartUpload) -> io::Result<String> { fn mpu_post(backend: &Backend, paste: MultipartUpload) -> Result<String> {
backend.upload_multipart(paste) backend.upload_multipart(paste)
} }
#[put("/", data = "<paste>")] #[put("/", data = "<paste>")]
fn upload_put(backend: &Backend, paste: Data) -> io::Result<String> { fn upload_put(backend: &Backend, paste: Data) -> Result<String> {
backend.upload(paste, None) backend.upload(paste, None)
} }
#[post("/", data = "<paste>")] #[post("/", data = "<paste>")]
fn upload_post(backend: &Backend, paste: Data) -> io::Result<String> { fn upload_post(backend: &Backend, paste: Data) -> Result<String> {
backend.upload(paste, None) backend.upload(paste, None)
} }
#[patch("/", data = "<paste>")] #[patch("/", data = "<paste>")]
fn upload_patch(backend: &Backend, paste: Data) -> io::Result<String> { fn upload_patch(backend: &Backend, paste: Data) -> Result<String> {
backend.upload(paste, None) backend.upload(paste, None)
} }
#[put("/<key>", data = "<paste>")] #[put("/<key>", data = "<paste>")]
fn upload_put_key(backend: &Backend, paste: Data, key: PasteID) -> io::Result<String> { fn upload_put_key(backend: &Backend, paste: Data, key: PasteID) -> Result<String> {
backend.upload(paste, Some(key)) backend.upload(paste, Some(key))
} }
#[post("/<key>", data = "<paste>")] #[post("/<key>", data = "<paste>")]
fn upload_post_key(backend: &Backend, paste: Data, key: PasteID) -> io::Result<String> { fn upload_post_key(backend: &Backend, paste: Data, key: PasteID) -> Result<String> {
backend.upload(paste, Some(key)) backend.upload(paste, Some(key))
} }
#[patch("/<key>", data = "<paste>")] #[patch("/<key>", data = "<paste>")]
fn upload_patch_key(backend: &Backend, paste: Data, key: PasteID) -> io::Result<String> { fn upload_patch_key(backend: &Backend, paste: Data, key: PasteID) -> Result<String> {
backend.upload(paste, Some(key)) backend.upload(paste, Some(key))
} }
@ -299,8 +333,9 @@ fn delete(id: PasteID, _key: PasteID) -> Option<content::Plain<File>> {
} }
#[get("/")] #[get("/")]
fn index() -> io::Result<NamedFile> { fn index() -> Result<NamedFile> {
NamedFile::open("static/index.html") let index = NamedFile::open("static/index.html")?;
Ok(index)
} }
#[get("/static/<file..>")] #[get("/static/<file..>")]
@ -320,7 +355,7 @@ fn rocket() -> rocket::Rocket {
} }
// adapted from io::copy // adapted from io::copy
fn copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> io::Result<u64> fn copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W, upload_max_size: u64) -> io::Result<u64>
where R: Read, W: Write where R: Read, W: Write
{ {
let mut buf : [u8; 8192] = [0; 8192]; let mut buf : [u8; 8192] = [0; 8192];
@ -335,40 +370,23 @@ fn copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> io::Result<u64>
}; };
writer.write_all(&buf[..len])?; writer.write_all(&buf[..len])?;
written += len as u64; written += len as u64;
if written > UPLOAD_MAX_SIZE { if written > upload_max_size {
return Err(std::io::Error::from(std::io::ErrorKind::InvalidData)) return Err(std::io::Error::from(std::io::ErrorKind::InvalidData))
} }
} }
} }
fn consume_paste(backend: &Backend, mut stream: TcpStream, timeout: Option<Duration>) -> io::Result<()>
{
let (filename, _info_filename, url) = backend.new_paste();
let mut paste_file = File::create(&filename)?;
stream.set_read_timeout(timeout)?;
stream.set_write_timeout(timeout)?;
stream.write(&url.into_bytes())?;
stream.flush()?;
copy(&mut stream, &mut paste_file)?;
Ok(())
}
fn run_tcp(){ fn run_tcp(){
// Bind the server's socket // Bind the server's socket
thread::spawn(|| { thread::spawn(|| {
let backend = &Backend::PlainFile; let backend = &Backend::PlainFile;
let timeout = Some(Duration::new(5, 0));
let listener = TcpListener::bind("127.0.0.1:12345").unwrap(); let listener = TcpListener::bind("127.0.0.1:12345").unwrap();
loop { loop {
match listener.accept() { match listener.accept() {
Ok((mut stream, _addr)) => { Ok((mut stream, _addr)) => {
thread::spawn(move || { thread::spawn(move || {
consume_paste(backend, stream, timeout).is_ok(); // again we don't care about this error backend.upload_tcp_stream(stream).is_ok(); // again we don't care about this error
}); });
}, },
Err(_e) => { Err(_e) => {