sharang
a3a96fe2cc
CI / Check (push) Has been skipped
CI / Detect Changes (push) Successful in 5s
CI / Deploy Agent (push) Successful in 8m13s
CI / Deploy Dashboard (push) Successful in 7m3s
CI / Deploy Docs (push) Has been skipped
CI / Deploy MCP (push) Successful in 1m50s
MCP server validates per-tenant bearer tokens on incoming calls and routes each tool to the caller's tenant DB. Closes the cross-tenant data leak in the MCP path identified in M7.3.
130 lines
4.3 KiB
Rust
130 lines
4.3 KiB
Rust
//! Bearer-token authentication for incoming MCP HTTP requests.
|
|
//!
|
|
//! LLM clients (Claude Desktop / Cursor / ChatGPT / etc.) can't run
|
|
//! Keycloak OIDC, so the MCP server uses opaque static tokens minted
|
|
//! per-tenant via the agent's `POST /api/v1/mcp-tokens` endpoint.
|
|
//!
|
|
//! Flow per request:
|
|
//! 1. Extract `Authorization: Bearer <token>`. Missing → 401.
|
|
//! 2. SHA-256 hash the token.
|
|
//! 3. Look up the hash in `<prefix>__admin.mcp_tokens`. Missing or
|
|
//! revoked → 401.
|
|
//! 4. Fire-and-forget update of `last_used_at` so the dashboard can
|
|
//! show staleness without blocking the handler.
|
|
//! 5. Stash the tenant_id in [`TENANT_ID`] (a `tokio::task_local`) so
|
|
//! the MCP tool handlers can read it without modifying rmcp's
|
|
//! handler signatures.
|
|
//!
|
|
//! The `task_local` is scoped around the inner service call via
|
|
//! [`bearer_auth`], so every handler invoked downstream sees the
|
|
//! tenant_id without us having to thread it through the macro-
|
|
//! generated tool router.
|
|
|
|
use axum::body::Body;
|
|
use axum::extract::{Request, State};
|
|
use axum::http::StatusCode;
|
|
use axum::middleware::Next;
|
|
use axum::response::{IntoResponse, Response};
|
|
use mongodb::bson::doc;
|
|
use sha2::{Digest, Sha256};
|
|
|
|
use crate::database::DatabasePool;
|
|
|
|
tokio::task_local! {
|
|
/// Tenant id resolved from the bearer for this request. Set by
|
|
/// [`bearer_auth`] before the inner service runs; read by the
|
|
/// MCP tool handlers via [`current_tenant_id`].
|
|
pub static TENANT_ID: String;
|
|
}
|
|
|
|
/// Mongo collection name in `<prefix>__admin`.
|
|
const COLLECTION: &str = "mcp_tokens";
|
|
|
|
/// Returns the tenant_id set by the auth middleware. `None` outside a
|
|
/// request scope (e.g. unit tests that bypass the middleware).
|
|
pub fn current_tenant_id() -> Option<String> {
|
|
TENANT_ID.try_with(|s| s.clone()).ok()
|
|
}
|
|
|
|
/// Axum middleware: validate bearer → set [`TENANT_ID`] → call inner.
|
|
pub async fn bearer_auth(
|
|
State(pool): State<DatabasePool>,
|
|
request: Request,
|
|
next: Next,
|
|
) -> Response {
|
|
let Some(token) = extract_bearer(&request) else {
|
|
return (StatusCode::UNAUTHORIZED, "Missing bearer token").into_response();
|
|
};
|
|
if !token.starts_with("mcpt_") {
|
|
return (StatusCode::UNAUTHORIZED, "Invalid token format").into_response();
|
|
}
|
|
let token_hash = sha256_hex(&token);
|
|
|
|
let col = pool.admin_db().collection::<TokenLookup>(COLLECTION);
|
|
let found = match col
|
|
.find_one(doc! { "token_hash": &token_hash, "revoked": false })
|
|
.await
|
|
{
|
|
Ok(Some(t)) => t,
|
|
Ok(None) => {
|
|
return (StatusCode::UNAUTHORIZED, "Invalid or revoked token").into_response();
|
|
}
|
|
Err(e) => {
|
|
tracing::error!("MCP token lookup failed: {e}");
|
|
return (StatusCode::INTERNAL_SERVER_ERROR, "Token lookup error").into_response();
|
|
}
|
|
};
|
|
|
|
// Fire-and-forget last_used_at update — never block the handler.
|
|
let col2 = pool.admin_db().collection::<TokenLookup>(COLLECTION);
|
|
let hash_for_update = token_hash.clone();
|
|
tokio::spawn(async move {
|
|
let _ = col2
|
|
.update_one(
|
|
doc! { "token_hash": &hash_for_update },
|
|
doc! { "$set": { "last_used_at": mongodb::bson::DateTime::now() } },
|
|
)
|
|
.await;
|
|
});
|
|
|
|
let tenant_id = found.tenant_id;
|
|
let inner = next.run(request);
|
|
TENANT_ID.scope(tenant_id, inner).await
|
|
}
|
|
|
|
/// Bare-bones projection — we don't need the whole `McpToken` here,
|
|
/// just enough to route and confirm validity.
|
|
#[derive(serde::Deserialize)]
|
|
struct TokenLookup {
|
|
tenant_id: String,
|
|
}
|
|
|
|
fn extract_bearer(req: &Request<Body>) -> Option<String> {
|
|
req.headers()
|
|
.get(axum::http::header::AUTHORIZATION)
|
|
.and_then(|v| v.to_str().ok())
|
|
.and_then(|s| s.strip_prefix("Bearer "))
|
|
.map(|s| s.trim().to_string())
|
|
.filter(|s| !s.is_empty())
|
|
}
|
|
|
|
fn sha256_hex(s: &str) -> String {
|
|
let mut h = Sha256::new();
|
|
h.update(s.as_bytes());
|
|
hex::encode(h.finalize())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn sha256_known_value() {
|
|
// python -c 'import hashlib; print(hashlib.sha256(b"mcpt_known").hexdigest())'
|
|
assert_eq!(
|
|
sha256_hex("mcpt_known"),
|
|
"27cf6cf678a44244106863c1c031be8e57b84c2b3019d742f755f8e7afa75dfd"
|
|
);
|
|
}
|
|
}
|