malice

pub mod idle;

use std::{future::Future, net::SocketAddr, time::Duration};

use axum::{

    Router,

    http::{HeaderName, HeaderValue, StatusCode, header},

    middleware,

    middleware::Next,

    response::{IntoResponse, Response},

    routing::any,

};

use listenfd::ListenFd;

use tower_http::{

    services::{ServeDir, ServeFile},

    trace::TraceLayer,

};

use crate::idle::IdleTracker;

fn bind_addr() -> SocketAddr {

    std::env::var("MALICE_BIND")

        .unwrap_or_else(|_| "127.0.0.1:3000".to_owned())

        .parse()

        .expect("MALICE_BIND must be host:port")

}

fn idle_exit_secs() -> Option<u64> {

    std::env::var("MALICE_IDLE_EXIT_SECS")

        .ok()

        .and_then(|value| value.trim().parse().ok())

}

/// Takes over a systemd-activated listener from fd 3 when present, so the

/// socket unit can hold the port before the process ever starts. Falls back

/// to a plain bind otherwise.

pub async fn bind_listener() -> tokio::net::TcpListener {

    let mut listenfd = ListenFd::from_env();

    if let Some(listener) = listenfd

        .take_tcp_listener(0)

        .expect("inspecting inherited sockets")

    {

        listener

            .set_nonblocking(true)

            .expect("setting inherited socket nonblocking");

        return tokio::net::TcpListener::from_std(listener)

            .expect("adopting inherited socket into tokio");

    }

    let addr = bind_addr();

    tracing::info!("listening on {addr}");

    tokio::net::TcpListener::bind(addr)

        .await

        .expect("binding the listener must succeed")

}

/// Serves `app` on `listener`, exiting on ctrl-c or, when

/// `MALICE_IDLE_EXIT_SECS` is set, after that long with no request in flight

/// or started. `on_shutdown` runs after the exit signal either way.

pub async fn serve(

    listener: tokio::net::TcpListener,

    app: Router,

    on_shutdown: impl Future<Output = ()> + Send + 'static,

) {

    let idle = idle_exit_secs().map(|secs| (IdleTracker::new(), secs));

    let app = match &idle {

        Some((tracker, _)) => {

            app.layer(middleware::from_fn_with_state(tracker.clone(), idle::track))

        }

        None => app,

    };

    axum::serve(listener, app)

        .with_graceful_shutdown(async move {

            match idle {

                Some((tracker, secs)) => {

                    tokio::select! {

                        _ = tokio::signal::ctrl_c() => {}

                        _ = tracker.wait_idle(Duration::from_secs(secs)) => {

                            tracing::info!("idle for {secs}s with no requests, exiting");

                        }

                    }

                }

                None => {

                    let _ = tokio::signal::ctrl_c().await;

                }

            }

            on_shutdown.await;

        })

        .await

        .expect("serving must succeed");

}

async fn api_not_found() -> StatusCode {

    StatusCode::NOT_FOUND

}

async fn missing_frontend() -> Response {

    (

        StatusCode::SERVICE_UNAVAILABLE,

        "frontend assets missing; set MALICE_WEB_DIR or use the Nix package",

    )

        .into_response()

}

fn is_no_cache_path(path: &str) -> bool {

    matches!(path, "/" | "/index.html" | "/sw.js")

}

/// index.html and sw.js must never be cached, or a stale service worker or

/// app shell survives a rebuild; the rest of the bundle is hashed and can be

/// cached freely.

async fn static_headers(request: axum::extract::Request, next: Next) -> Response {

    let no_cache = is_no_cache_path(request.uri().path());

    let mut response = next.run(request).await;

    if no_cache {

        let headers = response.headers_mut();

        headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-cache"));

        headers.insert(

            HeaderName::from_static("service-worker-allowed"),

            HeaderValue::from_static("/"),

        );

    }

    response

}

/// Serves only the built frontend: no `/api` routes, no session state, no

/// data dir. Used for the socket-activated on-device build.

pub fn static_router(web_dir: Option<String>) -> Router {

    let router = Router::new()

        .route("/api", any(api_not_found))

        .route("/api/{*rest}", any(api_not_found))

        .layer(TraceLayer::new_for_http());

    match web_dir {

        Some(dir) => {

            let files =

                ServeDir::new(&dir).not_found_service(ServeFile::new(format!("{dir}/index.html")));

            router

                .fallback_service(files)

                .layer(middleware::from_fn(static_headers))

        }

        None => router.fallback(missing_frontend),

    }

}