compliance-scanner-agent/compliance-agent/src/api/server.rs

use std::sync::Arc;

use axum::extract::Request;
use axum::http::HeaderValue;
use axum::middleware::Next;
use axum::response::Response;
use axum::routing::{delete, get};
use axum::{middleware, Extension, Router};
use tokio::sync::RwLock;
use tower_http::cors::CorsLayer;
use tower_http::set_header::SetResponseHeaderLayer;
use tower_http::trace::TraceLayer;

use compliance_core::auth::{require_jwt_auth, require_tenant_status, JwksState};
use compliance_core::{TenantContext, TenantStatus};

use crate::agent::ComplianceAgent;
use crate::api::handlers;
use crate::api::routes;
use crate::error::AgentError;

/// Synthetic tenant id used when Keycloak isn't configured (local dev,
/// `cargo run` against a bare Mongo). Lets the handler stack stay
/// uniformly tenant-scoped without the operator having to spin up KC
/// just to poke at the API. Override via `DEV_TENANT_ID`.
const DEFAULT_DEV_TENANT_ID: &str = "dev";

/// Inject a synthetic [`TenantContext`] for any request that lacks one.
/// Only mounted when Keycloak is NOT configured; with KC, the real
/// `require_jwt_auth` middleware owns this and we never reach here
/// without a context.
///
/// Public so the integration-test harness can mount it without
/// duplicating the synthetic-context shape.
pub async fn inject_dev_tenant(mut request: Request, next: Next) -> Response {
    if request.extensions().get::<TenantContext>().is_none() {
        let tenant_id =
            std::env::var("DEV_TENANT_ID").unwrap_or_else(|_| DEFAULT_DEV_TENANT_ID.to_string());
        let ctx = TenantContext {
            tenant_slug: tenant_id.clone(),
            tenant_id,
            org_roles: vec![],
            products: vec![],
            plan: "dev".to_string(),
            status: TenantStatus::Active,
            user_id: "dev-user".to_string(),
            user_name: None,
        };
        request.extensions_mut().insert(ctx);
    }
    next.run(request).await
}

pub async fn start_api_server(agent: ComplianceAgent, port: u16) -> Result<(), AgentError> {
    // Admin sub-router. Routes are only mounted when ADMIN_API_TOKEN is
    // configured — without it, the paths don't exist at all (404 rather
    // than 401), so an operator who hasn't opted in can't fingerprint
    // the surface area.
    let admin_router: Router = if agent.config.admin_api_token.is_some() {
        tracing::info!("Admin API enabled — /api/v1/admin/* mounted behind ADMIN_API_TOKEN bearer");
        Router::new()
            .route(
                "/api/v1/admin/tenants",
                get(handlers::admin::list_tenant_dbs),
            )
            .route(
                "/api/v1/admin/tenants/{tenant_id}",
                delete(handlers::admin::drop_tenant_db),
            )
            .layer(middleware::from_fn(handlers::admin::require_admin_token))
    } else {
        Router::new()
    };

    let mut app = routes::build_router()
        .merge(admin_router)
        .layer(Extension(Arc::new(agent.clone())))
        .layer(CorsLayer::permissive())
        .layer(TraceLayer::new_for_http())
        // Security headers (defense-in-depth, primary enforcement via Traefik)
        .layer(SetResponseHeaderLayer::overriding(
            axum::http::header::STRICT_TRANSPORT_SECURITY,
            HeaderValue::from_static("max-age=31536000; includeSubDomains"),
        ))
        .layer(SetResponseHeaderLayer::overriding(
            axum::http::header::X_FRAME_OPTIONS,
            HeaderValue::from_static("DENY"),
        ))
        .layer(SetResponseHeaderLayer::overriding(
            axum::http::header::X_CONTENT_TYPE_OPTIONS,
            HeaderValue::from_static("nosniff"),
        ))
        .layer(SetResponseHeaderLayer::overriding(
            axum::http::header::REFERRER_POLICY,
            HeaderValue::from_static("strict-origin-when-cross-origin"),
        ));

    if let (Some(kc_url), Some(kc_realm)) =
        (&agent.config.keycloak_url, &agent.config.keycloak_realm)
    {
        let jwks_url = format!("{kc_url}/realms/{kc_realm}/protocol/openid-connect/certs");
        let jwks_state = JwksState {
            jwks: Arc::new(RwLock::new(None)),
            jwks_url,
        };
        tracing::info!("Keycloak JWT auth enabled for realm '{kc_realm}'");
        // Layers execute outermost-first. Extension(jwks_state) must run
        // before require_jwt_auth so the middleware can read it; the
        // status gate runs after JWT so TenantContext is in extensions.
        app = app
            .layer(middleware::from_fn(require_tenant_status))
            .layer(middleware::from_fn(require_jwt_auth))
            .layer(Extension(jwks_state));
    } else {
        let tenant_id =
            std::env::var("DEV_TENANT_ID").unwrap_or_else(|_| DEFAULT_DEV_TENANT_ID.to_string());
        tracing::warn!(
            tenant_id = %tenant_id,
            "Keycloak not configured — running unauthenticated against the dev tenant. \
             DO NOT use in any environment with real customer data."
        );
        app = app.layer(middleware::from_fn(inject_dev_tenant));
    }

    let addr = format!("0.0.0.0:{port}");
    let listener = tokio::net::TcpListener::bind(&addr)
        .await
        .map_err(|e| AgentError::Other(format!("Failed to bind to {addr}: {e}")))?;

    tracing::info!("REST API listening on {addr}");
    axum::serve(listener, app)
        .await
        .map_err(|e| AgentError::Other(format!("API server error: {e}")))?;

    Ok(())
}