malice

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

use axum::{Router, middleware};

use listenfd::ListenFd;

use malice::{

    idle::IdleTracker,

    routes::{router, static_router},

    state::{AppState, ensure_dirs, load_or_create_master_key},

};

fn env_flag(name: &str, default: bool) -> bool {

    match std::env::var(name) {

        Ok(value) => {

            let value = value.trim().to_ascii_lowercase();

            match value.as_str() {

                "1" | "true" | "yes" | "on" => true,

                "0" | "false" | "no" | "off" => false,

                _ => default,

            }

        }

        Err(_) => default,

    }

}

fn data_dir() -> PathBuf {

    std::env::var("MALICE_DATA_DIR")

        .map(PathBuf::from)

        .unwrap_or_else(|_| PathBuf::from("local-data"))

}

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())

}

fn web_dir() -> Option<String> {

    if let Ok(dir) = std::env::var("MALICE_WEB_DIR") {

        assert!(

            std::path::Path::new(&dir).join("index.html").is_file(),

            "MALICE_WEB_DIR must point to built frontend assets containing index.html"

        );

        return Some(dir);

    }

    let local = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("web/dist");

    local

        .join("index.html")

        .is_file()

        .then(|| local.to_string_lossy().into_owned())

}

/// 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.

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.

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(),

            malice::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 run_hosted() {

    let data_dir = data_dir();

    ensure_dirs(&data_dir);

    let master_key =

        load_or_create_master_key(&data_dir).expect("loading the master key must succeed");

    let cookie_secure = env_flag("MALICE_SESSION_COOKIE_SECURE", true);

    let allow_http_homeserver = env_flag("MALICE_ALLOW_HTTP_HOMESERVER", false);

    let state = AppState::new(data_dir, master_key, cookie_secure, allow_http_homeserver);

    let app = router(state.clone(), web_dir());

    let listener = bind_listener().await;

    serve(

        listener,

        app,

        async move { state.shutdown_gracefully().await },

    )

    .await;

}

async fn run_static_only() {

    let app = static_router(web_dir());

    let listener = bind_listener().await;

    serve(listener, app, async {}).await;

}

#[tokio::main]

async fn main() {

    tracing_subscriber::fmt()

        .with_env_filter(

            tracing_subscriber::EnvFilter::try_from_default_env()

                .unwrap_or_else(|_| "malice=info,tower_http=info".into()),

        )

        .with_target(false)

        .compact()

        .init();

    if env_flag("MALICE_STATIC_ONLY", false) {

        run_static_only().await;

    } else {

        run_hosted().await;

    }

}