Don't use image crate to validate gifs

Cargo.lock

@@ -1381,6 +1381,7 @@ dependencies = [
  "bytes",
  "env_logger",
  "futures",
+ "gif",
  "image",
  "log",
  "mime",

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"

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=png_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=png_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=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()

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 {
+pub(crate) enum UploadError {
-    #[error("Invalid content type provided, {0}")]
-    ContentType(mime::Mime),
-
     #[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<actix_web::client::SendRequestError> 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,

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;
 

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<bytes::Bytes, UploadError>
-        })
-        .await?;
-
-        Ok((bytes, content_type))
-    }
-
     // 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();
@@ -606,13 +574,3 @@ fn variant_key_bounds(hash: &[u8]) -> (Vec<u8>, Vec<u8>) {
 
     (start, end)
 }
-
-fn valid_format(format: image::ImageFormat) -> Result<mime::Mime, UploadError> {
-    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),
-    }
-}

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<Format>,
+) -> 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<Bytes, UploadError> {
+    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<Bytes, UploadError> {
+    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<Bytes, UploadError> {
+    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<Bytes, GifError> {
+    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(())
+}