Talon/src/icons.rs

use std::{
    io::{BufReader, Read, Seek},
    path::PathBuf,
};

use image::io::Reader as ImageReader;
use poem::{
    error::ResponseError,
    handler,
    http::StatusCode,
    web::{Data, Path, StaticFileRequest, StaticFileResponse},
    FromRequest, IntoResponse, Request, Response,
};

use crate::Talon;

pub struct Icons {
    path: PathBuf,
}

#[derive(thiserror::Error, Debug)]
pub enum ImagesError {
    #[error("io error: {0}")]
    Io(#[from] std::io::Error),
    #[error("image error: {0}")]
    Image(#[from] image::ImageError),
    #[error("image for `{0}` not found")]
    NotFound(String),
}

impl ResponseError for ImagesError {
    fn status(&self) -> StatusCode {
        match self {
            ImagesError::NotFound(_) => StatusCode::NOT_FOUND,
            ImagesError::Io(_) | ImagesError::Image(_) => StatusCode::INTERNAL_SERVER_ERROR,
        }
    }
}

const IMAGE_SIZE: u32 = 48;
const MAX_IMAGE_SIZE: u32 = 4000;

type Result<T> = std::result::Result<T, ImagesError>;

impl Icons {
    /// Create a new image storage
    pub fn new<P: Into<PathBuf>>(path: P) -> Self {
        Self { path: path.into() }
    }

    fn icon_path(&self, subdomain: &str) -> PathBuf {
        self.path.join(format!("{subdomain}.png"))
    }

    pub fn insert_icon(&self, reader: impl Read + Seek, subdomain: &str) -> Result<()> {
        let mut img_reader = ImageReader::new(BufReader::new(reader));
        let mut limits = image::io::Limits::default();
        limits.max_image_height = Some(MAX_IMAGE_SIZE);
        limits.max_image_width = Some(MAX_IMAGE_SIZE);
        img_reader.limits(limits);

        let img = img_reader.with_guessed_format()?.decode()?;
        let img = img.resize(
            IMAGE_SIZE,
            IMAGE_SIZE,
            image::imageops::FilterType::Lanczos3,
        );
        img.save(self.icon_path(subdomain))?;
        Ok(())
    }

    pub fn delete_icon(&self, subdomain: &str) -> Result<()> {
        let path = self.get_icon(subdomain)?;
        std::fs::remove_file(path)?;
        Ok(())
    }

    pub fn get_icon(&self, subdomain: &str) -> Result<PathBuf> {
        let path = self.icon_path(subdomain);
        if path.is_file() {
            Ok(path)
        } else {
            Err(ImagesError::NotFound(subdomain.to_owned()))
        }
    }

    pub async fn get_icon_response(
        &self,
        subdomain: &str,
        request: &Request,
    ) -> poem::Result<StaticFileResponse> {
        let path = self.icon_path(subdomain);
        Ok(StaticFileRequest::from_request_without_body(request)
            .await?
            .create_response(path, false)?)
    }
}

#[handler]
pub async fn icon(
    request: &Request,
    talon: Data<&Talon>,
    subdomain: Path<String>,
) -> poem::Result<Response> {
    talon
        .icons
        .get_icon_response(&subdomain, request)
        .await
        .map(|r| r.into_response())
}