feat(server): allow to override filename when using random_url (#233)

* feat(server): allow to override filename when using random_url

* docs(README): remove line from features

* refactor(header): make const private

README.md

@@ -52,6 +52,7 @@ Here you can read the blog post about how it is deployed on Shuttle: [https://bl
     - [Paste file from remote URL](#paste-file-from-remote-url)
     - [Cleaning up expired files](#cleaning-up-expired-files)
     - [Delete file from server](#delete-file-from-server)
+    - [Override the filename when using `random_url`](#override-the-filename-when-using-random_url)
   - [Server](#server)
     - [List endpoint](#list-endpoint)
     - [HTML Form](#html-form)
@@ -254,6 +255,14 @@ $ curl -H "Authorization: <auth_token>" -X DELETE "<server_address>/file.txt"
 
 > The `DELETE` endpoint will not be exposed and will return `404` error if `delete_tokens` are not set.
 
+#### Override the filename when using `random_url`
+
+The generation of a random filename can be overridden by sending a header called `filename`:
+
+```sh
+curl -F "[email protected]" -H "filename: <file_name>" "<server_address>"
+```
+
 ### Server
 
 To start the server:

fixtures/test-file-upload-override-filename/config.toml

@@ -0,0 +1,9 @@
+[server]
+address = "127.0.0.1:8000"
+max_content_length = "10MB"
+upload_path = "./upload"
+
+[paste]
+random_url = { type = "alphanumeric", length = "4", suffix_mode = true }
+default_extension = "txt"
+duplicate_files = true

fixtures/test-file-upload-override-filename/test.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+
+content="test data"
+
+setup() {
+  echo "$content" > file
+}
+
+run_test() {
+  file_url=$(curl -s -F "file=@file" -H "filename:fn_from_header.txt" localhost:8000)
+  test "$file_url" = "http://localhost:8000/fn_from_header.txt"
+  test "$content" = "$(cat upload/fn_from_header.txt)"
+  test "$content" = "$(curl -s $file_url)"
+}
+
+teardown() {
+  rm file
+  rm -r upload
+}

src/header.rs

@@ -7,6 +7,9 @@ use std::time::Duration;
 /// Custom HTTP header for expiry dates.
 pub const EXPIRE: &str = "expire";
 
+/// Custom HTTP header to override filename.
+const FILENAME: &str = "filename";
+
 /// Parses the expiry date from the [`custom HTTP header`](EXPIRE).
 pub fn parse_expiry_date(headers: &HeaderMap, time: Duration) -> Result<Option<u128>, ActixError> {
     if let Some(expire_time) = headers.get(EXPIRE).and_then(|v| v.to_str().ok()) {
@@ -18,6 +21,15 @@ pub fn parse_expiry_date(headers: &HeaderMap, time: Duration) -> Result<Option<u
     }
 }
 
+/// Parses the filename from the header.
+pub fn parse_header_filename(headers: &HeaderMap) -> Result<Option<String>, ActixError> {
+    if let Some(file_name) = headers.get(FILENAME).and_then(|v| v.to_str().ok()) {
+        Ok(Some(file_name.to_string()))
+    } else {
+        Ok(None)
+    }
+}
+
 /// Wrapper for Actix content disposition header.
 ///
 /// Aims to parse the file data from multipart body.

src/paste.rs

@@ -88,13 +88,15 @@ impl Paste {
     /// - If `file_name` does not have an extension, it is replaced with [`default_extension`].
     /// - If `file_name` is "-", it is replaced with "stdin".
     /// - If [`random_url.enabled`] is `true`, `file_name` is replaced with a pet name or random string.
+    /// - If `header_filename` is set, it will override the filename.
     ///
     /// [`default_extension`]: crate::config::PasteConfig::default_extension
     /// [`random_url.enabled`]: crate::random::RandomURLConfig::enabled
     pub fn store_file(
         &self,
         file_name: &str,
         expiry_date: Option<u128>,
+        header_filename: Option<String>,
         config: &Config,
     ) -> IoResult<String> {
         let file_type = infer::get(&self.data);
@@ -166,6 +168,10 @@ impl Paste {
         }
         path.set_file_name(file_name);
         path.set_extension(extension);
+        if let Some(header_filename) = header_filename {
+            file_name = header_filename;
+            path.set_file_name(file_name);
+        }
         let file_name = path
             .file_name()
             .map(|v| v.to_string_lossy())
@@ -235,7 +241,7 @@ impl Paste {
                     .to_string());
             }
         }
-        Ok(self.store_file(file_name, expiry_date, &config)?)
+        Ok(self.store_file(file_name, expiry_date, None, &config)?)
     }
 
     /// Writes an URL to a file in upload directory.
@@ -295,7 +301,7 @@ mod tests {
             data: vec![65, 66, 67],
             type_: PasteType::File,
         };
-        let file_name = paste.store_file("test.txt", None, &config)?;
+        let file_name = paste.store_file("test.txt", None, None, &config)?;
         assert_eq!("ABC", fs::read_to_string(&file_name)?);
         assert_eq!(
             Some("txt"),
@@ -315,7 +321,7 @@ mod tests {
             data: vec![116, 101, 115, 115, 117, 115],
             type_: PasteType::File,
         };
-        let file_name = paste.store_file("foo.tar.gz", None, &config)?;
+        let file_name = paste.store_file("foo.tar.gz", None, None, &config)?;
         assert_eq!("tessus", fs::read_to_string(&file_name)?);
         assert!(file_name.ends_with(".tar.gz"));
         assert!(file_name.starts_with("foo."));
@@ -331,7 +337,7 @@ mod tests {
             data: vec![116, 101, 115, 115, 117, 115],
             type_: PasteType::File,
         };
-        let file_name = paste.store_file(".foo.tar.gz", None, &config)?;
+        let file_name = paste.store_file(".foo.tar.gz", None, None, &config)?;
         assert_eq!("tessus", fs::read_to_string(&file_name)?);
         assert!(file_name.ends_with(".tar.gz"));
         assert!(file_name.starts_with(".foo."));
@@ -347,7 +353,7 @@ mod tests {
             data: vec![116, 101, 115, 115, 117, 115],
             type_: PasteType::File,
         };
-        let file_name = paste.store_file("foo.tar.gz", None, &config)?;
+        let file_name = paste.store_file("foo.tar.gz", None, None, &config)?;
         assert_eq!("tessus", fs::read_to_string(&file_name)?);
         assert!(file_name.ends_with(".tar.gz"));
         fs::remove_file(file_name)?;
@@ -358,7 +364,7 @@ mod tests {
             data: vec![120, 121, 122],
             type_: PasteType::File,
         };
-        let file_name = paste.store_file(".foo", None, &config)?;
+        let file_name = paste.store_file(".foo", None, None, &config)?;
         assert_eq!("xyz", fs::read_to_string(&file_name)?);
         assert_eq!(".foo.txt", file_name);
         fs::remove_file(file_name)?;
@@ -373,7 +379,7 @@ mod tests {
             data: vec![120, 121, 122],
             type_: PasteType::File,
         };
-        let file_name = paste.store_file("random", None, &config)?;
+        let file_name = paste.store_file("random", None, None, &config)?;
         assert_eq!("xyz", fs::read_to_string(&file_name)?);
         assert_eq!(
             Some("bin"),
@@ -383,6 +389,46 @@ mod tests {
         );
         fs::remove_file(file_name)?;
 
+        config.paste.random_url = Some(RandomURLConfig {
+            length: Some(4),
+            type_: RandomURLType::Alphanumeric,
+            suffix_mode: Some(true),
+            ..RandomURLConfig::default()
+        });
+        let paste = Paste {
+            data: vec![116, 101, 115, 115, 117, 115],
+            type_: PasteType::File,
+        };
+        let file_name = paste.store_file(
+            "filename.txt",
+            None,
+            Some("fn_from_header.txt".to_string()),
+            &config,
+        )?;
+        assert_eq!("tessus", fs::read_to_string(&file_name)?);
+        assert_eq!("fn_from_header.txt", file_name);
+        fs::remove_file(file_name)?;
+
+        config.paste.random_url = Some(RandomURLConfig {
+            length: Some(4),
+            type_: RandomURLType::Alphanumeric,
+            suffix_mode: Some(true),
+            ..RandomURLConfig::default()
+        });
+        let paste = Paste {
+            data: vec![116, 101, 115, 115, 117, 115],
+            type_: PasteType::File,
+        };
+        let file_name = paste.store_file(
+            "filename.txt",
+            None,
+            Some("fn_from_header".to_string()),
+            &config,
+        )?;
+        assert_eq!("tessus", fs::read_to_string(&file_name)?);
+        assert_eq!("fn_from_header", file_name);
+        fs::remove_file(file_name)?;
+
         for paste_type in &[PasteType::Url, PasteType::Oneshot] {
             fs::create_dir_all(paste_type.get_path(&config.server.upload_path))?;
         }
@@ -393,7 +439,7 @@ mod tests {
             type_: PasteType::Oneshot,
         };
         let expiry_date = util::get_system_time()?.as_millis() + 100;
-        let file_name = paste.store_file("test.file", Some(expiry_date), &config)?;
+        let file_name = paste.store_file("test.file", Some(expiry_date), None, &config)?;
         let file_path = PasteType::Oneshot
             .get_path(&config.server.upload_path)
             .join(format!("{file_name}.{expiry_date}"));

src/server.rs

@@ -225,6 +225,7 @@ async fn upload(
     }
     let mut urls: Vec<String> = Vec::new();
     while let Some(item) = payload.next().await {
+        let header_filename = header::parse_header_filename(request.headers())?;
         let mut field = item?;
         let content = ContentDisposition::from(field.content_disposition().clone());
         if let Ok(paste_type) = PasteType::try_from(&content) {
@@ -274,7 +275,12 @@ async fn upload(
                     let config = config
                         .read()
                         .map_err(|_| error::ErrorInternalServerError("cannot acquire config"))?;
-                    paste.store_file(content.get_file_name()?, expiry_date, &config)?
+                    paste.store_file(
+                        content.get_file_name()?,
+                        expiry_date,
+                        header_filename,
+                        &config,
+                    )?
                 }
                 PasteType::RemoteFile => {
                     paste
@@ -853,6 +859,56 @@ mod tests {
         Ok(())
     }
 
+    #[actix_web::test]
+    async fn test_upload_file_override_filename() -> Result<(), Error> {
+        let mut config = Config::default();
+        config.server.upload_path = env::current_dir()?;
+
+        let app = test::init_service(
+            App::new()
+                .app_data(Data::new(RwLock::new(config)))
+                .app_data(Data::new(Client::default()))
+                .configure(configure_routes),
+        )
+        .await;
+
+        let file_name = "test_file.txt";
+        let header_filename = "fn_from_header.txt";
+        let timestamp = util::get_system_time()?.as_secs().to_string();
+        let response = test::call_service(
+            &app,
+            get_multipart_request(&timestamp, "file", file_name)
+                .insert_header((
+                    header::HeaderName::from_static("filename"),
+                    header::HeaderValue::from_static("fn_from_header.txt"),
+                ))
+                .to_request(),
+        )
+        .await;
+        assert_eq!(StatusCode::OK, response.status());
+        assert_body(
+            response.into_body(),
+            &format!("http://localhost:8080/{header_filename}\n"),
+        )
+        .await?;
+
+        let serve_request = TestRequest::get()
+            .uri(&format!("/{header_filename}"))
+            .to_request();
+        let response = test::call_service(&app, serve_request).await;
+        assert_eq!(StatusCode::OK, response.status());
+        assert_body(response.into_body(), &timestamp).await?;
+
+        fs::remove_file(header_filename)?;
+        let serve_request = TestRequest::get()
+            .uri(&format!("/{header_filename}"))
+            .to_request();
+        let response = test::call_service(&app, serve_request).await;
+        assert_eq!(StatusCode::NOT_FOUND, response.status());
+
+        Ok(())
+    }
+
     #[actix_web::test]
     #[allow(deprecated)]
     async fn test_upload_duplicate_file() -> Result<(), Error> {