pict-rs-proxy/src/main.rs

use actix_web::{
    body::BodyStream,
    client::Client,
    http::{
        header::{CacheControl, CacheDirective, ContentType, LOCATION},
        StatusCode,
    },
    middleware::Logger,
    web, App, HttpRequest, HttpResponse, HttpServer,
};
use once_cell::sync::Lazy;
use std::{io::Cursor, net::SocketAddr};
use structopt::StructOpt;
use url::Url;

include!(concat!(env!("OUT_DIR"), "/templates.rs"));

const HOURS: u32 = 60 * 60;

#[derive(Clone, Debug, StructOpt)]
struct Config {
    #[structopt(
        short,
        long,
        env = "PICTRS_PROXY_ADDR",
        default_value = "0.0.0.0:8081",
        help = "The address and port the server binds to"
    )]
    addr: SocketAddr,

    #[structopt(
        short,
        long,
        env = "PICTRS_PROXY_UPSTREAM",
        default_value = "http://localhost:8080",
        help = "The url of the upstream pict-rs server"
    )]
    upstream: Url,

    #[structopt(
        short,
        long,
        env = "PICTRS_PROXY_DOMAIN",
        default_value = "http://localhost:8081",
        help = "The scheme, domain, and optional port of the pict-rs proxy server"
    )]
    domain: Url,
}

impl Config {
    fn upstream_upload_url(&self) -> String {
        let mut url = self.upstream.clone();
        url.set_path("image");

        url.to_string()
    }

    fn upstream_image_url(&self, name: &str) -> String {
        let mut url = self.upstream.clone();
        url.set_path(&format!("image/{}", name));

        url.to_string()
    }

    fn upstream_thumbnail_url(&self, size: u64, name: &str) -> String {
        let mut url = self.upstream.clone();
        url.set_path(&format!("image/thumbnail{}/{}", size, name));

        url.to_string()
    }

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

        url.to_string()
    }

    fn image_url(&self, name: &str) -> String {
        let mut url = self.domain.clone();
        url.set_path(&format!("image/{}", name));

        url.to_string()
    }

    fn thumbnail_url(&self, size: u64, name: &str) -> String {
        let mut url = self.domain.clone();
        url.set_path(&format!("thumb/{}/{}", size, name));

        url.to_string()
    }

    fn delete_url(&self, token: &str, name: &str) -> String {
        let mut url = self.domain.clone();
        url.set_path(&format!("delete/{}/{}", token, name));

        url.to_string()
    }
}

static CONFIG: Lazy<Config> = Lazy::new(|| Config::from_args());

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

impl Images {
    fn files(&self) -> Option<&[Image]> {
        self.files.as_ref().map(|v| v.as_ref())
    }

    fn msg(&self) -> &str {
        &self.msg
    }

    fn is_ok(&self) -> bool {
        self.files().is_some()
    }
}

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

impl Image {
    fn filename(&self) -> &str {
        &self.file
    }

    fn link(&self) -> String {
        CONFIG.image_url(&self.file)
    }

    fn thumb(&self, size: u64) -> String {
        CONFIG.thumbnail_url(size, &self.file)
    }

    fn delete(&self) -> String {
        CONFIG.delete_url(&self.delete_token, &self.file)
    }
}

fn statics(file: &str) -> String {
    format!("/static/{}", file)
}

async fn index() -> Result<HttpResponse, actix_web::Error> {
    let mut cursor = Cursor::new(vec![]);
    self::templates::index(&mut cursor, "/upload", "images[]")?;
    Ok(HttpResponse::Ok()
        .content_type(mime::TEXT_HTML.essence_str())
        .body(cursor.into_inner()))
}

async fn upload(
    req: HttpRequest,
    body: web::Payload,
    client: web::Data<Client>,
) -> Result<HttpResponse, actix_web::Error> {
    let mut res = client
        .request_from(CONFIG.upstream_upload_url(), req.head())
        .if_some(req.head().peer_addr, |addr, req| {
            req.header("X-Forwarded-For", addr.to_string())
        })
        .send_stream(body)
        .await?;

    let images = res.json::<Images>().await?;

    let mut cursor = Cursor::new(vec![]);
    self::templates::images(&mut cursor, images)?;

    Ok(HttpResponse::build(res.status())
        .content_type(mime::TEXT_HTML.essence_str())
        .body(cursor.into_inner()))
}

async fn image(
    url: String,
    req: HttpRequest,
    client: web::Data<Client>,
) -> Result<HttpResponse, actix_web::Error> {
    let res = client
        .request_from(url, req.head())
        .if_some(req.head().peer_addr, |addr, req| {
            req.header("X-Forwarded-For", addr.to_string())
        })
        .no_decompress()
        .send()
        .await?;

    if res.status() == StatusCode::NOT_FOUND {
        return Ok(to_404());
    }

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

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

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

async fn thumbnail(
    parts: web::Path<(u64, String)>,
    req: HttpRequest,
    client: web::Data<Client>,
) -> Result<HttpResponse, actix_web::Error> {
    let (size, file) = parts.into_inner();

    if (size % 100) == 0 {
        let url = CONFIG.upstream_thumbnail_url(size, &file);

        return image(url, req, client).await;
    }

    Ok(to_404())
}

async fn full_res(
    filename: web::Path<String>,
    req: HttpRequest,
    client: web::Data<Client>,
) -> Result<HttpResponse, actix_web::Error> {
    let url = CONFIG.upstream_image_url(&filename.into_inner());

    image(url, req, client).await
}

async fn static_files(filename: web::Path<String>) -> HttpResponse {
    let filename = filename.into_inner();

    if let Some(data) = self::templates::statics::StaticFile::get(&filename) {
        return HttpResponse::Ok()
            .set(CacheControl(vec![
                CacheDirective::Public,
                CacheDirective::MaxAge(24 * HOURS),
                CacheDirective::Extension("immutable".to_owned(), None),
            ]))
            .set(ContentType(data.mime.clone()))
            .body(data.content);
    }

    to_404()
}

async fn delete(
    components: web::Path<(String, String)>,
    req: HttpRequest,
    client: web::Data<Client>,
) -> Result<HttpResponse, actix_web::Error> {
    let (token, file) = components.into_inner();

    let url = CONFIG.upstream_delete_url(&token, &file);
    client
        .request_from(url, req.head())
        .if_some(req.head().peer_addr, |addr, req| {
            req.header("X-Forwarded-For", addr.to_string())
        })
        .no_decompress()
        .send()
        .await?;

    Ok(go_home().await)
}

fn to_404() -> HttpResponse {
    HttpResponse::TemporaryRedirect()
        .header(LOCATION, "/404")
        .finish()
}

async fn not_found() -> Result<HttpResponse, actix_web::Error> {
    let mut cursor = Cursor::new(vec![]);

    self::templates::not_found(&mut cursor)?;

    Ok(HttpResponse::NotFound().body(cursor.into_inner()))
}

async fn go_home() -> HttpResponse {
    HttpResponse::TemporaryRedirect()
        .header(LOCATION, "/")
        .finish()
}

#[actix_rt::main]
async fn main() -> Result<(), anyhow::Error> {
    dotenv::dotenv().ok();

    if std::env::var("RUST_LOG").is_err() {
        std::env::set_var("RUST_LOG", "info");
    }

    env_logger::init();

    HttpServer::new(move || {
        let client = Client::build()
            .header("User-Agent", "pict-rs-frontend, v0.1.0")
            .finish();

        App::new()
            .data(client)
            .wrap(Logger::default())
            .service(web::resource("/").route(web::get().to(index)))
            .service(web::resource("/upload").route(web::post().to(upload)))
            .service(web::resource("/image/{filename}").route(web::get().to(full_res)))
            .service(web::resource("/thumb/{size}/{filename}").route(web::get().to(thumbnail)))
            .service(web::resource("/static/{filename}").route(web::get().to(static_files)))
            .service(web::resource("/delete/{token}/{filename}").route(web::get().to(delete)))
            .service(web::resource("/404").route(web::get().to(not_found)))
            .default_service(web::get().to(go_home))
    })
    .bind(CONFIG.addr)?
    .run()
    .await?;

    Ok(())
}