site-backend/src/main.rs

use std::{future::ready, net::SocketAddr};

use axum::{
    extract::Extension,
    http::StatusCode,
    middleware,
    routing::{get, get_service},
    Router,
};
use tower_http::{auth::RequireAuthorizationLayer, services::ServeDir, trace::TraceLayer};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

mod api;
mod matrix;
mod metrics;

/// Name of the directory where static sites are stored inside the data directory
const STATIC_DIR_NAME: &str = "static";
/// Name of the subdir of STATIC_DIR_NAME where the default (fallback) site is located
const DEFAULT_STATIC_SITE: &str = "default";

#[tokio::main]
async fn main() {
    // Enable tracing
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::new(
            std::env::var("RUST_LOG").unwrap_or_else(|_| "site=debug,tower_http=debug".into()),
        ))
        .with(tracing_subscriber::fmt::layer())
        .init();

    // Get required variables from env vars
    let api_key = std::env::var("API_KEY").expect("No API_KEY was provided.");
    let data_dir = std::env::var("DATA_DIR").expect("No DATA_DIR was provided.");
    let static_dir = format!("{}/{}", data_dir, STATIC_DIR_NAME);

    std::fs::create_dir_all(&static_dir);

    // Initialize metrics
    let recorder_handle = metrics::setup_metrics_recorder();

    let mut app = Router::new()
        // Handle Matrix .well-known files
        .nest("/", matrix::router())
        // Routes under /api path
        .nest(
            "/api",
            api::router().layer(RequireAuthorizationLayer::bearer(&api_key)),
        )
        .route("/metrics", get(move || ready(recorder_handle.render())));

    // Each static site gets mounted explicitely so that the default site can be used as fallback
    // Each entry is of the form (route, static dir name)
    let sites = [
        ("/docs/vieter", "docs-vieter"),
        ("/api-docs/vieter", "api-docs-vieter"),
    ];

    for (path, dir) in sites {
        let full_path = format!("{}/{}", static_dir, dir);

        app = app.nest(
            path,
            get_service(ServeDir::new(full_path)).handle_error(
                |error: std::io::Error| async move {
                    (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        format!("Unhandled internal error: {}", error),
                    )
                },
            ),
        );
    }

    app = app
        // The fallback option is to serve the actual static files
        .fallback(
            get_service(ServeDir::new(format!(
                "{}/{}",
                static_dir, DEFAULT_STATIC_SITE
            )))
            .handle_error(|error: std::io::Error| async move {
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Unhandled internal error: {}", error),
                )
            }),
        )
        .layer(middleware::from_fn(metrics::track_metrics))
        .layer(Extension(data_dir))
        .layer(TraceLayer::new_for_http());

    let addr = SocketAddr::from(([0, 0, 0, 0], 3000));
    tracing::debug!("listening on {}", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}