datatrash/src/main.rs

mod config;
mod db;
mod deleter;
mod download;
mod mime_relations;
mod multipart;
mod rate_limit;
mod template;
mod upload;

use crate::rate_limit::ForwardedPeerIpKeyExtractor;
use actix_files::Files;
use actix_governor::{Governor, GovernorConfigBuilder};
use actix_web::{
    http::header::{HeaderName, HeaderValue, CONTENT_SECURITY_POLICY, X_CONTENT_TYPE_OPTIONS},
    middleware::{self, DefaultHeaders, Logger},
    web::{self, Data},
    App, Error, HttpResponse, HttpServer,
};
use env_logger::Env;
use sqlx::postgres::PgPool;
use std::env;
use tokio::{sync::mpsc::channel, task};

const DEFAULT_CSP: (HeaderName, &str) = (
    CONTENT_SECURITY_POLICY,
    "default-src 'none'; connect-src 'self'; img-src 'self'; media-src 'self'; font-src 'self'; script-src 'self'; style-src 'self'; object-src 'none'; base-uri 'self'; frame-src 'none'; frame-ancestors 'none'; form-action 'self';"
);

async fn not_found() -> Result<HttpResponse, Error> {
    Ok(HttpResponse::NotFound()
        .content_type("text/plain")
        .body("not found"))
}

#[tokio::main]
async fn main() -> std::io::Result<()> {
    env_logger::Builder::from_env(Env::default().default_filter_or("info,sqlx=warn")).init();

    let pool: PgPool = db::setup().await;
    let config = config::from_env().await;
    let (sender, receiver) = channel(8);

    log::info!("omnomnom");

    let db = web::Data::new(pool.clone());
    let expiry_watch_sender = web::Data::new(sender);
    let bind_address = env::var("BIND_ADDRESS").unwrap_or_else(|_| "0.0.0.0:8000".to_owned());

    task::spawn(deleter::delete_old_files(
        receiver,
        pool,
        config.files_dir.clone(),
    ));

    template::write_prefillable_templates(&config).await;
    let config = Data::new(config);

    let governor_conf = GovernorConfigBuilder::default()
        .per_second(config.rate_limit_replenish_seconds)
        .burst_size(config.rate_limit_burst)
        .key_extractor(ForwardedPeerIpKeyExtractor {
            proxied: config.proxied,
        })
        .use_headers()
        .finish()
        .unwrap();

    HttpServer::new({
        move || {
            let app = App::new()
                .wrap(Logger::new(r#"%{r}a "%r" =%s %bbytes %Tsec"#))
                .wrap(
                    DefaultHeaders::new()
                        .add(DEFAULT_CSP)
                        .add((X_CONTENT_TYPE_OPTIONS, HeaderValue::from_static("nosniff"))),
                )
                .wrap(middleware::Compress::default())
                .app_data(db.clone())
                .app_data(expiry_watch_sender.clone())
                .app_data(config.clone())
                .service(web::resource("/").route(web::get().to(upload::index)))
                .service(web::resource("/upload").route(web::post().to(upload::upload)))
                .service(
                    web::resource(["/upload/{id}", "/upload/{id}/{name}"])
                        .route(web::get().to(upload::uploaded)),
                )
                .service(Files::new("/static", "static").disable_content_disposition())
                .default_service(web::route().to(not_found));
            if config.enable_rate_limit {
                app.service(
                    web::resource([
                        "/{id:[a-z0-9]{5}}",
                        "/{id:[a-z0-9]{5}}/",
                        "/{id:[a-z0-9]{5}}/{name}",
                    ])
                    .wrap(Governor::new(&governor_conf))
                    .route(web::get().to(download::download)),
                )
            } else {
                app.service(
                    web::resource([
                        "/{id:[a-z0-9]{5}}",
                        "/{id:[a-z0-9]{5}}/",
                        "/{id:[a-z0-9]{5}}/{name}",
                    ])
                    .route(web::get().to(download::download)),
                )
            }
        }
    })
    .bind(bind_address)?
    .run()
    .await
}