tetratto/crates/app/src/routes/api/v1/auth/images.rs

use axum::{
    Extension, Json,
    body::Body,
    extract::{Path, Query},
    response::IntoResponse,
};
use axum_extra::extract::CookieJar;
use pathbufd::{PathBufD, pathd};
use serde::Deserialize;
use std::{
    fs::{File, exists},
    io::Read,
};
use tetratto_core::model::{permissions::FinePermission, ApiReturn, Error};

use crate::{
    State,
    image::{Image, save_buffer},
    get_user_from_token,
};

pub fn read_image(path: PathBufD) -> Vec<u8> {
    let mut bytes = Vec::new();

    for byte in File::open(path).unwrap().bytes() {
        bytes.push(byte.unwrap())
    }

    bytes
}

#[derive(Deserialize, PartialEq, Eq)]
pub enum AvatarSelectorType {
    #[serde(alias = "username")]
    Username,
    #[serde(alias = "id")]
    Id,
}

#[derive(Deserialize)]
pub struct AvatarSelectorQuery {
    pub selector_type: AvatarSelectorType,
}

/// Get a profile's avatar image
/// `/api/v1/auth/user/{id}/avatar`
pub async fn avatar_request(
    Path(selector): Path<String>,
    Extension(data): Extension<State>,
    Query(req): Query<AvatarSelectorQuery>,
) -> impl IntoResponse {
    let data = &(data.read().await).0;

    let user = match if req.selector_type == AvatarSelectorType::Id {
        data.get_user_by_id(match selector.parse::<usize>() {
            Ok(d) => d,
            Err(_) => {
                return Err((
                    [("Content-Type", "image/svg+xml")],
                    Body::from(read_image(PathBufD::current().extend(&[
                        data.0.dirs.media.as_str(),
                        "images",
                        "default-avatar.svg",
                    ]))),
                ));
            }
        })
        .await
    } else {
        data.get_user_by_username(&selector).await
    } {
        Ok(ua) => ua,
        Err(_) => {
            return Err((
                [("Content-Type", "image/svg+xml")],
                Body::from(read_image(PathBufD::current().extend(&[
                    data.0.dirs.media.as_str(),
                    "images",
                    "default-avatar.svg",
                ]))),
            ));
        }
    };

    let path = PathBufD::current().extend(&[
        data.0.dirs.media.as_str(),
        "avatars",
        &format!(
            "{}.{}",
            &(user.id as i64),
            user.settings.avatar_mime.replace("image/", "")
        ),
    ]);

    if !exists(&path).unwrap() {
        return Err((
            [("Content-Type", "image/svg+xml")],
            Body::from(read_image(PathBufD::current().extend(&[
                data.0.dirs.media.as_str(),
                "images",
                "default-avatar.svg",
            ]))),
        ));
    }

    Ok((
        [(
            "Content-Type".to_string(),
            user.settings.avatar_mime.clone(),
        )],
        Body::from(read_image(path)),
    ))
}

/// Get a profile's banner image
/// `/api/v1/auth/user/{id}/banner`
pub async fn banner_request(
    Path(username): Path<String>,
    Extension(data): Extension<State>,
) -> impl IntoResponse {
    let data = &(data.read().await).0;

    let user = match data.get_user_by_username(&username).await {
        Ok(ua) => ua,
        Err(_) => {
            return Err((
                [("Content-Type", "image/svg+xml")],
                Body::from(read_image(PathBufD::current().extend(&[
                    data.0.dirs.media.as_str(),
                    "images",
                    "default-banner.svg",
                ]))),
            ));
        }
    };

    let path = PathBufD::current().extend(&[
        data.0.dirs.media.as_str(),
        "banners",
        &format!(
            "{}.{}",
            &(user.id as i64),
            user.settings.banner_mime.replace("image/", "")
        ),
    ]);

    if !exists(&path).unwrap() {
        return Err((
            [("Content-Type", "image/svg+xml")],
            Body::from(read_image(PathBufD::current().extend(&[
                data.0.dirs.media.as_str(),
                "images",
                "default-banner.svg",
            ]))),
        ));
    }

    Ok((
        [(
            "Content-Type".to_string(),
            user.settings.banner_mime.clone(),
        )],
        Body::from(read_image(path)),
    ))
}

pub static MAXIUMUM_FILE_SIZE: usize = 8388608;
pub static MAXIUMUM_GIF_FILE_SIZE: usize = 2097152;

/// Upload avatar
pub async fn upload_avatar_request(
    jar: CookieJar,
    Extension(data): Extension<State>,
    img: Image,
) -> impl IntoResponse {
    // get user from token
    let data = &(data.read().await).0;
    let mut auth_user = match get_user_from_token!(jar, data) {
        Some(ua) => ua,
        None => return Json(Error::NotAllowed.into()),
    };

    if img.1 == "image/gif" && !auth_user.permissions.check(FinePermission::SUPPORTER) {
        return Json(Error::RequiresSupporter.into());
    }

    let mime = if img.1 == "image/gif" {
        "image/gif"
    } else {
        "image/avif"
    };

    if auth_user.settings.avatar_mime != mime {
        // mime changed; delete old image
        let path = pathd!(
            "{}/avatars/{}.{}",
            data.0.dirs.media,
            &auth_user.id,
            auth_user.settings.avatar_mime.replace("image/", "")
        );

        if std::fs::exists(&path).unwrap() {
            std::fs::remove_file(path).unwrap();
        }
    }

    let path = pathd!(
        "{}/avatars/{}.{}",
        data.0.dirs.media,
        &auth_user.id,
        mime.replace("image/", "")
    );

    // update user settings
    auth_user.settings.avatar_mime = mime.to_string();
    if let Err(e) = data
        .update_user_settings(auth_user.id, auth_user.settings)
        .await
    {
        return Json(e.into());
    }

    // upload image (gif)
    if mime == "image/gif" {
        // gif image, don't encode
        if img.0.len() > MAXIUMUM_GIF_FILE_SIZE {
            return Json(Error::DataTooLong("gif".to_string()).into());
        }

        std::fs::write(&path, img.0).unwrap();
        return Json(ApiReturn {
            ok: true,
            message: "Avatar uploaded. It might take a bit to update".to_string(),
            payload: (),
        });
    }

    // check file size
    if img.0.len() > MAXIUMUM_FILE_SIZE {
        return Json(Error::DataTooLong("image".to_string()).into());
    }

    // upload image
    let mut bytes = Vec::new();

    for byte in img.0 {
        bytes.push(byte);
    }

    match save_buffer(
        &path,
        bytes,
        if mime == "image/gif" {
            image::ImageFormat::Gif
        } else {
            image::ImageFormat::Avif
        },
    ) {
        Ok(_) => Json(ApiReturn {
            ok: true,
            message: "Avatar uploaded. It might take a bit to update".to_string(),
            payload: (),
        }),
        Err(e) => Json(Error::MiscError(e.to_string()).into()),
    }
}

/// Upload banner
pub async fn upload_banner_request(
    jar: CookieJar,
    Extension(data): Extension<State>,
    img: Image,
) -> impl IntoResponse {
    // get user from token
    let data = &(data.read().await).0;
    let mut auth_user = match get_user_from_token!(jar, data) {
        Some(ua) => ua,
        None => return Json(Error::NotAllowed.into()),
    };

    if img.1 == "image/gif" && !auth_user.permissions.check(FinePermission::SUPPORTER) {
        return Json(Error::RequiresSupporter.into());
    }

    let mime = if img.1 == "image/gif" {
        "image/gif"
    } else {
        "image/avif"
    };

    if auth_user.settings.banner_mime != mime {
        // mime changed; delete old image
        let path = pathd!(
            "{}/banners/{}.{}",
            data.0.dirs.media,
            &auth_user.id,
            auth_user.settings.banner_mime.replace("image/", "")
        );

        if std::fs::exists(&path).unwrap() {
            std::fs::remove_file(path).unwrap();
        }
    }

    let path = pathd!(
        "{}/banners/{}.{}",
        data.0.dirs.media,
        &auth_user.id,
        mime.replace("image/", "")
    );

    // update user settings
    auth_user.settings.banner_mime = mime.to_string();
    if let Err(e) = data
        .update_user_settings(auth_user.id, auth_user.settings)
        .await
    {
        return Json(e.into());
    }

    // upload image (gif)
    if mime == "image/gif" {
        // gif image, don't encode
        if img.0.len() > MAXIUMUM_GIF_FILE_SIZE {
            return Json(Error::DataTooLong("gif".to_string()).into());
        }

        std::fs::write(&path, img.0).unwrap();
        return Json(ApiReturn {
            ok: true,
            message: "Banner uploaded. It might take a bit to update".to_string(),
            payload: (),
        });
    }

    // check file size
    if img.0.len() > MAXIUMUM_FILE_SIZE {
        return Json(Error::DataTooLong("image".to_string()).into());
    }

    // upload image
    let mut bytes = Vec::new();

    for byte in img.0 {
        bytes.push(byte);
    }

    match save_buffer(
        &path,
        bytes,
        if mime == "image/gif" {
            image::ImageFormat::Gif
        } else {
            image::ImageFormat::Avif
        },
    ) {
        Ok(_) => Json(ApiReturn {
            ok: true,
            message: "Banner uploaded. It might take a bit to update".to_string(),
            payload: (),
        }),
        Err(e) => Json(Error::MiscError(e.to_string()).into()),
    }
}