Menu
Public repository / master
AkurAI-Drive / backend/src/mcp.rs
Latest change f891f45338fe5e5d14370df77cc8001dceb06426 - Drive MCP: recover /mcp tools and make it an IdP OAuth protected resource by AkurAI Agent
//! Streamable-HTTP MCP endpoint (`POST /mcp`), an OAuth 2.0 protected resource
//! for the AkurAI IdP.
//!
//! Recovered from the deployed akurai-drive-api (its `src/mcp.rs` was never
//! committed): JSON-RPC `initialize` / `tools/list` / `tools/call` with the nine
//! `drive_*` tools. Hooked up to the IdP per RFC 9728:
//! - metadata at `/.well-known/oauth-protected-resource[/mcp]`;
//! - `401` + `WWW-Authenticate: Bearer resource_metadata=...` without a valid token;
//! - accepts IdP EdDSA access tokens whose `aud` contains `MCP_AUDIENCE`
//! (see `oauth.rs`) and, as before, Drive device tokens from `/auth/device`.
//!
//! Every tool is ownership-gated on the caller's `sub`, exactly like `/api/*`.
use axum::{
Json,
body::Bytes,
http::{HeaderMap, HeaderValue, StatusCode, header},
response::{IntoResponse, Response},
};
use serde_json::{Value, json};
use crate::api::{owned_buckets, storage_access};
use crate::config;
pub const PROTOCOL_VERSION: &str = "2024-11-05";
const SERVER_NAME: &str = "akurai-drive-api";
const RESOURCE_NAME: &str = "AkurAI Drive MCP";
const SCOPES: [&str; 5] = ["openid", "profile", "email", "groups", "offline_access"];
const PRESIGN_SECS: u64 = 900;
const MAX_LIMIT: u64 = 1000;
// ---------------------------------------------------------------------------
// Protected resource metadata and challenges
// ---------------------------------------------------------------------------
fn base(base_url: &str) -> &str {
base_url.trim_end_matches('/')
}
pub fn resource_url(base_url: &str) -> String {
format!("{}/mcp", base(base_url))
}
pub fn resource_metadata_url(base_url: &str) -> String {
format!(
"{}/.well-known/oauth-protected-resource/mcp",
base(base_url)
)
}
pub fn protected_resource_metadata(base_url: &str, issuer: &str) -> Value {
json!({
"resource": resource_url(base_url),
"authorization_servers": [issuer.trim_end_matches('/')],
"scopes_supported": SCOPES,
"bearer_methods_supported": ["header"],
"resource_name": RESOURCE_NAME,
})
}
/// `WWW-Authenticate` value; `invalid` marks a presented-but-rejected token.
pub fn www_authenticate(base_url: &str, invalid: bool) -> String {
let url = resource_metadata_url(base_url);
if invalid {
format!("Bearer error=\"invalid_token\", resource_metadata=\"{url}\"")
} else {
format!("Bearer resource_metadata=\"{url}\"")
}
}
/// Browsers must be same-origin; non-browser MCP clients send no Origin.
pub fn origin_allowed(origin: Option<&str>, base_url: &str) -> bool {
origin.is_none_or(|o| o == base(base_url))
}
/// `GET /.well-known/oauth-protected-resource[/mcp]`
pub async fn metadata() -> Response {
let cfg = config::get();
Json(protected_resource_metadata(
&cfg.drive_public_url,
&cfg.idp_issuer,
))
.into_response()
}
fn bearer(headers: &HeaderMap) -> Option<&str> {
let raw = headers.get(header::AUTHORIZATION)?.to_str().ok()?;
let token = raw
.strip_prefix("Bearer ")
.or_else(|| raw.strip_prefix("bearer "))?
.trim();
(!token.is_empty()).then_some(token)
}
fn challenge(invalid: bool) -> Response {
let cfg = config::get();
let value = www_authenticate(&cfg.drive_public_url, invalid);
let mut resp = (
StatusCode::UNAUTHORIZED,
Json(json!({"error": if invalid { "invalid bearer token" } else { "bearer authentication required" }})),
)
.into_response();
if let Ok(v) = HeaderValue::from_str(&value) {
resp.headers_mut().insert(header::WWW_AUTHENTICATE, v);
}
resp
}
/// Resolve the bearer to a Drive owner `sub`: an IdP JWT access token, or a
/// legacy opaque device token.
async fn caller_sub(headers: &HeaderMap, token: &str) -> Option<String> {
if token.split('.').count() == 3 {
let token = token.to_string();
return tokio::task::spawn_blocking(move || crate::oauth::authenticate(&token))
.await
.ok()?
.map(|grant| grant.sub)
.ok();
}
crate::device::lookup_bearer(headers).map(|u| u.sub)
}
/// `POST /mcp`
pub async fn handle(headers: HeaderMap, body: Bytes) -> Response {
let cfg = config::get();
let origin = headers.get(header::ORIGIN).and_then(|v| v.to_str().ok());
if !origin_allowed(origin, &cfg.drive_public_url) {
return (
StatusCode::FORBIDDEN,
Json(json!({"error": "cross_origin_request"})),
)
.into_response();
}
let Some(token) = bearer(&headers) else {
return challenge(false);
};
let Some(sub) = caller_sub(&headers, token).await else {
return challenge(true);
};
let Ok(request) = serde_json::from_slice::<Value>(&body) else {
return Json(rpc_error(Value::Null, -32700, "parse error")).into_response();
};
match dispatch(&sub, &request).await {
Some(reply) => Json(reply).into_response(),
None => StatusCode::ACCEPTED.into_response(),
}
}
// ---------------------------------------------------------------------------
// JSON-RPC
// ---------------------------------------------------------------------------
fn rpc_result(id: Value, result: Value) -> Value {
json!({"jsonrpc": "2.0", "id": id, "result": result})
}
fn rpc_error(id: Value, code: i64, message: &str) -> Value {
json!({"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message}})
}
/// Handle one JSON-RPC message; `None` for notifications (no `id`).
pub async fn dispatch(sub: &str, request: &Value) -> Option<Value> {
let id = request.get("id").cloned()?;
let method = request.get("method").and_then(Value::as_str).unwrap_or("");
let params = request.get("params").cloned().unwrap_or(json!({}));
Some(match method {
"initialize" => rpc_result(id, initialize_result()),
"ping" => rpc_result(id, json!({})),
"tools/list" => rpc_result(id, json!({"tools": tool_definitions()})),
"tools/call" => {
let name = params.get("name").and_then(Value::as_str).unwrap_or("");
let args = params.get("arguments").cloned().unwrap_or(json!({}));
rpc_result(id, tool_result(call_tool(sub, name, args).await))
}
_ => rpc_error(id, -32601, "method not found"),
})
}
pub fn initialize_result() -> Value {
json!({
"protocolVersion": PROTOCOL_VERSION,
"capabilities": {"tools": {}},
"serverInfo": {"name": SERVER_NAME, "version": env!("CARGO_PKG_VERSION")},
})
}
fn tool_result(outcome: Result<Value, String>) -> Value {
match outcome {
Ok(v) => json!({"content": [{"type": "text", "text": v.to_string()}], "isError": false}),
Err(e) => json!({"content": [{"type": "text", "text": e}], "isError": true}),
}
}
// ---------------------------------------------------------------------------
// Tools
// ---------------------------------------------------------------------------
fn tool(name: &str, description: &str, properties: Value, required: &[&str]) -> Value {
json!({
"name": name,
"description": description,
"inputSchema": {
"type": "object",
"properties": properties,
"required": required,
"additionalProperties": false,
},
})
}
pub fn tool_definitions() -> Vec<Value> {
let s = json!({"type": "string"});
let overwrite = json!({"type": "boolean", "default": false});
vec![
tool(
"drive_buckets",
"List the caller's Drive buckets",
json!({}),
&[],
),
tool(
"drive_list",
"List folders and files with bounded pagination",
json!({"bucket": s, "prefix": s, "cursor": s,
"limit": {"type": "integer", "minimum": 1, "maximum": MAX_LIMIT, "default": 100}}),
&["bucket"],
),
tool(
"drive_search",
"Search object keys within a bucket",
json!({"bucket": s, "query": s,
"limit": {"type": "integer", "minimum": 1, "maximum": MAX_LIMIT, "default": 100}}),
&["bucket", "query"],
),
tool(
"drive_create_folder",
"Create a virtual folder marker",
json!({"bucket": s, "path": s}),
&["bucket", "path"],
),
tool(
"drive_presign",
"Create a 15-minute upload or download URL",
json!({"bucket": s, "key": s,
"operation": {"type": "string", "enum": ["upload", "download"]}}),
&["bucket", "key", "operation"],
),
tool(
"drive_copy",
"Copy a file with collision protection",
json!({"bucket": s, "key": s, "destination": s, "overwrite": overwrite}),
&["bucket", "key", "destination"],
),
tool(
"drive_move",
"Move a file with collision protection",
json!({"bucket": s, "key": s, "destination": s, "overwrite": overwrite}),
&["bucket", "key", "destination"],
),
tool(
"drive_delete",
"Delete a file",
json!({"bucket": s, "key": s}),
&["bucket", "key"],
),
tool(
"drive_integrity",
"Read size, ETag, SHA-256 checksum, and version metadata",
json!({"bucket": s, "key": s}),
&["bucket", "key"],
),
]
}
fn arg(args: &Value, name: &str) -> Result<String, String> {
args.get(name)
.and_then(Value::as_str)
.map(str::trim)
.filter(|v| !v.is_empty())
.map(str::to_string)
.ok_or_else(|| format!("{name} is required"))
}
fn limit(args: &Value) -> usize {
args.get("limit")
.and_then(Value::as_u64)
.unwrap_or(100)
.clamp(1, MAX_LIMIT) as usize
}
/// Split the flat listing under `prefix` into immediate sub-folders and files,
/// then page the files from offset `cursor`.
pub fn list_page(
files: Vec<akurai_drive_contract::RemoteFile>,
prefix: &str,
cursor: usize,
limit: usize,
) -> Value {
let mut folders: Vec<String> = Vec::new();
let mut direct = Vec::new();
for f in files {
let rest = f.key.strip_prefix(prefix).unwrap_or(&f.key);
if rest.is_empty() {
continue; // the folder marker of `prefix` itself
}
match rest.find('/') {
Some(i) => {
let folder = format!("{prefix}{}", &rest[..=i]);
if !folders.contains(&folder) {
folders.push(folder);
}
}
None => direct.push(f),
}
}
let total = direct.len();
let page: Vec<_> = direct.into_iter().skip(cursor).take(limit).collect();
let next = cursor + page.len();
json!({
"folders": folders,
"files": page,
"nextCursor": (next < total).then(|| next.to_string()),
})
}
/// Run blocking S3 work on the blocking pool.
async fn blocking<T: Send + 'static>(
f: impl FnOnce() -> Result<T, String> + Send + 'static,
) -> Result<T, String> {
tokio::task::spawn_blocking(f)
.await
.map_err(|_| "storage task failed".to_string())?
}
pub async fn call_tool(sub: &str, name: &str, args: Value) -> Result<Value, String> {
if name == "drive_buckets" {
return Ok(json!({"buckets": owned_buckets(sub)}));
}
if !tool_definitions().iter().any(|t| t["name"] == name) {
return Err("unknown Drive tool".to_string());
}
let bucket = arg(&args, "bucket")?;
let region = storage_access(sub, &bucket)
.ok_or("bucket not found")?
.region;
match name {
"drive_list" => {
let prefix = args
.get("prefix")
.and_then(Value::as_str)
.unwrap_or("")
.to_string();
let cursor = args
.get("cursor")
.and_then(Value::as_str)
.and_then(|c| c.parse().ok())
.unwrap_or(0);
let limit = limit(&args);
let p = prefix.clone();
let files = blocking(move || {
crate::util::list_objects(
&bucket,
®ion,
Some(p.as_str()).filter(|p| !p.is_empty()),
)
})
.await?;
Ok(list_page(files, &prefix, cursor, limit))
}
"drive_search" => {
let query = arg(&args, "query")?.to_lowercase();
let limit = limit(&args);
let files = blocking(move || crate::util::list_objects(&bucket, ®ion, None)).await?;
let hits: Vec<_> = files
.into_iter()
.filter(|f| f.key.to_lowercase().contains(&query))
.take(limit)
.collect();
Ok(json!({"results": hits}))
}
"drive_create_folder" => {
let path = format!("{}/", arg(&args, "path")?.trim_matches('/'));
if path == "/" {
return Err("folder path is required".to_string());
}
let key = path.clone();
blocking(move || crate::util::put_empty_object(&bucket, ®ion, &key)).await?;
Ok(json!({"folder": path}))
}
"drive_presign" => presign(&bucket, ®ion, &arg(&args, "key")?, &args),
"drive_copy" | "drive_move" => {
let key = arg(&args, "key")?;
let dst = arg(&args, "destination")?;
if key == dst {
return Err("source and destination must be distinct non-empty keys".to_string());
}
let overwrite = args
.get("overwrite")
.and_then(Value::as_bool)
.unwrap_or(false);
let is_move = name == "drive_move";
let (k, d) = (key.clone(), dst.clone());
blocking(move || {
if !overwrite && crate::util::head_object(&bucket, ®ion, &d)?.is_some() {
return Err("destination already exists".to_string());
}
crate::util::copy_object(&bucket, ®ion, &k, &d)?;
if is_move {
crate::util::delete_object(&bucket, ®ion, &k)?;
}
Ok(())
})
.await?;
Ok(json!({"key": key, "destination": dst, "moved": is_move}))
}
"drive_delete" => {
let key = arg(&args, "key")?;
let k = key.clone();
blocking(move || crate::util::delete_object(&bucket, ®ion, &k)).await?;
Ok(json!({"deleted": key}))
}
"drive_integrity" => {
let key = arg(&args, "key")?;
let k = key.clone();
let head = blocking(move || crate::util::head_object(&bucket, ®ion, &k))
.await?
.ok_or("file not found")?;
Ok(json!({
"key": key,
"size": head.get("ContentLength"),
"etag": head.get("ETag").and_then(Value::as_str).map(|e| e.trim_matches('"')),
"checksumSha256": head.get("ChecksumSHA256"),
"versionId": head.get("VersionId"),
"lastModified": head.get("LastModified"),
}))
}
_ => Err("unknown Drive tool".to_string()),
}
}
fn presign(bucket: &str, region: &str, key: &str, args: &Value) -> Result<Value, String> {
let method = match args.get("operation").and_then(Value::as_str) {
Some("upload") => "PUT",
Some("download") => "GET",
_ => return Err("operation must be upload or download".to_string()),
};
let cfg = config::get();
if cfg.aws_access_key_id.is_empty() || cfg.aws_secret_access_key.is_empty() {
return Err("signing credentials are unavailable".to_string());
}
let url = crate::sigv4::presign(&crate::sigv4::Sigv4Params {
access_key: &cfg.aws_access_key_id,
secret_key: &cfg.aws_secret_access_key,
region,
bucket,
key,
method,
expires: PRESIGN_SECS,
now: chrono::Utc::now(),
});
Ok(json!({"url": url, "method": method, "expiresIn": PRESIGN_SECS}))
}
#[cfg(test)]
mod tests {
use super::*;
use akurai_drive_contract::RemoteFile;
const BASE: &str = "https://akurai-drive.olibuijr.com/";
#[test]
fn metadata_names_resource_and_idp() {
let m = protected_resource_metadata(BASE, "https://auth.olibuijr.com/");
assert_eq!(m["resource"], "https://akurai-drive.olibuijr.com/mcp");
assert_eq!(
m["authorization_servers"],
json!(["https://auth.olibuijr.com"])
);
assert_eq!(m["bearer_methods_supported"], json!(["header"]));
}
#[test]
fn challenge_points_at_metadata() {
let url = "https://akurai-drive.olibuijr.com/.well-known/oauth-protected-resource/mcp";
assert_eq!(
www_authenticate(BASE, false),
format!("Bearer resource_metadata=\"{url}\"")
);
assert_eq!(
www_authenticate(BASE, true),
format!("Bearer error=\"invalid_token\", resource_metadata=\"{url}\"")
);
}
#[test]
fn origin_check() {
assert!(origin_allowed(None, BASE));
assert!(origin_allowed(
Some("https://akurai-drive.olibuijr.com"),
BASE
));
assert!(!origin_allowed(Some("https://evil.example"), BASE));
}
#[test]
fn lists_the_nine_drive_tools() {
let names: Vec<_> = tool_definitions()
.iter()
.map(|t| t["name"].as_str().unwrap().to_string())
.collect();
assert_eq!(
names,
[
"drive_buckets",
"drive_list",
"drive_search",
"drive_create_folder",
"drive_presign",
"drive_copy",
"drive_move",
"drive_delete",
"drive_integrity"
]
);
}
#[tokio::test]
async fn rpc_initialize_unknown_method_and_notifications() {
let init = dispatch(
"u",
&json!({"jsonrpc": "2.0", "id": 1, "method": "initialize"}),
)
.await
.unwrap();
assert_eq!(init["result"]["protocolVersion"], PROTOCOL_VERSION);
let bad = dispatch("u", &json!({"jsonrpc": "2.0", "id": 2, "method": "nope"}))
.await
.unwrap();
assert_eq!(bad["error"]["code"], -32601);
let note = json!({"jsonrpc": "2.0", "method": "notifications/initialized"});
assert!(dispatch("u", ¬e).await.is_none());
}
#[tokio::test]
async fn unknown_tool_is_a_tool_error() {
assert_eq!(
call_tool("u", "drive_nope", json!({})).await,
Err("unknown Drive tool".to_string())
);
}
#[tokio::test]
async fn missing_bearer_gets_401_challenge() {
let resp = handle(HeaderMap::new(), Bytes::from_static(b"{}")).await;
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
let v = resp.headers()[header::WWW_AUTHENTICATE].to_str().unwrap();
assert!(v.starts_with("Bearer resource_metadata="));
}
#[tokio::test]
async fn foreign_origin_is_forbidden() {
let mut h = HeaderMap::new();
h.insert(
header::ORIGIN,
HeaderValue::from_static("https://evil.example"),
);
let resp = handle(h, Bytes::from_static(b"{}")).await;
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
}
fn file(key: &str) -> RemoteFile {
RemoteFile {
key: key.to_string(),
size: 1,
etag: String::new(),
last_modified: String::new(),
}
}
#[test]
fn list_page_splits_folders_and_pages_files() {
let files = [
"docs/a.txt",
"docs/b.txt",
"docs/sub/c.txt",
"docs/empty/",
"docs/c.txt",
]
.map(file)
.to_vec();
let page = list_page(files, "docs/", 0, 2);
assert_eq!(page["folders"], json!(["docs/sub/", "docs/empty/"]));
assert_eq!(page["files"].as_array().unwrap().len(), 2);
assert_eq!(page["nextCursor"], "2");
let files = ["docs/a.txt", "docs/b.txt", "docs/c.txt"]
.map(file)
.to_vec();
let last = list_page(files, "docs/", 2, 2);
assert_eq!(last["files"][0]["key"], "docs/c.txt");
assert!(last["nextCursor"].is_null());
}
}