From 2b62a7e1da8a6197aaebb7b0119a6b3803d69b27 Mon Sep 17 00:00:00 2001 From: moparisthebest Date: Thu, 16 Aug 2018 23:33:37 -0400 Subject: [PATCH] More refactoring --- src/main.rs | 98 +++++++++++++++++++++++++++++++---------------------- 1 file changed, 58 insertions(+), 40 deletions(-) diff --git a/src/main.rs b/src/main.rs index 7d8ac7a..5cf994e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -32,7 +32,7 @@ use mpu::MultipartUpload; const HOST: &'static str = "http://localhost:8000"; -const UPLOAD_MAX_SIZE: u64 = 8 * 1024 * 1024; +type Result = ::std::result::Result; #[derive(Debug)] pub enum Error { @@ -67,7 +67,7 @@ impl From for Error { } } -#[derive(Deserialize, Serialize, Default, Debug)] +#[derive(Deserialize, Serialize, Debug)] struct PasteInfo<'a> { // these never change key: Option>, // 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 } +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> { - fn read>(path: P) -> ::std::result::Result,Error> { + fn read>(path: P) -> Result> { let mut f = File::open(path)?; let mut input = String::new(); f.read_to_string(&mut input)?; @@ -88,7 +101,7 @@ impl<'a> PasteInfo<'a> { Ok(paste_info) } - fn write>(&self, path: P) -> ::std::result::Result<(),Error> { + fn write>(&self, path: P) -> Result<()> { let toml = toml::to_string(&self)?; fs::write(path, toml)?; 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) } - fn mark_viewed_and_write>(&mut self, path: P) -> ::std::result::Result<(),Error> { + fn mark_viewed_and_write>(&mut self, path: P) -> Result<()> { let mut must_write = false; if self.delete_after_num_views.is_some() { 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)) } - fn upload(&self, paste: Data, _key: Option) -> io::Result { + fn upload(&self, paste: Data, _key: Option) -> Result { let (filename, info_filename, url) = self.new_paste(); - + PasteInfo::default().write(info_filename)?; paste.stream_to_file(Path::new(&filename))?; Ok(url) } - fn upload_multipart(&self, paste: MultipartUpload) -> io::Result { + fn upload_multipart(&self, paste: MultipartUpload) -> Result { let (filename, info_filename, url) = self.new_paste(); + PasteInfo::default().write(info_filename)?; paste.stream_to_file(Path::new(&filename))?; Ok(url) } - fn upload_string(&self, paste: &PasteForm) -> ::std::result::Result { + fn upload_string(&self, paste: &PasteForm) -> Result { let (filename, info_filename, url) = self.new_paste(); PasteInfo::from(paste).write(info_filename)?; fs::write(filename, &paste.content)?; Ok(url) } - fn get(&self, id: PasteID) -> ::std::result::Result,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> { let (filename, info_filename, foldername) = self.file_paths(id); let mut paste_info = PasteInfo::read(&info_filename)?; // 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 #[post("/w", format = "application/x-www-form-urlencoded", data = "")] -fn web_post(backend: &Backend, paste: LenientForm) -> ::std::result::Result { +fn web_post(backend: &Backend, paste: LenientForm) -> Result { backend.upload_string(paste.get()) } // todo: change /w to /, shouldn't conflict because of format, but it does currently #[post("/m", format = "multipart/form-data", data = "")] -fn mpu_post(backend: &Backend, paste: MultipartUpload) -> io::Result { +fn mpu_post(backend: &Backend, paste: MultipartUpload) -> Result { backend.upload_multipart(paste) } #[put("/", data = "")] -fn upload_put(backend: &Backend, paste: Data) -> io::Result { +fn upload_put(backend: &Backend, paste: Data) -> Result { backend.upload(paste, None) } #[post("/", data = "")] -fn upload_post(backend: &Backend, paste: Data) -> io::Result { +fn upload_post(backend: &Backend, paste: Data) -> Result { backend.upload(paste, None) } #[patch("/", data = "")] -fn upload_patch(backend: &Backend, paste: Data) -> io::Result { +fn upload_patch(backend: &Backend, paste: Data) -> Result { backend.upload(paste, None) } #[put("/", data = "")] -fn upload_put_key(backend: &Backend, paste: Data, key: PasteID) -> io::Result { +fn upload_put_key(backend: &Backend, paste: Data, key: PasteID) -> Result { backend.upload(paste, Some(key)) } #[post("/", data = "")] -fn upload_post_key(backend: &Backend, paste: Data, key: PasteID) -> io::Result { +fn upload_post_key(backend: &Backend, paste: Data, key: PasteID) -> Result { backend.upload(paste, Some(key)) } #[patch("/", data = "")] -fn upload_patch_key(backend: &Backend, paste: Data, key: PasteID) -> io::Result { +fn upload_patch_key(backend: &Backend, paste: Data, key: PasteID) -> Result { backend.upload(paste, Some(key)) } @@ -299,8 +333,9 @@ fn delete(id: PasteID, _key: PasteID) -> Option> { } #[get("/")] -fn index() -> io::Result { - NamedFile::open("static/index.html") +fn index() -> Result { + let index = NamedFile::open("static/index.html")?; + Ok(index) } #[get("/static/")] @@ -320,7 +355,7 @@ fn rocket() -> rocket::Rocket { } // adapted from io::copy -fn copy(reader: &mut R, writer: &mut W) -> io::Result +fn copy(reader: &mut R, writer: &mut W, upload_max_size: u64) -> io::Result where R: Read, W: Write { let mut buf : [u8; 8192] = [0; 8192]; @@ -335,40 +370,23 @@ fn copy(reader: &mut R, writer: &mut W) -> io::Result }; writer.write_all(&buf[..len])?; written += len as u64; - if written > UPLOAD_MAX_SIZE { + if written > upload_max_size { return Err(std::io::Error::from(std::io::ErrorKind::InvalidData)) } } } -fn consume_paste(backend: &Backend, mut stream: TcpStream, timeout: Option) -> 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(){ // Bind the server's socket thread::spawn(|| { let backend = &Backend::PlainFile; - let timeout = Some(Duration::new(5, 0)); let listener = TcpListener::bind("127.0.0.1:12345").unwrap(); loop { match listener.accept() { Ok((mut stream, _addr)) => { 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) => {