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(
62            "/v1/memories/{id}/provenance",
63            get(handlers::get_provenance_handler),
64        )
65        .route(
66            "/v1/provenance/principal/{principal}",
67            get(handlers::provenance_by_principal_handler),
68        )
69        .route(
70            "/v1/provenance/session/{session_id}",
71            get(handlers::provenance_by_session_handler),
72        )
73        .route(
74            "/v1/provenance/verify",
75            get(handlers::verify_provenance_handler),
76        )
77        .route(
78            "/v1/provenance/forget",
79            post(handlers::forget_by_provenance_handler),
80        )
81        .route("/v1/ingest/otlp", post(handlers::otlp_ingest_handler))
82        .route("/v1/health", get(handlers::health_handler))
83        .layer(DefaultBodyLimit::max(2 * 1024 * 1024)) // 2 MB max request body
84        .layer(cors)
85        .layer(tower_http::trace::TraceLayer::new_for_http());
86
87    // Bearer-token gate (outermost so it runs before handlers). When unset,
88    // run open but log loudly — never silently serve an unauthenticated
89    // memory database without surfacing it.
90    let app = match auth_token {
91        Some(token) if !token.is_empty() => {
92            tracing::info!(
93                "REST bearer-token auth ENABLED (Authorization: Bearer <MNEMO_AUTH_TOKEN>)"
94            );
95            app.layer(middleware::from_fn_with_state(
96                Arc::new(token),
97                require_bearer,
98            ))
99        }
100        _ => {
101            tracing::warn!(
102                "REST API running WITHOUT authentication — set MNEMO_AUTH_TOKEN to require a \
103                 bearer token. Do not expose an unauthenticated memory server."
104            );
105            app
106        }
107    };
108
109    app.with_state(engine)
110}
111
112/// Axum middleware: require `Authorization: Bearer <expected>` on every request
113/// except `/v1/health` and CORS preflight (`OPTIONS`). Returns `401` otherwise.
114async fn require_bearer(
115    State(expected): State<Arc<String>>,
116    req: Request,
117    next: Next,
118) -> Result<Response, StatusCode> {
119    // Liveness probes and CORS preflight must not require the secret.
120    if req.method() == Method::OPTIONS || req.uri().path() == "/v1/health" {
121        return Ok(next.run(req).await);
122    }
123    let provided = req
124        .headers()
125        .get(header::AUTHORIZATION)
126        .and_then(|v| v.to_str().ok());
127    if mnemo_core::auth::bearer_token_matches(provided, &expected) {
128        Ok(next.run(req).await)
129    } else {
130        Err(StatusCode::UNAUTHORIZED)
131    }
132}
133
134fn build_cors_layer() -> CorsLayer {
135    use axum::http::{HeaderName, Method};
136
137    let base = CorsLayer::new()
138        .allow_methods([Method::GET, Method::POST, Method::DELETE, Method::OPTIONS])
139        .allow_headers([
140            HeaderName::from_static("content-type"),
141            HeaderName::from_static("authorization"),
142        ])
143        .max_age(std::time::Duration::from_secs(3600));
144
145    match std::env::var("MNEMO_CORS_ORIGINS") {
146        Ok(val) if val == "*" => base.allow_origin(AllowOrigin::any()),
147        Ok(val) => {
148            let origins: Vec<_> = val
149                .split(',')
150                .filter_map(|s| s.trim().parse().ok())
151                .collect();
152            base.allow_origin(origins)
153        }
154        Err(_) => {
155            // Default: localhost only
156            let origins: Vec<_> = [
157                "http://localhost:3000",
158                "http://localhost:8080",
159                "http://127.0.0.1:3000",
160                "http://127.0.0.1:8080",
161            ]
162            .iter()
163            .filter_map(|s| s.parse().ok())
164            .collect();
165            base.allow_origin(origins)
166        }
167    }
168}