added file expiration and fleshed out the API a bit

src/upload.rs

@@ -6,10 +6,10 @@ use actix_web_actors::ws::{self, CloseCode};
 use bytes::Bytes;
 use log::{debug, error, info, trace};
 use rand::distributions::{Alphanumeric, DistString};
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 use time::OffsetDateTime;
 
-use crate::{storage_dir, DownloadableFile, UploadedFile};
+use crate::store::{storage_dir, StoredFile, self};
 
 const MAX_FILES: usize = 256;
 const FILENAME_DATE_FORMAT: &[time::format_description::FormatItem] =
@@ -31,6 +31,10 @@ enum Error {
     NoFiles,
     #[error("Number of files submitted by client exceeded the maximum limit")]
     TooManyFiles,
+    #[error("Requested lifetime was too long")]
+    TooLong,
+    #[error("Upload size was too large, can be at most {0} bytes")]
+    TooBig(u64),
     #[error("Websocket was closed by client before completing transfer")]
     ClosedEarly(Option<CloseReason>),
     #[error("Client sent more data than they were supposed to")]
@@ -40,15 +44,17 @@ enum Error {
 impl Error {
     fn close_code(&self) -> CloseCode {
         match self {
-            Self::Parse(_) => CloseCode::Invalid,
-            Self::Storage(_) => CloseCode::Error,
-            Self::TimeFormat(_) => CloseCode::Error,
-            Self::DuplicateFilename => CloseCode::Policy,
-            Self::UnexpectedMessageType => CloseCode::Invalid,
-            Self::NoFiles => CloseCode::Policy,
-            Self::TooManyFiles => CloseCode::Policy,
-            Self::ClosedEarly(_) => CloseCode::Invalid,
-            Self::UnexpectedExtraData => CloseCode::Invalid,
+            Self::Storage(_)
+                | Self::TimeFormat(_) => CloseCode::Error,
+            Self::Parse(_)
+                | Self::UnexpectedMessageType
+                | Self::ClosedEarly(_)
+                | Self::UnexpectedExtraData => CloseCode::Invalid,
+            Self::DuplicateFilename
+                | Self::NoFiles
+                | Self::TooManyFiles
+                | Self::TooLong
+                | Self::TooBig(_) => CloseCode::Policy,
         }
     }
 }
@@ -64,7 +70,7 @@ impl Uploader {
     pub(crate) fn new(app_data: super::AppData) -> Self {
         Self {
             writer: None,
-            storage_filename: String::new(),
+            storage_filename: Alphanumeric.sample_string(&mut rand::thread_rng(), 8),
             app_data,
             bytes_remaining: 0,
         }
@@ -75,6 +81,22 @@ impl Actor for Uploader {
     type Context = ws::WebsocketContext<Self>;
 }
 
+pub struct UploadedFile {
+    pub name: String,
+    pub size: u64,
+    pub modtime: OffsetDateTime,
+}
+
+impl UploadedFile {
+    fn new(name: &str, size: u64, modtime: OffsetDateTime) -> Self {
+        Self {
+            name: sanitise_file_name::sanitise(name),
+            size,
+            modtime,
+        }
+    }
+}
+
 #[derive(Debug, Deserialize)]
 struct RawUploadedFile {
     name: String,
@@ -93,6 +115,31 @@ impl RawUploadedFile {
     }
 }
 
+#[derive(Deserialize)]
+struct UploadManifest {
+    files: Vec<RawUploadedFile>,
+    lifetime: u32,
+}
+
+#[derive(Serialize)]
+#[serde(rename_all = "snake_case", tag = "type")]
+enum ServerMessage {
+    Ready { code: String },
+    TooBig { max_size: u64 },
+    TooLong { max_days: u32 },
+    Error { details: String },
+}
+
+impl From<&Error> for ServerMessage {
+    fn from(e: &Error) -> Self {
+        match e {
+            Error::TooBig(max_size) => ServerMessage::TooBig { max_size: *max_size },
+            Error::TooLong => ServerMessage::TooLong { max_days: store::max_lifetime() },
+            _ => ServerMessage::Error { details: e.to_string() },
+        }
+    }
+}
+
 fn stop_and_flush<T>(_: T, u: &mut Uploader, ctx: &mut <Uploader as Actor>::Context) {
     ctx.stop();
     if let Some(w) = u.writer.as_mut() {
@@ -115,12 +162,7 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for Uploader {
 
         match self.handle_message(msg, ctx) {
             Err(e) => {
-                error!("{}", e);
-                ctx.close(Some(ws::CloseReason {
-                    code: e.close_code(),
-                    description: Some(e.to_string()),
-                }));
-                self.cleanup_after_error(ctx);
+                self.notify_error_and_cleanup(e, ctx);
             }
             Ok(true) => {
                 info!("Finished uploading data");
@@ -140,6 +182,16 @@ fn ack(ctx: &mut <Uploader as Actor>::Context) {
 }
 
 impl Uploader {
+    fn notify_error_and_cleanup(&mut self, e: Error, ctx: &mut <Self as Actor>::Context) {
+        error!("{}", e);
+        ctx.text(serde_json::to_string(&ServerMessage::from(&e)).unwrap());
+        ctx.close(Some(ws::CloseReason {
+            code: e.close_code(),
+            description: Some(e.to_string()),
+        }));
+        self.cleanup_after_error(ctx);
+    }
+
     fn handle_message(
         &mut self,
         msg: ws::Message,
@@ -151,12 +203,18 @@ impl Uploader {
                 if self.writer.is_some() {
                     return Err(Error::UnexpectedMessageType);
                 }
-                let raw_files: Vec<RawUploadedFile> = serde_json::from_slice(text.as_bytes())?;
+                let UploadManifest { files: raw_files, lifetime, } = serde_json::from_slice(text.as_bytes())?;
+                if lifetime > store::max_lifetime() {
+                    return Err(Error::TooLong);
+                }
                 info!("Received file list: {} files", raw_files.len());
                 debug!("{:?}", raw_files);
                 if raw_files.len() > MAX_FILES {
                     return Err(Error::TooManyFiles);
                 }
+                if raw_files.is_empty() {
+                    return Err(Error::NoFiles);
+                }
                 let mut filenames: HashSet<String> = HashSet::new();
                 let mut files = Vec::new();
                 for raw_file in raw_files.iter() {
@@ -172,18 +230,13 @@ impl Uploader {
                     self.bytes_remaining += file.size;
                     files.push(file);
                 }
-                if files.is_empty() {
-                    return Err(Error::NoFiles);
-                }
-                let storage_filename = Alphanumeric.sample_string(&mut rand::thread_rng(), 8);
-                self.storage_filename = storage_filename.clone();
-                let storage_path = storage_dir().join(storage_filename.clone());
+                let storage_path = storage_dir().join(self.storage_filename.clone());
                 info!("storing to: {:?}", storage_path);
                 let writer = File::options()
                     .write(true)
                     .create_new(true)
                     .open(&storage_path)?;
-                let (writer, downloadable_file): (Box<dyn Write>, _) = if files.len() > 1 {
+                let (writer, name, size, modtime): (Box<dyn Write>,_,_,_) = if files.len() > 1 {
                     info!("Wrapping in zipfile generator");
                     let now = OffsetDateTime::now_utc();
                     let zip_writer = super::zip::ZipGenerator::new(files, writer);
@@ -192,33 +245,40 @@ impl Uploader {
                         super::APP_NAME.to_owned() + &now.format(FILENAME_DATE_FORMAT)? + ".zip";
                     (
                         Box::new(zip_writer),
-                        DownloadableFile {
-                            name: download_filename,
-                            size,
-                            modtime: now,
-                        },
+                        download_filename,
+                        size,
+                        now,
                     )
                 } else {
                     (
                         Box::new(writer),
-                        DownloadableFile {
-                            name: files[0].name.clone(),
-                            size: files[0].size,
-                            modtime: files[0].modtime,
-                        },
+                        files[0].name.clone(),
+                        files[0].size,
+                        files[0].modtime,
                     )
                 };
                 self.writer = Some(writer);
+                let stored_file = StoredFile {
+                    name,
+                    size,
+                    modtime,
+                    expiry: OffsetDateTime::now_utc() + lifetime*time::Duration::DAY,
+                };
                 let data = self.app_data.clone();
+                let storage_filename = self.storage_filename.clone();
                 ctx.spawn(actix::fut::wrap_future(async move {
                     debug!("Spawned future to add entry {} to state", storage_filename);
                     data.write()
                         .await
-                        .add_file(storage_filename, downloadable_file)
+                        .add_file(storage_filename, stored_file)
                         .await
-                        .unwrap();
+                }).map(|res, u: &mut Self, ctx: &mut <Self as Actor>::Context| {
+                    match res {
+                        Ok(Ok(())) => ctx.text(serde_json::to_string(&ServerMessage::Ready { code: u.storage_filename.clone() }).unwrap()),
+                        Ok(Err(size)) => u.notify_error_and_cleanup(Error::TooBig(size), ctx),
+                        Err(e) => u.notify_error_and_cleanup(Error::from(e), ctx)
+                    }
                 }));
-                ctx.text(self.storage_filename.as_str());
             }
             ws::Message::Binary(data) | ws::Message::Continuation(Item::Last(data)) => {
                 let result = self.handle_data(data)?;