pict-rs/src/validate.rs

use crate::{config::Format, error::UploadError, upload_manager::tmp_file};
use actix_web::web;
use magick_rust::MagickWand;
use rexiv2::{MediaType, Metadata};
use std::{
    fs::File,
    io::{BufReader, BufWriter, Write},
    path::PathBuf,
};
use tracing::{debug, error, instrument, trace, warn, Span};

pub(crate) trait Op {
    fn op<F, T>(&self, f: F) -> Result<T, UploadError>
    where
        F: Fn(&Self) -> Result<T, &'static str>;

    fn op_mut<F, T>(&mut self, f: F) -> Result<T, UploadError>
    where
        F: Fn(&mut Self) -> Result<T, &'static str>;
}

impl Op for MagickWand {
    fn op<F, T>(&self, f: F) -> Result<T, UploadError>
    where
        F: Fn(&Self) -> Result<T, &'static str>,
    {
        match f(self) {
            Ok(t) => Ok(t),
            Err(e) => {
                if let Ok(e) = self.get_exception() {
                    error!("WandError: {}", e.0);
                    Err(UploadError::Wand(e.0.to_owned()))
                } else {
                    Err(UploadError::Wand(e.to_owned()))
                }
            }
        }
    }

    fn op_mut<F, T>(&mut self, f: F) -> Result<T, UploadError>
    where
        F: Fn(&mut Self) -> Result<T, &'static str>,
    {
        match f(self) {
            Ok(t) => Ok(t),
            Err(e) => {
                if let Ok(e) = self.get_exception() {
                    error!("WandError: {}", e.0);
                    Err(UploadError::Wand(e.0.to_owned()))
                } else {
                    Err(UploadError::Wand(e.to_owned()))
                }
            }
        }
    }
}

#[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),
}

pub(crate) fn image_webp() -> mime::Mime {
    "image/webp".parse().unwrap()
}

pub(crate) fn ptos(p: &PathBuf) -> Result<String, UploadError> {
    Ok(p.to_str().ok_or(UploadError::Path)?.to_owned())
}

fn validate_format(file: &str, format: &str) -> Result<(), UploadError> {
    let wand = MagickWand::new();
    debug!("reading");
    wand.op(|w| w.read_image(file))?;

    if wand.op(|w| w.get_image_format())? != format {
        return Err(UploadError::UnsupportedFormat);
    }

    Ok(())
}

// import & export image using the image crate
#[instrument]
pub(crate) async fn validate_image(
    tmpfile: PathBuf,
    prescribed_format: Option<Format>,
) -> Result<mime::Mime, UploadError> {
    let tmpfile_str = ptos(&tmpfile)?;
    let span = Span::current();

    let content_type = web::block(move || {
        let entered = span.enter();

        let meta = Metadata::new_from_path(&tmpfile)?;

        let content_type = match (prescribed_format, meta.get_media_type()?) {
            (_, MediaType::Gif) => {
                let newfile = tmp_file();
                validate_gif(&tmpfile, &newfile)?;

                mime::IMAGE_GIF
            }
            (Some(Format::Jpeg), MediaType::Jpeg) | (None, MediaType::Jpeg) => {
                validate_format(&tmpfile_str, "JPEG")?;

                meta.clear();
                meta.save_to_file(&tmpfile)?;

                mime::IMAGE_JPEG
            }
            (Some(Format::Png), MediaType::Png) | (None, MediaType::Png) => {
                validate_format(&tmpfile_str, "PNG")?;

                meta.clear();
                meta.save_to_file(&tmpfile)?;

                mime::IMAGE_PNG
            }
            (Some(Format::Webp), MediaType::Other(webp)) | (None, MediaType::Other(webp))
                if webp == "image/webp" =>
            {
                let newfile = tmp_file();
                let newfile_str = ptos(&newfile)?;
                // clean metadata by writing new webp, since exiv2 doesn't support webp yet
                {
                    let wand = MagickWand::new();

                    debug!("reading");
                    wand.op(|w| w.read_image(&tmpfile_str))?;

                    if wand.op(|w| w.get_image_format())? != "WEBP" {
                        return Err(UploadError::UnsupportedFormat);
                    }

                    wand.op(|w| w.write_image(&newfile_str))?;
                }

                std::fs::rename(&newfile, &tmpfile)?;

                image_webp()
            }
            (Some(format), _) => {
                let newfile = tmp_file();
                let newfile_str = ptos(&newfile)?;
                {
                    let mut wand = MagickWand::new();

                    debug!("reading: {}", tmpfile_str);
                    wand.op(|w| w.read_image(&tmpfile_str))?;

                    wand.op_mut(|w| w.set_image_format(format.to_magick_format()))?;

                    debug!("writing: {}", newfile_str);
                    wand.op(|w| w.write_image(&newfile_str))?;
                }

                std::fs::rename(&newfile, &tmpfile)?;

                format.to_mime()
            }
            (_, media_type) => {
                warn!("Unsupported media type, {}", media_type);
                return Err(UploadError::UnsupportedFormat);
            }
        };

        drop(entered);
        Ok(content_type) as Result<mime::Mime, UploadError>
    })
    .await?;

    Ok(content_type)
}

#[instrument]
fn validate_gif(from: &PathBuf, to: &PathBuf) -> Result<(), GifError> {
    debug!("Transmuting GIF");
    use gif::{Parameter, SetParameter};

    let mut decoder = gif::Decoder::new(BufReader::new(File::open(from)?));

    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 writer = BufWriter::new(File::create(to)?);
    let mut encoder = gif::Encoder::new(&mut writer, width, height, global_palette)?;

    gif::Repeat::Infinite.set_param(&mut encoder)?;

    while let Some(frame) = reader.read_next_frame()? {
        trace!("Writing frame");
        encoder.write_frame(frame)?;
    }

    drop(encoder);
    writer.flush()?;

    std::fs::rename(to, from)?;
    Ok(())
}