magnetar/fe_calckey/src/main.rs

mod frontend_api;
mod frontend_render;
mod manifest;
mod static_serve;
mod summary_proxy;

use crate::frontend_api::create_frontend_api;
use crate::frontend_render::{new_frontend_render_router, FrontendConfig};
use crate::manifest::handle_manifest;
use crate::static_serve::{static_serve, static_serve_svg, static_serve_sw};
use crate::summary_proxy::generate_summary;
use axum::routing::{any, get};
use axum::Router;
use axum_extra::TypedHeader;
use dotenvy::dotenv;
use headers::CacheControl;
use hyper::StatusCode;
use magnetar_common::config::MagnetarConfig;
use miette::{miette, IntoDiagnostic};
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tera::Tera;
use thiserror::Error;
use tokio::net::TcpListener;
use tokio::signal;
use tower_http::services::ServeFile;
use tower_http::trace::TraceLayer;
use tracing::error;
use tracing::log::info;
use tracing_subscriber::EnvFilter;

#[derive(Debug, Error)]
pub enum RouterSetupError {
    #[error("Failed to load template: {0}")]
    TemplateLoadError(#[from] tera::Error),
}

fn new_calckey_fe_router(config: &'static MagnetarConfig) -> Result<Router, RouterSetupError> {
    let tera = Tera::new("fe_calckey/frontend/assets-be/template/**/*.html")?;
    let frontend_renderer_config = FrontendConfig {
        magnetar_config: config,
        templater: Arc::new(tera),
    };

    Ok(Router::new()
        .nest_service(
            "/favicon.ico",
            ServeFile::new("fe_calckey/frontend/assets/favicon.ico"),
        )
        .nest_service(
            "/favicon.png",
            ServeFile::new("fe_calckey/frontend/assets/favicon.png"),
        )
        .nest_service(
            "/favicon.svg",
            ServeFile::new("fe_calckey/frontend/assets/favicon.svg"),
        )
        .route("/manifest.json", get(handle_manifest).with_state(config))
        .nest(
            "/sw.js",
            static_serve_sw("fe_calckey/frontend/built/_sw_dist_/sw.js"),
        )
        .nest("/static-assets", static_serve("fe_calckey/frontend/assets"))
        .nest(
            "/client-assets",
            static_serve("fe_calckey/frontend/client/assets"),
        )
        .nest(
            "/assets",
            static_serve("fe_calckey/frontend/built/_client_dist_"),
        )
        .nest(
            "/twemoji",
            static_serve_svg("fe_calckey/frontend/assets-be/twemoji"),
        )
        .nest("/fe-api", create_frontend_api("fe_calckey/frontend/assets"))
        .route("/url", get(generate_summary))
        .route(
            "/streaming",
            any(|| async {
                (
                    StatusCode::SERVICE_UNAVAILABLE,
                    TypedHeader(
                        CacheControl::new()
                            .with_max_age(Duration::from_secs(0))
                            .with_private(),
                    ),
                )
            }),
        )
        .nest("", new_frontend_render_router(frontend_renderer_config)))
}

#[tokio::main]
async fn main() -> miette::Result<()> {
    dotenv().ok();

    let filter_layer = EnvFilter::try_from_default_env()
        .or_else(|_| EnvFilter::try_new("info"))
        .unwrap();

    tracing_subscriber::fmt()
        .with_env_filter(filter_layer)
        .with_test_writer()
        .init();

    let config = &*Box::leak::<'static>(Box::new(
        magnetar_common::config::load_config().map_err(|e| miette!(e))?,
    ));
    info!("Loaded configuration: {config:#?}");

    let app = Router::new()
        .layer(TraceLayer::new_for_http())
        .nest("", new_calckey_fe_router(config).into_diagnostic()?);

    let addr = SocketAddr::from((
        config.calckey_frontend.bind_addr,
        config.calckey_frontend.port,
    ));
    tracing::info!("Binding to: {addr}");
    let listener = TcpListener::bind(addr).await.into_diagnostic()?;
    tracing::info!("Serving...");
    axum::serve(listener, app.into_make_service())
        .await
        .map_err(|e| miette!("Error running server: {}", e))
}

// FIXME: Plug this back in when Axum reimplements graceful shutdown

async fn shutdown_signal() {
    let ctrl_c = async {
        if let Err(e) = signal::ctrl_c().await {
            error!("Ctrl+C signal handler error: {}", e);
        }
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("SIGTERM handler error")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }

    info!("Shutting down...");
}