diff --git a/Cargo.lock b/Cargo.lock index e263b07..11c2553 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1381,6 +1381,7 @@ dependencies = [ "bytes", "env_logger", "futures", + "gif", "image", "log", "mime", diff --git a/Cargo.toml b/Cargo.toml index cab2d84..b5f4bbf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,7 @@ anyhow = "1.0" bytes = "0.5" env_logger = "0.7" futures = "0.3.4" +gif = "0.10.3" image = "0.23.4" log = "0.4" mime = "0.3.1" diff --git a/client-examples/earth.gif b/client-examples/earth.gif new file mode 100644 index 0000000..9df6477 Binary files /dev/null and b/client-examples/earth.gif differ diff --git a/client-examples/python/client.py b/client-examples/python/client.py index 7e3cf08..6d46abd 100755 --- a/client-examples/python/client.py +++ b/client-examples/python/client.py @@ -7,7 +7,8 @@ import asyncio import aiofiles import aiohttp -file_name = '../test.png' +png_name = '../test.png' +gif_name = '../earth.gif' url = 'http://localhost:8080/image' async def file_sender(file_name=None): @@ -21,9 +22,10 @@ async def file_sender(file_name=None): async def req(): async with aiohttp.ClientSession() as session: data = aiohttp.FormData(quote_fields=False) - data.add_field("images[]", file_sender(file_name=file_name), filename="image1.png", content_type="image/png") - data.add_field("images[]", file_sender(file_name=file_name), filename="image2.png", content_type="image/png") - data.add_field("images[]", file_sender(file_name=file_name), filename="image3.png", content_type="image/png") + data.add_field("images[]", file_sender(file_name=png_name), filename="image1.png", content_type="image/png") + data.add_field("images[]", file_sender(file_name=png_name), filename="image2.png", content_type="image/png") + data.add_field("images[]", file_sender(file_name=gif_name), filename="image1.gif", content_type="image/gif") + data.add_field("images[]", file_sender(file_name=gif_name), filename="image2.gif", content_type="image/gif") async with session.post(url, data=data) as resp: text = await resp.text() diff --git a/src/error.rs b/src/error.rs index d55c38a..d0bc1dd 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,10 +1,8 @@ +use crate::validate::GifError; use actix_web::{http::StatusCode, HttpResponse, ResponseError}; #[derive(Debug, thiserror::Error)] -pub enum UploadError { - #[error("Invalid content type provided, {0}")] - ContentType(mime::Mime), - +pub(crate) enum UploadError { #[error("Couln't upload file, {0}")] Upload(String), @@ -64,6 +62,9 @@ pub enum UploadError { #[error("Tried to save an image with an already-taken name")] DuplicateAlias, + + #[error("Error validating Gif file, {0}")] + Gif(#[from] GifError), } impl From for UploadError { @@ -102,9 +103,9 @@ where impl ResponseError for UploadError { fn status_code(&self) -> StatusCode { match self { - UploadError::DuplicateAlias + UploadError::Gif(_) + | UploadError::DuplicateAlias | UploadError::NoFiles - | UploadError::ContentType(_) | UploadError::Upload(_) => StatusCode::BAD_REQUEST, UploadError::MissingAlias | UploadError::MissingFilename => StatusCode::NOT_FOUND, UploadError::InvalidToken => StatusCode::FORBIDDEN, diff --git a/src/main.rs b/src/main.rs index a30c791..738fc00 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,16 +15,10 @@ mod config; mod error; mod processor; mod upload_manager; +mod validate; use self::{config::Config, error::UploadError, upload_manager::UploadManager}; -const ACCEPTED_MIMES: &[mime::Mime] = &[ - mime::IMAGE_BMP, - mime::IMAGE_GIF, - mime::IMAGE_JPEG, - mime::IMAGE_PNG, -]; - const MEGABYTES: usize = 1024 * 1024; const HOURS: u32 = 60 * 60; diff --git a/src/upload_manager.rs b/src/upload_manager.rs index 72ec76f..68b0663 100644 --- a/src/upload_manager.rs +++ b/src/upload_manager.rs @@ -1,4 +1,4 @@ -use crate::{config::Format, error::UploadError, safe_save_file, to_ext, ACCEPTED_MIMES}; +use crate::{config::Format, error::UploadError, safe_save_file, to_ext, validate::validate_image}; use actix_web::web; use futures::stream::{Stream, StreamExt}; use log::{error, warn}; @@ -206,7 +206,8 @@ impl UploadManager { let bytes = read_stream(stream).await?; let (bytes, content_type) = if validate { - self.validate_image(bytes).await? + let format = self.inner.format.clone(); + validate_image(bytes, format).await? } else { (bytes, content_type) }; @@ -233,7 +234,8 @@ impl UploadManager { let bytes = read_stream(stream).await?; // -- VALIDATE IMAGE -- - let (bytes, content_type) = self.validate_image(bytes).await?; + let format = self.inner.format.clone(); + let (bytes, content_type) = validate_image(bytes, format).await?; // -- DUPLICATE CHECKS -- @@ -331,40 +333,6 @@ impl UploadManager { Ok(()) } - // import & export image using the image crate - async fn validate_image( - &self, - bytes: bytes::Bytes, - ) -> Result<(bytes::Bytes, mime::Mime), UploadError> { - let (img, format) = 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 - .format - .as_ref() - .map(|f| (f.to_image_format(), f.to_mime())) - .unwrap_or((format.clone(), valid_format(format)?)); - - if ACCEPTED_MIMES.iter().all(|valid| *valid != content_type) { - return Err(UploadError::ContentType(content_type)); - } - - let bytes: bytes::Bytes = web::block(move || { - let mut bytes = std::io::Cursor::new(vec![]); - img.write_to(&mut bytes, format)?; - Ok(bytes::Bytes::from(bytes.into_inner())) as Result - }) - .await?; - - Ok((bytes, content_type)) - } - // produce a sh256sum of the uploaded file async fn hash(&self, bytes: bytes::Bytes) -> Result, UploadError> { let mut hasher = self.inner.hasher.clone(); @@ -606,13 +574,3 @@ fn variant_key_bounds(hash: &[u8]) -> (Vec, Vec) { (start, end) } - -fn valid_format(format: image::ImageFormat) -> Result { - match format { - image::ImageFormat::Jpeg => Ok(mime::IMAGE_JPEG), - image::ImageFormat::Png => Ok(mime::IMAGE_PNG), - image::ImageFormat::Gif => Ok(mime::IMAGE_GIF), - image::ImageFormat::Bmp => Ok(mime::IMAGE_BMP), - _ => Err(UploadError::UnsupportedFormat), - } -} diff --git a/src/validate.rs b/src/validate.rs new file mode 100644 index 0000000..4108e50 --- /dev/null +++ b/src/validate.rs @@ -0,0 +1,118 @@ +use crate::{config::Format, error::UploadError}; +use actix_web::web; +use bytes::Bytes; +use image::{ImageDecoder, ImageEncoder, ImageFormat}; +use std::io::Cursor; + +#[derive(Debug, thiserror::Error)] +pub(crate) enum GifError { + #[error("Error decoding gif")] + Decode(#[from] gif::DecodingError), + + #[error("Error reading bytes")] + Io(#[from] std::io::Error), +} + +// import & export image using the image crate +pub(crate) async fn validate_image( + bytes: Bytes, + prescribed_format: Option, +) -> Result<(Bytes, mime::Mime), UploadError> { + let tup = web::block(move || { + if let Some(prescribed) = prescribed_format { + let img = image::load_from_memory(&bytes).map_err(UploadError::InvalidImage)?; + + let mime = prescribed.to_mime(); + + let mut bytes = Cursor::new(vec![]); + img.write_to(&mut bytes, prescribed.to_image_format())?; + return Ok((Bytes::from(bytes.into_inner()), mime)); + } + + let format = image::guess_format(&bytes).map_err(UploadError::InvalidImage)?; + + match format { + ImageFormat::Png => Ok((validate_png(bytes)?, mime::IMAGE_PNG)), + ImageFormat::Jpeg => Ok((validate_jpg(bytes)?, mime::IMAGE_JPEG)), + ImageFormat::Bmp => Ok((validate_bmp(bytes)?, mime::IMAGE_BMP)), + ImageFormat::Gif => Ok((validate_gif(bytes)?, mime::IMAGE_GIF)), + _ => Err(UploadError::UnsupportedFormat), + } + }) + .await?; + + Ok(tup) +} + +fn validate_png(bytes: Bytes) -> Result { + let decoder = image::png::PngDecoder::new(Cursor::new(&bytes))?; + + let mut bytes = Cursor::new(vec![]); + let encoder = image::png::PNGEncoder::new(&mut bytes); + validate_still_image(decoder, encoder)?; + + Ok(Bytes::from(bytes.into_inner())) +} + +fn validate_jpg(bytes: Bytes) -> Result { + let decoder = image::jpeg::JpegDecoder::new(Cursor::new(&bytes))?; + + let mut bytes = Cursor::new(vec![]); + let encoder = image::jpeg::JPEGEncoder::new(&mut bytes); + validate_still_image(decoder, encoder)?; + + Ok(Bytes::from(bytes.into_inner())) +} + +fn validate_bmp(bytes: Bytes) -> Result { + let decoder = image::bmp::BmpDecoder::new(Cursor::new(&bytes))?; + + let mut bytes = Cursor::new(vec![]); + let encoder = image::bmp::BMPEncoder::new(&mut bytes); + validate_still_image(decoder, encoder)?; + + Ok(Bytes::from(bytes.into_inner())) +} + +fn validate_gif(bytes: Bytes) -> Result { + use gif::{Parameter, SetParameter}; + + let mut decoder = gif::Decoder::new(Cursor::new(&bytes)); + + decoder.set(gif::ColorOutput::Indexed); + + let mut reader = decoder.read_info()?; + + let width = reader.width(); + let height = reader.height(); + let global_palette = reader.global_palette().unwrap_or(&[]); + + let mut bytes = Cursor::new(vec![]); + { + let mut encoder = gif::Encoder::new(&mut bytes, width, height, global_palette)?; + + gif::Repeat::Infinite.set_param(&mut encoder)?; + + while let Some(frame) = reader.read_next_frame()? { + encoder.write_frame(frame)?; + } + } + + Ok(Bytes::from(bytes.into_inner())) +} + +fn validate_still_image<'a, D, E>(decoder: D, encoder: E) -> Result<(), UploadError> +where + D: ImageDecoder<'a>, + E: ImageEncoder, +{ + let (width, height) = decoder.dimensions(); + let color_type = decoder.color_type(); + let total_bytes = decoder.total_bytes(); + let mut decoded_bytes = vec![0u8; total_bytes as usize]; + decoder.read_image(&mut decoded_bytes)?; + + encoder.write_image(&decoded_bytes, width, height, color_type)?; + + Ok(()) +}