Add import API for filename preservation

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "pict-rs"
 description = "A simple image hosting service"
-version = "0.1.2"
+version = "0.1.3"
 authors = ["asonix <asonix@asonix.dog>"]
 license = "AGPL-3.0"
 readme = "README.md"

README.md

@@ -4,21 +4,26 @@ _a simple image hosting service_
 ## Usage
 ### Running
 ```
-pict-rs 0.1.0
+pict-rs 0.1.3
 
 USAGE:
-    pict-rs [OPTIONS] --addr <addr> --path <path>
+    pict-rs [FLAGS] [OPTIONS] --path <path>
 
 FLAGS:
-    -h, --help       Prints help information
+    -h, --help                     Prints help information
-    -V, --version    Prints version information
+    -s, --skip-validate-imports    Whether to skip validating images uploaded via the internal import API
+    -V, --version                  Prints version information
 
 OPTIONS:
-    -a, --addr <addr>                 The address and port the server binds to, e.g. 127.0.0.1:80
+    -a, --addr <addr>                      The address and port the server binds to. Default: 0.0.0.0:8080 [env:
-    -f, --format <format>             An image format to convert all uploaded files into, supports 'jpg' and 'png'
+                                           PICTRS_ADDR=]  [default: 0.0.0.0:8080]
-    -p, --path <path>                 The path to the data directory, e.g. data/
+    -f, --format <format>                  An optional image format to convert all uploaded files into, supports 'jpg'
-    -w, --whitelist <whitelist>...    An optional list of filters to whitelist, supports 'identity', 'thumbnail', and
+                                           and 'png' [env: PICTRS_FORMAT=]
-                                      'blur'
+    -m, --max-file-size <max-file-size>    Specify the maximum allowed uploaded file size (in Megabytes) [env:
+                                           PICTRS_MAX_FILE_SIZE=]  [default: 40]
+    -p, --path <path>                      The path to the data directory, e.g. data/ [env: PICTRS_PATH=]
+    -w, --whitelist <whitelist>...         An optional list of filters to whitelist, supports 'identity', 'thumbnail',
+                                           and 'blur' [env: PICTRS_FILTER_WHITELIST=]
 ```
 
 #### Example:
@@ -80,6 +85,9 @@ pict-rs offers four endpoints:
         "msg": "ok"
     }
     ```
+- `POST /import` for uploading an image while preserving the filename. This should not be exposed to
+    the public internet, as it can cause naming conflicts with saved files. The upload format and
+    response format are the same as the `POST /image` endpoint.
 - `GET /image/download?url=...` Download an image from a remote server, returning the same JSON
     payload as the `POST` endpoint
 - `GET /image/{file}` for getting a full-resolution image. `file` here is the `file` key from the

src/config.rs

@@ -2,6 +2,13 @@ use std::{collections::HashSet, net::SocketAddr, path::PathBuf};
 
 #[derive(Clone, Debug, structopt::StructOpt)]
 pub(crate) struct Config {
+    #[structopt(
+        short,
+        long,
+        help = "Whether to skip validating images uploaded via the internal import API"
+    )]
+    skip_validate_imports: bool,
+
     #[structopt(
         short,
         long,
@@ -12,10 +19,11 @@ pub(crate) struct Config {
     addr: SocketAddr,
 
     #[structopt(
-      short, 
+        short,
-      long, 
+        long,
-      env = "PICTRS_PATH",
+        env = "PICTRS_PATH",
-      help = "The path to the data directory, e.g. data/")]
+        help = "The path to the data directory, e.g. data/"
+    )]
     path: PathBuf,
 
     #[structopt(
@@ -33,6 +41,15 @@ pub(crate) struct Config {
         help = "An optional list of filters to whitelist, supports 'identity', 'thumbnail', and 'blur'"
     )]
     whitelist: Option<Vec<String>>,
+
+    #[structopt(
+        short,
+        long,
+        env = "PICTRS_MAX_FILE_SIZE",
+        help = "Specify the maximum allowed uploaded file size (in Megabytes)",
+        default_value = "40"
+    )]
+    max_file_size: usize,
 }
 
 impl Config {
@@ -53,6 +70,14 @@ impl Config {
             .as_ref()
             .map(|wl| wl.iter().cloned().collect())
     }
+
+    pub(crate) fn validate_imports(&self) -> bool {
+        !self.skip_validate_imports
+    }
+
+    pub(crate) fn max_file_size(&self) -> usize {
+        self.max_file_size
+    }
 }
 
 #[derive(Debug, thiserror::Error)]

src/error.rs

@@ -61,6 +61,9 @@ pub enum UploadError {
 
     #[error("Error converting Path to String")]
     Path,
+
+    #[error("Tried to save an image with an already-taken name")]
+    DuplicateAlias,
 }
 
 impl From<actix_web::client::SendRequestError> for UploadError {
@@ -99,9 +102,10 @@ where
 impl ResponseError for UploadError {
     fn status_code(&self) -> StatusCode {
         match self {
-            UploadError::NoFiles | UploadError::ContentType(_) | UploadError::Upload(_) => {
+            UploadError::DuplicateAlias
-                StatusCode::BAD_REQUEST
+            | UploadError::NoFiles
-            }
+            | UploadError::ContentType(_)
+            | UploadError::Upload(_) => StatusCode::BAD_REQUEST,
             UploadError::MissingAlias | UploadError::MissingFilename => StatusCode::NOT_FOUND,
             UploadError::InvalidToken => StatusCode::FORBIDDEN,
             _ => StatusCode::INTERNAL_SERVER_ERROR,

src/main.rs

@@ -270,7 +270,7 @@ async fn main() -> Result<(), anyhow::Error> {
     let manager2 = manager.clone();
     let form = Form::new()
         .max_files(10)
-        .max_file_size(40 * MEGABYTES)
+        .max_file_size(config.max_file_size() * MEGABYTES)
         .field(
             "images",
             Field::array(Field::file(move |_, _, stream| {
@@ -286,6 +286,32 @@ async fn main() -> Result<(), anyhow::Error> {
             })),
         );
 
+    // Create a new Multipart Form validator for internal imports
+    //
+    // This form is expecting a single array field, 'images' with at most 10 files in it
+    let validate_imports = config.validate_imports();
+    let manager2 = manager.clone();
+    let import_form = Form::new()
+        .max_files(10)
+        .max_file_size(config.max_file_size() * MEGABYTES)
+        .field(
+            "images",
+            Field::array(Field::file(move |filename, content_type, stream| {
+                let manager = manager2.clone();
+
+                async move {
+                    manager
+                        .import(filename, content_type, validate_imports, stream)
+                        .await
+                        .map(|alias| {
+                            let mut path = PathBuf::new();
+                            path.push(alias);
+                            Some(path)
+                        })
+                }
+            })),
+        );
+
     let config2 = config.clone();
     HttpServer::new(move || {
         let client = Client::build()
@@ -316,6 +342,11 @@ async fn main() -> Result<(), anyhow::Error> {
                     )
                     .service(web::resource("/{tail:.*}").route(web::get().to(serve))),
             )
+            .service(
+                web::resource("/import")
+                    .wrap(import_form.clone())
+                    .route(web::post().to(upload)),
+            )
     })
     .bind(config.bind_address())?
     .run()

src/upload_manager.rs

@@ -68,6 +68,7 @@ impl UploadManager {
         })
     }
 
+    /// Store the path to a generated image variant so we can easily clean it up later
     pub(crate) async fn store_variant(&self, path: PathBuf) -> Result<(), UploadError> {
         let filename = path
             .file_name()
@@ -88,6 +89,7 @@ impl UploadManager {
         Ok(())
     }
 
+    /// Delete the alias, and the file & variants if no more aliases exist
     pub(crate) async fn delete(&self, alias: String, token: String) -> Result<(), UploadError> {
         use sled::Transactional;
         let db = self.inner.db.clone();
@@ -159,6 +161,111 @@ impl UploadManager {
         Ok(())
     }
 
+    /// Generate a delete token for an alias
+    pub(crate) async fn delete_token(&self, alias: String) -> Result<String, UploadError> {
+        use rand::distributions::{Alphanumeric, Distribution};
+        let rng = rand::thread_rng();
+        let s: String = Alphanumeric.sample_iter(rng).take(10).collect();
+        let delete_token = s.clone();
+
+        let alias_tree = self.inner.alias_tree.clone();
+        let key = delete_key(&alias);
+        let res = web::block(move || {
+            alias_tree.compare_and_swap(
+                key.as_bytes(),
+                None as Option<sled::IVec>,
+                Some(s.as_bytes()),
+            )
+        })
+        .await?;
+
+        if let Err(sled::CompareAndSwapError {
+            current: Some(ivec),
+            ..
+        }) = res
+        {
+            let s = String::from_utf8(ivec.to_vec())?;
+
+            return Ok(s);
+        }
+
+        Ok(delete_token)
+    }
+
+    /// Upload the file while preserving the filename, optionally validating the uploaded image
+    pub(crate) async fn import<E>(
+        &self,
+        alias: String,
+        content_type: mime::Mime,
+        validate: bool,
+        stream: UploadStream<E>,
+    ) -> Result<String, UploadError>
+    where
+        UploadError: From<E>,
+    {
+        let bytes = read_stream(stream).await?;
+
+        let (bytes, content_type) = if validate {
+            self.validate_image(bytes).await?
+        } else {
+            (bytes, content_type)
+        };
+
+        // -- DUPLICATE CHECKS --
+
+        // Cloning bytes is fine because it's actually a pointer
+        let hash = self.hash(bytes.clone()).await?;
+
+        self.add_existing_alias(&hash, &alias).await?;
+
+        self.save_upload(bytes, hash, content_type).await?;
+
+        // Return alias to file
+        Ok(alias)
+    }
+
+    /// Upload the file, discarding bytes if it's already present, or saving if it's new
+    pub(crate) async fn upload<E>(&self, stream: UploadStream<E>) -> Result<String, UploadError>
+    where
+        UploadError: From<E>,
+    {
+        // -- READ IN BYTES FROM CLIENT --
+        let bytes = read_stream(stream).await?;
+
+        // -- VALIDATE IMAGE --
+        let (bytes, content_type) = self.validate_image(bytes).await?;
+
+        // -- DUPLICATE CHECKS --
+
+        // Cloning bytes is fine because it's actually a pointer
+        let hash = self.hash(bytes.clone()).await?;
+
+        let alias = self.add_alias(&hash, content_type.clone()).await?;
+
+        self.save_upload(bytes, hash, content_type).await?;
+
+        // Return alias to file
+        Ok(alias)
+    }
+
+    /// Fetch the real on-disk filename given an alias
+    pub(crate) async fn from_alias(&self, alias: String) -> Result<String, UploadError> {
+        let tree = self.inner.alias_tree.clone();
+        let hash = web::block(move || tree.get(alias.as_bytes()))
+            .await?
+            .ok_or(UploadError::MissingAlias)?;
+
+        let db = self.inner.db.clone();
+        let filename = web::block(move || db.get(hash))
+            .await?
+            .ok_or(UploadError::MissingFile)?;
+
+        let filename = String::from_utf8(filename.to_vec())?;
+
+        Ok(filename)
+    }
+
+    // Find image variants and remove them from the DB and the disk
     async fn cleanup_files(&self, filename: sled::IVec) -> Result<(), UploadError> {
         let mut path = self.image_dir();
         let fname = String::from_utf8(filename.to_vec())?;
@@ -201,61 +308,41 @@ impl UploadManager {
         Ok(())
     }
 
-    /// Generate a delete token for an alias
+    // check duplicates & store image if new
-    pub(crate) async fn delete_token(&self, alias: String) -> Result<String, UploadError> {
+    async fn save_upload(
-        use rand::distributions::{Alphanumeric, Distribution};
+        &self,
-        let rng = rand::thread_rng();
+        bytes: bytes::Bytes,
-        let s: String = Alphanumeric.sample_iter(rng).take(10).collect();
+        hash: Vec<u8>,
-        let delete_token = s.clone();
+        content_type: mime::Mime,
+    ) -> Result<(), UploadError> {
+        let (dup, name) = self.check_duplicate(hash, content_type).await?;
 
-        let alias_tree = self.inner.alias_tree.clone();
+        // bail early with alias to existing file if this is a duplicate
-        let key = delete_key(&alias);
+        if dup.exists() {
-        let res = web::block(move || {
+            return Ok(());
-            alias_tree.compare_and_swap(
-                key.as_bytes(),
-                None as Option<sled::IVec>,
-                Some(s.as_bytes()),
-            )
-        })
-        .await?;
-
-        if let Err(sled::CompareAndSwapError {
-            current: Some(ivec),
-            ..
-        }) = res
-        {
-            let s = String::from_utf8(ivec.to_vec())?;
-
-            return Ok(s);
         }
 
-        Ok(delete_token)
+        // -- WRITE NEW FILE --
+        let mut real_path = self.image_dir();
+        real_path.push(name);
+
+        safe_save_file(real_path, bytes).await?;
+
+        Ok(())
     }
 
-    /// Upload the file, discarding bytes if it's already present, or saving if it's new
+    // import & export image using the image crate
-    pub(crate) async fn upload<E>(&self, mut stream: UploadStream<E>) -> Result<String, UploadError>
+    async fn validate_image(
-    where
+        &self,
-        UploadError: From<E>,
+        bytes: bytes::Bytes,
-    {
+    ) -> Result<(bytes::Bytes, mime::Mime), UploadError> {
-        let (img, format) = {
+        let (img, format) = web::block(move || {
-            // -- READ IN BYTES FROM CLIENT --
+            let format = image::guess_format(&bytes).map_err(UploadError::InvalidImage)?;
-            let mut bytes = bytes::BytesMut::new();
+            let img = image::load_from_memory(&bytes).map_err(UploadError::InvalidImage)?;
 
-            while let Some(res) = stream.next().await {
+            Ok((img, format)) as Result<(image::DynamicImage, image::ImageFormat), UploadError>
-                bytes.extend(res?);
+        })
-            }
+        .await?;
-
-            let bytes = bytes.freeze();
-
-            // -- VALIDATE IMAGE --
-            web::block(move || {
-                let format = image::guess_format(&bytes).map_err(UploadError::InvalidImage)?;
-                let img = image::load_from_memory(&bytes).map_err(UploadError::InvalidImage)?;
-
-                Ok((img, format)) as Result<(image::DynamicImage, image::ImageFormat), UploadError>
-            })
-            .await?
-        };
 
         let (format, content_type) = self
             .inner
@@ -275,43 +362,7 @@ impl UploadManager {
         })
         .await?;
 
-        // -- DUPLICATE CHECKS --
+        Ok((bytes, content_type))
-
-        // Cloning bytes is fine because it's actually a pointer
-        let hash = self.hash(bytes.clone()).await?;
-
-        let alias = self.add_alias(&hash, content_type.clone()).await?;
-        let (dup, name) = self.check_duplicate(hash, content_type).await?;
-
-        // bail early with alias to existing file if this is a duplicate
-        if dup.exists() {
-            return Ok(alias);
-        }
-
-        // -- WRITE NEW FILE --
-        let mut real_path = self.image_dir();
-        real_path.push(name);
-
-        safe_save_file(real_path, bytes).await?;
-
-        // Return alias to file
-        Ok(alias)
-    }
-
-    pub(crate) async fn from_alias(&self, alias: String) -> Result<String, UploadError> {
-        let tree = self.inner.alias_tree.clone();
-        let hash = web::block(move || tree.get(alias.as_bytes()))
-            .await?
-            .ok_or(UploadError::MissingAlias)?;
-
-        let db = self.inner.db.clone();
-        let filename = web::block(move || db.get(hash))
-            .await?
-            .ok_or(UploadError::MissingFile)?;
-
-        let filename = String::from_utf8(filename.to_vec())?;
-
-        Ok(filename)
     }
 
     // produce a sh256sum of the uploaded file
@@ -387,6 +438,14 @@ impl UploadManager {
         }
     }
 
+    async fn add_existing_alias(&self, hash: &[u8], alias: &str) -> Result<(), UploadError> {
+        self.save_alias(hash, alias).await??;
+
+        self.store_alias(hash, alias).await?;
+
+        Ok(())
+    }
+
     // Add an alias to an existing file
     //
     // This will help if multiple 'users' upload the same file, and one of them wants to delete it
@@ -397,6 +456,16 @@ impl UploadManager {
     ) -> Result<String, UploadError> {
         let alias = self.next_alias(hash, content_type).await?;
 
+        self.store_alias(hash, &alias).await?;
+
+        Ok(alias)
+    }
+
+    // Add a pre-defined alias to an existin file
+    //
+    // DANGER: this can cause BAD BAD BAD conflicts if the same alias is used for multiple files
+    async fn store_alias(&self, hash: &[u8], alias: &str) -> Result<(), UploadError> {
+        let alias = alias.to_string();
         loop {
             let db = self.inner.db.clone();
             let id = web::block(move || db.generate_id()).await?.to_string();
@@ -418,7 +487,7 @@ impl UploadManager {
             }
         }
 
-        Ok(alias)
+        Ok(())
     }
 
     // Generate an alias to the file
@@ -430,26 +499,54 @@ impl UploadManager {
         use rand::distributions::{Alphanumeric, Distribution};
         let mut limit: usize = 10;
         let rng = rand::thread_rng();
-        let hvec = hash.to_vec();
         loop {
             let s: String = Alphanumeric.sample_iter(rng).take(limit).collect();
-            let filename = file_name(s, content_type.clone());
+            let alias = file_name(s, content_type.clone());
 
-            let tree = self.inner.alias_tree.clone();
+            let res = self.save_alias(hash, &alias).await?;
-            let vec = hvec.clone();
-            let filename2 = filename.clone();
-            let res = web::block(move || {
-                tree.compare_and_swap(filename2.as_bytes(), None as Option<sled::IVec>, Some(vec))
-            })
-            .await?;
 
             if res.is_ok() {
-                return Ok(filename);
+                return Ok(alias);
             }
 
             limit += 1;
         }
     }
+
+    // Save an alias to the database
+    async fn save_alias(
+        &self,
+        hash: &[u8],
+        alias: &str,
+    ) -> Result<Result<(), UploadError>, UploadError> {
+        let tree = self.inner.alias_tree.clone();
+        let vec = hash.to_vec();
+        let alias = alias.to_string();
+
+        let res = web::block(move || {
+            tree.compare_and_swap(alias.as_bytes(), None as Option<sled::IVec>, Some(vec))
+        })
+        .await?;
+
+        if res.is_err() {
+            return Ok(Err(UploadError::DuplicateAlias));
+        }
+
+        return Ok(Ok(()));
+    }
+}
+
+async fn read_stream<E>(mut stream: UploadStream<E>) -> Result<bytes::Bytes, UploadError>
+where
+    UploadError: From<E>,
+{
+    let mut bytes = bytes::BytesMut::new();
+
+    while let Some(res) = stream.next().await {
+        bytes.extend(res?);
+    }
+
+    Ok(bytes.freeze())
 }
 
 async fn remove_path(path: sled::IVec) -> Result<(), UploadError> {