tetratto/crates/app/src/main.rs

#![doc = include_str!("../../../README.md")]
#![doc(html_favicon_url = "/public/favicon.svg")]
#![doc(html_logo_url = "/public/tetratto_bunny.webp")]
mod assets;
mod cookie;
mod image;
mod macros;
mod routes;
mod sanitize;

use assets::{init_dirs, write_assets};
use stripe::Client as StripeClient;
use tetratto_core::model::{
    permissions::{FinePermission, SecondaryPermission},
    uploads::CustomEmoji,
};
pub use tetratto_core::*;

use axum::{
    http::{HeaderName, HeaderValue},
    Extension, Router,
};
use reqwest::Client;
use tera::{Tera, Value};
use tower_http::{
    catch_panic::CatchPanicLayer,
    set_header::SetResponseHeaderLayer,
    trace::{self, TraceLayer},
};
use tracing::{Level, info};

use std::{collections::HashMap, env::var, net::SocketAddr, process::exit, sync::Arc};
use tokio::sync::RwLock;

pub(crate) type InnerState = (DataManager, Tera, Client, Option<StripeClient>);
pub(crate) type State = Arc<RwLock<InnerState>>;

fn render_markdown(value: &Value, _: &HashMap<String, Value>) -> tera::Result<Value> {
    Ok(
        tetratto_shared::markdown::render_markdown(&CustomEmoji::replace(value.as_str().unwrap()))
            .replace("\\@", "@")
            .replace("%5C@", "@")
            .into(),
    )
}

fn render_emojis(value: &Value, _: &HashMap<String, Value>) -> tera::Result<Value> {
    Ok(CustomEmoji::replace(value.as_str().unwrap()).into())
}

fn color_escape(value: &Value, _: &HashMap<String, Value>) -> tera::Result<Value> {
    Ok(sanitize::color_escape(value.as_str().unwrap()).into())
}

fn check_supporter(value: &Value, _: &HashMap<String, Value>) -> tera::Result<Value> {
    Ok(FinePermission::from_bits(value.as_u64().unwrap() as u32)
        .unwrap()
        .check(FinePermission::SUPPORTER)
        .into())
}

fn check_dev_pass(value: &Value, _: &HashMap<String, Value>) -> tera::Result<Value> {
    Ok(
        SecondaryPermission::from_bits(value.as_u64().unwrap() as u32)
            .unwrap()
            .check(SecondaryPermission::DEVELOPER_PASS)
            .into(),
    )
}

fn check_staff_badge(value: &Value, _: &HashMap<String, Value>) -> tera::Result<Value> {
    Ok(FinePermission::from_bits(value.as_u64().unwrap() as u32)
        .unwrap()
        .check(FinePermission::STAFF_BADGE)
        .into())
}

fn check_banned(value: &Value, _: &HashMap<String, Value>) -> tera::Result<Value> {
    Ok(FinePermission::from_bits(value.as_u64().unwrap() as u32)
        .unwrap()
        .check_banned()
        .into())
}

fn remove_script_tags(value: &Value, _: &HashMap<String, Value>) -> tera::Result<Value> {
    Ok(value
        .as_str()
        .unwrap()
        .replace("</script>", "&lt;/script&gt;")
        .into())
}

#[tokio::main(flavor = "multi_thread")]
async fn main() {
    tracing_subscriber::fmt()
        .with_target(false)
        .compact()
        .init();

    let mut config = config::Config::get_config();
    if let Ok(port) = var("PORT") {
        let port = port.parse::<u16>().expect("port should be a u16");
        config.port = port;
    }

    // init
    init_dirs(&config).await;
    let html_path = write_assets(&config).await;

    // ...
    let database = DataManager::new(config.clone()).await.unwrap();
    database.init().await.unwrap();

    let mut tera = match Tera::new(&format!("{html_path}/**/*")) {
        Ok(t) => t,
        Err(e) => {
            println!("{e}");
            exit(1);
        }
    };

    tera.register_filter("markdown", render_markdown);
    tera.register_filter("color", color_escape);
    tera.register_filter("has_supporter", check_supporter);
    tera.register_filter("has_dev_pass", check_dev_pass);
    tera.register_filter("has_staff_badge", check_staff_badge);
    tera.register_filter("has_banned", check_banned);
    tera.register_filter("remove_script_tags", remove_script_tags);
    tera.register_filter("emojis", render_emojis);

    let client = Client::new();
    let mut app = Router::new();

    // create stripe client
    let stripe_client = if let Some(ref stripe) = config.stripe {
        Some(StripeClient::new(stripe.secret.clone()))
    } else {
        None
    };

    // add correct routes
    if var("LITTLEWEB").is_ok() {
        app = app.merge(routes::lw_routes());
    } else {
        app = app
            .merge(routes::routes(&config))
            .layer(SetResponseHeaderLayer::if_not_present(
                HeaderName::from_static("content-security-policy"),
                HeaderValue::from_static("default-src 'self' *.spotify.com musicbrainz.org; img-src * data:; media-src *; font-src *; style-src 'unsafe-inline' 'self' *; script-src 'self' 'unsafe-inline' *; worker-src * blob:; object-src 'self' *; upgrade-insecure-requests; connect-src * localhost; frame-src 'self' blob: *; frame-ancestors 'self'"),
            ));
    }

    // add junk
    app = app
        .layer(Extension(Arc::new(RwLock::new((
            database,
            tera,
            client,
            stripe_client,
        )))))
        .layer(axum::extract::DefaultBodyLimit::max(
            var("BODY_LIMIT")
                .unwrap_or("8388608".to_string())
                .parse::<usize>()
                .unwrap(),
        ))
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
                .on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
        )
        .layer(CatchPanicLayer::new());

    // ...
    let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{}", config.port))
        .await
        .unwrap();

    info!("🐇 tetratto.");
    info!("listening on http://0.0.0.0:{}", config.port);
    axum::serve(
        listener,
        app.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .await
    .unwrap();
}