xecute-cli 0.1.1

use std::{net::SocketAddr, sync::Arc, time::Instant};
use axum::{
    extract::{ConnectInfo, DefaultBodyLimit, MatchedPath, Request, State},
    http::{header, HeaderValue, Method, StatusCode},
    middleware::{self, Next},
    response::{IntoResponse, Response}, routing::get, Json, Router,
};
use tokio::sync::Semaphore;
use tower_http::{catch_panic::CatchPanicLayer, cors::CorsLayer, limit::RequestBodyLimitLayer};
use tracing::Instrument;
use crate::{config::Config, error, rate_limit::RateLimiter};

#[derive(Clone)]
pub struct RequestId(pub String);

struct GuardState {
    limiter: RateLimiter,
    permits: Arc<Semaphore>,
    timeout: std::time::Duration,
}

pub fn router(app: Router, config: &Config) -> Router {
    let state = Arc::new(GuardState {
        limiter: RateLimiter::new(config.rate_limit_requests, config.rate_limit_window, config.rate_limit_max_clients),
        permits: Arc::new(Semaphore::new(config.max_in_flight)),
        timeout: config.request_timeout,
    });
    let cors = CorsLayer::new()
        .allow_origin(config.cors_origins.clone())
        .allow_methods([Method::GET, Method::POST, Method::PUT, Method::PATCH, Method::DELETE, Method::HEAD, Method::OPTIONS])
        .allow_headers([header::CONTENT_TYPE, header::AUTHORIZATION])
        .expose_headers([header::HeaderName::from_static("x-request-id"), header::RETRY_AFTER]);
    app.route("/health/live", get(health))
        .route("/health/ready", get(health))
        .fallback(|| async { StatusCode::NOT_FOUND })
        .layer(DefaultBodyLimit::max(config.body_limit))
        .layer(RequestBodyLimitLayer::new(config.body_limit))
        .layer(CatchPanicLayer::custom(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response()))
        .layer(middleware::from_fn_with_state(state, guard))
        .layer(cors)
        // Outermost: rejection responses and preflights also receive IDs and logs.
        .layer(middleware::from_fn(observe))
}

async fn health() -> Json<serde_json::Value> {
    Json(serde_json::json!({ "status": "ok" }))
}

async fn guard(State(state): State<Arc<GuardState>>, request: Request, next: Next) -> Response {
    let healthy = request.method() == Method::GET && matches!(request.uri().path(), "/health/live" | "/health/ready");
    let request_id = request.extensions().get::<RequestId>().map(|id| id.0.as_str()).unwrap_or("");
    let permit = if healthy { None } else {
        // Use the actual socket peer. Never trust arbitrary X-Forwarded-For headers.
        let Some(ConnectInfo(peer)) = request.extensions().get::<ConnectInfo<SocketAddr>>() else {
            tracing::error!("socket peer missing; use into_make_service_with_connect_info");
            return error::response(StatusCode::INTERNAL_SERVER_ERROR, request_id);
        };
        if let Err(retry) = state.limiter.check(peer.ip()) {
            let mut response = error::response(StatusCode::TOO_MANY_REQUESTS, request_id);
            response.headers_mut().insert(header::RETRY_AFTER, HeaderValue::from_str(&retry.to_string()).unwrap());
            return response;
        }
        match state.permits.clone().try_acquire_owned() {
            Ok(permit) => Some(permit),
            Err(_) => {
                let mut response = error::response(StatusCode::SERVICE_UNAVAILABLE, request_id);
                response.headers_mut().insert(header::RETRY_AFTER, HeaderValue::from_static("1"));
                return response;
            }
        }
    };
    // The permit covers handler work up to response headers, not streaming bodies.
    let id = request_id.to_owned();
    let response = match tokio::time::timeout(state.timeout, next.run(request)).await {
        Ok(response) => response,
        Err(_) => error::response(StatusCode::GATEWAY_TIMEOUT, &id),
    };
    drop(permit);
    response
}

async fn observe(mut request: Request, next: Next) -> Response {
    let id = uuid::Uuid::new_v4().to_string();
    request.extensions_mut().insert(RequestId(id.clone()));
    let route = request.extensions().get::<MatchedPath>().map(|path| path.as_str()).unwrap_or("<unmatched>");
    // Log matched route templates, never headers, query strings, or payloads.
    let span = tracing::info_span!("http_request", request_id = %id, method = %request.method(), route);
    async move {
        let started = Instant::now();
        tracing::info!("request received");
        let mut response = next.run(request).await;
        // Normalize framework errors while retaining protocol/CORS/retry headers.
        if response.status().is_client_error() || response.status().is_server_error() {
            let normalized = error::response(response.status(), &id);
            *response.body_mut() = normalized.into_body();
            response.headers_mut().remove(header::CONTENT_LENGTH);
            response.headers_mut().insert(header::CONTENT_TYPE, HeaderValue::from_static("application/json"));
        }
        response.headers_mut().insert("x-request-id", id.parse().unwrap());
        response.headers_mut().insert("x-content-type-options", HeaderValue::from_static("nosniff"));
        response.headers_mut().insert("x-frame-options", HeaderValue::from_static("DENY"));
        response.headers_mut().insert("referrer-policy", HeaderValue::from_static("no-referrer"));
        response.headers_mut().insert("content-security-policy", HeaderValue::from_static("default-src 'none'; frame-ancestors 'none'"));
        let status = response.status().as_u16();
        let latency_ms = started.elapsed().as_millis() as u64;
        if status >= 500 { tracing::error!(status, latency_ms, "response sent"); }
        else if status >= 400 { tracing::warn!(status, latency_ms, "response sent"); }
        else { tracing::info!(status, latency_ms, "response sent"); }
        response
    }.instrument(span).await
}

#[cfg(test)]
mod tests;