From 459db42a88c12d55b6ca3f898514187e4d978215 Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Sat, 6 Jun 2020 20:44:26 -0500
Subject: [PATCH] Add optional image format coersion

---
 src/config.rs         | 51 ++++++++++++++++++++++++++++++++++++++++++-
 src/error.rs          |  3 +++
 src/main.rs           |  2 +-
 src/upload_manager.rs | 49 ++++++++++++++++++++++++++++++++++-------
 4 files changed, 95 insertions(+), 10 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 420f9e5..1a60f49 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -1,7 +1,7 @@
 use std::{net::SocketAddr, path::PathBuf};
 
 #[derive(structopt::StructOpt)]
-pub struct Config {
+pub(crate) struct Config {
     #[structopt(
         short,
         long,
@@ -11,6 +11,13 @@ pub struct Config {
 
     #[structopt(short, long, help = "The path to the data directory, e.g. data/")]
     path: PathBuf,
+
+    #[structopt(
+        short,
+        long,
+        help = "An image format to convert all uploaded files into, supports 'jpg' and 'png'"
+    )]
+    format: Option<Format>,
 }
 
 impl Config {
@@ -21,4 +28,46 @@ impl Config {
     pub(crate) fn data_dir(&self) -> PathBuf {
         self.path.clone()
     }
+
+    pub(crate) fn format(&self) -> Option<Format> {
+        self.format.clone()
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("Invalid format supplied, {0}")]
+pub(crate) struct FormatError(String);
+
+#[derive(Clone, Debug)]
+pub(crate) enum Format {
+    Jpeg,
+    Png,
+}
+
+impl Format {
+    pub(crate) fn to_image_format(&self) -> image::ImageFormat {
+        match self {
+            Format::Jpeg => image::ImageFormat::Jpeg,
+            Format::Png => image::ImageFormat::Png,
+        }
+    }
+
+    pub(crate) fn to_mime(&self) -> mime::Mime {
+        match self {
+            Format::Jpeg => mime::IMAGE_JPEG,
+            Format::Png => mime::IMAGE_PNG,
+        }
+    }
+}
+
+impl std::str::FromStr for Format {
+    type Err = FormatError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "png" => Ok(Format::Png),
+            "jpg" => Ok(Format::Jpeg),
+            other => Err(FormatError(other.to_string())),
+        }
+    }
 }
diff --git a/src/error.rs b/src/error.rs
index 58c5216..46060d1 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -40,6 +40,9 @@ pub enum UploadError {
 
     #[error("Provided token did not match expected token")]
     InvalidToken,
+
+    #[error("Uploaded content could not be validated as an image")]
+    InvalidImage(image::error::ImageError),
 }
 
 impl From<sled::transaction::TransactionError<UploadError>> for UploadError {
diff --git a/src/main.rs b/src/main.rs
index 038107f..7a06b0b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -239,7 +239,7 @@ async fn main() -> Result<(), anyhow::Error> {
     let config = Config::from_args();
     std::env::set_var("RUST_LOG", "info");
     env_logger::init();
-    let manager = UploadManager::new(config.data_dir()).await?;
+    let manager = UploadManager::new(config.data_dir(), config.format()).await?;
 
     // Create a new Multipart Form validator
     //
diff --git a/src/upload_manager.rs b/src/upload_manager.rs
index 1481f09..afc6f6f 100644
--- a/src/upload_manager.rs
+++ b/src/upload_manager.rs
@@ -1,4 +1,4 @@
-use crate::{error::UploadError, safe_save_file, to_ext, ACCEPTED_MIMES};
+use crate::{config::Format, error::UploadError, safe_save_file, to_ext, ACCEPTED_MIMES};
 use actix_web::web;
 use futures::stream::{Stream, StreamExt};
 use log::{error, warn};
@@ -11,6 +11,7 @@ pub struct UploadManager {
 }
 
 struct UploadManagerInner {
+    format: Option<Format>,
     hasher: sha2::Sha256,
     image_dir: PathBuf,
     alias_tree: sled::Tree,
@@ -40,7 +41,10 @@ impl UploadManager {
     }
 
     /// Create a new UploadManager
-    pub(crate) async fn new(mut root_dir: PathBuf) -> Result<Self, UploadError> {
+    pub(crate) async fn new(
+        mut root_dir: PathBuf,
+        format: Option<Format>,
+    ) -> Result<Self, UploadError> {
         let mut sled_dir = root_dir.clone();
         sled_dir.push("db");
         // sled automatically creates it's own directories
@@ -53,6 +57,7 @@ impl UploadManager {
 
         Ok(UploadManager {
             inner: Arc::new(UploadManagerInner {
+                format,
                 hasher: sha2::Sha256::new(),
                 image_dir: root_dir,
                 alias_tree: db.open_tree("alias")?,
@@ -181,14 +186,42 @@ impl UploadManager {
             return Err(UploadError::ContentType(content_type));
         }
 
-        // -- READ IN BYTES FROM CLIENT --
-        let mut bytes = bytes::BytesMut::new();
+        let (img, format) = {
+            // -- READ IN BYTES FROM CLIENT --
+            let mut bytes = bytes::BytesMut::new();
 
-        while let Some(res) = stream.next().await {
-            bytes.extend(res?);
-        }
+            while let Some(res) = stream.next().await {
+                bytes.extend(res?);
+            }
 
-        let bytes = bytes.freeze();
+            let bytes = bytes.freeze();
+
+            // -- VALIDATE IMAGE --
+            let format = image::guess_format(&bytes).map_err(UploadError::InvalidImage)?;
+            let img = image::load_from_memory(&bytes).map_err(UploadError::InvalidImage)?;
+
+            (img, format)
+        };
+
+        let format = self
+            .inner
+            .format
+            .as_ref()
+            .map(|f| f.to_image_format())
+            .unwrap_or(format);
+
+        let content_type = self
+            .inner
+            .format
+            .as_ref()
+            .map(|f| f.to_mime())
+            .unwrap_or(content_type);
+
+        let bytes: bytes::Bytes = {
+            let mut bytes = std::io::Cursor::new(vec![]);
+            img.write_to(&mut bytes, format)?;
+            bytes::Bytes::from(bytes.into_inner())
+        };
 
         // -- DUPLICATE CHECKS --