From 6b2bf4fab527d4a979af306029741f436e16d412 Mon Sep 17 00:00:00 2001 From: lr90 Date: Fri, 4 Sep 2026 20:39:33 +0800 Subject: [PATCH 1/4] feat(auth): add scoped keys for Astra --- memoria/crates/memoria-api/src/auth.rs | 147 ++++++++++++-- memoria/crates/memoria-api/src/lib.rs | 1 + memoria/crates/memoria-api/src/routes/auth.rs | 190 ++++++++++++++++-- memoria/crates/memoria-api/src/routes/mcp.rs | 60 +++++- memoria/crates/memoria-api/src/state.rs | 18 +- memoria/crates/memoria-api/tests/api_e2e.rs | 96 +++++++++ memoria/crates/memoria-storage/src/store.rs | 21 ++ 7 files changed, 496 insertions(+), 37 deletions(-) diff --git a/memoria/crates/memoria-api/src/auth.rs b/memoria/crates/memoria-api/src/auth.rs index 5a3db694..32729fa8 100644 --- a/memoria/crates/memoria-api/src/auth.rs +++ b/memoria/crates/memoria-api/src/auth.rs @@ -21,6 +21,63 @@ use tracing::warn; use crate::state::{AppState, CachedApiKeyPrincipal}; +pub const SCOPE_IDENTITY_READ: &str = "identity:read"; +pub const SCOPE_MEMORY_READ: &str = "memory:read"; +pub const SCOPE_MEMORY_WRITE: &str = "memory:write"; +pub const SCOPE_KEYS_MANAGE: &str = "keys:manage"; +pub const DEFAULT_API_KEY_SCOPES: &str = "identity:read,memory:read,memory:write,keys:manage"; + +pub fn parse_scopes(value: &str) -> Vec { + value + .split(',') + .map(str::trim) + .filter(|scope| !scope.is_empty()) + .map(str::to_string) + .collect() +} + +fn required_scope_for_request(method: &axum::http::Method, path: &str) -> Option<&'static str> { + let memory_path = [ + "/v1/memories", + "/v1/profiles", + "/v1/feedback", + "/v1/retrieval-params", + "/v1/governance", + "/v1/consolidate", + "/v1/reflect", + "/v1/extract-entities", + "/v1/entities", + "/v1/snapshots", + "/v1/branches", + "/v1/tasks", + "/v1/observe", + "/v1/sessions", + "/v1/pipeline", + "/v1/tool-usage", + ] + .iter() + .any(|prefix| path == *prefix || path.starts_with(&format!("{prefix}/"))); + + if !memory_path { + return None; + } + if method == axum::http::Method::GET + || (method == axum::http::Method::POST + && matches!( + path, + "/v1/memories/query" + | "/v1/memories/fulltext-search" + | "/v1/memories/retrieve" + | "/v1/memories/search" + )) + { + Some(SCOPE_MEMORY_READ) + } else { + Some(SCOPE_MEMORY_WRITE) + } +} + +#[derive(Clone)] pub struct AuthUser { pub user_id: String, /// Routing scope: equals `user_id` in personal mode, `group_id` (e.g. `grp_xxx`) @@ -29,6 +86,9 @@ pub struct AuthUser { pub scope_id: String, pub group_id: Option, pub is_master: bool, + pub key_id: Option, + pub key_prefix: Option, + pub scopes: Vec, } impl AuthUser { @@ -47,6 +107,21 @@ impl AuthUser { pub fn is_group_scoped(&self) -> bool { self.group_id.is_some() } + + pub fn has_scope(&self, scope: &str) -> bool { + self.is_master || self.scopes.iter().any(|granted| granted == scope) + } + + pub fn require_scope(&self, scope: &str) -> Result<(), (StatusCode, String)> { + if self.has_scope(scope) { + Ok(()) + } else { + Err(( + StatusCode::FORBIDDEN, + format!("API key missing required scope: {scope}"), + )) + } + } } async fn cached_or_db_principal(token: &str, state: &AppState) -> Option { @@ -61,7 +136,7 @@ async fn cached_or_db_principal(token: &str, state: &AppState) -> Option NOW(6))", ) @@ -72,14 +147,13 @@ async fn cached_or_db_principal(token: &str, state: &AppState) -> Option("scopes").ok()?), }; - state.api_key_cache.insert( - key_hash, - principal.user_id.clone(), - principal.group_id.clone(), - ); + state.api_key_cache.insert(key_hash, principal.clone()); Some(principal) } @@ -988,7 +1062,19 @@ impl FromRequestParts for AuthUser { // fall through } // 2) API key — user_id resolved from DB, never master - else if let Some((uid, group_id)) = validate_api_key(token, state).await { + else if let Some(principal) = validate_api_key(token, state).await { + if let Some(required_scope) = + required_scope_for_request(&parts.method, parts.uri.path()) + { + if !principal.scopes.iter().any(|scope| scope == required_scope) { + return Err(( + StatusCode::FORBIDDEN, + format!("API key missing required scope: {required_scope}"), + )); + } + } + let uid = principal.user_id.clone(); + let group_id = principal.group_id.clone(); if let Some(tool) = tool_name { state.tool_usage_batcher.mark_used(uid.clone(), tool); } @@ -1006,6 +1092,9 @@ impl FromRequestParts for AuthUser { scope_id, group_id, is_master: false, + key_id: Some(principal.key_id), + key_prefix: Some(principal.key_prefix), + scopes: principal.scopes, }); } else { crate::metrics::registry().security.auth_failures.inc(); @@ -1052,6 +1141,9 @@ impl FromRequestParts for AuthUser { group_id: None, user_id, is_master: true, + key_id: None, + key_prefix: None, + scopes: parse_scopes(DEFAULT_API_KEY_SCOPES), }) } } @@ -1062,7 +1154,7 @@ impl FromRequestParts for AuthUser { /// Uses a dedicated auth connection pool so that auth validation is never /// blocked by slow business queries on the main pool. /// `last_used_at` is updated via batched writes (see [`LastUsedBatcher`]). -async fn validate_api_key(token: &str, state: &AppState) -> Option<(String, Option)> { +async fn validate_api_key(token: &str, state: &AppState) -> Option { state.service.sql_store.as_ref()?; let key_hash = format!("{:x}", Sha256::digest(token.as_bytes())); @@ -1080,7 +1172,7 @@ async fn validate_api_key(token: &str, state: &AppState) -> Option<(String, Opti if let Some(principal) = state.api_key_cache.get(&key_hash) { // Still enqueue last_used_at update (batched, no DB pressure) state.last_used_batcher.mark_used(key_hash); - return Some((principal.user_id, principal.group_id)); + return Some(principal); } let Some(pool) = state.auth_pool.as_ref() else { @@ -1089,7 +1181,7 @@ async fn validate_api_key(token: &str, state: &AppState) -> Option<(String, Opti }; let row = sqlx::query( - "SELECT user_id, group_id FROM mem_api_keys \ + "SELECT key_id, user_id, group_id, key_prefix, scopes FROM mem_api_keys \ WHERE key_hash = ? AND is_active = 1 \ AND (expires_at IS NULL OR expires_at > NOW(6))", ) @@ -1101,6 +1193,13 @@ async fn validate_api_key(token: &str, state: &AppState) -> Option<(String, Opti let user_id: String = row.try_get("user_id").ok()?; let group_id: Option = row.try_get("group_id").ok().flatten(); + let principal = CachedApiKeyPrincipal { + key_id: row.try_get("key_id").ok()?, + user_id: user_id.clone(), + group_id: group_id.clone(), + key_prefix: row.try_get("key_prefix").ok()?, + scopes: parse_scopes(&row.try_get::("scopes").ok()?), + }; // Enforce real-time group membership: even if the key references a group, // the user must still be an active member in `mem_group_members` and the @@ -1133,17 +1232,41 @@ async fn validate_api_key(token: &str, state: &AppState) -> Option<(String, Opti // Cache the result (TTL 5 min) state .api_key_cache - .insert(key_hash.clone(), user_id.clone(), group_id.clone()); + .insert(key_hash.clone(), principal.clone()); // Enqueue batched last_used_at update — zero DB pressure on hot path state.last_used_batcher.mark_used(key_hash); - Some((user_id, group_id)) + Some(principal) } #[cfg(test)] mod tests { use super::*; + #[test] + fn classifies_memory_read_and_write_routes() { + assert_eq!( + required_scope_for_request(&axum::http::Method::POST, "/v1/memories/retrieve"), + Some(SCOPE_MEMORY_READ) + ); + assert_eq!( + required_scope_for_request(&axum::http::Method::GET, "/v1/memories/abc"), + Some(SCOPE_MEMORY_READ) + ); + assert_eq!( + required_scope_for_request(&axum::http::Method::POST, "/v1/memories"), + Some(SCOPE_MEMORY_WRITE) + ); + assert_eq!( + required_scope_for_request(&axum::http::Method::DELETE, "/v1/snapshots/one"), + Some(SCOPE_MEMORY_WRITE) + ); + assert_eq!( + required_scope_for_request(&axum::http::Method::GET, "/auth/whoami"), + None + ); + } + #[test] fn test_tool_usage_mark_and_query() { let b = ToolUsageBatcher::new(); diff --git a/memoria/crates/memoria-api/src/lib.rs b/memoria/crates/memoria-api/src/lib.rs index 69010d6c..1db7475a 100644 --- a/memoria/crates/memoria-api/src/lib.rs +++ b/memoria/crates/memoria-api/src/lib.rs @@ -339,6 +339,7 @@ pub fn build_router(state: AppState) -> Router { // Sessions .route("/v1/tasks/:task_id", get(routes::sessions::get_task_status)) // API key management + .route("/auth/whoami", get(routes::auth::whoami)) .route("/auth/keys", post(routes::auth::create_key)) .route("/auth/keys", get(routes::auth::list_keys)) .route("/auth/keys/:id", get(routes::auth::get_key)) diff --git a/memoria/crates/memoria-api/src/routes/auth.rs b/memoria/crates/memoria-api/src/routes/auth.rs index d7c839d9..6c5dcf78 100644 --- a/memoria/crates/memoria-api/src/routes/auth.rs +++ b/memoria/crates/memoria-api/src/routes/auth.rs @@ -9,7 +9,14 @@ use axum::{ use serde::{Deserialize, Serialize}; use sqlx::Row; -use crate::{auth::AuthUser, routes::memory::api_err, state::AppState}; +use crate::{ + auth::{ + parse_scopes, AuthUser, DEFAULT_API_KEY_SCOPES, SCOPE_IDENTITY_READ, SCOPE_KEYS_MANAGE, + SCOPE_MEMORY_READ, SCOPE_MEMORY_WRITE, + }, + routes::memory::api_err, + state::AppState, +}; fn auth_pool(state: &AppState) -> Result<&sqlx::MySqlPool, (StatusCode, String)> { state @@ -34,6 +41,47 @@ fn generate_key() -> (String, String, String) { (raw, hash, prefix) } +fn normalize_scopes(requested: Option>) -> Result, (StatusCode, String)> { + let requested = requested.unwrap_or_else(|| parse_scopes(DEFAULT_API_KEY_SCOPES)); + let supported = [ + SCOPE_IDENTITY_READ, + SCOPE_MEMORY_READ, + SCOPE_MEMORY_WRITE, + SCOPE_KEYS_MANAGE, + ]; + if let Some(scope) = requested + .iter() + .map(|scope| scope.trim()) + .find(|scope| !supported.contains(scope)) + { + return Err(( + StatusCode::BAD_REQUEST, + format!("Unsupported API key scope: {scope}"), + )); + } + + let scopes: Vec = supported + .iter() + .filter(|scope| requested.iter().any(|item| item.trim() == **scope)) + .map(|scope| (*scope).to_string()) + .collect(); + if !scopes.iter().any(|scope| scope == SCOPE_IDENTITY_READ) { + return Err(( + StatusCode::BAD_REQUEST, + format!("API key scope {SCOPE_IDENTITY_READ} is required"), + )); + } + if scopes.iter().any(|scope| scope == SCOPE_MEMORY_WRITE) + && !scopes.iter().any(|scope| scope == SCOPE_MEMORY_READ) + { + return Err(( + StatusCode::BAD_REQUEST, + format!("API key scope {SCOPE_MEMORY_WRITE} requires {SCOPE_MEMORY_READ}"), + )); + } + Ok(scopes) +} + // ── Request / Response ──────────────────────────────────────────────────────── #[derive(Deserialize)] @@ -42,6 +90,7 @@ pub struct CreateKeyRequest { pub name: String, pub expires_at: Option, pub group_id: Option, + pub scopes: Option>, } #[derive(Serialize)] @@ -51,6 +100,7 @@ pub struct KeyResponse { pub group_id: Option, pub name: String, pub key_prefix: String, + pub scopes: Vec, pub created_at: String, pub expires_at: Option, pub last_used_at: Option, @@ -58,6 +108,26 @@ pub struct KeyResponse { pub raw_key: Option, } +#[derive(Serialize)] +pub struct WhoAmIScope { + #[serde(rename = "type")] + pub kind: &'static str, + pub id: String, +} + +#[derive(Serialize)] +pub struct WhoAmIResponse { + pub user_id: String, + pub key_id: Option, + pub key_prefix: Option, + pub scope: WhoAmIScope, + pub granted_scopes: Vec, + pub api_version: &'static str, + pub capabilities: [&'static str; 2], + pub is_active: bool, + pub is_master: bool, +} + async fn ensure_group_membership( pool: &sqlx::MySqlPool, user_id: &str, @@ -85,6 +155,26 @@ async fn ensure_group_membership( // ── Handlers ────────────────────────────────────────────────────────────────── +/// GET /auth/whoami — resolve the authenticated principal and granted scopes. +pub async fn whoami(auth: AuthUser) -> Result, (StatusCode, String)> { + auth.require_scope(SCOPE_IDENTITY_READ)?; + let (kind, id) = match auth.group_id.clone() { + Some(group_id) => ("group", group_id), + None => ("personal", auth.user_id.clone()), + }; + Ok(Json(WhoAmIResponse { + user_id: auth.user_id, + key_id: auth.key_id, + key_prefix: auth.key_prefix, + scope: WhoAmIScope { kind, id }, + granted_scopes: auth.scopes, + api_version: "1", + capabilities: ["api_key_scopes", "memory_filters_v1"], + is_active: true, + is_master: auth.is_master, + })) +} + /// POST /auth/keys — create API key /// /// Access: master key can create any key. Group owners can create keys @@ -94,7 +184,10 @@ pub async fn create_key( auth: AuthUser, Json(req): Json, ) -> Result<(StatusCode, Json), (StatusCode, String)> { + auth.require_scope(SCOPE_KEYS_MANAGE)?; let pool = auth_pool(&state)?; + let scopes = normalize_scopes(req.scopes.clone())?; + let scopes_csv = scopes.join(","); match req.group_id.as_deref() { Some(group_id) => { @@ -131,11 +224,11 @@ pub async fn create_key( let now = chrono::Utc::now().naive_utc(); sqlx::query( - "INSERT INTO mem_api_keys (key_id, user_id, group_id, name, key_hash, key_prefix, is_active, created_at, expires_at) \ - VALUES (?,?,?,?,?,?,1,?,?)" + "INSERT INTO mem_api_keys (key_id, user_id, group_id, name, key_hash, key_prefix, scopes, is_active, created_at, expires_at) \ + VALUES (?,?,?,?,?,?,?,1,?,?)" ) .bind(&key_id).bind(&req.user_id).bind(req.group_id.as_deref()).bind(&req.name) - .bind(&key_hash).bind(&key_prefix).bind(now) + .bind(&key_hash).bind(&key_prefix).bind(&scopes_csv).bind(now) .bind(req.expires_at.as_deref()) .execute(pool).await.map_err(api_err)?; @@ -147,6 +240,7 @@ pub async fn create_key( group_id: req.group_id, name: req.name, key_prefix, + scopes, created_at: now.to_string(), expires_at: req.expires_at, last_used_at: None, @@ -158,12 +252,14 @@ pub async fn create_key( /// GET /auth/keys — list keys for current user pub async fn list_keys( State(state): State, - AuthUser { user_id, .. }: AuthUser, + auth: AuthUser, ) -> Result>, (StatusCode, String)> { + auth.require_scope(SCOPE_KEYS_MANAGE)?; + let user_id = auth.user_id; let pool = auth_pool(&state)?; let rows = sqlx::query( - "SELECT key_id, user_id, group_id, name, key_prefix, created_at, expires_at, last_used_at \ + "SELECT key_id, user_id, group_id, name, key_prefix, scopes, created_at, expires_at, last_used_at \ FROM mem_api_keys WHERE user_id = ? AND is_active = 1 ORDER BY created_at DESC", ) .bind(&user_id) @@ -179,6 +275,10 @@ pub async fn list_keys( group_id: r.try_get("group_id").ok(), name: r.try_get("name").unwrap_or_default(), key_prefix: r.try_get("key_prefix").unwrap_or_default(), + scopes: parse_scopes( + &r.try_get::("scopes") + .unwrap_or_else(|_| DEFAULT_API_KEY_SCOPES.to_string()), + ), created_at: r .try_get::("created_at") .map(|d| d.to_string()) @@ -203,15 +303,16 @@ pub async fn list_keys( /// GET /auth/keys/:id — get a single API key by ID pub async fn get_key( State(state): State, - AuthUser { - user_id, is_master, .. - }: AuthUser, + auth: AuthUser, Path(key_id): Path, ) -> Result, (StatusCode, String)> { + auth.require_scope(SCOPE_KEYS_MANAGE)?; + let user_id = auth.user_id; + let is_master = auth.is_master; let pool = auth_pool(&state)?; let row = sqlx::query( - "SELECT key_id, user_id, group_id, name, key_prefix, created_at, expires_at, last_used_at \ + "SELECT key_id, user_id, group_id, name, key_prefix, scopes, created_at, expires_at, last_used_at \ FROM mem_api_keys WHERE key_id = ? AND is_active = 1", ) .bind(&key_id) @@ -230,6 +331,10 @@ pub async fn get_key( group_id: r.try_get("group_id").ok(), name: r.try_get("name").unwrap_or_default(), key_prefix: r.try_get("key_prefix").unwrap_or_default(), + scopes: parse_scopes( + &r.try_get::("scopes") + .unwrap_or_else(|_| DEFAULT_API_KEY_SCOPES.to_string()), + ), created_at: r .try_get::("created_at") .map(|d| d.to_string()) @@ -251,15 +356,16 @@ pub async fn get_key( /// PUT /auth/keys/:id/rotate — revoke old key, issue new one pub async fn rotate_key( State(state): State, - AuthUser { - user_id, is_master, .. - }: AuthUser, + auth: AuthUser, Path(key_id): Path, ) -> Result<(StatusCode, Json), (StatusCode, String)> { + auth.require_scope(SCOPE_KEYS_MANAGE)?; + let user_id = auth.user_id; + let is_master = auth.is_master; let pool = auth_pool(&state)?; let old = sqlx::query( - "SELECT user_id, group_id, name, expires_at, key_hash FROM mem_api_keys WHERE key_id = ? AND is_active = 1", + "SELECT user_id, group_id, name, scopes, expires_at, key_hash FROM mem_api_keys WHERE key_id = ? AND is_active = 1", ) .bind(&key_id) .fetch_optional(pool) @@ -275,6 +381,10 @@ pub async fn rotate_key( let name: String = old.try_get("name").map_err(api_err)?; let group_id: Option = old.try_get("group_id").ok(); let expires_at: Option = old.try_get("expires_at").ok().flatten(); + let scopes_csv: String = old + .try_get("scopes") + .unwrap_or_else(|_| DEFAULT_API_KEY_SCOPES.to_string()); + let scopes = parse_scopes(&scopes_csv); // Invalidate cache before DB update if let Ok(key_hash) = old.try_get::("key_hash") { @@ -294,11 +404,11 @@ pub async fn rotate_key( let now = chrono::Utc::now().naive_utc(); sqlx::query( - "INSERT INTO mem_api_keys (key_id, user_id, group_id, name, key_hash, key_prefix, is_active, created_at, expires_at) \ - VALUES (?,?,?,?,?,?,1,?,?)" + "INSERT INTO mem_api_keys (key_id, user_id, group_id, name, key_hash, key_prefix, scopes, is_active, created_at, expires_at) \ + VALUES (?,?,?,?,?,?,?,1,?,?)" ) .bind(&new_id).bind(&old_user).bind(group_id.as_deref()).bind(&name) - .bind(&key_hash).bind(&key_prefix).bind(now) + .bind(&key_hash).bind(&key_prefix).bind(&scopes_csv).bind(now) .bind(expires_at) .execute(pool).await.map_err(api_err)?; @@ -310,6 +420,7 @@ pub async fn rotate_key( group_id, name, key_prefix, + scopes, created_at: now.to_string(), expires_at: expires_at.map(|d| d.to_string()), last_used_at: None, @@ -324,11 +435,12 @@ pub async fn rotate_key( /// group owner can revoke any key scoped to their group. pub async fn revoke_key( State(state): State, - AuthUser { - user_id, is_master, .. - }: AuthUser, + auth: AuthUser, Path(key_id): Path, ) -> Result { + auth.require_scope(SCOPE_KEYS_MANAGE)?; + let user_id = auth.user_id; + let is_master = auth.is_master; let pool = auth_pool(&state)?; let row = sqlx::query("SELECT user_id, group_id, key_hash FROM mem_api_keys WHERE key_id = ?") @@ -374,3 +486,41 @@ pub async fn revoke_key( Ok(StatusCode::NO_CONTENT) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn scopes_are_canonicalized_and_deduplicated() { + let scopes = normalize_scopes(Some(vec![ + SCOPE_MEMORY_READ.to_string(), + SCOPE_IDENTITY_READ.to_string(), + SCOPE_MEMORY_READ.to_string(), + ])) + .unwrap(); + assert_eq!( + scopes, + vec![ + SCOPE_IDENTITY_READ.to_string(), + SCOPE_MEMORY_READ.to_string() + ] + ); + } + + #[test] + fn write_scope_requires_read_scope() { + let err = normalize_scopes(Some(vec![ + SCOPE_IDENTITY_READ.to_string(), + SCOPE_MEMORY_WRITE.to_string(), + ])) + .unwrap_err(); + assert_eq!(err.0, StatusCode::BAD_REQUEST); + } + + #[test] + fn identity_scope_is_required() { + let err = normalize_scopes(Some(vec![SCOPE_MEMORY_READ.to_string()])).unwrap_err(); + assert_eq!(err.0, StatusCode::BAD_REQUEST); + } +} diff --git a/memoria/crates/memoria-api/src/routes/mcp.rs b/memoria/crates/memoria-api/src/routes/mcp.rs index f9d0379c..bb652c68 100644 --- a/memoria/crates/memoria-api/src/routes/mcp.rs +++ b/memoria/crates/memoria-api/src/routes/mcp.rs @@ -22,7 +22,7 @@ use axum::{extract::State, http::StatusCode, response::IntoResponse, Json}; use serde_json::json; use crate::{ - auth::{AuthUser, RpcMeta}, + auth::{AuthUser, RpcMeta, SCOPE_MEMORY_READ, SCOPE_MEMORY_WRITE}, state::AppState, }; @@ -135,6 +135,14 @@ fn mcp_tool_dirty_mask(tool: &str) -> Option } } +fn mcp_tool_required_scope(tool: &str) -> &'static str { + if mcp_tool_dirty_mask(tool).is_some() { + SCOPE_MEMORY_WRITE + } else { + SCOPE_MEMORY_READ + } +} + fn spawn_metrics_dirty_mark( state: AppState, user_id: String, @@ -276,6 +284,10 @@ pub async fn mcp_handler( } } }; + let missing_scope = tracked_tool + .as_deref() + .map(mcp_tool_required_scope) + .filter(|required| !auth.has_scope(required)); // ── Group main-write guard (computed once, shared by both code paths) ───── // Resolved here — before the Notification early-return — so that @@ -344,6 +356,18 @@ pub async fn mcp_handler( // JSON-RPC 2.0: a Notification is a *valid* Request without an "id" member. // The server MUST NOT reply to Notifications. if req.get("id").is_none() { + if missing_scope.is_some() { + report_stats(&track_path, false); + state.call_log_batcher.record_rpc( + user_id, + "POST".to_string(), + track_path, + 204, + t.elapsed().as_millis() as u32, + RpcMeta::err(-32003), + ); + return StatusCode::NO_CONTENT.into_response(); + } // Write guard: per JSON-RPC 2.0 the server MUST NOT reply to Notifications, // so we silently drop blocked writes without dispatching. if blocked_tool.is_some() { @@ -391,6 +415,27 @@ pub async fn mcp_handler( let id = req["id"].clone(); + if let Some(required_scope) = missing_scope { + let err_body = Json(json!({ + "jsonrpc": "2.0", + "id": id, + "error": { + "code": -32003, + "message": format!("API key missing required scope: {required_scope}") + } + })); + report_stats(&track_path, false); + state.call_log_batcher.record_rpc( + user_id, + "POST".to_string(), + track_path, + 200, + t.elapsed().as_millis() as u32, + RpcMeta::err(-32003), + ); + return err_body.into_response(); + } + // Use the pre-computed write-guard decision (see above). if let Some(tool) = &blocked_tool { let err_body = Json(json!({ @@ -473,7 +518,8 @@ pub async fn mcp_handler( #[cfg(test)] mod tests { - use super::{mcp_tool_dirty_mask, tracking_path}; + use super::{mcp_tool_dirty_mask, mcp_tool_required_scope, tracking_path}; + use crate::auth::{SCOPE_MEMORY_READ, SCOPE_MEMORY_WRITE}; use serde_json::json; // ── tools/call — happy path ─────────────────────────────────────────────── @@ -511,6 +557,16 @@ mod tests { assert!(mcp_tool_dirty_mask("memory_branches").is_none()); } + #[test] + fn memory_tools_require_matching_scopes() { + assert_eq!(mcp_tool_required_scope("memory_search"), SCOPE_MEMORY_READ); + assert_eq!(mcp_tool_required_scope("memory_store"), SCOPE_MEMORY_WRITE); + assert_eq!( + mcp_tool_required_scope("memory_checkout"), + SCOPE_MEMORY_WRITE + ); + } + // ── tools/call — missing / malformed name ───────────────────────────────── #[test] diff --git a/memoria/crates/memoria-api/src/state.rs b/memoria/crates/memoria-api/src/state.rs index 297bc554..02c210d9 100644 --- a/memoria/crates/memoria-api/src/state.rs +++ b/memoria/crates/memoria-api/src/state.rs @@ -26,15 +26,21 @@ pub struct CachedMetrics { } struct ApiKeyCacheEntry { + key_id: String, user_id: String, group_id: Option, + key_prefix: String, + scopes: Vec, cached_at: Instant, } #[derive(Clone)] pub struct CachedApiKeyPrincipal { + pub key_id: String, pub user_id: String, pub group_id: Option, + pub key_prefix: String, + pub scopes: Vec, } #[derive(Clone)] @@ -57,8 +63,11 @@ impl ApiKeyCache { if let Some(entry) = cache.get(key_hash) { if now.duration_since(entry.cached_at) < self.ttl { return Some(CachedApiKeyPrincipal { + key_id: entry.key_id.clone(), user_id: entry.user_id.clone(), group_id: entry.group_id.clone(), + key_prefix: entry.key_prefix.clone(), + scopes: entry.scopes.clone(), }); } } @@ -68,13 +77,16 @@ impl ApiKeyCache { None } - pub fn insert(&self, key_hash: String, user_id: String, group_id: Option) { + pub fn insert(&self, key_hash: String, principal: CachedApiKeyPrincipal) { if let Ok(mut cache) = self.inner.write() { cache.insert( key_hash, ApiKeyCacheEntry { - user_id, - group_id, + key_id: principal.key_id, + user_id: principal.user_id, + group_id: principal.group_id, + key_prefix: principal.key_prefix, + scopes: principal.scopes, cached_at: Instant::now(), }, ); diff --git a/memoria/crates/memoria-api/tests/api_e2e.rs b/memoria/crates/memoria-api/tests/api_e2e.rs index eb9918bc..5a76ff4c 100644 --- a/memoria/crates/memoria-api/tests/api_e2e.rs +++ b/memoria/crates/memoria-api/tests/api_e2e.rs @@ -1897,6 +1897,102 @@ async fn test_api_key_crud() { println!("✅ revoke nonexistent → 404"); } +#[tokio::test] +async fn test_scoped_api_key_whoami_and_memory_authorization() { + let mk = "test-master-key-scopes"; + let (base, client, _server) = spawn_server_with_master_key(mk).await; + let auth = format!("Bearer {mk}"); + let user_id = uid(); + + let r = client + .post(format!("{base}/auth/keys")) + .header("Authorization", &auth) + .json(&json!({ + "user_id": user_id, + "name": "astra-read-only", + "scopes": ["identity:read", "memory:read"] + })) + .send() + .await + .unwrap(); + assert_eq!(r.status(), 201); + let key_body: Value = r.json().await.unwrap(); + let raw_key = key_body["raw_key"].as_str().unwrap(); + assert_eq!(key_body["scopes"], json!(["identity:read", "memory:read"])); + + let bearer = format!("Bearer {raw_key}"); + let r = client + .get(format!("{base}/auth/whoami")) + .header("Authorization", &bearer) + .send() + .await + .unwrap(); + assert_eq!(r.status(), 200); + let whoami: Value = r.json().await.unwrap(); + assert_eq!(whoami["user_id"], user_id); + assert_eq!(whoami["scope"]["type"], "personal"); + assert_eq!(whoami["scope"]["id"], user_id); + assert_eq!( + whoami["granted_scopes"], + json!(["identity:read", "memory:read"]) + ); + assert_eq!( + whoami["capabilities"], + json!(["api_key_scopes", "memory_filters_v1"]) + ); + + let r = client + .get(format!("{base}/v1/memories")) + .header("Authorization", &bearer) + .send() + .await + .unwrap(); + assert_eq!(r.status(), 200, "read-only key must be able to read"); + + let r = client + .post(format!("{base}/v1/memories")) + .header("Authorization", &bearer) + .json(&json!({"content": "must not be stored"})) + .send() + .await + .unwrap(); + assert_eq!(r.status(), 403, "read-only key must not be able to write"); + assert!(r.text().await.unwrap().contains("memory:write")); + + let r = client + .get(format!("{base}/auth/keys")) + .header("Authorization", &bearer) + .send() + .await + .unwrap(); + assert_eq!( + r.status(), + 403, + "integration keys must not be able to mint or manage credentials" + ); + assert!(r.text().await.unwrap().contains("keys:manage")); + + let r = client + .post(format!("{base}/mcp")) + .header("Authorization", &bearer) + .json(&json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "memory_store", "arguments": {"content": "blocked"}} + })) + .send() + .await + .unwrap(); + assert_eq!(r.status(), 200); + let body: Value = r.json().await.unwrap(); + assert_eq!(body["error"]["code"], -32003); + assert!(body["error"]["message"] + .as_str() + .unwrap() + .contains("memory:write")); +} + #[tokio::test] async fn test_api_key_cannot_get_other_users_memory() { let mk = "test-master-key-memory-read"; diff --git a/memoria/crates/memoria-storage/src/store.rs b/memoria/crates/memoria-storage/src/store.rs index f5de83b1..371df9ee 100644 --- a/memoria/crates/memoria-storage/src/store.rs +++ b/memoria/crates/memoria-storage/src/store.rs @@ -2172,6 +2172,7 @@ impl SqlMemoryStore { name VARCHAR(100) NOT NULL, key_hash VARCHAR(64) NOT NULL, key_prefix VARCHAR(12) NOT NULL, + scopes VARCHAR(512) NOT NULL DEFAULT 'identity:read,memory:read,memory:write,keys:manage', is_active TINYINT(1) NOT NULL DEFAULT 1, created_at DATETIME(6) NOT NULL, expires_at DATETIME(6) DEFAULT NULL, @@ -2209,6 +2210,26 @@ impl SqlMemoryStore { let _ = alter_idx.execute(&mut *conn).await; } + let key_scopes_col_exists: Option = sqlx::query_scalar( + "SELECT 1 FROM information_schema.columns \ + WHERE table_schema = DATABASE() AND table_name = 'mem_api_keys' AND column_name = 'scopes' \ + LIMIT 1", + ) + .fetch_optional(&mut *conn) + .await + .map_err(db_err)?; + if key_scopes_col_exists.is_none() { + let alter = sqlx::query( + "ALTER TABLE mem_api_keys ADD COLUMN scopes VARCHAR(512) NOT NULL \ + DEFAULT 'identity:read,memory:read,memory:write,keys:manage' AFTER key_prefix", + ); + if let Err(e) = alter.execute(&mut *conn).await { + if !is_duplicate_column(&e) { + return Err(db_err(e)); + } + } + } + sqlx::query( r#"CREATE TABLE IF NOT EXISTS mem_groups ( group_id VARCHAR(64) PRIMARY KEY, From 835f8f648acc879e50e62b13a6aacc7901488ee7 Mon Sep 17 00:00:00 2001 From: lr90 Date: Mon, 7 Sep 2026 21:12:58 +0800 Subject: [PATCH 2/4] fix(auth): enforce scoped routes and fresh identity revocation --- memoria/crates/memoria-api/src/auth.rs | 110 +++++++++++++++++--- memoria/crates/memoria-api/tests/api_e2e.rs | 70 +++++++++++++ 2 files changed, 166 insertions(+), 14 deletions(-) diff --git a/memoria/crates/memoria-api/src/auth.rs b/memoria/crates/memoria-api/src/auth.rs index 32729fa8..51f97167 100644 --- a/memoria/crates/memoria-api/src/auth.rs +++ b/memoria/crates/memoria-api/src/auth.rs @@ -37,6 +37,18 @@ pub fn parse_scopes(value: &str) -> Vec { } fn required_scope_for_request(method: &axum::http::Method, path: &str) -> Option<&'static str> { + let under = |prefix: &str| path == prefix || path.starts_with(&format!("{prefix}/")); + if (path == "/auth/whoami" && method == axum::http::Method::GET) + || (path == "/mcp" && method == axum::http::Method::POST) + { + // MCP additionally authorizes each tool in its handler. + return Some(SCOPE_IDENTITY_READ); + } + if under("/auth/keys") || under("/v1/groups") { + // Group administration can copy personal memories, grant access to + // other accounts and delete databases. Memory scopes never grant it. + return Some(SCOPE_KEYS_MANAGE); + } let memory_path = [ "/v1/memories", "/v1/profiles", @@ -54,6 +66,7 @@ fn required_scope_for_request(method: &axum::http::Method, path: &str) -> Option "/v1/sessions", "/v1/pipeline", "/v1/tool-usage", + "/v1/health", ] .iter() .any(|prefix| path == *prefix || path.starts_with(&format!("{prefix}/"))); @@ -62,6 +75,7 @@ fn required_scope_for_request(method: &axum::http::Method, path: &str) -> Option return None; } if method == axum::http::Method::GET + || method == axum::http::Method::HEAD || (method == axum::http::Method::POST && matches!( path, @@ -1062,17 +1076,10 @@ impl FromRequestParts for AuthUser { // fall through } // 2) API key — user_id resolved from DB, never master - else if let Some(principal) = validate_api_key(token, state).await { - if let Some(required_scope) = - required_scope_for_request(&parts.method, parts.uri.path()) - { - if !principal.scopes.iter().any(|scope| scope == required_scope) { - return Err(( - StatusCode::FORBIDDEN, - format!("API key missing required scope: {required_scope}"), - )); - } - } + else if let Some(principal) = + validate_api_key(token, state, parts.uri.path() == "/auth/whoami").await + { + authorize_api_key_route(&parts.method, parts.uri.path(), &principal.scopes)?; let uid = principal.user_id.clone(); let group_id = principal.group_id.clone(); if let Some(tool) = tool_name { @@ -1154,7 +1161,11 @@ impl FromRequestParts for AuthUser { /// Uses a dedicated auth connection pool so that auth validation is never /// blocked by slow business queries on the main pool. /// `last_used_at` is updated via batched writes (see [`LastUsedBatcher`]). -async fn validate_api_key(token: &str, state: &AppState) -> Option { +async fn validate_api_key( + token: &str, + state: &AppState, + fresh: bool, +) -> Option { state.service.sql_store.as_ref()?; let key_hash = format!("{:x}", Sha256::digest(token.as_bytes())); @@ -1169,7 +1180,12 @@ async fn validate_api_key(token: &str, state: &AppState) -> Option Option Result<(), (StatusCode, String)> { + let required = required_scope_for_request(method, path).ok_or(( + StatusCode::FORBIDDEN, + "API key access is not enabled for this route".to_string(), + ))?; + if scopes.iter().any(|scope| scope == required) { + Ok(()) + } else { + Err(( + StatusCode::FORBIDDEN, + format!("API key missing required scope: {required}"), + )) + } +} #[cfg(test)] mod tests { use super::*; @@ -1263,10 +1298,57 @@ mod tests { ); assert_eq!( required_scope_for_request(&axum::http::Method::GET, "/auth/whoami"), - None + Some(SCOPE_IDENTITY_READ) ); } + #[test] + fn restricted_keys_cannot_administer_groups_or_use_unclassified_routes() { + use axum::http::Method; + for scopes in [ + "identity:read", + "identity:read,memory:read", + "identity:read,memory:read,memory:write", + ] { + let scopes = parse_scopes(scopes); + for (method, path) in [ + (Method::GET, "/v1/groups"), + (Method::POST, "/v1/groups"), + (Method::POST, "/v1/groups/group/members/another-user"), + (Method::DELETE, "/v1/groups/group"), + (Method::DELETE, "/v1/groups/group/members/another-user"), + (Method::GET, "/unclassified-sensitive-route"), + (Method::POST, "/admin/users"), + ] { + assert_eq!( + authorize_api_key_route(&method, path, &scopes) + .unwrap_err() + .0, + StatusCode::FORBIDDEN + ); + } + assert!(authorize_api_key_route(&Method::GET, "/auth/whoami", &scopes).is_ok()); + } + assert!(authorize_api_key_route( + &Method::POST, + "/v1/groups", + &parse_scopes(DEFAULT_API_KEY_SCOPES) + ) + .is_ok()); + assert!(authorize_api_key_route( + &Method::GET, + "/unclassified-sensitive-route", + &parse_scopes(DEFAULT_API_KEY_SCOPES) + ) + .is_err()); + assert!(authorize_api_key_route( + &Method::GET, + "/v1/health/analyze", + &parse_scopes("identity:read") + ) + .is_err()); + } + #[test] fn test_tool_usage_mark_and_query() { let b = ToolUsageBatcher::new(); diff --git a/memoria/crates/memoria-api/tests/api_e2e.rs b/memoria/crates/memoria-api/tests/api_e2e.rs index 5a76ff4c..525fa885 100644 --- a/memoria/crates/memoria-api/tests/api_e2e.rs +++ b/memoria/crates/memoria-api/tests/api_e2e.rs @@ -1993,6 +1993,76 @@ async fn test_scoped_api_key_whoami_and_memory_authorization() { .contains("memory:write")); } +#[tokio::test] +async fn test_scoped_keys_deny_groups_and_whoami_observes_uncached_revocation() { + let (base, client, server) = spawn_server_with_master_key("review-master").await; + let user = uid(); + for scopes in [ + json!(["identity:read"]), + json!(["identity:read", "memory:read"]), + json!(["identity:read", "memory:read", "memory:write"]), + ] { + let response = client + .post(format!("{base}/auth/keys")) + .bearer_auth("review-master") + .json(&json!({"user_id":user, "name":"review-scoped-key", "scopes":scopes})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 201); + let key: Value = response.json().await.unwrap(); + let raw = key["raw_key"].as_str().unwrap(); + for (method, path) in [ + (reqwest::Method::GET, "/v1/groups"), + (reqwest::Method::POST, "/v1/groups"), + ( + reqwest::Method::POST, + "/v1/groups/owned-group/members/other-user", + ), + (reqwest::Method::DELETE, "/v1/groups/owned-group"), + ] { + let response = client.request(method, format!("{base}{path}")).bearer_auth(raw) + .json(&json!({"group_name":"must-not-be-created", "seed":{"db_name":"personal", "mode":"active_only"}})) + .send().await.unwrap(); + assert_eq!( + response.status(), + 403, + "scope check must precede group lookup/mutation: {path}" + ); + } + assert_eq!( + client + .get(format!("{base}/auth/whoami")) + .bearer_auth(raw) + .send() + .await + .unwrap() + .status(), + 200 + ); + // Simulate revocation through a different replica: do not invalidate + // this server's warm principal cache. + sqlx::query(&format!( + "UPDATE {} SET is_active = 0 WHERE key_id = ?", + server.shared_table("mem_api_keys") + )) + .bind(key["key_id"].as_str().unwrap()) + .execute(&server.shared_pool()) + .await + .unwrap(); + assert_eq!( + client + .get(format!("{base}/auth/whoami")) + .bearer_auth(raw) + .send() + .await + .unwrap() + .status(), + 401 + ); + } +} + #[tokio::test] async fn test_api_key_cannot_get_other_users_memory() { let mk = "test-master-key-memory-read"; From 58744ff16ca9cc6c5a7bd6697f17b665ed229dc3 Mon Sep 17 00:00:00 2001 From: lr90 Date: Tue, 8 Sep 2026 01:10:04 +0800 Subject: [PATCH 3/4] fix(auth): classify MCP tool permissions explicitly and fail closed --- memoria/crates/memoria-api/src/routes/mcp.rs | 98 ++++++++-- memoria/crates/memoria-api/tests/api_e2e.rs | 195 +++++++++++++++++++ 2 files changed, 278 insertions(+), 15 deletions(-) diff --git a/memoria/crates/memoria-api/src/routes/mcp.rs b/memoria/crates/memoria-api/src/routes/mcp.rs index bb652c68..0b0801e1 100644 --- a/memoria/crates/memoria-api/src/routes/mcp.rs +++ b/memoria/crates/memoria-api/src/routes/mcp.rs @@ -135,11 +135,41 @@ fn mcp_tool_dirty_mask(tool: &str) -> Option } } -fn mcp_tool_required_scope(tool: &str) -> &'static str { - if mcp_tool_dirty_mask(tool).is_some() { - SCOPE_MEMORY_WRITE - } else { - SCOPE_MEMORY_READ +/// Authorization is independent of metrics invalidation. Include callable tools +/// that are not advertised by tools/list, and fail closed for unclassified names. +fn mcp_tool_required_scope(tool: &str) -> Option<&'static str> { + match tool { + "memory_retrieve" + | "memory_search" + | "memory_profile" + | "memory_list" + | "memory_capabilities" + | "memory_get_retrieval_params" + | "memory_snapshots" + | "memory_branches" + | "memory_diff" => Some(SCOPE_MEMORY_READ), + "memory_store" + | "memory_correct" + | "memory_purge" + | "memory_observe" + | "memory_governance" + | "memory_rebuild_index" + | "memory_consolidate" + | "memory_reflect" + | "memory_extract_entities" + | "memory_link_entities" + | "memory_feedback" + | "memory_tune_params" + | "memory_snapshot" + | "memory_snapshot_delete" + | "memory_rollback" + | "memory_branch" + | "memory_checkout" + | "memory_merge" + | "memory_pick" + | "memory_branch_delete" + | "memory_apply" => Some(SCOPE_MEMORY_WRITE), + _ => None, } } @@ -284,10 +314,23 @@ pub async fn mcp_handler( } } }; - let missing_scope = tracked_tool - .as_deref() - .map(mcp_tool_required_scope) - .filter(|required| !auth.has_scope(required)); + // Use the exact dispatch name, never a sanitized/truncated metrics label. + let authorization_error = if method == "tools/call" { + let name = params + .as_ref() + .and_then(|p| p.get("name")) + .and_then(|n| n.as_str()) + .unwrap_or(""); + match mcp_tool_required_scope(name) { + Some(scope) if !auth.has_scope(scope) => { + Some(format!("API key missing required scope: {scope}")) + } + Some(_) => None, + None => Some("MCP tool is not authorized: unclassified tool".to_string()), + } + } else { + None + }; // ── Group main-write guard (computed once, shared by both code paths) ───── // Resolved here — before the Notification early-return — so that @@ -356,7 +399,7 @@ pub async fn mcp_handler( // JSON-RPC 2.0: a Notification is a *valid* Request without an "id" member. // The server MUST NOT reply to Notifications. if req.get("id").is_none() { - if missing_scope.is_some() { + if authorization_error.is_some() { report_stats(&track_path, false); state.call_log_batcher.record_rpc( user_id, @@ -415,13 +458,13 @@ pub async fn mcp_handler( let id = req["id"].clone(); - if let Some(required_scope) = missing_scope { + if let Some(message) = authorization_error { let err_body = Json(json!({ "jsonrpc": "2.0", "id": id, "error": { "code": -32003, - "message": format!("API key missing required scope: {required_scope}") + "message": message } })); report_stats(&track_path, false); @@ -559,12 +602,37 @@ mod tests { #[test] fn memory_tools_require_matching_scopes() { - assert_eq!(mcp_tool_required_scope("memory_search"), SCOPE_MEMORY_READ); - assert_eq!(mcp_tool_required_scope("memory_store"), SCOPE_MEMORY_WRITE); + assert_eq!( + mcp_tool_required_scope("memory_search"), + Some(SCOPE_MEMORY_READ) + ); + assert_eq!( + mcp_tool_required_scope("memory_store"), + Some(SCOPE_MEMORY_WRITE) + ); assert_eq!( mcp_tool_required_scope("memory_checkout"), - SCOPE_MEMORY_WRITE + Some(SCOPE_MEMORY_WRITE) ); + for name in ["memory_apply", "memory_rebuild_index", "memory_tune_params"] { + assert_eq!(mcp_tool_required_scope(name), Some(SCOPE_MEMORY_WRITE)); + } + for name in ["", "memory_future_tool", "memory_search/", " memory_search"] { + assert_eq!(mcp_tool_required_scope(name), None); + } + } + + #[test] + fn all_advertised_tools_have_explicit_authorization() { + let mut tools = memoria_mcp::tools::list().as_array().unwrap().clone(); + tools.extend(memoria_mcp::git_tools::list().as_array().unwrap().clone()); + for tool in tools { + let name = tool["name"].as_str().unwrap(); + assert!( + mcp_tool_required_scope(name).is_some(), + "unclassified advertised tool: {name}" + ); + } } // ── tools/call — missing / malformed name ───────────────────────────────── diff --git a/memoria/crates/memoria-api/tests/api_e2e.rs b/memoria/crates/memoria-api/tests/api_e2e.rs index 525fa885..fb273dd5 100644 --- a/memoria/crates/memoria-api/tests/api_e2e.rs +++ b/memoria/crates/memoria-api/tests/api_e2e.rs @@ -1993,6 +1993,201 @@ async fn test_scoped_api_key_whoami_and_memory_authorization() { .contains("memory:write")); } +#[tokio::test] +async fn test_scoped_mcp_write_authorization_is_independent_of_metrics() { + let master = "mcp-scope-regression-master"; + let (base, client, server) = spawn_server_with_master_key(master).await; + let user = uid(); + let mut keys = Vec::new(); + for scopes in [ + json!(["identity:read", "memory:read"]), + json!(["identity:read", "memory:read", "memory:write"]), + json!(["identity:read"]), + ] { + let response = client + .post(format!("{base}/auth/keys")) + .bearer_auth(master) + .json(&json!({"user_id": user, "name": "mcp-scope-regression", "scopes": scopes})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 201); + let body: Value = response.json().await.unwrap(); + keys.push(body["raw_key"].as_str().unwrap().to_owned()); + } + let read_key = &keys[0]; + let write_key = &keys[1]; + let identity_key = &keys[2]; + + // Seed enough real feedback that an accidentally dispatched tune call would + // persist a change, including when the caller uses a notification (no id). + let response = client + .post(format!("{base}/v1/memories")) + .bearer_auth(write_key) + .json(&json!({"content": "MCP authorization regression fixture"})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 201); + let memory: Value = response.json().await.unwrap(); + let store = server.user_store(&user).await; + for _ in 0..10 { + store + .record_feedback(&user, memory["memory_id"].as_str().unwrap(), "useful", None) + .await + .unwrap(); + } + let before = store + .get_user_retrieval_params(&user) + .await + .unwrap() + .feedback_weight; + + for (name, arguments) in [ + ( + "memory_apply", + json!({"source": "review-branch", "removes": [memory["memory_id"]]}), + ), + ("memory_rebuild_index", json!({"table": "mem_memories"})), + ("memory_tune_params", json!({})), + ] { + for notification in [false, true] { + let mut request = json!({"jsonrpc": "2.0", "method": "tools/call", + "params": {"name": name, "arguments": arguments}}); + if !notification { + request["id"] = json!(1); + } + let response = client + .post(format!("{base}/mcp")) + .bearer_auth(read_key) + .json(&request) + .send() + .await + .unwrap(); + if notification { + assert_eq!(response.status(), 204, "{name}"); + assert!(response.bytes().await.unwrap().is_empty()); + } else { + assert_eq!(response.status(), 200, "{name}"); + let body: Value = response.json().await.unwrap(); + assert_eq!(body["error"]["code"], -32003, "{name}: {body}"); + assert!(body["error"]["message"] + .as_str() + .unwrap() + .contains("memory:write")); + } + } + } + assert_eq!( + store + .get_user_retrieval_params(&user) + .await + .unwrap() + .feedback_weight, + before, + "denied requests and notifications must not tune persisted parameters" + ); + + // Positive controls: scope checks allow writes into dispatch. Avoid an + // expensive index rebuild and branch mutation by using handler validation. + for (name, arguments, expected) in [ + ( + "memory_apply", + json!({"source": "main"}), + "Cannot apply from main", + ), + ( + "memory_rebuild_index", + json!({"table": "invalid-table"}), + "Invalid table", + ), + ("memory_tune_params", json!({}), "Parameters tuned"), + ] { + let response = client + .post(format!("{base}/mcp")) + .bearer_auth(write_key) + .json(&json!({"jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": {"name": name, "arguments": arguments}})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200); + let body: Value = response.json().await.unwrap(); + assert!(body.get("error").is_none(), "{name}: {body}"); + assert!( + body["result"]["content"][0]["text"] + .as_str() + .unwrap() + .contains(expected), + "{body}" + ); + } + assert!( + store + .get_user_retrieval_params(&user) + .await + .unwrap() + .feedback_weight + > before, + "the same write-authorized tune call must change persisted parameters" + ); + + for (key, allowed) in [(read_key, true), (write_key, true), (identity_key, false)] { + let response = client + .post(format!("{base}/mcp")) + .bearer_auth(key) + .json(&json!({"jsonrpc": "2.0", "id": 3, "method": "tools/call", + "params": {"name": "memory_get_retrieval_params", "arguments": {}}})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200); + let body: Value = response.json().await.unwrap(); + if allowed { + assert!(body.get("error").is_none(), "{body}"); + assert!(body["result"]["content"][0]["text"] + .as_str() + .unwrap() + .contains("feedback_weight")); + } else { + assert_eq!(body["error"]["code"], -32003, "{body}"); + assert!(body["error"]["message"] + .as_str() + .unwrap() + .contains("memory:read")); + } + } + + // Full-access credentials must not bypass classification, and metrics-name + // sanitization must not turn malformed names into authorized read tools. + for key in [read_key.as_str(), master] { + for name in [ + Value::Null, + json!(42), + json!(""), + json!("memory_future_tool"), + json!("memory_search/"), + json!(" memory_search"), + ] { + let response = client + .post(format!("{base}/mcp")) + .bearer_auth(key) + .json(&json!({"jsonrpc": "2.0", "id": 4, "method": "tools/call", + "params": {"name": name, "arguments": {}}})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200); + let body: Value = response.json().await.unwrap(); + assert_eq!(body["error"]["code"], -32003, "{name}: {body}"); + assert!(body["error"]["message"] + .as_str() + .unwrap() + .contains("unclassified tool")); + } + } +} + #[tokio::test] async fn test_scoped_keys_deny_groups_and_whoami_observes_uncached_revocation() { let (base, client, server) = spawn_server_with_master_key("review-master").await; From 40f7dbd6706233a750fe7250734642b96c3a4f28 Mon Sep 17 00:00:00 2001 From: lr90 Date: Tue, 8 Sep 2026 12:08:26 +0800 Subject: [PATCH 4/4] fix(auth): enforce cached key expiry and gate MCP storage side effects Check absolute API key expiry on every cache read and reject unauthorized MCP calls before per-user telemetry or group storage access. Add expiry and database-provisioning regressions and document the authorization boundary. --- docs/per-user-database-architecture.md | 8 + memoria/crates/memoria-api/src/auth.rs | 57 ++- memoria/crates/memoria-api/src/routes/mcp.rs | 122 +++---- memoria/crates/memoria-api/src/state.rs | 73 +++- memoria/crates/memoria-api/tests/api_e2e.rs | 354 +++++++++++++++++++ 5 files changed, 518 insertions(+), 96 deletions(-) diff --git a/docs/per-user-database-architecture.md b/docs/per-user-database-architecture.md index e13bb85a..79f7e7db 100644 --- a/docs/per-user-database-architecture.md +++ b/docs/per-user-database-architecture.md @@ -107,6 +107,14 @@ per-user DB 放两类数据： 注意：这些“用户级控制元数据”虽然在用户库里，但它们不是 rollback 的业务恢复目标；它们要保持当前态，否则会出现 branch / snapshot 注册表与底层事实不一致。 +### 3.4 鉴权与首次建库的边界 + +API Key 校验只依赖 shared DB。缓存命中仍须检查 Key 的 `expires_at`（UTC）；缓存 TTL 不能延长 Key 的有效期。无到期时间的 Key 继续遵守正常缓存 TTL。 + +MCP 的工具权限检查必须先于群组存储访问、工具使用统计和个人调用日志入队。因 scope 不足而拒绝的请求（包括 notification），以及格式错误的请求，不得为了记录拒绝而创建个人记忆库。这些事件保留在运行日志中，并在启用 shared stats reporter 时计入共享汇总统计。 + +仅有 `identity:read` 的 Key 调用身份或 MCP 元数据接口时，也不写入可能触发建库的个人统计；即使带有 `X-Memoria-Tool` 或 `X-Tool-Name` Header，行为不变。获准的记忆工具调用仍正常记录使用情况，并可按原有流程首次建库。以上边界不需要新增表或数据库迁移。 + --- ## 4. 最终实现选择：global user pool + qualified tables diff --git a/memoria/crates/memoria-api/src/auth.rs b/memoria/crates/memoria-api/src/auth.rs index 51f97167..ba56aa5f 100644 --- a/memoria/crates/memoria-api/src/auth.rs +++ b/memoria/crates/memoria-api/src/auth.rs @@ -36,6 +36,17 @@ pub fn parse_scopes(value: &str) -> Vec { .collect() } +/// Agent labels are recorded only after the request's scope admission. MCP +/// defers this until its tool-specific admission, not merely bearer validation. +pub(crate) fn request_tool_name(headers: &axum::http::HeaderMap) -> Option { + headers + .get("X-Memoria-Tool") + .or_else(|| headers.get("X-Tool-Name")) + .and_then(|v| v.to_str().ok()) + .filter(|v| !v.is_empty()) + .map(String::from) +} + fn required_scope_for_request(method: &axum::http::Method, path: &str) -> Option<&'static str> { let under = |prefix: &str| path == prefix || path.starts_with(&format!("{prefix}/")); if (path == "/auth/whoami" && method == axum::http::Method::GET) @@ -150,7 +161,7 @@ async fn cached_or_db_principal(token: &str, state: &AppState) -> Option NOW(6))", ) @@ -166,6 +177,7 @@ async fn cached_or_db_principal(token: &str, state: &AppState) -> Option("scopes").ok()?), + expires_at: row.try_get("expires_at").ok()?, }; state.api_key_cache.insert(key_hash, principal.clone()); Some(principal) @@ -252,6 +264,11 @@ pub async fn group_main_write_guard( .await .filter(|p| p.group_id.is_some()) { + if let Err(rejection) = + authorize_api_key_route(req.method(), req.uri().path(), &p.scopes) + { + return rejection.into_response(); + } let gid = p.group_id.as_ref().unwrap(); // Set task-local so active_branch_name resolves per-member state let user_id = p.user_id.clone(); @@ -1052,13 +1069,11 @@ impl FromRequestParts for AuthUser { // Agents send X-Memoria-Tool with their name: cursor / kiro / claude / codex / openclaw. // Fall back to X-Tool-Name for backwards compatibility with older clients. // Any non-empty value is accepted — no whitelist, so new agents work automatically. - let tool_name = parts - .headers - .get("X-Memoria-Tool") - .or_else(|| parts.headers.get("X-Tool-Name")) - .and_then(|v| v.to_str().ok()) - .filter(|v| !v.is_empty()) - .map(String::from); + let tool_name = if parts.uri.path() == "/mcp" { + None + } else { + request_tool_name(&parts.headers) + }; let bearer = parts .headers @@ -1082,15 +1097,20 @@ impl FromRequestParts for AuthUser { authorize_api_key_route(&parts.method, parts.uri.path(), &principal.scopes)?; let uid = principal.user_id.clone(); let group_id = principal.group_id.clone(); - if let Some(tool) = tool_name { - state.tool_usage_batcher.mark_used(uid.clone(), tool); - } - // Notify call-log middleware (if present) of the resolved user_id. - // The middleware inserted CallLogContext into extensions before calling next; - // we fill in the user_id so it can record the call after the handler returns. - if let Some(ctx) = parts.extensions.get::() { - if let Ok(mut guard) = ctx.0.lock() { - *guard = Some(uid.clone()); + let memory_telemetry = principal + .scopes + .iter() + .any(|scope| matches!(scope.as_str(), SCOPE_MEMORY_READ | SCOPE_MEMORY_WRITE)); + if memory_telemetry { + if let Some(tool) = tool_name { + state.tool_usage_batcher.mark_used(uid.clone(), tool); + } + // Only admitted memory-capable requests may enqueue logs + // whose persistence can provision a personal memory DB. + if let Some(ctx) = parts.extensions.get::() { + if let Ok(mut guard) = ctx.0.lock() { + *guard = Some(uid.clone()); + } } } let scope_id = group_id.clone().unwrap_or_else(|| uid.clone()); @@ -1197,7 +1217,7 @@ async fn validate_api_key( }; let row = sqlx::query( - "SELECT key_id, user_id, group_id, key_prefix, scopes FROM mem_api_keys \ + "SELECT key_id, user_id, group_id, key_prefix, scopes, expires_at FROM mem_api_keys \ WHERE key_hash = ? AND is_active = 1 \ AND (expires_at IS NULL OR expires_at > NOW(6))", ) @@ -1215,6 +1235,7 @@ async fn validate_api_key( group_id: group_id.clone(), key_prefix: row.try_get("key_prefix").ok()?, scopes: parse_scopes(&row.try_get::("scopes").ok()?), + expires_at: row.try_get("expires_at").ok()?, }; // Enforce real-time group membership: even if the key references a group, diff --git a/memoria/crates/memoria-api/src/routes/mcp.rs b/memoria/crates/memoria-api/src/routes/mcp.rs index 0b0801e1..cdb6c5c5 100644 --- a/memoria/crates/memoria-api/src/routes/mcp.rs +++ b/memoria/crates/memoria-api/src/routes/mcp.rs @@ -192,6 +192,7 @@ fn spawn_metrics_dirty_mark( pub async fn mcp_handler( State(state): State, auth: AuthUser, + headers: axum::http::HeaderMap, body: String, ) -> impl IntoResponse { // Start timing after auth. Billable MCP requests are recorded below; transport @@ -202,14 +203,9 @@ pub async fn mcp_handler( // Uses underscore-prefixed paths so they never collide with real tool names. macro_rules! validation_err { ($path:expr, $code:expr, $body:expr) => {{ - state.call_log_batcher.record_rpc( - auth.user_id.clone(), - "POST".to_string(), - $path.to_string(), - 200, - t.elapsed().as_millis() as u32, - RpcMeta::err($code), - ); + // No tool has been admitted. Per-user log flushing can provision a + // memory DB, so malformed requests use non-memory telemetry only. + tracing::warn!(user_id = %auth.user_id, path = $path, rpc_code = $code, "invalid MCP request"); if let Some(reporter) = &state.stats_reporter { reporter.report(memoria_service::stats_reporter::StatsEvent::ApiCallLogged { user_id: auth.user_id.clone(), @@ -294,9 +290,6 @@ pub async fn mcp_handler( }; let user_id = auth.user_id.clone(); let scope_id = auth.scope_id.clone(); - if let Some(tool) = tracked_tool.clone() { - state.tool_usage_batcher.mark_used(user_id.clone(), tool); - } // Single reporting point for MCP call stats, shared by both the // notification path and the regular-request path below. @@ -332,6 +325,44 @@ pub async fn mcp_handler( None }; + // Reject before *any* per-user instrumentation or group storage lookup. + // A missing memory grant is not authorization to create a database merely + // to log the refusal. Shared stats and tracing retain the denial evidence. + if let Some(message) = authorization_error { + report_stats(&track_path, false); + tracing::warn!(user_id = %user_id, path = %track_path, rpc_code = -32003, "MCP scope admission denied"); + if req.get("id").is_none() { + return StatusCode::NO_CONTENT.into_response(); + } + return Json(json!({"jsonrpc": "2.0", "id": req["id"], + "error": {"code": -32003, "message": message}})) + .into_response(); + } + + // Identity-only metadata calls (initialize, tools/list, etc.) must not + // implicitly enable memory storage either. AuthUser defers MCP headers here. + let memory_telemetry = auth.has_scope(SCOPE_MEMORY_READ) || auth.has_scope(SCOPE_MEMORY_WRITE); + if memory_telemetry { + if let Some(tool) = tracked_tool.clone() { + state.tool_usage_batcher.mark_used(user_id.clone(), tool); + } + if let Some(agent) = crate::auth::request_tool_name(&headers) { + state.tool_usage_batcher.mark_used(user_id.clone(), agent); + } + } + let record_call = |status_code, rpc| { + if memory_telemetry { + state.call_log_batcher.record_rpc( + user_id.clone(), + "POST".to_string(), + track_path.clone(), + status_code, + t.elapsed().as_millis() as u32, + rpc, + ); + } + }; + // ── Group main-write guard (computed once, shared by both code paths) ───── // Resolved here — before the Notification early-return — so that // Notification-form write calls to `main` are also blocked instead of @@ -399,30 +430,11 @@ pub async fn mcp_handler( // JSON-RPC 2.0: a Notification is a *valid* Request without an "id" member. // The server MUST NOT reply to Notifications. if req.get("id").is_none() { - if authorization_error.is_some() { - report_stats(&track_path, false); - state.call_log_batcher.record_rpc( - user_id, - "POST".to_string(), - track_path, - 204, - t.elapsed().as_millis() as u32, - RpcMeta::err(-32003), - ); - return StatusCode::NO_CONTENT.into_response(); - } // Write guard: per JSON-RPC 2.0 the server MUST NOT reply to Notifications, // so we silently drop blocked writes without dispatching. if blocked_tool.is_some() { report_stats(&track_path, false); - state.call_log_batcher.record_rpc( - user_id, - "POST".to_string(), - track_path, - 204, - t.elapsed().as_millis() as u32, - RpcMeta::err(-32001), - ); + record_call(204, RpcMeta::err(-32001)); return StatusCode::NO_CONTENT.into_response(); } let dispatch_result = memoria_mcp::dispatch_http( @@ -445,40 +457,12 @@ pub async fn mcp_handler( // Report accurate ops metrics using the real RPC path and success flag // (JSON-RPC errors still return HTTP 200, so is_success must come from rpc.success). report_stats(&track_path, rpc.success); - state.call_log_batcher.record_rpc( - user_id, - "POST".to_string(), - track_path, - 204, // HTTP 204 No Content — correct for notifications - t.elapsed().as_millis() as u32, - rpc, - ); + record_call(204, rpc); return StatusCode::NO_CONTENT.into_response(); } let id = req["id"].clone(); - if let Some(message) = authorization_error { - let err_body = Json(json!({ - "jsonrpc": "2.0", - "id": id, - "error": { - "code": -32003, - "message": message - } - })); - report_stats(&track_path, false); - state.call_log_batcher.record_rpc( - user_id, - "POST".to_string(), - track_path, - 200, - t.elapsed().as_millis() as u32, - RpcMeta::err(-32003), - ); - return err_body.into_response(); - } - // Use the pre-computed write-guard decision (see above). if let Some(tool) = &blocked_tool { let err_body = Json(json!({ @@ -498,14 +482,7 @@ pub async fn mcp_handler( } })); report_stats(&track_path, false); - state.call_log_batcher.record_rpc( - user_id, - "POST".to_string(), - track_path, - 200, - t.elapsed().as_millis() as u32, - RpcMeta::err(-32001), - ); + record_call(200, RpcMeta::err(-32001)); return err_body.into_response(); } @@ -547,14 +524,7 @@ pub async fn mcp_handler( // Report accurate ops metrics using the real RPC path and success flag // (JSON-RPC errors still return HTTP 200, so is_success must come from rpc.success). report_stats(&track_path, rpc.success); - state.call_log_batcher.record_rpc( - user_id, - "POST".to_string(), - track_path, - 200, // HTTP 200 — always correct for JSON-RPC responses - t.elapsed().as_millis() as u32, - rpc, - ); + record_call(200, rpc); response } diff --git a/memoria/crates/memoria-api/src/state.rs b/memoria/crates/memoria-api/src/state.rs index 02c210d9..f2c3b5de 100644 --- a/memoria/crates/memoria-api/src/state.rs +++ b/memoria/crates/memoria-api/src/state.rs @@ -31,6 +31,7 @@ struct ApiKeyCacheEntry { group_id: Option, key_prefix: String, scopes: Vec, + expires_at: Option, cached_at: Instant, } @@ -41,6 +42,7 @@ pub struct CachedApiKeyPrincipal { pub group_id: Option, pub key_prefix: String, pub scopes: Vec, + pub expires_at: Option, } #[derive(Clone)] @@ -58,16 +60,27 @@ impl ApiKeyCache { } pub fn get(&self, key_hash: &str) -> Option { - let now = Instant::now(); + self.get_at(key_hash, Instant::now(), chrono::Utc::now().naive_utc()) + } + + fn get_at( + &self, + key_hash: &str, + now: Instant, + wall_now: chrono::NaiveDateTime, + ) -> Option { if let Ok(cache) = self.inner.read() { if let Some(entry) = cache.get(key_hash) { - if now.duration_since(entry.cached_at) < self.ttl { + if now.duration_since(entry.cached_at) < self.ttl + && entry.expires_at.is_none_or(|expiry| wall_now < expiry) + { return Some(CachedApiKeyPrincipal { key_id: entry.key_id.clone(), user_id: entry.user_id.clone(), group_id: entry.group_id.clone(), key_prefix: entry.key_prefix.clone(), scopes: entry.scopes.clone(), + expires_at: entry.expires_at, }); } } @@ -87,6 +100,7 @@ impl ApiKeyCache { group_id: principal.group_id, key_prefix: principal.key_prefix, scopes: principal.scopes, + expires_at: principal.expires_at, cached_at: Instant::now(), }, ); @@ -366,6 +380,61 @@ impl AppState { mod tests { use super::*; + fn cached_principal(expires_at: Option) -> CachedApiKeyPrincipal { + CachedApiKeyPrincipal { + key_id: "test-key".into(), + user_id: "test-user".into(), + group_id: None, + key_prefix: "sk-test".into(), + scopes: vec!["identity:read".into()], + expires_at, + } + } + + #[test] + fn api_key_cache_enforces_exact_key_expiry_before_cache_ttl() { + let cache = ApiKeyCache::new(Duration::from_secs(300)); + let expiry = chrono::Utc::now().naive_utc() + chrono::Duration::seconds(10); + cache.insert("hash".into(), cached_principal(Some(expiry))); + let now = Instant::now(); + assert!(cache + .get_at("hash", now, expiry - chrono::Duration::microseconds(1)) + .is_some()); + assert!(cache.get_at("hash", now, expiry).is_none()); + assert!( + !cache.inner.read().unwrap().contains_key("hash"), + "expired entry evicted" + ); + } + + #[test] + fn api_key_cache_keeps_nonexpiring_keys_but_still_enforces_cache_ttl() { + let cache = ApiKeyCache::new(Duration::from_secs(300)); + cache.insert("hash".into(), cached_principal(None)); + let now = Instant::now(); + let wall_now = chrono::Utc::now().naive_utc(); + assert!(cache.get_at("hash", now, wall_now).is_some()); + assert!(cache + .get_at("hash", now + Duration::from_secs(301), wall_now) + .is_none()); + } + + #[test] + fn api_key_cache_rejects_expired_entries_on_every_public_read() { + let cache = ApiKeyCache::new(Duration::from_secs(300)); + let expired = chrono::Utc::now().naive_utc() - chrono::Duration::seconds(1); + cache.insert("hash".into(), cached_principal(Some(expired))); + assert!(cache.get("hash").is_none()); + // A concurrent stale lookup must not extend the key's absolute lifetime. + cache.insert("hash".into(), cached_principal(Some(expired))); + assert!(cache.get("hash").is_none()); + cache.insert( + "hash".into(), + cached_principal(Some(expired + chrono::Duration::hours(1))), + ); + assert!(cache.get("hash").is_some()); + } + #[tokio::test] async fn test_metrics_cache_hit() { let cache: Arc>> = Arc::new(RwLock::new(None)); diff --git a/memoria/crates/memoria-api/tests/api_e2e.rs b/memoria/crates/memoria-api/tests/api_e2e.rs index fb273dd5..b136f2ea 100644 --- a/memoria/crates/memoria-api/tests/api_e2e.rs +++ b/memoria/crates/memoria-api/tests/api_e2e.rs @@ -1993,6 +1993,360 @@ async fn test_scoped_api_key_whoami_and_memory_authorization() { .contains("memory:write")); } +async fn assert_scoped_user_has_no_database(server: &support::multi_db::ApiTestServer, user: &str) { + let registry: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM mem_user_registry WHERE user_id = ?") + .bind(user) + .fetch_one(&server.shared_pool()) + .await + .unwrap(); + assert_eq!(registry, 0, "denied request registered {user}"); + let db_name = memoria_storage::DbRouter::user_db_name_for_id(user); + let databases: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = ?", + ) + .bind(db_name) + .fetch_one(&server.shared_pool()) + .await + .unwrap(); + assert_eq!(databases, 0, "denied request created a database for {user}"); +} + +#[tokio::test] +async fn test_scoped_mcp_denials_do_not_provision_storage_even_after_flush() { + let master = "denied-telemetry-test-master"; + let (base, client, server) = spawn_server_with_master_key(master).await; + let mut users = Vec::new(); + let mut write_key = String::new(); + for (scopes, denied_tools) in [ + ( + json!(["identity:read"]), + vec!["memory_store", "memory_search"], + ), + ( + json!(["identity:read", "memory:read"]), + vec!["memory_store", "memory_tune_params"], + ), + ( + json!(["identity:read", "memory:read", "memory:write"]), + vec!["memory_unclassified"], + ), + ] { + let user = uid(); + let response = client + .post(format!("{base}/auth/keys")) + .bearer_auth(master) + .json(&json!({"user_id":user,"name":"denial-test","scopes":scopes})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 201); + let key = response.json::().await.unwrap()["raw_key"] + .as_str() + .unwrap() + .to_owned(); + assert_scoped_user_has_no_database(&server, &user).await; + for tool in denied_tools { + for notification in [false, true] { + for header in ["X-Memoria-Tool", "X-Tool-Name"] { + let mut request = json!({"jsonrpc":"2.0","method":"tools/call", + "params":{"name":tool,"arguments":{"content":"must not be stored"}}}); + if !notification { + request["id"] = json!(1); + } + let response = client + .post(format!("{base}/mcp")) + .bearer_auth(&key) + .header(header, "denied-test-agent") + .json(&request) + .send() + .await + .unwrap(); + if notification { + assert_eq!(response.status(), 204); + assert!(response.bytes().await.unwrap().is_empty()); + } else { + assert_eq!(response.status(), 200); + let result: Value = response.json().await.unwrap(); + assert_eq!(result["error"]["code"], -32003, "{result}"); + } + } + } + } + // Parse failures have not admitted any tool either. + for invalid in [ + "not json", + "[]", + r#"{"jsonrpc":"1.0","method":"tools/call"}"#, + ] { + let response = client + .post(format!("{base}/mcp")) + .bearer_auth(&key) + .header("X-Memoria-Tool", "denied-test-agent") + .body(invalid) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200); + assert!(response + .json::() + .await + .unwrap() + .get("error") + .is_some()); + } + if scopes.as_array().unwrap().len() == 1 { + let response = client + .get(format!("{base}/auth/whoami")) + .bearer_auth(&key) + .header("X-Memoria-Tool", "identity-test-agent") + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200); + for method in [ + "initialize", + "tools/list", + "ping", + "notifications/initialized", + ] { + let response = client + .post(format!("{base}/mcp")) + .bearer_auth(&key) + .header("X-Tool-Name", "identity-test-agent") + .json(&json!({"jsonrpc":"2.0","id":2,"method":method,"params":{}})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200); + } + } + assert!(server + .state() + .tool_usage_batcher + .get_user_tool_usage(&user) + .is_empty()); + users.push(user); + write_key = key; + } + server + .state() + .tool_usage_batcher + .flush(&server.service()) + .await; + server + .state() + .call_log_batcher + .flush(&server.service()) + .await; + for user in &users { + assert_scoped_user_has_no_database(&server, user).await; + } + server.state().drain_flushers().await; + for user in &users { + assert_scoped_user_has_no_database(&server, user).await; + } + + // Positive control: an admitted write still provisions and tracks usage. + let response = client + .post(format!("{base}/mcp")) + .bearer_auth(&write_key) + .header("X-Memoria-Tool", "allowed-test-agent") + .json(&json!({"jsonrpc":"2.0","id":3,"method":"tools/call", + "params":{"name":"memory_store","arguments":{"content":"authorized fixture"}}})) + .send() + .await + .unwrap(); + let result: Value = response.json().await.unwrap(); + assert!(result.get("error").is_none(), "{result}"); + assert_ne!(result["result"]["isError"], true, "{result}"); + let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM mem_user_registry WHERE user_id = ?") + .bind(users.last().unwrap()) + .fetch_one(&server.shared_pool()) + .await + .unwrap(); + assert_eq!(count, 1); + assert!(!server + .state() + .tool_usage_batcher + .get_user_tool_usage(users.last().unwrap()) + .is_empty()); + server + .state() + .tool_usage_batcher + .flush(&server.service()) + .await; + server + .state() + .call_log_batcher + .flush(&server.service()) + .await; +} + +#[tokio::test] +async fn test_scoped_key_warm_cache_expires_for_rest_and_mcp() { + use sha2::{Digest, Sha256}; + let master = "expiry-test-master"; + let (base, client, server) = spawn_server_with_master_key(master).await; + // DATETIME(6) stores microseconds; Linux clocks can expose nanoseconds. + let expiry = chrono::DateTime::from_timestamp_micros(chrono::Utc::now().timestamp_micros()) + .unwrap() + .naive_utc() + + chrono::Duration::seconds(5); + let mut keys = Vec::new(); + for expires_at in [ + Some(expiry), + Some(expiry), + None, + Some(expiry + chrono::Duration::hours(1)), + ] { + let response = client.post(format!("{base}/auth/keys")).bearer_auth(master) + .json(&json!({"user_id":uid(),"name":"expiry-test","expires_at":expires_at.map(|ts|ts.to_string()), + "scopes":["identity:read","memory:read","memory:write"]})).send().await.unwrap(); + assert_eq!(response.status(), 201); + let key = response.json::().await.unwrap()["raw_key"] + .as_str() + .unwrap() + .to_owned(); + let response = client + .post(format!("{base}/mcp")) + .bearer_auth(&key) + .json(&json!({"jsonrpc":"2.0","id":1,"method":"ping"})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200, "must warm before expiry"); + let hash = format!("{:x}", Sha256::digest(key.as_bytes())); + let cached = server + .state() + .api_key_cache + .get(&hash) + .expect("principal cached"); + assert_eq!(cached.expires_at, expires_at); + keys.push(key); + } + let wait = (expiry - chrono::Utc::now().naive_utc()) + .to_std() + .unwrap_or_default(); + tokio::time::sleep(wait + std::time::Duration::from_millis(100)).await; + // Separate warmed entries: REST rejection must not evict the MCP entry + // before its own first post-expiry authentication attempt. + let rest = client + .get(format!("{base}/v1/memories")) + .bearer_auth(&keys[0]) + .send() + .await + .unwrap(); + assert_eq!(rest.status(), 401); + let mcp = client + .post(format!("{base}/mcp")) + .bearer_auth(&keys[1]) + .json(&json!({"jsonrpc":"2.0","id":2,"method":"tools/call", + "params":{"name":"memory_store","arguments":{"content":"expired denial"}}})) + .send() + .await + .unwrap(); + assert_eq!(mcp.status(), 401); + for key in &keys[..2] { + let response = client + .get(format!("{base}/auth/whoami")) + .bearer_auth(key) + .send() + .await + .unwrap(); + assert_eq!( + response.status(), + 401, + "fresh DB validation agrees with cache expiry" + ); + } + for key in &keys[2..] { + let response = client + .get(format!("{base}/v1/memories")) + .bearer_auth(key) + .send() + .await + .unwrap(); + assert_eq!( + response.status(), + 200, + "non-expiring/future keys remain valid" + ); + let response = client + .post(format!("{base}/mcp")) + .bearer_auth(key) + .json(&json!({"jsonrpc":"2.0","id":3,"method":"ping"})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200); + } +} + +#[tokio::test] +async fn test_scoped_group_denial_precedes_memory_storage_lookup() { + let master = "group-denial-test-master"; + let (base, client, server) = spawn_server_with_master_key(master).await; + let user = uid(); + let group = format!("grp_{}", uuid::Uuid::new_v4().simple()); + let group_db = memoria_storage::DbRouter::user_db_name_for_id(&group); + // Seed only control-plane metadata; neither member nor group has a memory DB. + sqlx::query("INSERT INTO mem_groups (group_id,group_name,db_name,owner_user_id,status,created_at,updated_at) VALUES (?, 'denial fixture', ?, ?, 'active', NOW(6), NOW(6))") + .bind(&group).bind(&group_db).bind(&user).execute(&server.shared_pool()).await.unwrap(); + sqlx::query("INSERT INTO mem_group_members (group_id,user_id,role,is_active,joined_at) VALUES (?, ?, 'owner', 1, NOW(6))") + .bind(&group).bind(&user).execute(&server.shared_pool()).await.unwrap(); + let response = client.post(format!("{base}/auth/keys")).bearer_auth(master) + .json(&json!({"user_id":user,"group_id":group,"name":"group-denial","scopes":["identity:read","memory:read"]})).send().await.unwrap(); + assert_eq!(response.status(), 201); + let key = response.json::().await.unwrap()["raw_key"] + .as_str() + .unwrap() + .to_owned(); + for notification in [false, true] { + let mut body = json!({"jsonrpc":"2.0","method":"tools/call","params":{"name":"memory_store","arguments":{"content":"denied"}}}); + if !notification { + body["id"] = json!(1); + } + let response = client + .post(format!("{base}/mcp")) + .bearer_auth(&key) + .header("X-Memoria-Tool", "denied-group-agent") + .json(&body) + .send() + .await + .unwrap(); + if notification { + assert_eq!(response.status(), 204); + } else { + assert_eq!( + response.json::().await.unwrap()["error"]["code"], + -32003 + ); + } + } + let response = client + .post(format!("{base}/v1/memories")) + .bearer_auth(&key) + .json(&json!({"content":"denied"})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 403); + server.state().drain_flushers().await; + assert_scoped_user_has_no_database(&server, &user).await; + let databases: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = ?", + ) + .bind(group_db) + .fetch_one(&server.shared_pool()) + .await + .unwrap(); + assert_eq!( + databases, 0, + "unauthorized guard lookup provisioned the group DB" + ); +} + #[tokio::test] async fn test_scoped_mcp_write_authorization_is_independent_of_metrics() { let master = "mcp-scope-regression-master";