ARCADIA-Studio/crates/oa-server/src/lib.rs

//! Headless REST service over the shared Arcadia model and command boundary.
pub mod error;
mod routes;
pub mod store;

#[cfg(test)]
mod preparation_http_tests;

use axum::{
    Router,
    extract::{DefaultBodyLimit, Request, State},
    http::{StatusCode, header},
    middleware::{self, Next},
    response::{IntoResponse, Response},
};
use error::{ApiError, ApiResult};
use std::{net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};
use store::Store;
use subtle::ConstantTimeEq;
use tokio::sync::Semaphore;
use tower_http::limit::RequestBodyLimitLayer;

#[derive(Clone)]
pub(crate) struct AppState {
    store: Arc<Store>,
    token: Option<Arc<str>>,
    requests: Arc<Semaphore>,
    workers: Arc<Semaphore>,
}

/// Authentication is optional on loopback; it is mandatory on any other bind.
pub struct ServerConfig {
    pub bind: SocketAddr,
    pub data_dir: PathBuf,
    pub bearer_token: Option<String>,
}

impl ServerConfig {
    pub fn validate(&self) -> Result<(), String> {
        if let Some(token) = &self.bearer_token
            && (token.len() < 32 || !token.bytes().all(|byte| byte.is_ascii_graphic()))
        {
            return Err("Bearer token must contain at least 32 printable ASCII characters without whitespace.".into());
        }
        if !self.bind.ip().is_loopback() && self.bearer_token.is_none() {
            return Err(
                "A non-loopback bind requires --token-file (or ARCADIA_SERVER_TOKEN).".into(),
            );
        }
        Ok(())
    }
}

/// Build a router for an already opened, exclusively owned store.
/// Call `ServerConfig::validate` before binding a network listener.
pub fn app(store: Arc<Store>, bearer_token: Option<String>) -> Router {
    let state = AppState {
        store,
        token: bearer_token.map(Arc::from),
        requests: Arc::new(Semaphore::new(8)),
        workers: Arc::new(Semaphore::new(8)),
    };
    routes::router()
        .with_state(state.clone())
        .layer(DefaultBodyLimit::max(store::MAX_MODEL_BYTES))
        .layer(RequestBodyLimitLayer::new(store::MAX_MODEL_BYTES))
        .layer(middleware::from_fn_with_state(state, guard))
}

async fn guard(State(state): State<AppState>, request: Request, next: Next) -> Response {
    let request_id = oa_core::Id::new_v4().to_string();
    let public = matches!(
        request.uri().path(),
        "/health" | "/health/live" | "/health/ready" | "/openapi.json" | "/api/v1/openapi.json"
    );
    let mut response = if public {
        next.run(request).await
    } else if request.headers().contains_key(header::ORIGIN) {
        ApiError::new(
            StatusCode::FORBIDDEN,
            "browser_origin_denied",
            "Browser origins are disabled; use a native client.",
        )
        .into_response()
    } else if !authorized(&state, request.headers()) {
        ApiError::new(
            StatusCode::UNAUTHORIZED,
            "unauthorized",
            "A valid bearer token is required.",
        )
        .into_response()
    } else if state.token.is_none() && !loopback_host(request.headers()) {
        ApiError::new(
            StatusCode::FORBIDDEN,
            "host_not_allowed",
            "Unauthenticated loopback access requires a localhost or loopback-IP Host header.",
        )
        .into_response()
    } else if request.uri().query().is_some() && !allows_query(&request) {
        ApiError::bad_request(
            "This route does not accept query parameters; historical revisions are not retained.",
        )
        .into_response()
    } else if request.headers().contains_key(header::CONTENT_ENCODING) {
        ApiError::new(
            StatusCode::UNSUPPORTED_MEDIA_TYPE,
            "content_encoding_unsupported",
            "Compressed or encoded request bodies are not supported.",
        )
        .into_response()
    } else {
        match state.requests.clone().try_acquire_owned() {
            Err(_) => ApiError::new(StatusCode::SERVICE_UNAVAILABLE, "overloaded", "All request slots are busy. Retry later.").into_response(),
            Ok(_permit) => match tokio::time::timeout(Duration::from_secs(30), next.run(request)).await {
                Ok(response) => response,
                Err(_) => ApiError::new(StatusCode::GATEWAY_TIMEOUT, "outcome_unknown", "Request deadline exceeded. An admitted write may still complete; read the project before retrying.").into_response(),
            },
        }
    };
    // Normalize middleware/framework rejections too, including body limits.
    if response.status().is_client_error() || response.status().is_server_error() {
        let is_problem = response
            .headers()
            .get(header::CONTENT_TYPE)
            .is_some_and(|value| value == "application/problem+json");
        if !is_problem {
            response = ApiError::new(
                response.status(),
                "request_rejected",
                "The HTTP request was rejected.",
            )
            .into_response();
        }
    }
    if response.status() == StatusCode::SERVICE_UNAVAILABLE {
        response
            .headers_mut()
            .insert(header::RETRY_AFTER, header::HeaderValue::from_static("1"));
    }
    response.headers_mut().insert(
        "x-request-id",
        request_id.parse().expect("UUID is a header value"),
    );
    response.headers_mut().insert(
        header::CACHE_CONTROL,
        header::HeaderValue::from_static("no-cache"),
    );
    response
}

fn loopback_host(headers: &header::HeaderMap) -> bool {
    let mut hosts = headers.get_all(header::HOST).iter();
    let authority = hosts
        .next()
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.parse::<axum::http::uri::Authority>().ok());
    if hosts.next().is_some() {
        return false;
    }
    authority.is_some_and(|authority| {
        let host = authority.host().trim_matches(['[', ']']);
        host.eq_ignore_ascii_case("localhost")
            || host
                .parse::<std::net::IpAddr>()
                .is_ok_and(|ip| ip.is_loopback())
    })
}

fn allows_query(request: &Request) -> bool {
    request.method() == axum::http::Method::GET
        && (request.uri().path() == "/api/v1/projects"
            || ["/elements", "/relations", "/artifact-links", "/traces"]
                .iter()
                .any(|suffix| request.uri().path().ends_with(suffix)))
}

fn authorized(state: &AppState, headers: &header::HeaderMap) -> bool {
    let Some(expected) = &state.token else {
        return true;
    };
    let mut values = headers.get_all(header::AUTHORIZATION).iter();
    let supplied = values
        .next()
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.strip_prefix("Bearer "));
    if values.next().is_some() {
        return false;
    }
    supplied.is_some_and(|supplied| bool::from(supplied.as_bytes().ct_eq(expected.as_bytes())))
}

pub(crate) async fn blocking<R: Send + 'static>(
    state: AppState,
    work: impl FnOnce(&Store) -> ApiResult<R> + Send + 'static,
) -> ApiResult<R> {
    let permit = state.workers.clone().try_acquire_owned().map_err(|_| {
        ApiError::new(
            StatusCode::SERVICE_UNAVAILABLE,
            "overloaded",
            "All model workers are busy. Retry later.",
        )
    })?;
    tokio::task::spawn_blocking(move || {
        // Retain the work permit even if the HTTP future is cancelled.
        let _permit = permit;
        // Parsing, queries and validation are worker work, not catalog-lock
        // work. Store accessors capture snapshots and recheck writes at commit.
        work(&state.store)
    })
    .await
    .map_err(|_| ApiError::internal("Model worker failed."))?
}