Split into multiple files

src/error.rs

@@ -0,0 +1,64 @@
+use actix_web::{http::StatusCode, HttpResponse, ResponseError};
+
+#[derive(Debug, thiserror::Error)]
+pub enum UploadError {
+    #[error("Invalid content type provided, {0}")]
+    ContentType(mime::Mime),
+
+    #[error("Couln't upload file, {0}")]
+    Upload(String),
+
+    #[error("Couldn't save file, {0}")]
+    Save(#[from] actix_fs::Error),
+
+    #[error("Error in DB, {0}")]
+    Db(#[from] sled::Error),
+
+    #[error("Error parsing string, {0}")]
+    ParseString(#[from] std::string::FromUtf8Error),
+
+    #[error("Error processing image, {0}")]
+    Image(#[from] image::error::ImageError),
+
+    #[error("Panic in blocking operation")]
+    Canceled,
+
+    #[error("No files present in upload")]
+    NoFiles,
+
+    #[error("Uploaded image could not be served, extension is missing")]
+    MissingExtension,
+}
+
+impl From<actix_form_data::Error> for UploadError {
+    fn from(e: actix_form_data::Error) -> Self {
+        UploadError::Upload(e.to_string())
+    }
+}
+
+impl<T> From<actix_web::error::BlockingError<T>> for UploadError
+where
+    T: Into<UploadError> + std::fmt::Debug,
+{
+    fn from(e: actix_web::error::BlockingError<T>) -> Self {
+        match e {
+            actix_web::error::BlockingError::Error(e) => e.into(),
+            _ => UploadError::Canceled,
+        }
+    }
+}
+
+impl ResponseError for UploadError {
+    fn status_code(&self) -> StatusCode {
+        match self {
+            UploadError::NoFiles | UploadError::ContentType(_) | UploadError::Upload(_) => {
+                StatusCode::BAD_REQUEST
+            }
+            _ => StatusCode::INTERNAL_SERVER_ERROR,
+        }
+    }
+
+    fn error_response(&self) -> HttpResponse {
+        HttpResponse::build(self.status_code()).json(serde_json::json!({ "msg": self.to_string() }))
+    }
+}

src/main.rs

@@ -1,17 +1,18 @@
 use actix_form_data::{Field, Form, Value};
 use actix_web::{
     guard,
-    http::{
+    http::header::{CacheControl, CacheDirective},
-        header::{CacheControl, CacheDirective},
-        StatusCode,
-    },
     middleware::Logger,
-    web, App, HttpResponse, HttpServer, ResponseError,
+    web, App, HttpResponse, HttpServer,
 };
-use futures::stream::{Stream, StreamExt, TryStreamExt};
+use futures::stream::{Stream, TryStreamExt};
 use log::{error, info};
-use sha2::Digest;
+use std::path::PathBuf;
-use std::{path::PathBuf, pin::Pin, sync::Arc};
+
+mod error;
+mod upload_manager;
+
+use self::{error::UploadError, upload_manager::UploadManager};
 
 const ACCEPTED_MIMES: &[mime::Mime] = &[
     mime::IMAGE_BMP,
@@ -23,216 +24,34 @@ const ACCEPTED_MIMES: &[mime::Mime] = &[
 const MEGABYTES: usize = 1024 * 1024;
 const HOURS: u32 = 60 * 60;
 
-#[derive(Clone)]
+// Try writing to a file
-struct UploadManager {
+async fn safe_save_file(path: PathBuf, bytes: bytes::Bytes) -> Result<(), UploadError> {
-    inner: Arc<UploadManagerInner>,
+    if let Some(path) = path.parent() {
-}
+        // create the directory for the file
-
+        actix_fs::create_dir_all(path.to_owned()).await?;
-struct UploadManagerInner {
-    hasher: sha2::Sha256,
-    base_dir: PathBuf,
-    db: sled::Db,
-}
-
-#[derive(Debug, thiserror::Error)]
-enum UploadError {
-    #[error("Invalid content type provided, {0}")]
-    ContentType(mime::Mime),
-
-    #[error("Couln't upload file, {0}")]
-    Upload(String),
-
-    #[error("Couldn't save file, {0}")]
-    Save(#[from] actix_fs::Error),
-
-    #[error("Error in DB, {0}")]
-    Db(#[from] sled::Error),
-
-    #[error("Error parsing string, {0}")]
-    ParseString(#[from] std::string::FromUtf8Error),
-
-    #[error("Error processing image, {0}")]
-    Image(#[from] image::error::ImageError),
-
-    #[error("Panic in blocking operation")]
-    Canceled,
-
-    #[error("No files present in upload")]
-    NoFiles,
-
-    #[error("Uploaded image could not be served, extension is missing")]
-    MissingExtension,
-}
-
-impl From<actix_form_data::Error> for UploadError {
-    fn from(e: actix_form_data::Error) -> Self {
-        UploadError::Upload(e.to_string())
-    }
-}
-
-impl<T> From<actix_web::error::BlockingError<T>> for UploadError
-where
-    T: Into<UploadError> + std::fmt::Debug,
-{
-    fn from(e: actix_web::error::BlockingError<T>) -> Self {
-        match e {
-            actix_web::error::BlockingError::Error(e) => e.into(),
-            _ => UploadError::Canceled,
-        }
-    }
-}
-
-impl ResponseError for UploadError {
-    fn status_code(&self) -> StatusCode {
-        match self {
-            UploadError::ContentType(_) | UploadError::Upload(_) => StatusCode::BAD_REQUEST,
-            _ => StatusCode::INTERNAL_SERVER_ERROR,
-        }
     }
 
-    fn error_response(&self) -> HttpResponse {
+    // Only write the file if it doesn't already exist
-        HttpResponse::build(self.status_code()).json(serde_json::json!({ "msg": self.to_string() }))
+    if let Err(e) = actix_fs::metadata(path.clone()).await {
+        if e.kind() != Some(std::io::ErrorKind::NotFound) {
+            return Err(e.into());
         }
-}
+    } else {
-
+        return Ok(());
-type UploadStream = Pin<Box<dyn Stream<Item = Result<bytes::Bytes, actix_form_data::Error>>>>;
-
-enum Dup {
-    Exists,
-    New,
-}
-
-impl Dup {
-    fn exists(&self) -> bool {
-        match self {
-            Dup::Exists => true,
-            _ => false,
-        }
-    }
-}
-
-impl UploadManager {
-    async fn new(mut root_dir: PathBuf) -> Result<Self, UploadError> {
-        let mut sled_dir = root_dir.clone();
-        sled_dir.push("db");
-        let db = sled::open(sled_dir)?;
-
-        root_dir.push("files");
-
-        actix_fs::create_dir_all(root_dir.clone()).await?;
-
-        Ok(UploadManager {
-            inner: Arc::new(UploadManagerInner {
-                hasher: sha2::Sha256::new(),
-                base_dir: root_dir,
-                db,
-            }),
-        })
-    }
-
-    async fn upload(
-        &self,
-        _filename: String,
-        content_type: mime::Mime,
-        mut stream: UploadStream,
-    ) -> Result<Option<PathBuf>, UploadError> {
-        if ACCEPTED_MIMES.iter().all(|valid| *valid != content_type) {
-            return Err(UploadError::ContentType(content_type));
-        }
-
-        let mut bytes = bytes::BytesMut::new();
-
-        while let Some(res) = stream.next().await {
-            bytes.extend(res?);
-        }
-
-        let bytes = bytes.freeze();
-
-        let hash = self.hash(bytes.clone()).await?;
-
-        let (dup, path) = self.check_duplicate(hash, content_type).await?;
-
-        if dup.exists() {
-            return Ok(Some(path));
     }
 
+    // Open the file for writing
     let file = actix_fs::file::create(path.clone()).await?;
 
+    // try writing
     if let Err(e) = actix_fs::file::write(file, bytes).await {
-            error!("Error saving file, {}", e);
+        error!("Error writing file, {}", e);
+        // remove file if writing failed before completion
         actix_fs::remove_file(path).await?;
         return Err(e.into());
     }
 
-        Ok(Some(path))
+    Ok(())
-    }
-
-    async fn hash(&self, bytes: bytes::Bytes) -> Result<Vec<u8>, UploadError> {
-        let mut hasher = self.inner.hasher.clone();
-        let hash = web::block(move || {
-            hasher.input(&bytes);
-            Ok(hasher.result().to_vec()) as Result<_, UploadError>
-        })
-        .await?;
-
-        Ok(hash)
-    }
-
-    async fn check_duplicate(
-        &self,
-        hash: Vec<u8>,
-        content_type: mime::Mime,
-    ) -> Result<(Dup, PathBuf), UploadError> {
-        let mut path = self.inner.base_dir.clone();
-        let db = self.inner.db.clone();
-
-        let filename = self.next_file(content_type).await?;
-
-        let filename2 = filename.clone();
-        let res = web::block(move || {
-            db.compare_and_swap(hash, None as Option<sled::IVec>, Some(filename2.as_bytes()))
-        })
-        .await?;
-
-        if let Err(sled::CompareAndSwapError {
-            current: Some(ivec),
-            ..
-        }) = res
-        {
-            let name = String::from_utf8(ivec.to_vec())?;
-            path.push(name);
-
-            return Ok((Dup::Exists, path));
-        }
-
-        path.push(filename);
-
-        Ok((Dup::New, path))
-    }
-
-    async fn next_file(&self, content_type: mime::Mime) -> Result<String, UploadError> {
-        let base_dir = self.inner.base_dir.clone();
-        use rand::distributions::{Alphanumeric, Distribution};
-        let mut limit: usize = 10;
-        let rng = rand::thread_rng();
-        loop {
-            let mut path = base_dir.clone();
-            let s: String = Alphanumeric.sample_iter(rng).take(limit).collect();
-
-            let filename = format!("{}{}", s, to_ext(content_type.clone()));
-
-            path.push(filename.clone());
-
-            if let Err(e) = actix_fs::metadata(path).await {
-                if e.kind() == Some(std::io::ErrorKind::NotFound) {
-                    return Ok(filename);
-                }
-                return Err(e.into());
-            }
-
-            limit += 1;
-        }
-    }
 }
 
 fn to_ext(mime: mime::Mime) -> &'static str {
@@ -256,6 +75,7 @@ fn from_ext(ext: std::ffi::OsString) -> mime::Mime {
     }
 }
 
+/// Handle responding to succesful uploads
 async fn upload(value: Value) -> Result<HttpResponse, UploadError> {
     let images = value
         .map()
@@ -279,12 +99,14 @@ async fn upload(value: Value) -> Result<HttpResponse, UploadError> {
     Ok(HttpResponse::Created().json(serde_json::json!({ "msg": "ok", "files": files })))
 }
 
+/// Serve original files
 async fn serve(
     manager: web::Data<UploadManager>,
     filename: web::Path<String>,
 ) -> Result<HttpResponse, UploadError> {
-    let mut path = manager.inner.base_dir.clone();
+    let mut path = manager.image_dir();
     path.push(filename.into_inner());
+
     let ext = path
         .extension()
         .ok_or(UploadError::MissingExtension)?
@@ -296,47 +118,54 @@ async fn serve(
     Ok(srv_response(stream, ext))
 }
 
+/// Serve resized files
 async fn serve_resized(
     manager: web::Data<UploadManager>,
     filename: web::Path<(u32, String)>,
 ) -> Result<HttpResponse, UploadError> {
     use image::GenericImageView;
 
-    let mut path = manager.inner.base_dir.clone();
+    let mut path = manager.image_dir();
 
     let (size, name) = filename.into_inner();
     path.push(size.to_string());
     path.push(name.clone());
+
     let ext = path
         .extension()
         .ok_or(UploadError::MissingExtension)?
         .to_owned();
     let ext = from_ext(ext);
 
+    // If the thumbnail doesn't exist, we need to create it
     if let Err(e) = actix_fs::metadata(path.clone()).await {
         if e.kind() != Some(std::io::ErrorKind::NotFound) {
             error!("Error looking up thumbnail, {}", e);
             return Err(e.into());
         }
 
-        let mut original_path = manager.inner.base_dir.clone();
+        let mut original_path = manager.image_dir();
         original_path.push(name.clone());
 
+        // Read the image file & produce a DynamicImage
+        //
+        // Drop bytes so we don't keep it around in memory longer than we need to
         let (img, format) = {
             let bytes = actix_fs::read(original_path.clone()).await?;
             let format = image::guess_format(&bytes)?;
-            let img = image::load_from_memory(&bytes)?;
+            let img = web::block(move || image::load_from_memory(&bytes)).await?;
 
             (img, format)
         };
 
-        if !img.in_bounds(size, size) {
         // return original image if resize target is larger
+        if !img.in_bounds(size, size) {
             drop(img);
             let stream = actix_fs::read_to_stream(original_path).await?;
             return Ok(srv_response(stream, ext));
         }
 
+        // perform thumbnail operation in a blocking thread
         let img_bytes: bytes::Bytes = web::block(move || {
             let mut bytes = std::io::Cursor::new(vec![]);
             img.thumbnail(size, size).write_to(&mut bytes, format)?;
@@ -347,21 +176,10 @@ async fn serve_resized(
         let path2 = path.clone();
         let img_bytes2 = img_bytes.clone();
 
+        // Save the file in another task, we want to return the thumbnail now
         actix_rt::spawn(async move {
-            if let Some(path) = path2.parent() {
+            if let Err(e) = safe_save_file(path2, img_bytes2).await {
-                if let Err(e) = actix_fs::create_dir_all(path.to_owned()).await {
+                error!("Error saving file, {}", e);
-                    error!("Couldn't create directory for thumbnail, {}", e);
-                }
-            }
-
-            if let Err(e) = actix_fs::metadata(path2.clone()).await {
-                if e.kind() == Some(std::io::ErrorKind::NotFound) {
-                    if let Err(e) = actix_fs::write(path2, img_bytes2).await {
-                        error!("Error saving image, {}", e);
-                    }
-                } else {
-                    error!("Error checking image, {}", e);
-                }
             }
         });
 
@@ -378,6 +196,7 @@ async fn serve_resized(
     Ok(srv_response(stream, ext))
 }
 
+// A helper method to produce responses with proper cache headers
 fn srv_response<S, E>(stream: S, ext: mime::Mime) -> HttpResponse
 where
     S: Stream<Item = Result<bytes::Bytes, E>> + Unpin + 'static,
@@ -398,6 +217,9 @@ async fn main() -> Result<(), anyhow::Error> {
     env_logger::init();
     let manager = UploadManager::new("data/".to_string().into()).await?;
 
+    // Create a new Multipart Form validator
+    //
+    // This form is expecting a single array field, 'images' with at most 10 files in it
     let manager2 = manager.clone();
     let form = Form::new()
         .max_files(10)

src/upload_manager.rs

@@ -0,0 +1,173 @@
+use crate::{error::UploadError, safe_save_file, to_ext, ACCEPTED_MIMES};
+use actix_web::web;
+use futures::stream::{Stream, StreamExt};
+use sha2::Digest;
+use std::{path::PathBuf, pin::Pin, sync::Arc};
+
+#[derive(Clone)]
+pub struct UploadManager {
+    inner: Arc<UploadManagerInner>,
+}
+
+struct UploadManagerInner {
+    hasher: sha2::Sha256,
+    image_dir: PathBuf,
+    db: sled::Db,
+}
+
+type UploadStream = Pin<Box<dyn Stream<Item = Result<bytes::Bytes, actix_form_data::Error>>>>;
+
+enum Dup {
+    Exists,
+    New,
+}
+
+impl Dup {
+    fn exists(&self) -> bool {
+        match self {
+            Dup::Exists => true,
+            _ => false,
+        }
+    }
+}
+
+impl UploadManager {
+    /// Get the image directory
+    pub(crate) fn image_dir(&self) -> PathBuf {
+        self.inner.image_dir.clone()
+    }
+
+    /// Create a new UploadManager
+    pub(crate) async fn new(mut root_dir: PathBuf) -> Result<Self, UploadError> {
+        let mut sled_dir = root_dir.clone();
+        sled_dir.push("db");
+        // sled automatically creates it's own directories
+        //
+        // This is technically a blocking operation but it's fine because it happens before we
+        // start handling requests
+        let db = sled::open(sled_dir)?;
+
+        root_dir.push("files");
+
+        // Ensure file dir exists
+        actix_fs::create_dir_all(root_dir.clone()).await?;
+
+        Ok(UploadManager {
+            inner: Arc::new(UploadManagerInner {
+                hasher: sha2::Sha256::new(),
+                image_dir: root_dir,
+                db,
+            }),
+        })
+    }
+
+    /// Upload the file, discarding bytes if it's already present, or saving if it's new
+    pub(crate) async fn upload(
+        &self,
+        _filename: String,
+        content_type: mime::Mime,
+        mut stream: UploadStream,
+    ) -> Result<Option<PathBuf>, UploadError> {
+        if ACCEPTED_MIMES.iter().all(|valid| *valid != content_type) {
+            return Err(UploadError::ContentType(content_type));
+        }
+
+        // -- READ IN BYTES FROM CLIENT --
+        let mut bytes = bytes::BytesMut::new();
+
+        while let Some(res) = stream.next().await {
+            bytes.extend(res?);
+        }
+
+        let bytes = bytes.freeze();
+
+        // -- DUPLICATE CHECKS --
+
+        // Cloning bytes is fine because it's actually a pointer
+        let hash = self.hash(bytes.clone()).await?;
+
+        let (dup, path) = self.check_duplicate(hash, content_type).await?;
+
+        // bail early with path to existing file if this is a duplicate
+        if dup.exists() {
+            return Ok(Some(path));
+        }
+
+        // TODO: validate image before saving
+
+        // -- WRITE NEW FILE --
+        safe_save_file(path.clone(), bytes).await?;
+
+        Ok(Some(path))
+    }
+
+    // produce a sh256sum of the uploaded file
+    async fn hash(&self, bytes: bytes::Bytes) -> Result<Vec<u8>, UploadError> {
+        let mut hasher = self.inner.hasher.clone();
+        let hash = web::block(move || {
+            hasher.input(&bytes);
+            Ok(hasher.result().to_vec()) as Result<_, UploadError>
+        })
+        .await?;
+
+        Ok(hash)
+    }
+
+    // check for an already-uploaded image with this hash, returning the path to the target file
+    async fn check_duplicate(
+        &self,
+        hash: Vec<u8>,
+        content_type: mime::Mime,
+    ) -> Result<(Dup, PathBuf), UploadError> {
+        let mut path = self.inner.image_dir.clone();
+        let db = self.inner.db.clone();
+
+        let filename = self.next_file(content_type).await?;
+
+        let filename2 = filename.clone();
+        let res = web::block(move || {
+            db.compare_and_swap(hash, None as Option<sled::IVec>, Some(filename2.as_bytes()))
+        })
+        .await?;
+
+        if let Err(sled::CompareAndSwapError {
+            current: Some(ivec),
+            ..
+        }) = res
+        {
+            let name = String::from_utf8(ivec.to_vec())?;
+            path.push(name);
+
+            return Ok((Dup::Exists, path));
+        }
+
+        path.push(filename);
+
+        Ok((Dup::New, path))
+    }
+
+    // generate a short filename that isn't already in-use
+    async fn next_file(&self, content_type: mime::Mime) -> Result<String, UploadError> {
+        let image_dir = self.inner.image_dir.clone();
+        use rand::distributions::{Alphanumeric, Distribution};
+        let mut limit: usize = 10;
+        let rng = rand::thread_rng();
+        loop {
+            let mut path = image_dir.clone();
+            let s: String = Alphanumeric.sample_iter(rng).take(limit).collect();
+
+            let filename = format!("{}{}", s, to_ext(content_type.clone()));
+
+            path.push(filename.clone());
+
+            if let Err(e) = actix_fs::metadata(path).await {
+                if e.kind() == Some(std::io::ErrorKind::NotFound) {
+                    return Ok(filename);
+                }
+                return Err(e.into());
+            }
+
+            limit += 1;
+        }
+    }
+}