mcp/src/http/mod.rs

use crate::{
    auth::{client_ip::ClientIpPolicy, ip_whitelist, AuthorizationService},
    config::Config,
    logs::{follower::LiveMessage, index::HistoryStore},
};
use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        ConnectInfo, Query, State,
    },
    http::{header, HeaderMap, StatusCode},
    response::{IntoResponse, Response},
    routing::get,
    Json, Router,
};
use serde::Deserialize;
use std::{
    net::{IpAddr, SocketAddr},
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
};
use tokio::sync::broadcast;
use tower_http::{
    catch_panic::CatchPanicLayer, limit::RequestBodyLimitLayer, set_header::SetResponseHeaderLayer,
    trace::TraceLayer,
};

#[derive(Clone)]
pub struct AppState {
    pub config: Arc<Config>,
    pub auth: AuthorizationService,
    pub history: HistoryStore,
    pub live_tx: broadcast::Sender<LiveMessage>,
    pub ip_policy: ClientIpPolicy,
    pub ws_count: Arc<AtomicUsize>,
}

pub fn router(state: AppState) -> Router {
    Router::new().route("/api/status",get(status)).route("/api/logs/recent",get(recent)).route("/api/logs/history",get(history)).route("/api/live",get(websocket)).fallback(crate::frontend::serve).with_state(state).layer(RequestBodyLimitLayer::new(16*1024)).layer(SetResponseHeaderLayer::if_not_present(header::X_CONTENT_TYPE_OPTIONS,header::HeaderValue::from_static("nosniff"))).layer(SetResponseHeaderLayer::if_not_present(header::REFERRER_POLICY,header::HeaderValue::from_static("no-referrer"))).layer(SetResponseHeaderLayer::if_not_present(header::CONTENT_SECURITY_POLICY,header::HeaderValue::from_static("default-src 'self'; connect-src 'self' ws: wss:; style-src 'self'; script-src 'self'; frame-ancestors 'none'"))).layer(CatchPanicLayer::new()).layer(TraceLayer::new_for_http())
}

async fn status() -> impl IntoResponse {
    Json(serde_json::json!({"status":"ok"}))
}
#[derive(Deserialize)]
struct PageQuery {
    before: Option<String>,
    limit: Option<usize>,
}
async fn recent(
    State(state): State<AppState>,
    ConnectInfo(peer): ConnectInfo<SocketAddr>,
    headers: HeaderMap,
    Query(q): Query<PageQuery>,
) -> Response {
    logs_response(state, peer, headers, None, q.limit).await
}
async fn history(
    State(state): State<AppState>,
    ConnectInfo(peer): ConnectInfo<SocketAddr>,
    headers: HeaderMap,
    Query(q): Query<PageQuery>,
) -> Response {
    logs_response(state, peer, headers, q.before.as_deref(), q.limit).await
}
async fn logs_response(
    state: AppState,
    peer: SocketAddr,
    headers: HeaderMap,
    before: Option<&str>,
    limit: Option<usize>,
) -> Response {
    if let Err(r) = authorize(&state, peer, &headers).await {
        return r;
    }
    let limit = limit
        .unwrap_or(state.config.initial_log_lines)
        .min(state.config.max_history_lines)
        .max(1);
    match state.history.page(before, limit).await {
        Ok(page) => Json(page).into_response(),
        Err(code) if code == "invalid_cursor" => error(StatusCode::BAD_REQUEST, "invalid_cursor"),
        Err(code) if code == "stale_cursor" => {
            tracing::debug!(event = "history_page_failed", reason = "stale_cursor");
            error(StatusCode::CONFLICT, "stale_cursor")
        }
        Err(code) => {
            tracing::warn!(event = "history_page_failed", reason = %code);
            error(StatusCode::UNPROCESSABLE_ENTITY, "history_unavailable")
        }
    }
}

async fn websocket(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    ConnectInfo(peer): ConnectInfo<SocketAddr>,
    headers: HeaderMap,
) -> Response {
    if let Err(r) = authorize(&state, peer, &headers).await {
        return r;
    }
    if state.ws_count.load(Ordering::Relaxed) >= state.config.max_ws_connections {
        return error(StatusCode::SERVICE_UNAVAILABLE, "connection_limit");
    }
    state.ws_count.fetch_add(1, Ordering::Relaxed);
    ws.on_upgrade(move |socket| serve_socket(socket, state))
        .into_response()
}
async fn serve_socket(mut socket: WebSocket, state: AppState) {
    let _guard = WsGuard(state.ws_count.clone());
    let mut rx = state.live_tx.subscribe();
    if let Ok(hello) = serde_json::to_string(&LiveMessage::Hello) {
        if socket.send(Message::Text(hello.into())).await.is_err() {
            return;
        }
    }
    loop {
        match rx.recv().await {
            Ok(event) => {
                if let Ok(json) = serde_json::to_string(&event) {
                    if socket.send(Message::Text(json.into())).await.is_err() {
                        break;
                    }
                }
            }
            Err(broadcast::error::RecvError::Lagged(_)) => {
                let _ = socket
                    .send(Message::Text(r#"{"type":"resync_required"}"#.into()))
                    .await;
                break;
            }
            Err(_) => break,
        }
    }
}
struct WsGuard(Arc<AtomicUsize>);
impl Drop for WsGuard {
    fn drop(&mut self) {
        self.0.fetch_sub(1, Ordering::Relaxed);
    }
}

async fn authorize(
    state: &AppState,
    peer: SocketAddr,
    headers: &HeaderMap,
) -> Result<IpAddr, Response> {
    validate_origin(&state.config.public_origin, headers).map_err(|_| {
        tracing::debug!(
            event = "authorization_decision",
            result = "denied",
            reason = "origin_denied",
            peer_ip = %peer.ip()
        );
        error(StatusCode::FORBIDDEN, "origin_denied")
    })?;
    let ip = state
        .ip_policy
        .extract(peer.ip(), headers)
        .map_err(|extract_error| {
            tracing::debug!(
                event = "authorization_decision",
                result = "denied",
                reason = "invalid_client_ip",
                peer_ip = %peer.ip(),
                error = %extract_error
            );
            error(StatusCode::BAD_REQUEST, "invalid_client_ip")
        })?;
    if state.auth.is_allowed(ip).await {
        let allowed_player_ips = state.auth.allowed_ips().await;
        tracing::debug!(
            event = "authorization_decision",
            result = "allowed",
            source = "player_join",
            client_ip = %ip,
            peer_ip = %peer.ip(),
            ?allowed_player_ips
        );
        return Ok(ip);
    }
    let allowed_player_ips = state.auth.allowed_ips().await;
    if let Some(path) = state.config.ip_whitelist_file.clone() {
        let log_path = path.clone();
        let result = tokio::task::spawn_blocking(move || ip_whitelist::check(&path, ip)).await;
        match result {
            Ok(Ok(check)) if check.matched => {
                tracing::debug!(
                    event = "authorization_decision",
                    result = "allowed",
                    source = "ip_whitelist_file",
                    client_ip = %ip,
                    peer_ip = %peer.ip(),
                    path = %log_path.display(),
                    allowed_override_networks = ?check.entries,
                    ?allowed_player_ips
                );
                return Ok(ip);
            }
            Ok(Ok(check)) => tracing::debug!(
                event = "ip_whitelist_checked",
                result = "no_match",
                client_ip = %ip,
                path = %log_path.display(),
                allowed_override_networks = ?check.entries,
                ?allowed_player_ips
            ),
            Ok(Err(error)) => tracing::warn!(
                event = "ip_whitelist_read_failed",
                client_ip = %ip,
                path = %log_path.display(),
                ?allowed_player_ips,
                %error
            ),
            Err(error) => tracing::warn!(
                event = "ip_whitelist_check_failed",
                client_ip = %ip,
                path = %log_path.display(),
                ?allowed_player_ips,
                %error
            ),
        }
    } else {
        tracing::debug!(
            event = "ip_whitelist_disabled",
            client_ip = %ip,
            ?allowed_player_ips
        );
    }
    tracing::debug!(
        event = "authorization_decision",
        result = "denied",
        reason = "no_authorization_match",
        client_ip = %ip,
        peer_ip = %peer.ip(),
        ?allowed_player_ips
    );
    Err((
        StatusCode::FORBIDDEN,
        Json(serde_json::json!({"error":"access_denied", "client_ip":ip})),
    )
        .into_response())
}
fn validate_origin(expected: &str, headers: &HeaderMap) -> Result<(), ()> {
    if let Some(origin) = headers.get(header::ORIGIN) {
        if origin.to_str().ok() != Some(expected) {
            return Err(());
        }
    }
    Ok(())
}
fn error(status: StatusCode, code: &str) -> Response {
    (status, Json(serde_json::json!({"error":code}))).into_response()
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn rejects_cross_origin() {
        let mut h = HeaderMap::new();
        h.insert(header::ORIGIN, "https://evil.example".parse().unwrap());
        assert!(validate_origin("https://logs.example", &h).is_err());
    }
    #[test]
    fn accepts_matching_or_non_browser_origin() {
        let mut h = HeaderMap::new();
        assert!(validate_origin("https://logs.example", &h).is_ok());
        h.insert(header::ORIGIN, "https://logs.example".parse().unwrap());
        assert!(validate_origin("https://logs.example", &h).is_ok());
    }
}