From 459db42a88c12d55b6ca3f898514187e4d978215 Mon Sep 17 00:00:00 2001 From: asonix 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, } impl Config { @@ -21,4 +28,46 @@ impl Config { pub(crate) fn data_dir(&self) -> PathBuf { self.path.clone() } + + pub(crate) fn format(&self) -> Option { + 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 { + 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> 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, 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 { + pub(crate) async fn new( + mut root_dir: PathBuf, + format: Option, + ) -> Result { 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 --