hyaenidae/profiles/src/pictrs.rs

use actix_web::{body::BodyStream, dev::Payload, HttpRequest, HttpResponse};
use awc::Client;
use std::{fmt, sync::Arc, time::Duration};
use url::Url;

#[derive(Debug, Clone)]
struct Serde<T> {
    inner: T,
}

impl<T> serde::Serialize for Serde<T>
where
    T: std::fmt::Display,
{
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let s = self.inner.to_string();
        serde::Serialize::serialize(s.as_str(), serializer)
    }
}

impl<'de, T> serde::Deserialize<'de> for Serde<T>
where
    T: std::str::FromStr,
    <T as std::str::FromStr>::Err: std::fmt::Display,
{
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let s: String = serde::Deserialize::deserialize(deserializer)?;
        let inner = s
            .parse::<T>()
            .map_err(|e| serde::de::Error::custom(e.to_string()))?;

        Ok(Serde { inner })
    }
}

#[derive(Debug, serde::Deserialize)]
struct Details {
    width: usize,
    height: usize,
    content_type: Serde<mime::Mime>,
}

#[derive(Debug, serde::Deserialize)]
pub struct Image {
    file: String,
    delete_token: String,
    details: Details,
}

impl Image {
    pub(crate) fn key(&self) -> &str {
        &self.file
    }

    pub(crate) fn token(&self) -> &str {
        &self.delete_token
    }

    pub(crate) fn width(&self) -> usize {
        self.details.width
    }

    pub(crate) fn height(&self) -> usize {
        self.details.height
    }

    pub(crate) fn mime(&self) -> mime::Mime {
        self.details.content_type.inner.clone()
    }
}

#[derive(Debug, serde::Deserialize)]
pub struct Images {
    msg: String,
    files: Option<Vec<Image>>,
}

impl Images {
    fn is_err(&self) -> bool {
        self.files.is_none()
    }

    pub(crate) fn images(&self) -> impl DoubleEndedIterator<Item = &Image> {
        self.files.iter().flat_map(|v| v.iter())
    }
}

#[derive(Debug, thiserror::Error)]
pub enum UploadError {
    #[error("Failed to send request")]
    Request,

    #[error("Failed to parse image json")]
    Json,

    #[error("Error in pict-rs: {0}")]
    Other(String),
}

#[derive(Debug, thiserror::Error)]
pub enum DeleteError {
    #[error("Failed to send request")]
    Request,

    #[error("Status was not 2XX")]
    Status,
}

pub trait ImageInfo {
    fn image_url(&self, key: &str) -> Url;
}

#[derive(Clone, Copy, Debug, serde::Deserialize, serde::Serialize)]
pub enum ImageType {
    #[serde(rename = "jpg")]
    Jpeg,

    #[serde(rename = "png")]
    Png,

    #[serde(rename = "webp")]
    Webp,
}

#[derive(Clone)]
pub struct State {
    upstream: Url,
    pub(super) info: Arc<dyn ImageInfo + Send + Sync>,
}

impl State {
    pub(super) fn new(upstream: Url, image_info: impl ImageInfo + Send + Sync + 'static) -> Self {
        State {
            upstream,
            info: Arc::new(image_info),
        }
    }

    pub async fn serve_thumbnail(
        &self,
        req: &HttpRequest,
        key: &str,
        kind: ImageType,
        size: u64,
        client: &Client,
    ) -> Result<HttpResponse, actix_web::Error> {
        self.serve_image(self.thumbnail(key, kind, size), req, client)
            .await
    }

    pub async fn serve_banner(
        &self,
        req: &HttpRequest,
        key: &str,
        kind: ImageType,
        size: u64,
        client: &Client,
    ) -> Result<HttpResponse, actix_web::Error> {
        self.serve_image(self.banner(key, kind, size), req, client)
            .await
    }

    pub async fn serve_icon(
        &self,
        req: &HttpRequest,
        key: &str,
        kind: ImageType,
        size: u64,
        client: &Client,
    ) -> Result<HttpResponse, actix_web::Error> {
        self.serve_image(self.icon(key, kind, size), req, client)
            .await
    }

    pub async fn serve_full(
        &self,
        req: &HttpRequest,
        key: &str,
        client: &Client,
    ) -> Result<HttpResponse, actix_web::Error> {
        self.serve_image(self.full_size(key), req, client).await
    }

    async fn serve_image(
        &self,
        url: Url,
        req: &HttpRequest,
        client: &Client,
    ) -> Result<HttpResponse, actix_web::Error> {
        let client_req = client
            .request_from(url.as_str(), req.head())
            .timeout(Duration::from_secs(30));
        let client_req = if let Some(addr) = req.head().peer_addr {
            client_req.insert_header(("X-Forwarded-For", addr.to_string()))
        } else {
            client_req
        };

        let res = client_req.no_decompress().send().await?;

        let mut client_res = HttpResponse::build(res.status());

        for (name, value) in res.headers().iter().filter(|(h, _)| *h != "connection") {
            client_res.insert_header((name.clone(), value.clone()));
        }

        Ok(client_res.body(BodyStream::new(res)))
    }

    pub(super) async fn delete_image(
        &self,
        key: &str,
        token: &str,
        client: &Client,
    ) -> Result<(), DeleteError> {
        let res = client
            .delete(self.delete(key, token).as_str())
            .send()
            .await
            .map_err(|e| {
                log::error!("{}", e);
                DeleteError::Request
            })?;

        if !res.status().is_success() {
            return Err(DeleteError::Status);
        }

        Ok(())
    }

    pub(super) async fn proxy_upload<'a>(
        &'a self,
        req: HttpRequest,
        body: Payload,
        client: &'a Client,
    ) -> Result<Images, UploadError> {
        let client_req = client
            .request_from(self.upload().as_str(), req.head())
            .timeout(Duration::from_secs(30));

        let client_req = if let Some(addr) = req.head().peer_addr {
            client_req.insert_header(("X-Forwarded-For", addr.to_string()))
        } else {
            client_req
        };

        let mut res = client_req
            .send_stream(body)
            .await
            .map_err(|_| UploadError::Request)?;

        let images = res.json::<Images>().await.map_err(|_| UploadError::Json)?;

        if images.is_err() {
            return Err(UploadError::Other(images.msg));
        }

        Ok(images)
    }

    pub(super) async fn download_image(
        &self,
        url: &Url,
        client: &Client,
    ) -> Result<Images, UploadError> {
        let mut res = client
            .get(self.download(url).as_str())
            .timeout(Duration::from_secs(30))
            .send()
            .await
            .map_err(|e| {
                log::error!("Error in download: {}", e);
                UploadError::Request
            })?;

        let images: Images = res.json().await.map_err(|e| {
            log::error!("Error in download: {}", e);
            UploadError::Json
        })?;

        if images.is_err() {
            return Err(UploadError::Other(images.msg));
        }

        Ok(images)
    }

    fn upload(&self) -> Url {
        let mut url = self.upstream.clone();
        url.set_path("/image");
        url
    }

    fn full_size(&self, key: &str) -> Url {
        let mut url = self.upstream.clone();
        url.set_path(&format!("/image/original/{}", key));
        url
    }

    fn process(&self, kind: ImageType) -> Url {
        let mut url = self.upstream.clone();
        url.set_path(&format!("/image/process.{}", kind));
        url
    }

    fn icon(&self, key: &str, kind: ImageType, size: u64) -> Url {
        let mut url = self.process(kind);
        url.set_query(Some(&format!("src={}&crop=1x1&resize={}", key, size)));
        url
    }

    fn banner(&self, key: &str, kind: ImageType, size: u64) -> Url {
        let mut url = self.process(kind);
        url.set_query(Some(&format!("src={}&crop=3x1&resize={}", key, size)));
        url
    }

    fn thumbnail(&self, key: &str, kind: ImageType, size: u64) -> Url {
        let mut url = self.process(kind);
        url.set_query(Some(&format!("src={}&resize={}", key, size)));
        url
    }

    fn download(&self, image_url: &Url) -> Url {
        let mut url = self.upstream.clone();
        url.set_path("/image/download");
        url.set_query(Some(&format!("url={}", image_url)));
        url
    }

    fn delete(&self, key: &str, token: &str) -> Url {
        let mut url = self.upstream.clone();
        url.set_path(&format!("/image/delete/{}/{}", token, key));
        url
    }
}

impl fmt::Display for ImageType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s = match self {
            Self::Png => "png",
            Self::Jpeg => "jpg",
            Self::Webp => "webp",
        };

        write!(f, "{}", s)
    }
}

impl fmt::Debug for State {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("State")
            .field("upstream", &self.upstream)
            .field("image_info", &"Rc<dyn ImageInfo>")
            .finish()
    }
}