tidaldb/tidal-server/src/router.rs

use std::{
    collections::HashMap,
    sync::{
        Arc,
        atomic::{AtomicU64, Ordering},
    },
    time::Duration,
};

use axum::{
    Json, Router,
    extract::{Query, Request, State},
    http::{StatusCode, header::AUTHORIZATION},
    middleware::{self, Next},
    response::{IntoResponse, Response},
    routing::{get, post},
};
use subtle::ConstantTimeEq;
use tidaldb::{
    query::{retrieve::Retrieve, search::Search},
    schema::EntityId,
};
use tower::{ServiceBuilder, limit::ConcurrencyLimitLayer};
use tower_http::{
    request_id::{MakeRequestId, PropagateRequestIdLayer, RequestId, SetRequestIdLayer},
    timeout::TimeoutLayer,
    trace::TraceLayer,
};

use crate::{
    dto::{
        EmbeddingRequest, FeedQuery, FeedResponse, ItemRequest, SearchQueryParams, SearchResponse,
        SignalRequest, VectorSearchRequest, VectorSearchResponse, feed_items, search_items,
        vector_matches,
    },
    error::{Result, ServerError},
    state::ServerState,
};

/// Maximum request body size. Requests exceeding this are rejected with 413
/// before any deserialization occurs.
///
/// Shared with the cluster router ([`crate::cluster::build_cluster_router`]) so
/// the standalone and cluster surfaces can never drift on the body cap (raising
/// one and forgetting the other).
pub(crate) const BODY_LIMIT_BYTES: usize = 2 * 1024 * 1024;

/// Body cap for the ONE internal node-to-node control payload that carries
/// state proportional to the corpus: `POST /cluster/reconcile/snapshot`
/// exchanges a whole-shard CRDT `StateSnapshot` (one entry per entity ×
/// signal type, each carrying per-node contributions).
///
/// It is NOT [`BODY_LIMIT_BYTES`]. The 2 MiB data-surface cap is sized for a
/// single client write; measured against the fleet's live 33k-document,
/// 13.3M-event cluster the snapshot is several MiB, so the shared cap made
/// divergence unhealable in production: `POST /cluster/reconcile` failed with
/// `503 region unreachable: reconcile peer returned 413 Payload Too Large`,
/// which reads like a network fault and sent the operator to TLS and
/// NetworkPolicy first.
///
/// This route is internal (`x-tidal-internal` marker), authenticated, and
/// driven only by an operator verb, so a large body here is a control-plane
/// cost, not an exposed DoS surface.
///
/// This is a CEILING, not a design: the snapshot grows with the corpus and at
/// ~1M entities it will outgrow this too. The durable fix is a chunked
/// reconcile (cursor over entity ranges, merge per chunk) — until then the
/// sender pre-checks its own snapshot size and says so plainly.
pub(crate) const RECONCILE_BODY_LIMIT_BYTES: usize = 64 * 1024 * 1024;

/// Maximum wall-clock time a single request may occupy. Exceeded requests
/// receive 408 Request Timeout.
///
/// Shared with the cluster router so both surfaces time out identically.
pub(crate) const REQUEST_TIMEOUT_SECS: u64 = 30;

/// Maximum number of requests processed concurrently. Additional requests
/// are queued by the concurrency layer until a slot opens.
///
/// Shared with the cluster router so both surfaces cap in-flight load
/// identically.
pub(crate) const MAX_CONCURRENCY: usize = 100;

/// Sequential request ID generator — assigns monotonically increasing IDs.
#[derive(Clone, Default)]
struct SequentialRequestId(Arc<AtomicU64>);

impl MakeRequestId for SequentialRequestId {
    fn make_request_id<B>(&mut self, _: &Request<B>) -> Option<RequestId> {
        // Relaxed ordering is sufficient: we only need unique IDs, not
        // cross-thread happens-before ordering.
        let id = self.0.fetch_add(1, Ordering::Relaxed);
        axum::http::HeaderValue::from_str(&id.to_string())
            .ok()
            .map(RequestId::new)
    }
}

/// Build the application router.
///
/// Routes are split into two groups:
/// - **Public** (`/health`) — never requires auth; safe for liveness/readiness probes.
/// - **Protected** — require `Authorization: Bearer <key>` when a bearer key is
///   configured. The key is read from [`ClusterCreds`] PER REQUEST (not captured
///   once), so a `TIDAL_API_KEY_FILE` rotation takes effect without a restart
///   (m11p7).
///
/// Global middleware stack (applied to all routes, outermost → innermost):
/// 1. `SetRequestIdLayer` — assigns sequential `x-request-id` before the span is created
/// 2. `PropagateRequestIdLayer` — echoes `x-request-id` into the response
/// 3. `TraceLayer` — creates a structured span; reads the ID set in step 1
///
/// Protected-only middleware (not applied to `/health`):
/// 4. `TimeoutLayer` — returns 408 for requests exceeding [`REQUEST_TIMEOUT_SECS`]
/// 5. `ConcurrencyLimitLayer` — queues beyond [`MAX_CONCURRENCY`] in-flight requests
///
/// Keeping `/health` outside the timeout/concurrency layers means health probes
/// are never queued or timed out under saturation, preventing false liveness failures.
pub fn build_router(
    state: Arc<ServerState>,
    creds: Arc<crate::cluster::security::ClusterCreds>,
) -> Router {
    // Public routes — exempt from auth so health probes always work. The
    // startup/live probes are shared with the cluster router via [`crate::health`]
    // so the two modes can never advertise a different probe contract.
    let public = Router::new()
        .route("/health", get(health))
        .route("/health/startup", get(crate::health::health_startup))
        .route("/health/live", get(crate::health::health_live))
        // The OpenAPI document describes the contract only (no data/secrets), so
        // it is served unauthenticated next to the probes — a client needs the
        // spec to learn how to authenticate the data routes.
        .route("/openapi.json", get(crate::openapi::serve_standalone))
        .with_state(Arc::clone(&state));

    // Protected routes — gated by Bearer token when a key is configured.
    let protected = Router::new()
        .route("/items", post(create_item))
        .route("/embeddings", post(write_embedding))
        .route("/signals", post(write_signal))
        .route("/feed", get(feed))
        .route("/search", get(search))
        .route("/vector_search", post(vector_search))
        .layer(axum::extract::DefaultBodyLimit::max(BODY_LIMIT_BYTES))
        .with_state(state);

    // Read the bearer key from `creds` PER REQUEST so a rotation (file swap)
    // takes effect with no restart. When no key is configured the layer passes
    // through (open) — the same posture as before, with the loud startup WARN
    // emitted once by `ClusterCreds::from_env`.
    let protected = protected.layer(middleware::from_fn(move |req: Request, next: Next| {
        let creds = Arc::clone(&creds);
        async move {
            if let Some(key) = creds.bearer()
                && !bearer_token_ok(req.headers(), &key)
            {
                return unauthorized_response();
            }
            // m11p7 per-principal rate limit (standalone principals are always
            // external — there is no inter-node plane here).
            let principal = creds.principal(req.headers());
            if let Err((retry_after_ms, limit)) = creds.check_rate(&principal) {
                return too_many_requests(retry_after_ms, limit);
            }
            next.run(req).await
        }
    }));

    // Timeout and concurrency apply only to protected routes so health probes
    // are never queued or dropped during overload.
    let protected = protected.layer(
        ServiceBuilder::new()
            .layer(TimeoutLayer::with_status_code(
                StatusCode::REQUEST_TIMEOUT,
                Duration::from_secs(REQUEST_TIMEOUT_SECS),
            ))
            .layer(ConcurrencyLimitLayer::new(MAX_CONCURRENCY)),
    );

    // SetRequestId must be outermost so the ID is in headers when TraceLayer
    // creates its span. In ServiceBuilder the first .layer() is outermost.
    with_request_id_tracing(public.merge(protected))
}

/// Wrap `router` with the shared observability layer stack (m11p8):
/// `SetRequestIdLayer` (outermost) assigns a sequential `x-request-id` when one
/// is absent, `PropagateRequestIdLayer` echoes it into the response, `TraceLayer`
/// opens a per-request span carrying the id, and the HTTP metrics layer counts
/// the response. Applied to the standalone router AND both cluster routers
/// (single-process and multi-process) so every HTTP surface correlates by
/// `x-request-id` and reports the same request/status series.
///
/// Because `SetRequestId` is a no-op when the header is already present, an
/// `x-request-id` forwarded from a gateway (see `cluster::forward`) survives the
/// hop: the leader's span shares the originating gateway's id.
///
/// The metrics layer sits INSIDE the request-id layers but OUTSIDE everything
/// else, so it observes the status actually returned to the client — the 401/403
/// from the auth gates, the 408 from the timeout layer, the 429 from the rate
/// limiter. Counting further in would miss precisely the failures worth
/// counting.
pub(crate) fn with_request_id_tracing(router: Router) -> Router {
    let router = router.layer(axum::middleware::from_fn(
        |req: Request<axum::body::Body>, next: axum::middleware::Next| {
            crate::http_metrics::track(crate::http_metrics::global(), req, next)
        },
    ));
    router.layer(
        ServiceBuilder::new()
            .layer(SetRequestIdLayer::x_request_id(
                SequentialRequestId::default(),
            ))
            .layer(PropagateRequestIdLayer::x_request_id())
            .layer(
                TraceLayer::new_for_http().make_span_with(|req: &Request<_>| {
                    // x-request-id is guaranteed present: SetRequestIdLayer runs first.
                    let id = req
                        .headers()
                        .get("x-request-id")
                        .and_then(|v| v.to_str().ok())
                        .unwrap_or("-");
                    tracing::info_span!(
                        "request",
                        method = %req.method(),
                        uri = %req.uri(),
                        request_id = %id,
                    )
                }),
            ),
    )
}

/// Validate an `Authorization: Bearer <token>` header.
///
/// Returns `401 Unauthorized` with a `WWW-Authenticate: Bearer` header if the
/// token is absent or does not match the expected key. The comparison uses
/// constant-time equality to prevent timing-based token reconstruction.
pub async fn bearer_auth(request: Request, next: Next, expected_key: &str) -> Response {
    if bearer_token_ok(request.headers(), expected_key) {
        next.run(request).await
    } else {
        unauthorized_response()
    }
}

/// Whether the request's `Authorization: Bearer <token>` matches `expected_key`
/// in constant time. The predicate behind [`bearer_auth`], extracted so the
/// cluster auth layer can compose it with the m11p7 marker-pinning gate without
/// running `next` twice.
#[must_use]
pub(crate) fn bearer_token_ok(headers: &axum::http::HeaderMap, expected_key: &str) -> bool {
    // RFC 7235 §2.1: auth-scheme tokens are case-insensitive.
    let token = headers
        .get(AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .and_then(|s| {
            // "Bearer " is 7 bytes; check case-insensitively then slice the token.
            if s.len() > 7 && s[..7].eq_ignore_ascii_case("bearer ") {
                Some(&s[7..])
            } else {
                None
            }
        });
    token.is_some_and(|t| {
        let a = t.as_bytes();
        let b = expected_key.as_bytes();
        // Length mismatch does not leak token content; short-circuit is safe.
        a.len() == b.len() && bool::from(a.ct_eq(b))
    })
}

/// The 401 returned when the bearer token is missing or invalid.
#[must_use]
pub(crate) fn unauthorized_response() -> Response {
    (
        StatusCode::UNAUTHORIZED,
        [("www-authenticate", "Bearer")],
        Json(serde_json::json!({"error": "missing or invalid api key"})),
    )
        .into_response()
}

/// The 403 returned when a valid data-plane bearer is presented to a DESTRUCTIVE
/// cluster verb but no cluster-admin credential (or verified sibling node token)
/// accompanies it.
///
/// 403, not 401: the caller authenticated successfully: it simply lacks operator
/// authority. Shared by both cluster routers so the two surfaces cannot drift.
#[must_use]
pub(crate) fn admin_forbidden_response() -> Response {
    (
        StatusCode::FORBIDDEN,
        Json(serde_json::json!({
            "error": "this cluster verb requires the cluster-admin credential \
                      (TIDAL_ADMIN_KEY) or a verified sibling node token; the \
                      data-plane bearer does not grant operator authority"
        })),
    )
        .into_response()
}

/// The 429 returned when a principal exceeds its per-principal rate limit
/// (m11p7), carrying a `Retry-After` header (seconds, ceil) and the limit + the
/// millisecond hint in the body so a client can back off precisely.
#[must_use]
pub(crate) fn too_many_requests(retry_after_ms: u64, limit: f64) -> Response {
    let retry_secs = retry_after_ms.div_ceil(1000).max(1);
    (
        StatusCode::TOO_MANY_REQUESTS,
        [("retry-after", retry_secs.to_string())],
        Json(serde_json::json!({
            "error": "rate limit exceeded",
            "retry_after_ms": retry_after_ms,
            "limit_per_second": limit,
        })),
    )
        .into_response()
}

#[utoipa::path(
    post,
    path = "/items",
    tag = "data",
    request_body = ItemRequest,
    responses(
        (status = 201, description = "Item created"),
        (status = 400, description = "Invalid request"),
        (status = 401, description = "Missing or invalid API key"),
    ),
    security(("bearerAuth" = [])),
)]
pub(crate) async fn create_item(
    State(state): State<Arc<ServerState>>,
    Json(req): Json<ItemRequest>,
) -> Result<StatusCode, AppError> {
    state
        .write_item(EntityId::new(req.entity_id), &req.metadata)
        .map_err(AppError)?;
    Ok(StatusCode::CREATED)
}

#[utoipa::path(
    post,
    path = "/embeddings",
    tag = "data",
    request_body = EmbeddingRequest,
    responses(
        (status = 204, description = "Embedding written"),
        (status = 400, description = "Invalid request"),
        (status = 401, description = "Missing or invalid API key"),
    ),
    security(("bearerAuth" = [])),
)]
pub(crate) async fn write_embedding(
    State(state): State<Arc<ServerState>>,
    Json(req): Json<EmbeddingRequest>,
) -> Result<StatusCode, AppError> {
    state
        .write_embedding(EntityId::new(req.entity_id), &req.values)
        .map_err(AppError)?;
    Ok(StatusCode::NO_CONTENT)
}

#[utoipa::path(
    post,
    path = "/signals",
    tag = "data",
    request_body = SignalRequest,
    responses(
        (status = 204, description = "Signal recorded"),
        (status = 400, description = "Invalid request"),
        (status = 401, description = "Missing or invalid API key"),
    ),
    security(("bearerAuth" = [])),
)]
pub(crate) async fn write_signal(
    State(state): State<Arc<ServerState>>,
    Json(req): Json<SignalRequest>,
) -> Result<StatusCode, AppError> {
    state
        .signal(
            &req.signal,
            EntityId::new(req.entity_id),
            req.weight,
            req.user_id,
            req.creator_id,
        )
        .map_err(AppError)?;
    Ok(StatusCode::NO_CONTENT)
}

#[utoipa::path(
    get,
    path = "/feed",
    tag = "data",
    params(FeedQuery),
    responses(
        (status = 200, description = "Ranked feed", body = FeedResponse),
        (status = 400, description = "Invalid request (e.g. region param standalone)"),
        (status = 401, description = "Missing or invalid API key"),
    ),
    security(("bearerAuth" = [])),
)]
pub(crate) async fn feed(
    State(state): State<Arc<ServerState>>,
    Query(query): Query<FeedQuery>,
) -> Result<Json<FeedResponse>, AppError> {
    // `limit` is client-supplied across the trust boundary; clamp it to
    // `dto::MAX_LIMIT` before it sizes the engine's candidate cap so an
    // oversized request cannot drive a memory-amplification DoS.
    let mut builder = Retrieve::builder()
        .profile(&query.profile)
        .limit(query.clamped_limit() as usize);

    if let Some(user_id) = query.user_id {
        builder = builder.for_user(user_id);
    }
    // m12p2: seed for "more like this" — `profile=related` resolves this item's
    // embedding and sources candidates by ANN over it.
    if let Some(seed) = query.similar_to {
        builder = builder.similar_to(EntityId::new(seed));
    }
    let retrieve = builder.build().map_err(|e| TidalErrorWrapper(e.into()))?;

    // `retrieve` is a synchronous, CPU-bound engine call (candidate scan +
    // scoring + diversity). Running it inline would pin this reactor worker for
    // the whole query, so a burst of feeds could starve every other request.
    // Offload to the blocking pool — the same treatment the cluster path gives
    // the identical work (see [`crate::offload`]).
    let offload_state = Arc::clone(&state);
    let region = query.region.clone();
    let result =
        crate::offload::offload_read(move || offload_state.retrieve(region.as_deref(), &retrieve))
            .await
            .map_err(AppError)?;

    Ok(Json(FeedResponse {
        items: feed_items(&result.items),
        total_candidates: result.total_candidates,
        region: query.region,
        unavailable_shards: None, // standalone: no shards, never degraded
    }))
}

#[utoipa::path(
    get,
    path = "/search",
    tag = "data",
    params(SearchQueryParams),
    responses(
        (status = 200, description = "Ranked search results", body = SearchResponse),
        (status = 400, description = "Invalid request (e.g. region param standalone)"),
        (status = 401, description = "Missing or invalid API key"),
    ),
    security(("bearerAuth" = [])),
)]
pub(crate) async fn search(
    State(state): State<Arc<ServerState>>,
    Query(query): Query<SearchQueryParams>,
) -> Result<Json<SearchResponse>, AppError> {
    // `limit` is client-supplied; clamp it to `dto::MAX_LIMIT` at the boundary
    // (see `feed`) before it reaches the search candidate cap.
    let mut builder = Search::builder()
        .query(&query.query)
        .limit(query.clamped_limit());
    if let Some(user_id) = query.user_id {
        builder = builder.for_user(user_id);
    }
    let search = builder.build().map_err(|e| TidalErrorWrapper(e.into()))?;

    // The text-index reload and the search are both synchronous, CPU/IO-bound
    // engine calls (tantivy reload + BM25/ANN/RRF pipeline + scoring +
    // diversity). Offload the pair together so the reactor stays free — see
    // [`crate::offload`].
    let offload_state = Arc::clone(&state);
    let region = query.region.clone();
    let result = crate::offload::offload_read(move || {
        offload_state.reload_text_index(region.as_deref())?;
        offload_state.search(region.as_deref(), &search)
    })
    .await
    .map_err(AppError)?;

    Ok(Json(SearchResponse {
        items: search_items(&result.items),
        total_candidates: result.total_candidates,
        region: query.region,
        unavailable_shards: None, // standalone: no shards, never degraded
    }))
}

#[utoipa::path(
    post,
    path = "/vector_search",
    tag = "data",
    request_body = VectorSearchRequest,
    responses(
        (status = 200, description = "Nearest items by vector distance, closest-first", body = VectorSearchResponse),
        (status = 400, description = "Empty/dimension-mismatched query vector, or no embedding slot"),
        (status = 401, description = "Missing or invalid API key"),
    ),
    security(("bearerAuth" = [])),
)]
pub(crate) async fn vector_search(
    State(state): State<Arc<ServerState>>,
    Json(req): Json<VectorSearchRequest>,
) -> Result<Json<VectorSearchResponse>, AppError> {
    // Reject an empty vector at the boundary with a clear 400 rather than letting
    // it reach the engine as a dimension mismatch.
    if req.vector.is_empty() {
        return Err(AppError(ServerError::BadRequest(
            "vector_search requires a non-empty query vector".into(),
        )));
    }
    let k = req.clamped_k();
    let ef_search = req.ef_search();
    let vector = req.vector;

    // Pure k-NN is a synchronous, CPU-bound index search; offload to the blocking
    // pool so a burst of recall probes cannot pin a reactor worker — the same
    // treatment `/feed` and `/search` give their engine calls.
    let offload_state = Arc::clone(&state);
    let result = crate::offload::offload_read(move || {
        offload_state.vector_search(None, &vector, k, ef_search)
    })
    .await
    .map_err(AppError)?;

    Ok(Json(VectorSearchResponse {
        items: vector_matches(&result),
        region: None,
        unavailable_shards: None, // standalone: no shards, never degraded
    }))
}

/// Readiness probe: 200 when ready, 503 when shutting down.
#[utoipa::path(
    get,
    path = "/health",
    tag = "health",
    responses(
        (status = 200, description = "Service ready"),
        (status = 503, description = "Service shutting down"),
    ),
)]
pub(crate) async fn health(
    State(state): State<Arc<ServerState>>,
    Query(query): Query<HashMap<String, String>>,
) -> std::result::Result<(StatusCode, Json<serde_json::Value>), AppError> {
    if state.is_shutting_down() {
        return Ok((
            StatusCode::SERVICE_UNAVAILABLE,
            Json(serde_json::json!({
                "ok": false,
                "service": "tidaldb",
                "cause": "shutting down"
            })),
        ));
    }

    let region = query.get("region").map(std::string::String::as_str);
    let items = state.item_count(region).map_err(AppError)?;

    Ok((
        StatusCode::OK,
        Json(serde_json::json!({
            "ok": true,
            "service": "tidaldb",
            "mode": "standalone",
            "items": items,
        })),
    ))
}

struct TidalErrorWrapper(tidaldb::TidalError);

impl From<TidalErrorWrapper> for AppError {
    fn from(value: TidalErrorWrapper) -> Self {
        Self(ServerError::Tidal(value.0))
    }
}

pub(crate) struct AppError(ServerError);

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        let status = status_from_error(&self.0);
        let body = serde_json::json!({
            "error": self.0.to_string()
        });
        (status, Json(body)).into_response()
    }
}

#[must_use]
pub const fn status_from_error(err: &ServerError) -> StatusCode {
    match err {
        ServerError::BadRequest(_)
        | ServerError::SchemaConfig(_)
        | ServerError::NotLocal { .. } => StatusCode::BAD_REQUEST,
        // A non-leader write degrades to 503 (retryable elsewhere) naming the
        // leader in the body; an unreachable leader/region forward degrades the
        // same way, naming the unreachable node and the transport cause.
        ServerError::Unavailable(_)
        | ServerError::NotLeader { .. }
        | ServerError::LeaderUnreachable { .. }
        | ServerError::RegionUnreachable { .. }
        | ServerError::QuorumTimeout { .. } => StatusCode::SERVICE_UNAVAILABLE,
        ServerError::Tidal(tidal_err) => match tidal_err {
            tidaldb::TidalError::NotFound { .. } => StatusCode::NOT_FOUND,
            tidaldb::TidalError::Schema(_) | tidaldb::TidalError::InvalidInput(_) => {
                StatusCode::BAD_REQUEST
            }
            tidaldb::TidalError::Backpressure { .. } | tidaldb::TidalError::RateLimited { .. } => {
                StatusCode::TOO_MANY_REQUESTS
            }
            tidaldb::TidalError::PolicyViolation { .. }
            | tidaldb::TidalError::SessionExpired { .. } => StatusCode::FORBIDDEN,
            _ => StatusCode::INTERNAL_SERVER_ERROR,
        },
        _ => StatusCode::INTERNAL_SERVER_ERROR,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A post-shutdown cluster access surfaces as [`ServerError::Unavailable`],
    /// which must map to 503 (clean degradation) rather than a 500 or a panic —
    /// the load-bearing behavior of the `cluster_ref`/`cluster_arc` Result change.
    #[test]
    fn unavailable_maps_to_503() {
        let err = ServerError::Unavailable("server shutting down".into());
        assert_eq!(status_from_error(&err), StatusCode::SERVICE_UNAVAILABLE);
    }

    #[test]
    fn bad_request_maps_to_400() {
        let err = ServerError::BadRequest("bad".into());
        assert_eq!(status_from_error(&err), StatusCode::BAD_REQUEST);
    }

    #[test]
    fn cluster_error_maps_to_500() {
        let err = ServerError::Cluster("boom".into());
        assert_eq!(status_from_error(&err), StatusCode::INTERNAL_SERVER_ERROR);
    }
}