kompress_eshop/src/media/mod.rs

use crate::shared::guard;
use axum::{
    body::Body,
    extract::{DefaultBodyLimit, Multipart},
    http::header,
};
use bytes::Bytes;
use loco_rs::{config::Config, prelude::*};
use serde::Serialize;
use std::path::{Path as StdPath, PathBuf};
use uuid::Uuid;

pub(crate) const IMAGE_MAX_BYTES: usize = 10 * 1024 * 1024;
pub const IMAGE_STORAGE_DIR: &str = "images";

#[derive(Debug, Serialize)]
struct UploadResponse {
    filename: String,
    url: String,
    size: usize,
}

pub fn uploads_root(config: &Config) -> Result<PathBuf> {
    config
        .settings
        .as_ref()
        .and_then(|settings| settings.get("uploads_root"))
        .and_then(|value| value.as_str())
        .filter(|value| !value.trim().is_empty())
        .map(PathBuf::from)
        .ok_or_else(|| Error::string("settings.uploads_root must be configured"))
}

fn safe_filename(filename: &str) -> Result<&str> {
    if filename.is_empty()
        || filename.contains('/')
        || filename.contains('\\')
        || filename.contains("..")
    {
        return Err(Error::BadRequest("invalid filename".to_string()));
    }
    Ok(filename)
}

fn image_content_type(extension: &str) -> &'static str {
    match extension {
        "gif" => "image/gif",
        "jpg" | "jpeg" => "image/jpeg",
        "png" => "image/png",
        "webp" => "image/webp",
        _ => "application/octet-stream",
    }
}

pub(crate) fn detect_image_extension(data: &[u8]) -> Result<&'static str> {
    if data.len() < 12 {
        return Err(Error::BadRequest("image file is too small".to_string()));
    }
    if data.starts_with(&[0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a]) {
        return Ok("png");
    }
    if data.starts_with(&[0xff, 0xd8, 0xff]) {
        return Ok("jpg");
    }
    if data.starts_with(b"RIFF") && &data[8..12] == b"WEBP" {
        return Ok("webp");
    }
    if data.starts_with(b"GIF87a") || data.starts_with(b"GIF89a") {
        return Ok("gif");
    }
    Err(Error::BadRequest("unsupported image format".to_string()))
}

async fn read_multipart_file(mut multipart: Multipart, max_bytes: usize) -> Result<Vec<u8>> {
    while let Some(mut field) = multipart
        .next_field()
        .await
        .map_err(|err| Error::BadRequest(format!("invalid multipart data: {err}")))?
    {
        if field.name() != Some("file") {
            continue;
        }

        let mut data = Vec::new();
        while let Some(chunk) = field
            .chunk()
            .await
            .map_err(|err| Error::BadRequest(format!("invalid multipart chunk: {err}")))?
        {
            data.extend_from_slice(&chunk);
            if data.len() > max_bytes {
                return Err(Error::BadRequest(format!(
                    "file is larger than {} MB",
                    max_bytes / 1024 / 1024
                )));
            }
        }

        if data.is_empty() {
            return Err(Error::BadRequest("empty file upload".to_string()));
        }
        return Ok(data);
    }

    Err(Error::BadRequest(
        "multipart field `file` is required".to_string(),
    ))
}

pub(crate) async fn store_upload(
    ctx: &AppContext,
    folder: &str,
    extension: &str,
    data: Vec<u8>,
) -> Result<String> {
    let filename = format!("{}.{}", Uuid::new_v4(), extension);
    let key = format!("{folder}/{filename}");
    ctx.storage
        .upload(StdPath::new(&key), &Bytes::from(data))
        .await?;
    Ok(filename)
}

#[debug_handler]
async fn image_upload(
    auth: auth::JWT,
    State(ctx): State<AppContext>,
    multipart: Multipart,
) -> Result<Response> {
    guard::current_admin(auth, &ctx).await?;
    let data = read_multipart_file(multipart, IMAGE_MAX_BYTES).await?;
    let extension = detect_image_extension(&data)?;
    let size = data.len();
    let filename = store_upload(&ctx, IMAGE_STORAGE_DIR, extension, data).await?;

    format::json(UploadResponse {
        url: format!("/images/{filename}"),
        filename,
        size,
    })
}

#[debug_handler]
async fn image_serve(
    Path(filename): Path<String>,
    State(ctx): State<AppContext>,
) -> Result<Response> {
    let filename = safe_filename(&filename)?;
    let extension = filename.rsplit('.').next().unwrap_or("");
    let key = format!("{IMAGE_STORAGE_DIR}/{filename}");
    let body: Vec<u8> = ctx.storage.download(StdPath::new(&key)).await?;
    Response::builder()
        .header(header::CONTENT_TYPE, image_content_type(extension))
        .header(header::CACHE_CONTROL, "public, max-age=31536000, immutable")
        .body(Body::from(body))
        .map_err(Error::from)
}

pub fn routes() -> Routes {
    Routes::new()
        .add(
            "/images/upload",
            post(image_upload).layer(DefaultBodyLimit::max(IMAGE_MAX_BYTES + 1024 * 1024)),
        )
        .add("/images/{filename}", get(image_serve))
}