AkurAI-IDP / src/main.rs

Latest change 8568faf1a542a58372fafd296dec8814018315de - AKURAI-IDP-4 remove legacy API health route by Ólafur Búi Ólafsson

#![allow(clippy::result_large_err)]
#![allow(clippy::too_many_arguments)]
#![allow(special_module_name)]

pub mod config;
pub mod db;
pub mod lib;
pub mod middleware;
pub mod routes;

use axum::{
    Extension, Json, Router,
    body::Body,
    extract::Request,
    http::{Response, StatusCode, header},
    middleware as axum_mw,
    response::{IntoResponse, Redirect},
    routing::get,
};
use serde_json::json;
use subtle::ConstantTimeEq;
use tower_http::trace::TraceLayer;

use crate::middleware::rate_limit::RateLimitState;

// ---------------------------------------------------------------------------
// Embedded static assets (compiled into the binary via include_bytes!)
// ---------------------------------------------------------------------------
const FAVICON_SVG: &[u8] = include_bytes!("../assets/favicon.svg");
const APPLE_TOUCH_ICON_PNG: &[u8] = include_bytes!("../assets/apple-touch-icon.png");
const THEME_JS: &[u8] = include_bytes!("../assets/theme.js");
const THEMES_CSS: &[u8] = include_bytes!("../assets/themes.css");
const THEMES_JSON: &[u8] = include_bytes!("../assets/themes.json");
const LANG_JS: &[u8] = include_bytes!("../assets/lang.js");
const AGENT_JS: &[u8] = include_bytes!("../assets/agent.js");

#[tokio::main]
async fn main() {
    // Init tracing
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();

    tracing::info!("akurai-idp starting up");

    // Ensure DB is initialized (triggers LazyLock)
    db::with_db(|_conn| {});

    // Ensure signing key exists
    ensure_signing_key();

    let cfg = config::get();

    // --- Build the router ---

    // Admin routes (token-auth'd)
    let admin_router = Router::new()
        .nest("/users", routes::admin::users::router())
        .nest("/groups", routes::admin::groups::router())
        .nest("/tenants", routes::admin::tenants::router())
        .nest("/clients", routes::admin::clients::router())
        .nest("/audit", routes::admin::audit::router())
        .layer(axum_mw::from_fn(admin_auth_middleware));

    // Build the full app using each module's router()
    let app = Router::new()
        // Auth pages (login/mfa/logout) with CSRF
        .merge(
            routes::auth_pages::router().layer(axum_mw::from_fn(middleware::csrf::csrf_protection)),
        )
        // OIDC discovery
        .merge(routes::well_known::router())
        // OIDC endpoints
        .merge(
            routes::authorize::router().layer(axum_mw::from_fn(middleware::auth::session_optional)),
        )
        .merge(routes::token::router())
        .merge(routes::userinfo::router().layer(axum_mw::from_fn(middleware::auth::bearer_auth)))
        .merge(routes::introspect::router())
        .merge(routes::revoke::router())
        // Authenticated agent console
        .merge(
            routes::agent::router()
                .layer(axum_mw::from_fn(middleware::auth::session_auth))
                .layer(axum_mw::from_fn(middleware::csrf::csrf_protection)),
        )
        // Account: the account router declares full `/account/...` paths, so it must
        // be merged at the root — nesting under "/account" double-prefixes every route
        // to `/account/account/...` and 404s the real ones (incl. the "Account" nav link).
        .merge(routes::account::router().layer(axum_mw::from_fn(middleware::auth::session_auth)))
        // Bare root: apps link to the issuer's bare domain for "Manage your account",
        // so send `/` to the account page (which itself redirects to /login if needed).
        .route("/", get(root_redirect))
        // Admin API
        .nest("/admin", admin_router)
        // Static assets
        .route("/favicon.svg", get(serve_favicon))
        .route("/apple-touch-icon.png", get(serve_apple_touch_icon))
        .route("/theme.js", get(serve_theme_js))
        .route("/themes.css", get(serve_themes_css))
        .route("/themes.json", get(serve_themes_json))
        .route("/lang.js", get(serve_lang_js))
        .route("/agent.js", get(serve_agent_js))
        // Health check
        .route("/health", get(health))
        // 404 catch-all
        .fallback(not_found)
        // Global middleware: secure headers + default rate limit + request logging
        .layer(axum_mw::from_fn(middleware::secure_headers::secure_headers))
        .layer(axum_mw::from_fn(middleware::rate_limit::rate_limit))
        .layer(TraceLayer::new_for_http())
        .layer(Extension(RateLimitState::new(60_000, 100)));

    tracing::info!("listening on {}", cfg.listen_addr);

    let listener = tokio::net::TcpListener::bind(&cfg.listen_addr)
        .await
        .expect("failed to bind listener");

    axum::serve(listener, app).await.expect("server error");
}

/// Admin auth middleware: requires Bearer token matching IDP_ADMIN_TOKEN (constant-time compare).
async fn admin_auth_middleware(req: Request, next: axum::middleware::Next) -> impl IntoResponse {
    let cfg = config::get();

    if cfg.admin_token.is_empty() {
        return (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(json!({"error": "admin API not configured (IDP_ADMIN_TOKEN not set)"})),
        )
            .into_response();
    }

    let auth_header = req
        .headers()
        .get(header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok());

    let token = match auth_header {
        Some(h) if h.starts_with("Bearer ") => &h[7..],
        _ => {
            return (
                StatusCode::UNAUTHORIZED,
                Json(json!({"error": "missing or invalid Authorization header"})),
            )
                .into_response();
        }
    };

    // Constant-time comparison
    let expected = cfg.admin_token.as_bytes();
    let provided = token.as_bytes();

    if expected.len() != provided.len() || expected.ct_eq(provided).unwrap_u8() != 1 {
        return (
            StatusCode::UNAUTHORIZED,
            Json(json!({"error": "invalid admin token"})),
        )
            .into_response();
    }

    next.run(req).await.into_response()
}

/// Ensure at least one active signing key exists. Generate one if the table is empty.
fn ensure_signing_key() {
    let has_active = db::with_db(|conn| {
        let count: i64 = conn
            .query_row(
                "SELECT COUNT(*) FROM signing_keys WHERE active = 1",
                [],
                |row| row.get(0),
            )
            .unwrap_or(0);
        count > 0
    });

    if !has_active {
        tracing::info!("no active signing key found, generating one");
        let (kid, pub_pem, priv_pem) = lib::jwt::generate_signing_key();
        let now = chrono::Utc::now().timestamp();

        db::with_db(|conn| {
            conn.execute(
                "INSERT INTO signing_keys (kid, alg, public_key, private_key_enc, active, created_at) VALUES (?1, ?2, ?3, ?4, 1, ?5)",
                rusqlite::params![kid, "EdDSA", pub_pem, priv_pem, now],
            )
            .expect("failed to insert signing key");
        });

        tracing::info!("signing key generated: kid={kid}");
    }
}

async fn serve_favicon() -> Response<Body> {
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "image/svg+xml")
        .header(header::CACHE_CONTROL, "public, max-age=86400")
        .body(Body::from(FAVICON_SVG.to_vec()))
        .unwrap()
}

async fn serve_apple_touch_icon() -> Response<Body> {
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "image/png")
        .header(header::CACHE_CONTROL, "public, max-age=86400")
        .body(Body::from(APPLE_TOUCH_ICON_PNG.to_vec()))
        .unwrap()
}

async fn serve_theme_js() -> Response<Body> {
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "text/javascript; charset=utf-8")
        .header(header::CACHE_CONTROL, "public, max-age=3600")
        .body(Body::from(THEME_JS.to_vec()))
        .unwrap()
}

async fn serve_themes_css() -> Response<Body> {
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "text/css; charset=utf-8")
        .header(header::CACHE_CONTROL, "public, max-age=3600")
        .body(Body::from(THEMES_CSS.to_vec()))
        .unwrap()
}

async fn serve_themes_json() -> Response<Body> {
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "application/json; charset=utf-8")
        .header(header::CACHE_CONTROL, "public, max-age=3600")
        .body(Body::from(THEMES_JSON.to_vec()))
        .unwrap()
}

async fn serve_lang_js() -> Response<Body> {
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "text/javascript; charset=utf-8")
        .header(header::CACHE_CONTROL, "public, max-age=3600")
        .body(Body::from(LANG_JS.to_vec()))
        .unwrap()
}

async fn serve_agent_js() -> Response<Body> {
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "text/javascript; charset=utf-8")
        .header(header::CACHE_CONTROL, "no-cache")
        .body(Body::from(AGENT_JS.to_vec()))
        .unwrap()
}

async fn health() -> impl IntoResponse {
    Json(json!({"ok": true, "service": "akurai-idp"}))
}

/// Bare root `/` → the account page. Relying-party apps point their "Manage your
/// account" links at the issuer's bare domain; without this the IDP 404s there.
async fn root_redirect() -> Redirect {
    Redirect::to("/account")
}

async fn not_found() -> impl IntoResponse {
    (StatusCode::NOT_FOUND, Json(json!({"error": "not found"})))
}