tetratto/crates/app/src/image.rs

use axum::{
    body::Bytes,
    extract::{FromRequest, Request},
    http::{StatusCode, header::CONTENT_TYPE},
};
use axum_extra::extract::Multipart;
use serde::de::DeserializeOwned;
use std::{fs::File, io::BufWriter};

/// An image extractor accepting:
/// * `multipart/form-data`
/// * `image/png`
/// * `image/jpeg`
/// * `image/avif`
/// * `image/webp`
pub struct Image(pub Bytes, pub String);

impl<S> FromRequest<S> for Image
where
    Bytes: FromRequest<S>,
    S: Send + Sync,
{
    type Rejection = StatusCode;

    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
        let Some(content_type) = req.headers().get(CONTENT_TYPE) else {
            return Err(StatusCode::BAD_REQUEST);
        };

        let content_type = content_type.to_str().unwrap();
        let content_type_string = content_type.to_string();

        let body = if content_type.starts_with("multipart/form-data") {
            let mut multipart = Multipart::from_request(req, state)
                .await
                .map_err(|_| StatusCode::BAD_REQUEST)?;

            let Ok(Some(field)) = multipart.next_field().await else {
                return Err(StatusCode::BAD_REQUEST);
            };

            field.bytes().await.map_err(|_| StatusCode::BAD_REQUEST)?
        } else if (content_type == "image/avif")
            | (content_type == "image/jpeg")
            | (content_type == "image/png")
            | (content_type == "image/webp")
            | (content_type == "image/gif")
        {
            Bytes::from_request(req, state)
                .await
                .map_err(|_| StatusCode::BAD_REQUEST)?
        } else {
            return Err(StatusCode::BAD_REQUEST);
        };

        Ok(Self(body, content_type_string))
    }
}

/// A file extractor accepting:
/// * `multipart/form-data`
///
/// Will also attempt to parse out the **last** field in the multipart upload
/// as the given struct from JSON. Every other field is put into a vector of bytes,
/// as they are seen as raw binary data.
pub struct JsonMultipart<T: DeserializeOwned>(pub Vec<Bytes>, pub T);

impl<S, T> FromRequest<S> for JsonMultipart<T>
where
    Bytes: FromRequest<S>,
    S: Send + Sync,
    T: DeserializeOwned,
{
    type Rejection = (StatusCode, String);

    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
        let Some(content_type) = req.headers().get(CONTENT_TYPE) else {
            return Err((
                StatusCode::BAD_REQUEST,
                "no content type header".to_string(),
            ));
        };

        let content_type = content_type.to_str().unwrap();

        if !content_type.starts_with("multipart/form-data") {
            return Err((
                StatusCode::BAD_REQUEST,
                "expected multipart/form-data".to_string(),
            ));
        }

        let mut multipart = Multipart::from_request(req, state).await.map_err(|_| {
            (
                StatusCode::BAD_REQUEST,
                "could not read multipart".to_string(),
            )
        })?;

        let mut body: Vec<Bytes> = {
            let mut out = Vec::new();

            while let Ok(Some(field)) = multipart.next_field().await {
                out.push(
                    field
                        .bytes()
                        .await
                        .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?,
                );
            }

            out
        };

        let last = match body.pop() {
            Some(b) => b,
            None => {
                return Err((
                    StatusCode::BAD_REQUEST,
                    "could not read json data".to_string(),
                ));
            }
        };

        let json: T = match serde_json::from_str(&match String::from_utf8(last.to_vec()) {
            Ok(s) => s,
            Err(_) => return Err((StatusCode::BAD_REQUEST, "json data isn't utf8".to_string())),
        }) {
            Ok(s) => s,
            Err(e) => {
                return Err((StatusCode::BAD_REQUEST, e.to_string()));
            }
        };

        Ok(Self(body, json))
    }
}

/// Create an image buffer given an input of `bytes`.
pub fn save_buffer(path: &str, bytes: Vec<u8>, format: image::ImageFormat) -> std::io::Result<()> {
    let pre_img_buffer = match image::load_from_memory(&bytes) {
        Ok(i) => i,
        Err(_) => {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "Image failed",
            ));
        }
    };

    let file = File::create(path)?;
    let mut writer = BufWriter::new(file);

    if pre_img_buffer.write_to(&mut writer, format).is_err() {
        return Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "Image conversion failed",
        ));
    };

    Ok(())
}

const WEBP_ENCODE_QUALITY: f32 = 85.0;

/// Create a WEBP image buffer given an input of `bytes`.
pub fn save_webp_buffer(path: &str, bytes: Vec<u8>, quality: Option<f32>) -> std::io::Result<()> {
    let img = match image::load_from_memory(&bytes) {
        Ok(i) => i,
        Err(_) => {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "Image failed",
            ));
        }
    };

    let encoder = match webp::Encoder::from_image(&img) {
        Ok(e) => e,
        Err(e) => {
            return Err(std::io::Error::new(
                std::io::ErrorKind::Other,
                e.to_string(),
            ));
        }
    };

    let mem = encoder.encode(match quality {
        Some(q) => q,
        None => WEBP_ENCODE_QUALITY,
    });

    if std::fs::write(path, &*mem).is_err() {
        return Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "Image conversion failed",
        ));
    };

    Ok(())
}