Skip to main content

mnemo_rest/
lib.rs

1pub mod handlers;
2
3use std::sync::Arc;
4
5use axum::Router;
6use axum::extract::{DefaultBodyLimit, Request, State};
7use axum::http::{Method, StatusCode, header};
8use axum::middleware::{self, Next};
9use axum::response::Response;
10use axum::routing::{get, post};
11use mnemo_core::query::MnemoEngine;
12use tower_http::cors::{AllowOrigin, CorsLayer};
13
14/// Construct the full Axum router for the Mnemo REST API, reading the
15/// bearer-token secret from the `MNEMO_AUTH_TOKEN` environment variable.
16///
17/// When `MNEMO_AUTH_TOKEN` is set (non-empty), every request except
18/// `/v1/health` and CORS preflight (`OPTIONS`) must carry a matching
19/// `Authorization: Bearer <token>` header or it is rejected with `401`. When
20/// the variable is unset, the server runs **open** and logs a warning — the
21/// floor for "don't run an unauthenticated memory server" is opt-in but loud.
22///
23/// All routes are nested under `/v1/` and the router carries
24/// `Arc<MnemoEngine>` as shared state. CORS is restrictive by default
25/// (localhost only); set `MNEMO_CORS_ORIGINS` to override.
26pub fn router(engine: Arc<MnemoEngine>) -> Router {
27    let token = std::env::var("MNEMO_AUTH_TOKEN")
28        .ok()
29        .filter(|s| !s.is_empty());
30    router_with_auth(engine, token)
31}
32
33/// Like [`router`] but with the bearer secret passed explicitly (so tests and
34/// embedders can configure auth without touching the process environment).
35/// `Some(token)` enables bearer auth; `None` runs open (with a warning).
36pub fn router_with_auth(engine: Arc<MnemoEngine>, auth_token: Option<String>) -> Router {
37    let cors = build_cors_layer();
38
39    let app = Router::new()
40        .route(
41            "/v1/memories",
42            post(handlers::remember_handler).get(handlers::recall_handler),
43        )
44        .route(
45            "/v1/memories/{id}",
46            get(handlers::get_memory_handler).delete(handlers::forget_handler),
47        )
48        .route("/v1/memories/{id}/share", post(handlers::share_handler))
49        .route("/v1/checkpoints", post(handlers::checkpoint_handler))
50        .route("/v1/consolidate", post(handlers::consolidate_handler))
51        .route("/v1/branches", post(handlers::branch_handler))
52        .route("/v1/merge", post(handlers::merge_handler))
53        .route("/v1/replay", post(handlers::replay_handler))
54        .route("/v1/verify", post(handlers::verify_handler))
55        .route(
56            "/v1/compliance/trajectory_audit",
57            post(handlers::trajectory_audit_handler),
58        )
59        .route("/v1/delegate", post(handlers::delegate_handler))
60        .route("/v1/forget_subject", post(handlers::forget_subject_handler))
61        .route("/v1/ingest/otlp", post(handlers::otlp_ingest_handler))
62        .route("/v1/health", get(handlers::health_handler))
63        .layer(DefaultBodyLimit::max(2 * 1024 * 1024)) // 2 MB max request body
64        .layer(cors)
65        .layer(tower_http::trace::TraceLayer::new_for_http());
66
67    // Bearer-token gate (outermost so it runs before handlers). When unset,
68    // run open but log loudly — never silently serve an unauthenticated
69    // memory database without surfacing it.
70    let app = match auth_token {
71        Some(token) if !token.is_empty() => {
72            tracing::info!(
73                "REST bearer-token auth ENABLED (Authorization: Bearer <MNEMO_AUTH_TOKEN>)"
74            );
75            app.layer(middleware::from_fn_with_state(
76                Arc::new(token),
77                require_bearer,
78            ))
79        }
80        _ => {
81            tracing::warn!(
82                "REST API running WITHOUT authentication — set MNEMO_AUTH_TOKEN to require a \
83                 bearer token. Do not expose an unauthenticated memory server."
84            );
85            app
86        }
87    };
88
89    app.with_state(engine)
90}
91
92/// Axum middleware: require `Authorization: Bearer <expected>` on every request
93/// except `/v1/health` and CORS preflight (`OPTIONS`). Returns `401` otherwise.
94async fn require_bearer(
95    State(expected): State<Arc<String>>,
96    req: Request,
97    next: Next,
98) -> Result<Response, StatusCode> {
99    // Liveness probes and CORS preflight must not require the secret.
100    if req.method() == Method::OPTIONS || req.uri().path() == "/v1/health" {
101        return Ok(next.run(req).await);
102    }
103    let provided = req
104        .headers()
105        .get(header::AUTHORIZATION)
106        .and_then(|v| v.to_str().ok());
107    if mnemo_core::auth::bearer_token_matches(provided, &expected) {
108        Ok(next.run(req).await)
109    } else {
110        Err(StatusCode::UNAUTHORIZED)
111    }
112}
113
114fn build_cors_layer() -> CorsLayer {
115    use axum::http::{HeaderName, Method};
116
117    let base = CorsLayer::new()
118        .allow_methods([Method::GET, Method::POST, Method::DELETE, Method::OPTIONS])
119        .allow_headers([
120            HeaderName::from_static("content-type"),
121            HeaderName::from_static("authorization"),
122        ])
123        .max_age(std::time::Duration::from_secs(3600));
124
125    match std::env::var("MNEMO_CORS_ORIGINS") {
126        Ok(val) if val == "*" => base.allow_origin(AllowOrigin::any()),
127        Ok(val) => {
128            let origins: Vec<_> = val
129                .split(',')
130                .filter_map(|s| s.trim().parse().ok())
131                .collect();
132            base.allow_origin(origins)
133        }
134        Err(_) => {
135            // Default: localhost only
136            let origins: Vec<_> = [
137                "http://localhost:3000",
138                "http://localhost:8080",
139                "http://127.0.0.1:3000",
140                "http://127.0.0.1:8080",
141            ]
142            .iter()
143            .filter_map(|s| s.parse().ok())
144            .collect();
145            base.allow_origin(origins)
146        }
147    }
148}