From e995ba800d8c02b0055f7a9984701ce3e0554cf7 Mon Sep 17 00:00:00 2001 From: Logic Date: Thu, 30 Jul 2026 15:36:18 +0800 Subject: [PATCH 1/5] maintenance: validate MCP OAuth exchanges --- mcp-servers/mcp-bash-server/Cargo.lock | 1 + mcp-servers/mcp-bash-server/Cargo.toml | 1 + mcp-servers/mcp-bash-server/README.md | 30 +- .../mcp-bash-server/src/common/oauth.rs | 307 ++++++++++++++++-- mcp-servers/mcp-bash-server/src/main.rs | 51 ++- .../templates/mcp_oauth_authorize.html | 29 +- 6 files changed, 374 insertions(+), 45 deletions(-) diff --git a/mcp-servers/mcp-bash-server/Cargo.lock b/mcp-servers/mcp-bash-server/Cargo.lock index f4a3287d686..d7e7e1758cf 100644 --- a/mcp-servers/mcp-bash-server/Cargo.lock +++ b/mcp-servers/mcp-bash-server/Cargo.lock @@ -1071,6 +1071,7 @@ dependencies = [ "serde", "serde_urlencoded", "shlex", + "subtle", "tempfile", "tokio", "tokio-stream", diff --git a/mcp-servers/mcp-bash-server/Cargo.toml b/mcp-servers/mcp-bash-server/Cargo.toml index f4e96235294..0f87a032b3f 100644 --- a/mcp-servers/mcp-bash-server/Cargo.toml +++ b/mcp-servers/mcp-bash-server/Cargo.toml @@ -43,6 +43,7 @@ oauth2 = "5.0" toml = "0.8" regex = "1.11.1" shlex = "1.3" +subtle = "2.6" [dev-dependencies] tokio-test = "0.4" diff --git a/mcp-servers/mcp-bash-server/README.md b/mcp-servers/mcp-bash-server/README.md index 06172f70c56..d1538f2416d 100644 --- a/mcp-servers/mcp-bash-server/README.md +++ b/mcp-servers/mcp-bash-server/README.md @@ -18,12 +18,17 @@ Version `1.88.0` can absolutely work, and we recommend using the latest version If you want to run this MCP server locally using the default settings provided by the project, simply run the following command in the project root directory: -```Rust +```shell +export MCP_OAUTH_APPROVAL_SECRET="" cargo run ``` This MCP server will be deployed at `http://127.0.0.1:4000/mcp`, and you can use the `modelcontextprotocol/inspector` tool to connect to and use this MCP server. +Production mode requires `MCP_OAUTH_APPROVAL_SECRET`. The operator enters this +secret on the OAuth approval page; it is separate from dynamically registered +client credentials. Development mode does not require it. + For information on how to use the modelcontextprotocol/inspector tool, refer to the [inspector documentation](https://github.com/modelcontextprotocol/inspector). If you encounter any issues while using Inspector, it is recommended to use version `v0.16.2`. Other versions may also work. @@ -47,23 +52,35 @@ docker build --build-arg HTTPS_PROXY= --build-arg HTTP_PROXY=< After building, use the following command to run it: ```shell -docker run -d --name mcp-bash-server -p 4000:4000 --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest +docker run -d --name mcp-bash-server -p 127.0.0.1:4000:4000 \ + -e MCP_OAUTH_APPROVAL_SECRET="" \ + --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest ``` -The MCP Server inside the container runs on 0.0.0.0:4000. On the host machine, use the inspector with URL `http://localhost:4000/mcp` to connect to the MCP Server inside the container. +The MCP Server inside the container runs on 0.0.0.0:4000, while the example +publishes it only on the host loopback interface. On the host machine, use the +inspector with URL `http://localhost:4000/mcp` to connect to the MCP Server +inside the container. Remote deployments must terminate TLS before forwarding +OAuth endpoints to the container. #### Use custom config in container Container's workdir is `/app` and it will run the `/app/mcp-bash-server` when it start, this program will read the `config.toml` at the same directory, so you can put the `config.toml` in the `/app` directory to cover the default config in image. Use the command below to do it. ```shell -docker run -d --name mcp-bash-server -p 4000:4000 -v `pwd`/config.toml:/app/config.toml --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest +docker run -d --name mcp-bash-server -p 127.0.0.1:4000:4000 \ + -e MCP_OAUTH_APPROVAL_SECRET="" \ + -v `pwd`/config.toml:/app/config.toml \ + --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest ``` If you are using SELinux, you may need to run the command instead to let the container access the file in host. ```shell -docker run -d --name mcp-bash-server -p 4000:4000 -v `pwd`/config.toml:/app/config.toml:Z --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest +docker run -d --name mcp-bash-server -p 127.0.0.1:4000:4000 \ + -e MCP_OAUTH_APPROVAL_SECRET="" \ + -v `pwd`/config.toml:/app/config.toml:Z \ + --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest ``` To check if the config.toml is used, do this @@ -231,7 +248,8 @@ The method for using OAuth verification and connection is as follows: 1. Click `Open Auth Settings` 2. Click `Quick OAuth Flow` -3. Click `Approve` on the pop-up webpage +3. Enter the server operator's `MCP_OAUTH_APPROVAL_SECRET`, then click `Approve` + on the pop-up webpage 4. Return to the MCP inspector, click on the Access Tokens under `Authentication Complete` in `OAuth Flow Progress`. Copy the `access_token` from there 5. Click `Authentication`, paste the previously copied token into the `Bearer Token` field, then click Connect diff --git a/mcp-servers/mcp-bash-server/src/common/oauth.rs b/mcp-servers/mcp-bash-server/src/common/oauth.rs index 5956688309a..dedd9e4ec77 100644 --- a/mcp-servers/mcp-bash-server/src/common/oauth.rs +++ b/mcp-servers/mcp-bash-server/src/common/oauth.rs @@ -43,6 +43,7 @@ use rmcp::transport::auth::{ AuthorizationMetadata, ClientRegistrationRequest, ClientRegistrationResponse, OAuthClientConfig, }; use serde::{Deserialize, Serialize}; +use subtle::ConstantTimeEq; use tokio::sync::RwLock; use tracing::{debug, error, info, warn}; use uuid::Uuid; @@ -61,11 +62,22 @@ pub struct McpOAuthStore { pub auth_sessions: Arc>>, /// Valid access tokens indexed by token string pub access_tokens: Arc>>, + /// Secret required for a resource owner to approve an authorization request + approval_secret: Arc, } impl McpOAuthStore { - /// Create a new OAuth store with a default client configuration + /// Create a new OAuth store with a random approval secret. + /// + /// Production callers should use [`Self::with_approval_secret`] so an + /// operator-controlled secret protects the approval step. + #[cfg(test)] pub fn new() -> Self { + Self::with_approval_secret(generate_random_string(32)) + } + + /// Create a new OAuth store with an operator-controlled approval secret. + pub fn with_approval_secret(approval_secret: String) -> Self { let mut clients = HashMap::new(); clients.insert( "mcp-client".to_string(), @@ -81,9 +93,15 @@ impl McpOAuthStore { clients: Arc::new(RwLock::new(clients)), auth_sessions: Arc::new(RwLock::new(HashMap::new())), access_tokens: Arc::new(RwLock::new(HashMap::new())), + approval_secret: Arc::new(approval_secret), } } + /// Validate the resource-owner credential used by the approval form. + pub fn validate_approval_secret(&self, candidate: &str) -> bool { + bool::from(self.approval_secret.as_bytes().ct_eq(candidate.as_bytes())) + } + /// Validate client credentials and redirect URI /// Returns Some(client_config) if valid, None otherwise pub async fn validate_client( @@ -104,6 +122,22 @@ impl McpOAuthStore { None } + /// Validate confidential client credentials and redirect URI. + pub async fn validate_client_credentials( + &self, + client_id: &str, + client_secret: &str, + redirect_uri: &str, + ) -> Option { + let client = self.validate_client(client_id, redirect_uri).await?; + let expected_secret = client.client_secret.as_deref()?; + if bool::from(expected_secret.as_bytes().ct_eq(client_secret.as_bytes())) { + Some(client) + } else { + None + } + } + /// Create a new authorization session for the OAuth flow /// Returns the session ID for tracking the auth process pub async fn create_auth_session( @@ -146,6 +180,7 @@ impl McpOAuthStore { /// Create a new MCP access token linked to an authorization session /// Returns the generated McpAccessToken on success + #[cfg(test)] pub async fn create_mcp_token(&self, session_id: &str) -> Result { let sessions = self.auth_sessions.read().await; if let Some(session) = sessions.get(session_id) { @@ -176,6 +211,43 @@ impl McpOAuthStore { } } + /// Exchange an authorization code once and bind it to the approved client. + pub async fn exchange_authorization_code( + &self, + session_id: &str, + client_id: &str, + ) -> Result { + let mut sessions = self.auth_sessions.write().await; + let session = sessions + .get(session_id) + .ok_or_else(|| "Authorization code not found or already used".to_string())?; + if session.client_id != client_id { + return Err("Authorization code does not belong to client".to_string()); + } + let session = sessions + .remove(session_id) + .ok_or_else(|| "Authorization code not found or already used".to_string())?; + let auth_token = session + .auth_token + .ok_or_else(|| "No third-party token available for session".to_string())?; + + let access_token = format!("mcp-token-{}", Uuid::new_v4()); + let token = McpAccessToken { + access_token: access_token.clone(), + token_type: "bearer".to_string(), + expires_in: Some(3600), + refresh_token: Some(format!("mcp-refresh-{}", Uuid::new_v4())), + scope: session.scope, + auth_token, + client_id: session.client_id, + }; + self.access_tokens + .write() + .await + .insert(access_token, token.clone()); + Ok(token) + } + /// Validate an access token and return the associated McpAccessToken if valid pub async fn validate_token(&self, token: &str) -> Option { self.access_tokens.read().await.get(token).cloned() @@ -268,6 +340,8 @@ pub struct ApprovalForm { pub scope: String, pub state: String, pub approved: String, + #[serde(default)] + pub approval_secret: String, } /// Generate a cryptographically secure random string @@ -321,6 +395,32 @@ pub async fn oauth_approve( State(state): State>, Form(form): Form, ) -> impl IntoResponse { + if state + .validate_client(&form.client_id, &form.redirect_uri) + .await + .is_none() + { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": "invalid_request", + "error_description": "invalid client id or redirect uri" + })), + ) + .into_response(); + } + if !state.validate_approval_secret(&form.approval_secret) { + warn!("Rejected OAuth approval with an invalid resource-owner credential"); + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({ + "error": "access_denied", + "error_description": "approval authentication failed" + })), + ) + .into_response(); + } + if form.approved != "true" { // user rejected the authorization request let redirect_url = format!( @@ -381,7 +481,7 @@ pub async fn oauth_approve( } ); - info!("authorization approved, redirecting to: {}", redirect_url); + info!("Authorization approved for client {}", form.client_id); Redirect::to(&redirect_url).into_response() } @@ -408,12 +508,12 @@ pub async fn oauth_token( } }; - let body_str = String::from_utf8_lossy(&bytes); - info!("request body: {}", body_str); - let token_req = match serde_urlencoded::from_bytes::(&bytes) { Ok(form) => { - info!("successfully parsed form data: {:?}", form); + debug!( + "Parsed token request for grant type {} and client {}", + form.grant_type, form.client_id + ); form } Err(e) => { @@ -453,7 +553,7 @@ pub async fn oauth_token( // get session_id from code if !token_req.code.starts_with("mcp-code-") { - info!("invalid authorization code: {}", token_req.code); + warn!("invalid authorization code format"); return ( StatusCode::BAD_REQUEST, Json(serde_json::json!({ @@ -464,24 +564,27 @@ pub async fn oauth_token( .into_response(); } - // handle empty client_id - let client_id = if token_req.client_id.is_empty() { - "mcp-client".to_string() - } else { - token_req.client_id.clone() - }; - - // validate client + // Validate the confidential client before consuming the authorization code. match state - .validate_client(&client_id, &token_req.redirect_uri) + .validate_client_credentials( + &token_req.client_id, + &token_req.client_secret, + &token_req.redirect_uri, + ) .await { Some(_) => { - let session_id = token_req.code.replace("mcp-code-", ""); - info!("got session id: {}", session_id); + let session_id = token_req.code.strip_prefix("mcp-code-").unwrap_or_default(); + debug!( + "Exchanging authorization code for client {}", + token_req.client_id + ); - // create mcp access token - match state.create_mcp_token(&session_id).await { + // Consume the authorization code and create an MCP access token. + match state + .exchange_authorization_code(session_id, &token_req.client_id) + .await + { Ok(token) => { info!("successfully created access token"); ( @@ -497,12 +600,12 @@ pub async fn oauth_token( .into_response() } Err(e) => { - error!("failed to create access token: {}", e); + warn!("failed to exchange authorization code: {}", e); ( - StatusCode::INTERNAL_SERVER_ERROR, + StatusCode::BAD_REQUEST, Json(serde_json::json!({ - "error": "server_error", - "error_description": format!("failed to create access token: {}", e) + "error": "invalid_grant", + "error_description": "authorization code is invalid or already used" })), ) .into_response() @@ -510,15 +613,12 @@ pub async fn oauth_token( } } None => { - info!( - "invalid client id or redirect uri: {} / {}", - client_id, token_req.redirect_uri - ); + warn!("invalid confidential client credentials"); ( - StatusCode::BAD_REQUEST, + StatusCode::UNAUTHORIZED, Json(serde_json::json!({ "error": "invalid_client", - "error_description": "invalid client id or redirect uri" + "error_description": "client authentication failed" })), ) .into_response() @@ -710,6 +810,153 @@ mod tests { assert!(result.is_none()); } + #[tokio::test] + async fn test_oauth_approve_rejects_direct_unauthenticated_post() { + let store = Arc::new(create_test_oauth_store()); + let form = ApprovalForm { + client_id: "mcp-client".to_string(), + redirect_uri: "http://localhost:8080/callback".to_string(), + scope: "profile".to_string(), + state: "state123".to_string(), + approved: "true".to_string(), + approval_secret: String::new(), + }; + + let response = oauth_approve(State(store.clone()), Form(form)) + .await + .into_response(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + assert!(store.auth_sessions.read().await.is_empty()); + } + + #[tokio::test] + async fn test_oauth_token_rejects_invalid_client_secret() { + let store = Arc::new(create_test_oauth_store()); + let session_id = store + .create_auth_session( + "mcp-client".to_string(), + Some("profile".to_string()), + None, + "secret-check-session".to_string(), + ) + .await; + let auth_token = AuthToken::new( + AccessToken::new("third-party-token".to_string()), + oauth2::basic::BasicTokenType::Bearer, + EmptyExtraTokenFields {}, + ); + store + .update_auth_session_token(&session_id, auth_token) + .await + .unwrap(); + let request_body = serde_urlencoded::to_string(TokenRequest { + grant_type: "authorization_code".to_string(), + code: format!("mcp-code-{session_id}"), + client_id: "mcp-client".to_string(), + client_secret: "wrong-secret".to_string(), + redirect_uri: "http://localhost:8080/callback".to_string(), + code_verifier: None, + refresh_token: String::new(), + }) + .unwrap(); + let request = Request::builder() + .header("content-type", "application/x-www-form-urlencoded") + .body(Body::from(request_body)) + .unwrap(); + + let response = oauth_token(State(store), request).await.into_response(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + } + + #[tokio::test] + async fn test_oauth_approve_requires_valid_client_and_resource_owner_secret() { + let store = Arc::new(McpOAuthStore::with_approval_secret( + "resource-owner-secret-with-32-characters".to_string(), + )); + let valid_form = ApprovalForm { + client_id: "mcp-client".to_string(), + redirect_uri: "http://localhost:8080/callback".to_string(), + scope: "profile".to_string(), + state: "state123".to_string(), + approved: "true".to_string(), + approval_secret: "resource-owner-secret-with-32-characters".to_string(), + }; + + let response = oauth_approve(State(store.clone()), Form(valid_form)) + .await + .into_response(); + + assert_eq!(response.status(), StatusCode::SEE_OTHER); + assert_eq!(store.auth_sessions.read().await.len(), 1); + + let invalid_redirect_form = ApprovalForm { + client_id: "mcp-client".to_string(), + redirect_uri: "http://attacker.invalid/callback".to_string(), + scope: "profile".to_string(), + state: "state123".to_string(), + approved: "true".to_string(), + approval_secret: "resource-owner-secret-with-32-characters".to_string(), + }; + let response = oauth_approve(State(store.clone()), Form(invalid_redirect_form)) + .await + .into_response(); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!(store.auth_sessions.read().await.len(), 1); + } + + #[tokio::test] + async fn test_oauth_token_consumes_authorization_code_once() { + let store = Arc::new(create_test_oauth_store()); + let session_id = store + .create_auth_session( + "mcp-client".to_string(), + Some("profile".to_string()), + None, + "single-use-session".to_string(), + ) + .await; + let auth_token = AuthToken::new( + AccessToken::new("third-party-token".to_string()), + oauth2::basic::BasicTokenType::Bearer, + EmptyExtraTokenFields {}, + ); + store + .update_auth_session_token(&session_id, auth_token) + .await + .unwrap(); + let request_body = serde_urlencoded::to_string(TokenRequest { + grant_type: "authorization_code".to_string(), + code: format!("mcp-code-{session_id}"), + client_id: "mcp-client".to_string(), + client_secret: "mcp-client-secret".to_string(), + redirect_uri: "http://localhost:8080/callback".to_string(), + code_verifier: None, + refresh_token: String::new(), + }) + .unwrap(); + + let first_request = Request::builder() + .header("content-type", "application/x-www-form-urlencoded") + .body(Body::from(request_body.clone())) + .unwrap(); + let first_response = oauth_token(State(store.clone()), first_request) + .await + .into_response(); + assert_eq!(first_response.status(), StatusCode::OK); + + let replay_request = Request::builder() + .header("content-type", "application/x-www-form-urlencoded") + .body(Body::from(request_body)) + .unwrap(); + let replay_response = oauth_token(State(store), replay_request) + .await + .into_response(); + assert_eq!(replay_response.status(), StatusCode::BAD_REQUEST); + } + #[tokio::test] async fn test_create_auth_session() { let store = create_test_oauth_store(); diff --git a/mcp-servers/mcp-bash-server/src/main.rs b/mcp-servers/mcp-bash-server/src/main.rs index 0611b72a55d..c715003d04f 100644 --- a/mcp-servers/mcp-bash-server/src/main.rs +++ b/mcp-servers/mcp-bash-server/src/main.rs @@ -30,7 +30,7 @@ use std::sync::OnceLock; use std::{net::SocketAddr, sync::Arc}; -use anyhow::Result; +use anyhow::{Context, Result, bail}; use axum::{ Router, body::Body, @@ -51,8 +51,8 @@ mod common; use common::bash_server::BashServer; use common::config; use common::oauth::{ - McpOAuthStore, oauth_approve, oauth_authorization_server, oauth_authorize, oauth_register, - oauth_token, validate_token_middleware, + McpOAuthStore, generate_random_string, oauth_approve, oauth_authorization_server, + oauth_authorize, oauth_register, oauth_token, validate_token_middleware, }; const INDEX_HTML: &str = include_str!("html/mcp_oauth_index.html"); @@ -85,8 +85,12 @@ async fn log_request(request: Request, next: Next) -> Response { // Log headers let headers = request.headers().clone(); let mut header_log = String::new(); - for (key, value) in headers.iter() { - let value_str = value.to_str().unwrap_or(""); + for (key, value) in &headers { + let value_str = if key == "authorization" || key == "cookie" { + "" + } else { + value.to_str().unwrap_or("") + }; header_log.push_str(&format!("\n {key}: {value_str}")); } @@ -116,6 +120,18 @@ async fn log_request(request: Request, next: Next) -> Response { response } +fn approval_secret_for_mode(is_dev: bool, configured: Option) -> Result { + if is_dev { + return Ok(configured.unwrap_or_else(|| generate_random_string(32))); + } + let approval_secret = configured + .context("MCP_OAUTH_APPROVAL_SECRET must be set when the server runs in production mode")?; + if approval_secret.len() < 32 { + bail!("MCP_OAUTH_APPROVAL_SECRET must contain at least 32 characters"); + } + Ok(approval_secret) +} + /// Main application entry point /// Sets up logging, OAuth store, HTTP server, and starts the MCP bash server #[tokio::main] @@ -142,8 +158,11 @@ async fn main() -> Result<()> { .unwrap_or_else(|| "production".to_string()); let is_dev = env_mode == "development"; - // Create the OAuth store - let oauth_store = Arc::new(McpOAuthStore::new()); + let approval_secret = + approval_secret_for_mode(is_dev, std::env::var("MCP_OAUTH_APPROVAL_SECRET").ok())?; + + // Create the OAuth store with an operator-controlled resource-owner secret. + let oauth_store = Arc::new(McpOAuthStore::with_approval_secret(approval_secret)); let host = config.settings.host.clone(); let port = config.settings.port; @@ -267,6 +286,24 @@ mod tests { assert!(INDEX_HTML.contains("html") || INDEX_HTML.contains("HTML")); } + #[test] + fn test_production_requires_strong_oauth_approval_secret() { + assert!(approval_secret_for_mode(false, None).is_err()); + assert!(approval_secret_for_mode(false, Some("too-short".to_string())).is_err()); + + let configured = "resource-owner-secret-with-32-characters".to_string(); + assert_eq!( + approval_secret_for_mode(false, Some(configured.clone())).unwrap(), + configured + ); + } + + #[test] + fn test_development_generates_oauth_approval_secret() { + let generated = approval_secret_for_mode(true, None).unwrap(); + assert_eq!(generated.len(), 32); + } + #[tokio::test] async fn test_log_request_middleware_functionality() { // Test basic properties of log_request function diff --git a/mcp-servers/mcp-bash-server/templates/mcp_oauth_authorize.html b/mcp-servers/mcp-bash-server/templates/mcp_oauth_authorize.html index c9a15936fb3..b51483b45b9 100644 --- a/mcp-servers/mcp-bash-server/templates/mcp_oauth_authorize.html +++ b/mcp-servers/mcp-bash-server/templates/mcp_oauth_authorize.html @@ -72,6 +72,25 @@ justify-content: center; } + .approval-secret { + margin-bottom: 1.5rem; + } + + .approval-secret label { + display: block; + font-weight: 600; + margin-bottom: 0.5rem; + } + + .approval-secret input { + box-sizing: border-box; + width: 100%; + padding: 0.75rem; + border: 1px solid var(--border-color); + border-radius: 6px; + font-size: 1rem; + } + .btn { padding: 0.75rem 1.5rem; border-radius: 6px; @@ -114,7 +133,13 @@

MCP OAuth

- + +
+ + +
+
@@ -122,4 +147,4 @@

MCP OAuth

- \ No newline at end of file + From ee866c9592dff36a71d7dcf56c17e16f51ab84af Mon Sep 17 00:00:00 2001 From: Logic Date: Thu, 30 Jul 2026 21:01:43 +0800 Subject: [PATCH 2/5] complete OAuth protocol validation --- mcp-servers/mcp-bash-server/Cargo.lock | 4 +- mcp-servers/mcp-bash-server/Cargo.toml | 4 +- mcp-servers/mcp-bash-server/README.md | 27 +- .../mcp-bash-server/src/common/oauth.rs | 2165 +++++++++-------- .../src/html/mcp_oauth_index.html | 33 +- mcp-servers/mcp-bash-server/src/main.rs | 574 +---- .../templates/mcp_oauth_authorize.html | 6 +- 7 files changed, 1283 insertions(+), 1530 deletions(-) diff --git a/mcp-servers/mcp-bash-server/Cargo.lock b/mcp-servers/mcp-bash-server/Cargo.lock index d7e7e1758cf..bbf0d6d796c 100644 --- a/mcp-servers/mcp-bash-server/Cargo.lock +++ b/mcp-servers/mcp-bash-server/Cargo.lock @@ -1062,14 +1062,15 @@ dependencies = [ "askama", "axum 0.8.4", "axum-test", + "base64", "chrono", "hyper", - "oauth2", "rand 0.8.5", "regex", "rmcp", "serde", "serde_urlencoded", + "sha2", "shlex", "subtle", "tempfile", @@ -1083,6 +1084,7 @@ dependencies = [ "tracing", "tracing-appender", "tracing-subscriber", + "url", "uuid", ] diff --git a/mcp-servers/mcp-bash-server/Cargo.toml b/mcp-servers/mcp-bash-server/Cargo.toml index 0f87a032b3f..5863483a5e1 100644 --- a/mcp-servers/mcp-bash-server/Cargo.toml +++ b/mcp-servers/mcp-bash-server/Cargo.toml @@ -36,14 +36,16 @@ axum = { version = "0.8", features = ["macros"] } chrono = "0.4" tower-http = { version = "0.6", features = ["cors"] } askama = { version = "0.14"} +base64 = "0.22" rand = { version = "0.8", features = ["std"] } uuid = { version = "1.6", features = ["v4", "serde"] } serde_urlencoded = "0.7" -oauth2 = "5.0" +sha2 = "0.10" toml = "0.8" regex = "1.11.1" shlex = "1.3" subtle = "2.6" +url = "2.5" [dev-dependencies] tokio-test = "0.4" diff --git a/mcp-servers/mcp-bash-server/README.md b/mcp-servers/mcp-bash-server/README.md index d1538f2416d..fd339514c83 100644 --- a/mcp-servers/mcp-bash-server/README.md +++ b/mcp-servers/mcp-bash-server/README.md @@ -20,6 +20,7 @@ If you want to run this MCP server locally using the default settings provided b ```shell export MCP_OAUTH_APPROVAL_SECRET="" +export MCP_OAUTH_PUBLIC_BASE_URL="https://mcp.example.com" cargo run ``` @@ -27,7 +28,17 @@ This MCP server will be deployed at `http://127.0.0.1:4000/mcp`, and you can use Production mode requires `MCP_OAUTH_APPROVAL_SECRET`. The operator enters this secret on the OAuth approval page; it is separate from dynamically registered -client credentials. Development mode does not require it. +client credentials. Production also requires an explicit HTTPS +`MCP_OAUTH_PUBLIC_BASE_URL`; OAuth metadata never derives public endpoints from +the request `Host` header. Development mode generates a temporary approval +secret and may derive a loopback HTTP URL from the local bind address. + +Dynamic registration supports public clients (`token_endpoint_auth_method: +none`) and confidential clients (`client_secret_post`). Authorization requests +must use PKCE S256. Authorization transactions and codes are one-time and +short-lived; access tokens expire after one hour, and refresh tokens expire +after one day and rotate on every use. OAuth form and JSON bodies are limited +to 16 KiB. For information on how to use the modelcontextprotocol/inspector tool, refer to the [inspector documentation](https://github.com/modelcontextprotocol/inspector). @@ -54,6 +65,7 @@ After building, use the following command to run it: ```shell docker run -d --name mcp-bash-server -p 127.0.0.1:4000:4000 \ -e MCP_OAUTH_APPROVAL_SECRET="" \ + -e MCP_OAUTH_PUBLIC_BASE_URL="https://mcp.example.com" \ --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest ``` @@ -70,6 +82,7 @@ Container's workdir is `/app` and it will run the `/app/mcp-bash-server` when it ```shell docker run -d --name mcp-bash-server -p 127.0.0.1:4000:4000 \ -e MCP_OAUTH_APPROVAL_SECRET="" \ + -e MCP_OAUTH_PUBLIC_BASE_URL="https://mcp.example.com" \ -v `pwd`/config.toml:/app/config.toml \ --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest ``` @@ -79,6 +92,7 @@ If you are using SELinux, you may need to run the command instead to let the con ```shell docker run -d --name mcp-bash-server -p 127.0.0.1:4000:4000 \ -e MCP_OAUTH_APPROVAL_SECRET="" \ + -e MCP_OAUTH_PUBLIC_BASE_URL="https://mcp.example.com" \ -v `pwd`/config.toml:/app/config.toml:Z \ --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest ``` @@ -115,14 +129,9 @@ Start the MCP Server in daemon mode, then add the settings to your Vscode Copilo } ``` -**Currently Vscode MCP OAuth can not automatically authorize this bash-server** -The vscode mcp OAuth flow is: - -1. GET /.well-known/oauth-authorization-server -2. GET /authorize with query-params -3. ... - -But we requires the client registration before accessing endpoint `/authorize` with query-params that contains invalid client-id. So we can only set the token manually now. +OAuth-capable MCP clients can discover the authorization metadata, dynamically +register a public client, and complete the PKCE flow. A manually configured +bearer token remains available for clients that do not implement MCP OAuth. ## Configuration diff --git a/mcp-servers/mcp-bash-server/src/common/oauth.rs b/mcp-servers/mcp-bash-server/src/common/oauth.rs index dedd9e4ec77..1acf7535284 100644 --- a/mcp-servers/mcp-bash-server/src/common/oauth.rs +++ b/mcp-servers/mcp-bash-server/src/common/oauth.rs @@ -1,298 +1,517 @@ /* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to You under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. */ -//! OAuth2 authentication implementation for MCP server +//! OAuth 2.0 authorization-code support for the MCP server. //! -//! This module provides OAuth2 authentication capabilities including: -//! - Client registration and validation -//! - Authorization code flow -//! - Token management and validation -//! - Session management for auth flows -//! - Middleware for request authentication - -use std::{collections::HashMap, sync::Arc}; +//! The implementation deliberately keeps the authorization server small, but +//! it still enforces the protocol properties on which bearer-token safety +//! depends: registered redirect URIs, PKCE S256, one-time authorization and +//! consent transactions, expiring codes and tokens, refresh-token rotation, +//! bounded request bodies, and a configured public issuer URL. + +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; use askama::Template; use axum::{ Json, body::Body, - extract::{Form, Query, State}, - http::{HeaderMap, Request, StatusCode}, + extract::{Query, State}, + http::{Request, StatusCode}, middleware::Next, response::{Html, IntoResponse, Redirect, Response}, }; -use chrono; -use oauth2::{AccessToken, EmptyExtraTokenFields, RefreshToken, StandardTokenResponse}; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use chrono::{DateTime, Duration, Utc}; use rand::{Rng, distributions::Alphanumeric}; use rmcp::serde_json::{self, Value}; use rmcp::transport::auth::{ - AuthorizationMetadata, ClientRegistrationRequest, ClientRegistrationResponse, OAuthClientConfig, + AuthorizationMetadata, ClientRegistrationRequest, ClientRegistrationResponse, }; use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; use tokio::sync::RwLock; -use tracing::{debug, error, info, warn}; +use tracing::{debug, info, warn}; +use url::Url; use uuid::Uuid; -/// Type alias for OAuth2 standard token response -/// Type alias for OAuth2 standard token response -pub type AuthToken = StandardTokenResponse; +pub const MAX_OAUTH_BODY_BYTES: usize = 16 * 1024; +const AUTH_TRANSACTION_TTL_SECONDS: i64 = 5 * 60; +const AUTHORIZATION_CODE_TTL_SECONDS: i64 = 2 * 60; +const ACCESS_TOKEN_TTL_SECONDS: i64 = 60 * 60; +const REFRESH_TOKEN_TTL_SECONDS: i64 = 24 * 60 * 60; +const MAX_CLIENTS: usize = 1_024; +const MAX_AUTH_TRANSACTIONS: usize = 4_096; +const MAX_AUTHORIZATION_CODES: usize = 4_096; +const MAX_ACCESS_TOKENS: usize = 4_096; +const MAX_REFRESH_TOKENS: usize = 4_096; +const SUPPORTED_SCOPES: [&str; 2] = ["profile", "email"]; + +#[derive(Clone, Debug, PartialEq, Eq)] +enum TokenEndpointAuthMethod { + None, + ClientSecretPost, +} + +impl TokenEndpointAuthMethod { + fn parse(value: &str) -> Option { + match value { + "none" => Some(Self::None), + "client_secret_post" => Some(Self::ClientSecretPost), + _ => None, + } + } + + fn as_str(&self) -> &'static str { + match self { + Self::None => "none", + Self::ClientSecretPost => "client_secret_post", + } + } +} + +#[derive(Clone, Debug)] +struct RegisteredClient { + client_id: String, + client_secret: Option, + redirect_uris: Vec, + token_endpoint_auth_method: TokenEndpointAuthMethod, +} + +#[derive(Clone, Debug)] +struct AuthorizationTransaction { + client_id: String, + redirect_uri: String, + scope: Option, + state: Option, + code_challenge: String, + consent_nonce: String, + expires_at: DateTime, +} + +#[derive(Clone, Debug)] +struct AuthorizationCode { + client_id: String, + redirect_uri: String, + scope: Option, + code_challenge: String, + expires_at: DateTime, +} + +#[derive(Clone, Debug)] +struct RefreshTokenRecord { + client_id: String, + scope: Option, + expires_at: DateTime, +} + +/// Access-token state retained by the authorization server. +#[derive(Clone, Debug, Serialize)] +pub struct McpAccessToken { + pub access_token: String, + pub token_type: String, + pub expires_in: u64, + pub refresh_token: Option, + pub scope: Option, + pub client_id: String, + #[serde(skip)] + pub issued_at: DateTime, + #[serde(skip)] + pub expires_at: DateTime, +} + +#[derive(Clone, Debug)] +struct IssuedTokens { + access: McpAccessToken, +} -/// Centralized OAuth store for managing clients, sessions, and tokens -/// Provides thread-safe access to OAuth-related data structures +/// Central OAuth state. Every remotely growable map has a hard cardinality +/// bound and is pruned by expiry on the operations that access it. #[derive(Clone, Debug)] pub struct McpOAuthStore { - /// Registered OAuth clients with their configurations - pub clients: Arc>>, - /// Active authorization sessions indexed by session ID - pub auth_sessions: Arc>>, - /// Valid access tokens indexed by token string - pub access_tokens: Arc>>, - /// Secret required for a resource owner to approve an authorization request + clients: Arc>>, + auth_transactions: Arc>>, + authorization_codes: Arc>>, + access_tokens: Arc>>, + refresh_tokens: Arc>>, approval_secret: Arc, + public_base_url: Arc, } impl McpOAuthStore { - /// Create a new OAuth store with a random approval secret. - /// - /// Production callers should use [`Self::with_approval_secret`] so an - /// operator-controlled secret protects the approval step. - #[cfg(test)] - pub fn new() -> Self { - Self::with_approval_secret(generate_random_string(32)) + /// Create the production store. Clients are registered dynamically; there + /// is intentionally no repository-known default confidential credential. + pub fn with_settings(approval_secret: String, public_base_url: Url) -> Self { + Self { + clients: Arc::new(RwLock::new(HashMap::new())), + auth_transactions: Arc::new(RwLock::new(HashMap::new())), + authorization_codes: Arc::new(RwLock::new(HashMap::new())), + access_tokens: Arc::new(RwLock::new(HashMap::new())), + refresh_tokens: Arc::new(RwLock::new(HashMap::new())), + approval_secret: Arc::new(approval_secret), + public_base_url: Arc::new(public_base_url), + } } - /// Create a new OAuth store with an operator-controlled approval secret. - pub fn with_approval_secret(approval_secret: String) -> Self { - let mut clients = HashMap::new(); + #[cfg(test)] + fn new() -> Self { + let mut store = Self::with_settings( + generate_random_string(32), + Url::parse("http://127.0.0.1:4000/").unwrap(), + ); + let clients = Arc::get_mut(&mut store.clients).unwrap().get_mut(); + clients.insert( + "public-test-client".to_string(), + RegisteredClient { + client_id: "public-test-client".to_string(), + client_secret: None, + redirect_uris: vec!["http://127.0.0.1:8080/callback".to_string()], + token_endpoint_auth_method: TokenEndpointAuthMethod::None, + }, + ); clients.insert( - "mcp-client".to_string(), - OAuthClientConfig { - client_id: "mcp-client".to_string(), - client_secret: Some("mcp-client-secret".to_string()), - scopes: vec!["profile".to_string(), "email".to_string()], - redirect_uri: "http://localhost:8080/callback".to_string(), + "confidential-test-client".to_string(), + RegisteredClient { + client_id: "confidential-test-client".to_string(), + client_secret: Some("test-only-confidential-secret".to_string()), + redirect_uris: vec!["https://client.example/callback".to_string()], + token_endpoint_auth_method: TokenEndpointAuthMethod::ClientSecretPost, }, ); + store + } - Self { - clients: Arc::new(RwLock::new(clients)), - auth_sessions: Arc::new(RwLock::new(HashMap::new())), - access_tokens: Arc::new(RwLock::new(HashMap::new())), - approval_secret: Arc::new(approval_secret), - } + pub fn public_base_url(&self) -> &Url { + self.public_base_url.as_ref() } - /// Validate the resource-owner credential used by the approval form. pub fn validate_approval_secret(&self, candidate: &str) -> bool { bool::from(self.approval_secret.as_bytes().ct_eq(candidate.as_bytes())) } - /// Validate client credentials and redirect URI - /// Returns Some(client_config) if valid, None otherwise - pub async fn validate_client( + async fn validate_client( &self, client_id: &str, redirect_uri: &str, - ) -> Option { + ) -> Option { let clients = self.clients.read().await; - if let Some(client) = clients.get(client_id) { - info!("client.redirect_uri: {}", client.redirect_uri); - info!("redirect_uri: {redirect_uri}"); - if client.redirect_uri == redirect_uri { - return Some(client.clone()); - } - } else { - error!("Invalid client_id: {client_id}"); - } - None + let client = clients.get(client_id)?; + client + .redirect_uris + .iter() + .any(|registered| registered == redirect_uri) + .then(|| client.clone()) } - /// Validate confidential client credentials and redirect URI. - pub async fn validate_client_credentials( + async fn validate_token_client( &self, client_id: &str, client_secret: &str, - redirect_uri: &str, - ) -> Option { - let client = self.validate_client(client_id, redirect_uri).await?; - let expected_secret = client.client_secret.as_deref()?; - if bool::from(expected_secret.as_bytes().ct_eq(client_secret.as_bytes())) { - Some(client) - } else { - None + ) -> Option { + let clients = self.clients.read().await; + let client = clients.get(client_id)?; + match client.token_endpoint_auth_method { + TokenEndpointAuthMethod::None => client_secret.is_empty().then(|| client.clone()), + TokenEndpointAuthMethod::ClientSecretPost => { + let expected = client.client_secret.as_deref()?; + bool::from(expected.as_bytes().ct_eq(client_secret.as_bytes())) + .then(|| client.clone()) + } } } - /// Create a new authorization session for the OAuth flow - /// Returns the session ID for tracking the auth process - pub async fn create_auth_session( + async fn create_authorization_transaction( &self, - client_id: String, - scope: Option, - state: Option, - session_id: String, - ) -> String { - let session = AuthSession { - client_id, + params: &AuthorizeQuery, + ) -> Result<(String, AuthorizationTransaction), OAuthError> { + if params.response_type != "code" { + return Err(OAuthError::InvalidRequest( + "response_type must be code".to_string(), + )); + } + if params.code_challenge_method.as_deref() != Some("S256") { + return Err(OAuthError::InvalidRequest( + "code_challenge_method must be S256".to_string(), + )); + } + let code_challenge = params + .code_challenge + .as_deref() + .ok_or_else(|| OAuthError::InvalidRequest("code_challenge is required".to_string()))?; + if !valid_pkce_challenge(code_challenge) { + return Err(OAuthError::InvalidRequest( + "code_challenge is not a valid S256 challenge".to_string(), + )); + } + let client = self + .validate_client(¶ms.client_id, ¶ms.redirect_uri) + .await + .ok_or_else(|| { + OAuthError::InvalidRequest("invalid client id or redirect uri".to_string()) + })?; + let scope = validate_scope(params.scope.as_deref())?; + debug!( + "Starting authorization transaction for client {} using {}", + client.client_id, + client.token_endpoint_auth_method.as_str() + ); + + let now = Utc::now(); + let mut transactions = self.auth_transactions.write().await; + transactions.retain(|_, transaction| transaction.expires_at > now); + if transactions.len() >= MAX_AUTH_TRANSACTIONS { + return Err(OAuthError::TemporarilyUnavailable); + } + + let transaction_id = random_prefixed("mcp-auth"); + let transaction = AuthorizationTransaction { + client_id: params.client_id.clone(), + redirect_uri: params.redirect_uri.clone(), scope, - _state: state, - _created_at: chrono::Utc::now(), - auth_token: None, + state: params.state.clone(), + code_challenge: code_challenge.to_string(), + consent_nonce: generate_random_string(48), + expires_at: now + Duration::seconds(AUTH_TRANSACTION_TTL_SECONDS), }; + transactions.insert(transaction_id.clone(), transaction.clone()); + Ok((transaction_id, transaction)) + } - self.auth_sessions - .write() - .await - .insert(session_id.clone(), session); - session_id + async fn consume_authorization_transaction( + &self, + transaction_id: &str, + consent_nonce: &str, + ) -> Result { + let now = Utc::now(); + let mut transactions = self.auth_transactions.write().await; + transactions.retain(|_, transaction| transaction.expires_at > now); + let transaction = transactions + .get(transaction_id) + .ok_or(OAuthError::InvalidTransaction)?; + if !bool::from( + transaction + .consent_nonce + .as_bytes() + .ct_eq(consent_nonce.as_bytes()), + ) { + return Err(OAuthError::InvalidTransaction); + } + transactions + .remove(transaction_id) + .ok_or(OAuthError::InvalidTransaction) } - /// Update an authorization session with a generated token - /// Links the OAuth token to the session for later retrieval - pub async fn update_auth_session_token( + async fn create_authorization_code( &self, - session_id: &str, - token: AuthToken, - ) -> Result<(), String> { - let mut sessions = self.auth_sessions.write().await; - if let Some(session) = sessions.get_mut(session_id) { - session.auth_token = Some(token); - Ok(()) - } else { - Err("Session not found".to_string()) + transaction: AuthorizationTransaction, + ) -> Result { + let now = Utc::now(); + let mut codes = self.authorization_codes.write().await; + codes.retain(|_, code| code.expires_at > now); + if codes.len() >= MAX_AUTHORIZATION_CODES { + return Err(OAuthError::TemporarilyUnavailable); } + let code_value = random_prefixed("mcp-code"); + codes.insert( + code_value.clone(), + AuthorizationCode { + client_id: transaction.client_id, + redirect_uri: transaction.redirect_uri, + scope: transaction.scope, + code_challenge: transaction.code_challenge, + expires_at: now + Duration::seconds(AUTHORIZATION_CODE_TTL_SECONDS), + }, + ); + Ok(code_value) } - /// Create a new MCP access token linked to an authorization session - /// Returns the generated McpAccessToken on success - #[cfg(test)] - pub async fn create_mcp_token(&self, session_id: &str) -> Result { - let sessions = self.auth_sessions.read().await; - if let Some(session) = sessions.get(session_id) { - if let Some(auth_token) = &session.auth_token { - let access_token = format!("mcp-token-{}", Uuid::new_v4()); - let refresh_token = format!("mcp-refresh-{}", Uuid::new_v4()); - - let token = McpAccessToken { - access_token: access_token.clone(), - token_type: "Bearer".to_string().to_lowercase(), - expires_in: Some(3600), - refresh_token: Some(refresh_token), - scope: session.scope.clone(), - auth_token: auth_token.clone(), - client_id: session.client_id.clone(), - }; - - self.access_tokens - .write() - .await - .insert(access_token.clone(), token.clone()); - Ok(token) - } else { - Err("No third-party token available for session".to_string()) - } - } else { - Err("Session not found".to_string()) + async fn exchange_authorization_code( + &self, + request: &TokenRequest, + ) -> Result { + self.validate_token_client(&request.client_id, &request.client_secret) + .await + .ok_or(OAuthError::InvalidClient)?; + + let now = Utc::now(); + let mut codes = self.authorization_codes.write().await; + codes.retain(|_, code| code.expires_at > now); + let code = codes + .get(&request.code) + .cloned() + .ok_or(OAuthError::InvalidGrant)?; + if code.client_id != request.client_id || code.redirect_uri != request.redirect_uri { + return Err(OAuthError::InvalidGrant); } + let verifier = request + .code_verifier + .as_deref() + .ok_or(OAuthError::InvalidGrant)?; + if !pkce_matches(verifier, &code.code_challenge) { + // A verifier failure consumes the code so it cannot become a + // brute-force oracle. + codes.remove(&request.code); + return Err(OAuthError::InvalidGrant); + } + let code = codes + .remove(&request.code) + .ok_or(OAuthError::InvalidGrant)?; + drop(codes); + + self.issue_tokens(code.client_id, code.scope).await } - /// Exchange an authorization code once and bind it to the approved client. - pub async fn exchange_authorization_code( + async fn exchange_refresh_token( &self, - session_id: &str, - client_id: &str, - ) -> Result { - let mut sessions = self.auth_sessions.write().await; - let session = sessions - .get(session_id) - .ok_or_else(|| "Authorization code not found or already used".to_string())?; - if session.client_id != client_id { - return Err("Authorization code does not belong to client".to_string()); + request: &TokenRequest, + ) -> Result { + self.validate_token_client(&request.client_id, &request.client_secret) + .await + .ok_or(OAuthError::InvalidClient)?; + + let now = Utc::now(); + let mut refresh_tokens = self.refresh_tokens.write().await; + refresh_tokens.retain(|_, token| token.expires_at > now); + let record = refresh_tokens + .get(&request.refresh_token) + .cloned() + .ok_or(OAuthError::InvalidGrant)?; + if record.client_id != request.client_id { + return Err(OAuthError::InvalidGrant); } - let session = sessions - .remove(session_id) - .ok_or_else(|| "Authorization code not found or already used".to_string())?; - let auth_token = session - .auth_token - .ok_or_else(|| "No third-party token available for session".to_string())?; - - let access_token = format!("mcp-token-{}", Uuid::new_v4()); - let token = McpAccessToken { - access_token: access_token.clone(), + let record = refresh_tokens + .remove(&request.refresh_token) + .ok_or(OAuthError::InvalidGrant)?; + drop(refresh_tokens); + + self.issue_tokens(record.client_id, record.scope).await + } + + async fn issue_tokens( + &self, + client_id: String, + scope: Option, + ) -> Result { + let now = Utc::now(); + let access_token_value = random_prefixed("mcp-token"); + let refresh_token_value = random_prefixed("mcp-refresh"); + let access = McpAccessToken { + access_token: access_token_value.clone(), token_type: "bearer".to_string(), - expires_in: Some(3600), - refresh_token: Some(format!("mcp-refresh-{}", Uuid::new_v4())), - scope: session.scope, - auth_token, - client_id: session.client_id, + expires_in: ACCESS_TOKEN_TTL_SECONDS as u64, + refresh_token: Some(refresh_token_value.clone()), + scope: scope.clone(), + client_id: client_id.clone(), + issued_at: now, + expires_at: now + Duration::seconds(ACCESS_TOKEN_TTL_SECONDS), }; - self.access_tokens - .write() - .await - .insert(access_token, token.clone()); - Ok(token) + + let mut access_tokens = self.access_tokens.write().await; + access_tokens.retain(|_, token| token.expires_at > now); + if access_tokens.len() >= MAX_ACCESS_TOKENS { + return Err(OAuthError::TemporarilyUnavailable); + } + let mut refresh_tokens = self.refresh_tokens.write().await; + refresh_tokens.retain(|_, token| token.expires_at > now); + if refresh_tokens.len() >= MAX_REFRESH_TOKENS { + return Err(OAuthError::TemporarilyUnavailable); + } + access_tokens.insert(access_token_value, access.clone()); + refresh_tokens.insert( + refresh_token_value, + RefreshTokenRecord { + client_id, + scope, + expires_at: now + Duration::seconds(REFRESH_TOKEN_TTL_SECONDS), + }, + ); + Ok(IssuedTokens { access }) } - /// Validate an access token and return the associated McpAccessToken if valid + /// Validate a bearer token and atomically discard it after expiry. pub async fn validate_token(&self, token: &str) -> Option { - self.access_tokens.read().await.get(token).cloned() + let now = Utc::now(); + let mut tokens = self.access_tokens.write().await; + let current = tokens.get(token)?.clone(); + if current.expires_at <= now || current.issued_at > now { + tokens.remove(token); + None + } else { + Some(current) + } } } -/// Authorization session data structure -/// Tracks ongoing OAuth authorization flows with client and state information -#[derive(Clone, Debug)] -pub struct AuthSession { - pub client_id: String, - pub scope: Option, - pub _state: Option, - pub _created_at: chrono::DateTime, - pub auth_token: Option, +#[derive(Debug)] +enum OAuthError { + InvalidRequest(String), + InvalidClient, + InvalidGrant, + InvalidTransaction, + TemporarilyUnavailable, } -/// MCP-specific access token structure -/// Wraps OAuth2 standard tokens with additional MCP metadata -#[derive(Clone, Debug, Serialize)] -pub struct McpAccessToken { - pub access_token: String, - pub token_type: String, - pub expires_in: Option, - pub refresh_token: Option, - pub scope: Option, - pub auth_token: AuthToken, - pub client_id: String, +impl OAuthError { + fn response(&self) -> Response { + match self { + Self::InvalidRequest(description) => { + oauth_json_error(StatusCode::BAD_REQUEST, "invalid_request", description) + } + Self::InvalidClient => oauth_json_error( + StatusCode::UNAUTHORIZED, + "invalid_client", + "client authentication failed", + ), + Self::InvalidGrant => oauth_json_error( + StatusCode::BAD_REQUEST, + "invalid_grant", + "authorization grant is invalid, expired, or already used", + ), + Self::InvalidTransaction => oauth_json_error( + StatusCode::BAD_REQUEST, + "invalid_request", + "authorization transaction is invalid, expired, or already used", + ), + Self::TemporarilyUnavailable => oauth_json_error( + StatusCode::SERVICE_UNAVAILABLE, + "temporarily_unavailable", + "authorization server capacity is temporarily exhausted", + ), + } + } } -/// OAuth authorization request parameters -/// Contains all required fields for initiating an OAuth authorization flow #[derive(Debug, Deserialize)] pub struct AuthorizeQuery { - #[allow(dead_code)] pub response_type: String, pub client_id: String, pub redirect_uri: String, pub scope: Option, pub state: Option, + pub code_challenge: Option, + pub code_challenge_method: Option, } -/// OAuth token request parameters -/// Used for exchanging authorization codes for access tokens -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub struct TokenRequest { pub grant_type: String, #[serde(default)] @@ -309,43 +528,24 @@ pub struct TokenRequest { pub refresh_token: String, } -/// User information structure for OAuth responses -/// Contains standard user profile data -#[derive(Debug, Deserialize, Serialize)] -pub struct UserInfo { - pub sub: String, - pub name: String, - pub email: String, - pub username: String, -} - -/// Template context for OAuth authorization page -/// Contains all data needed to render the authorization consent form #[derive(Template)] #[template(path = "mcp_oauth_authorize.html")] pub struct OAuthAuthorizeTemplate { pub client_id: String, - pub redirect_uri: String, - pub scope: String, - pub state: String, pub scopes: String, + pub transaction_id: String, + pub consent_nonce: String, } -/// Form data for user authorization approval -/// Contains user's decision and associated OAuth parameters -#[derive(Debug, Deserialize)] +#[derive(Debug, Deserialize, Serialize)] pub struct ApprovalForm { - pub client_id: String, - pub redirect_uri: String, - pub scope: String, - pub state: String, + pub transaction_id: String, + pub consent_nonce: String, pub approved: String, #[serde(default)] pub approval_secret: String, } -/// Generate a cryptographically secure random string -/// Used for creating client secrets and other security tokens pub fn generate_random_string(length: usize) -> String { rand::thread_rng() .sample_iter(&Alphanumeric) @@ -354,334 +554,302 @@ pub fn generate_random_string(length: usize) -> String { .collect() } -/// OAuth authorization endpoint handler -/// Displays the authorization consent page to users -pub async fn oauth_authorize( - Query(params): Query, - State(state): State>, -) -> impl IntoResponse { - debug!("doing oauth_authorize"); - if let Some(_client) = state - .validate_client(¶ms.client_id, ¶ms.redirect_uri) - .await +fn random_prefixed(prefix: &str) -> String { + format!("{prefix}-{}-{}", Uuid::new_v4(), generate_random_string(24)) +} + +fn validate_scope(scope: Option<&str>) -> Result, OAuthError> { + let Some(scope) = scope.map(str::trim).filter(|scope| !scope.is_empty()) else { + return Ok(None); + }; + let requested: HashSet<&str> = scope.split_whitespace().collect(); + if requested + .iter() + .any(|requested| !SUPPORTED_SCOPES.contains(requested)) { - let template = OAuthAuthorizeTemplate { - client_id: params.client_id, - redirect_uri: params.redirect_uri, - scope: params.scope.clone().unwrap_or_default(), - state: params.state.clone().unwrap_or_default(), - scopes: params - .scope - .clone() - .unwrap_or_else(|| "Basic scope".to_string()), - }; + return Err(OAuthError::InvalidRequest( + "requested scope is not supported".to_string(), + )); + } + let normalized = SUPPORTED_SCOPES + .iter() + .filter(|supported| requested.contains(**supported)) + .copied() + .collect::>() + .join(" "); + Ok(Some(normalized)) +} - Html(template.render().unwrap()).into_response() - } else { - ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": "invalid_request", - "error_description": "invalid client id or redirect uri" - })), - ) - .into_response() +fn valid_pkce_challenge(challenge: &str) -> bool { + challenge.len() == 43 + && challenge + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')) +} + +fn valid_pkce_verifier(verifier: &str) -> bool { + (43..=128).contains(&verifier.len()) + && verifier + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~')) +} + +fn pkce_matches(verifier: &str, expected_challenge: &str) -> bool { + if !valid_pkce_verifier(verifier) { + return false; } + let actual = URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes())); + bool::from(actual.as_bytes().ct_eq(expected_challenge.as_bytes())) } -/// Handle user approval/rejection of OAuth authorization -/// Processes the consent form and generates authorization codes -pub async fn oauth_approve( - State(state): State>, - Form(form): Form, -) -> impl IntoResponse { - if state - .validate_client(&form.client_id, &form.redirect_uri) - .await - .is_none() +fn redirect_with_params(redirect_uri: &str, params: &[(&str, &str)]) -> Result { + let mut url = Url::parse(redirect_uri) + .map_err(|_| OAuthError::InvalidRequest("invalid redirect uri".to_string()))?; { - return ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": "invalid_request", - "error_description": "invalid client id or redirect uri" - })), - ) - .into_response(); + let mut query = url.query_pairs_mut(); + for (key, value) in params { + if !value.is_empty() { + query.append_pair(key, value); + } + } } - if !state.validate_approval_secret(&form.approval_secret) { - warn!("Rejected OAuth approval with an invalid resource-owner credential"); - return ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({ - "error": "access_denied", - "error_description": "approval authentication failed" - })), - ) - .into_response(); + Ok(url.into()) +} + +pub fn validate_redirect_uri(value: &str) -> bool { + let Ok(url) = Url::parse(value) else { + return false; + }; + if !url.username().is_empty() + || url.password().is_some() + || url.fragment().is_some() + || url.host_str().is_none() + { + return false; } + if url.scheme() == "https" { + return true; + } + url.scheme() == "http" && matches!(url.host_str(), Some("localhost" | "127.0.0.1" | "::1")) +} - if form.approved != "true" { - // user rejected the authorization request - let redirect_url = format!( - "{}?error=access_denied&error_description={}{}", - form.redirect_uri, - "user rejected the authorization request", - if form.state.is_empty() { - "".to_string() - } else { - format!("&state={}", form.state) - } - ); - return Redirect::to(&redirect_url).into_response(); +pub fn validate_public_base_url(value: &str, require_https: bool) -> Result { + let mut url = Url::parse(value).map_err(|_| "public base URL is invalid".to_string())?; + if !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || url.host_str().is_none() + { + return Err("public base URL must not contain credentials, query, or fragment".to_string()); + } + if require_https && url.scheme() != "https" { + return Err("public base URL must use https in production".to_string()); + } + if !matches!(url.scheme(), "http" | "https") { + return Err("public base URL must use http or https".to_string()); } + if !matches!(url.path(), "" | "/") { + return Err("public base URL must not contain a path".to_string()); + } + url.set_path("/"); + Ok(url) +} - // user approved the authorization request, generate authorization code - let session_id = Uuid::new_v4().to_string(); - let auth_code = format!("mcp-code-{session_id}"); - - // create new session record authorization information - let session_id = state - .create_auth_session( - form.client_id.clone(), - Some(form.scope.clone()), - Some(form.state.clone()), - session_id.clone(), - ) - .await; +async fn read_limited_body(request: Request) -> Result, Response> { + match axum::body::to_bytes(request.into_body(), MAX_OAUTH_BODY_BYTES).await { + Ok(bytes) => Ok(bytes.to_vec()), + Err(_) => Err(oauth_json_error( + StatusCode::PAYLOAD_TOO_LARGE, + "invalid_request", + "request body exceeds the authorization endpoint limit", + )), + } +} - // create token using oauth2 standard - let access_token = AccessToken::new(format!("tp-token-{}", Uuid::new_v4())); - let refresh_token = RefreshToken::new(format!("tp-refresh-{}", Uuid::new_v4())); - let token_type = oauth2::basic::BasicTokenType::Bearer; +fn oauth_json_error(status: StatusCode, error: &str, description: &str) -> Response { + ( + status, + Json(serde_json::json!({ + "error": error, + "error_description": description + })), + ) + .into_response() +} - let mut created_token = - StandardTokenResponse::new(access_token, token_type, EmptyExtraTokenFields {}); - created_token.set_expires_in(Some(&std::time::Duration::from_secs(3600))); - created_token.set_refresh_token(Some(refresh_token)); - created_token.set_scopes(Some(vec![oauth2::Scope::new(form.scope.clone())])); +pub async fn oauth_authorize( + Query(params): Query, + State(state): State>, +) -> Response { + let (transaction_id, transaction) = match state.create_authorization_transaction(¶ms).await + { + Ok(created) => created, + Err(error) => return error.response(), + }; + let template = OAuthAuthorizeTemplate { + client_id: transaction.client_id, + scopes: transaction + .scope + .unwrap_or_else(|| "No additional scopes".to_string()), + transaction_id, + consent_nonce: transaction.consent_nonce, + }; + match template.render() { + Ok(html) => Html(html).into_response(), + Err(_) => oauth_json_error( + StatusCode::INTERNAL_SERVER_ERROR, + "server_error", + "authorization page could not be rendered", + ), + } +} - // update session token - if let Err(e) = state - .update_auth_session_token(&session_id, created_token) +pub async fn oauth_approve( + State(state): State>, + request: Request, +) -> Response { + let bytes = match read_limited_body(request).await { + Ok(bytes) => bytes, + Err(response) => return response, + }; + let form = match serde_urlencoded::from_bytes::(&bytes) { + Ok(form) => form, + Err(_) => { + return oauth_json_error( + StatusCode::BAD_REQUEST, + "invalid_request", + "approval form is invalid", + ); + } + }; + if !state.validate_approval_secret(&form.approval_secret) { + warn!("Rejected OAuth approval with an invalid resource-owner credential"); + return oauth_json_error( + StatusCode::UNAUTHORIZED, + "access_denied", + "approval authentication failed", + ); + } + let transaction = match state + .consume_authorization_transaction(&form.transaction_id, &form.consent_nonce) .await { - error!("Failed to update session token: {}", e); - } + Ok(transaction) => transaction, + Err(error) => return error.response(), + }; - // redirect back to client, with authorization code - let redirect_url = format!( - "{}?code={}{}", - form.redirect_uri, - auth_code, - if form.state.is_empty() { - "".to_string() - } else { - format!("&state={}", form.state) + if form.approved != "true" { + let mut params = vec![ + ("error", "access_denied"), + ( + "error_description", + "user rejected the authorization request", + ), + ]; + if let Some(state) = transaction.state.as_deref() { + params.push(("state", state)); } - ); + return match redirect_with_params(&transaction.redirect_uri, ¶ms) { + Ok(url) => Redirect::to(&url).into_response(), + Err(error) => error.response(), + }; + } - info!("Authorization approved for client {}", form.client_id); - Redirect::to(&redirect_url).into_response() + let state_value = transaction.state.clone(); + let redirect_uri = transaction.redirect_uri.clone(); + let client_id = transaction.client_id.clone(); + let code = match state.create_authorization_code(transaction).await { + Ok(code) => code, + Err(error) => return error.response(), + }; + let mut params = vec![("code", code.as_str())]; + if let Some(value) = state_value.as_deref() { + params.push(("state", value)); + } + info!("Authorization approved for client {}", client_id); + match redirect_with_params(&redirect_uri, ¶ms) { + Ok(url) => Redirect::to(&url).into_response(), + Err(error) => error.response(), + } } -/// OAuth token endpoint handler -/// Exchanges authorization codes for access tokens pub async fn oauth_token( State(state): State>, - request: axum::http::Request, -) -> impl IntoResponse { - info!("Received token request"); - - let bytes = match axum::body::to_bytes(request.into_body(), usize::MAX).await { + request: Request, +) -> Response { + let bytes = match read_limited_body(request).await { Ok(bytes) => bytes, - Err(e) => { - error!("can't read request body: {}", e); - return ( + Err(response) => return response, + }; + let request = match serde_urlencoded::from_bytes::(&bytes) { + Ok(request) => request, + Err(_) => { + return oauth_json_error( StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": "invalid_request", - "error_description": "can't read request body" - })), - ) - .into_response(); + "invalid_request", + "token request is invalid", + ); } }; - - let token_req = match serde_urlencoded::from_bytes::(&bytes) { - Ok(form) => { - debug!( - "Parsed token request for grant type {} and client {}", - form.grant_type, form.client_id + if request.client_id.is_empty() { + return OAuthError::InvalidClient.response(); + } + let issued = match request.grant_type.as_str() { + "authorization_code" => state.exchange_authorization_code(&request).await, + "refresh_token" => state.exchange_refresh_token(&request).await, + _ => { + return oauth_json_error( + StatusCode::BAD_REQUEST, + "unsupported_grant_type", + "authorization_code and refresh_token are supported", ); - form - } - Err(e) => { - error!("can't parse form data: {}", e); - return ( - StatusCode::UNPROCESSABLE_ENTITY, - Json(serde_json::json!({ - "error": "invalid_request", - "error_description": format!("can't parse form data: {}", e) - })), - ) - .into_response(); } }; - if token_req.grant_type == "refresh_token" { - warn!("this easy server only support authorization_code now"); - return ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": "unsupported_grant_type", - "error_description": "only authorization_code is supported" - })), - ) - .into_response(); - } - if token_req.grant_type != "authorization_code" { - info!("unsupported grant type: {}", token_req.grant_type); - return ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": "unsupported_grant_type", - "error_description": "only authorization_code is supported" - })), - ) - .into_response(); - } - - // get session_id from code - if !token_req.code.starts_with("mcp-code-") { - warn!("invalid authorization code format"); - return ( - StatusCode::BAD_REQUEST, + match issued { + Ok(issued) => ( + StatusCode::OK, Json(serde_json::json!({ - "error": "invalid_grant", - "error_description": "invalid authorization code" + "access_token": issued.access.access_token, + "token_type": issued.access.token_type, + "expires_in": issued.access.expires_in, + "refresh_token": issued.access.refresh_token, + "scope": issued.access.scope, })), ) - .into_response(); - } - - // Validate the confidential client before consuming the authorization code. - match state - .validate_client_credentials( - &token_req.client_id, - &token_req.client_secret, - &token_req.redirect_uri, - ) - .await - { - Some(_) => { - let session_id = token_req.code.strip_prefix("mcp-code-").unwrap_or_default(); - debug!( - "Exchanging authorization code for client {}", - token_req.client_id - ); - - // Consume the authorization code and create an MCP access token. - match state - .exchange_authorization_code(session_id, &token_req.client_id) - .await - { - Ok(token) => { - info!("successfully created access token"); - ( - StatusCode::OK, - Json(serde_json::json!({ - "access_token": token.access_token, - "token_type": token.token_type, - "expires_in": token.expires_in, - "refresh_token": token.refresh_token, - "scope": token.scope, - })), - ) - .into_response() - } - Err(e) => { - warn!("failed to exchange authorization code: {}", e); - ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": "invalid_grant", - "error_description": "authorization code is invalid or already used" - })), - ) - .into_response() - } - } - } - None => { - warn!("invalid confidential client credentials"); - ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({ - "error": "invalid_client", - "error_description": "client authentication failed" - })), - ) - .into_response() - } + .into_response(), + Err(error) => error.response(), } } -/// Authentication middleware for validating Bearer tokens -/// Intercepts requests and validates access tokens before allowing access pub async fn validate_token_middleware( State(token_store): State>, - request: Request, + request: Request, next: Next, ) -> Response { - debug!("validate_token_middleware"); - // Extract the access token from the Authorization header - let auth_header = request.headers().get("Authorization"); - let token = match auth_header { - Some(header) => { - let header_str = header.to_str().unwrap_or(""); - if let Some(stripped) = header_str.strip_prefix("Bearer ") { - stripped.to_string() - } else { - return StatusCode::UNAUTHORIZED.into_response(); - } - } - None => { - return StatusCode::UNAUTHORIZED.into_response(); - } + let Some(token) = request + .headers() + .get("Authorization") + .and_then(|header| header.to_str().ok()) + .and_then(|header| header.strip_prefix("Bearer ")) + else { + return StatusCode::UNAUTHORIZED.into_response(); }; - - // Validate the token - match token_store.validate_token(&token).await { + match token_store.validate_token(token).await { Some(_) => next.run(request).await, None => StatusCode::UNAUTHORIZED.into_response(), } } -/// Get the actual IP address to use for endpoints -/// Returns the host from request headers if bind_address is 0.0.0.0, otherwise returns the original address -fn get_endpoint_address(bind_address: &str, host_header: Option<&str>) -> String { - if bind_address.starts_with("0.0.0.0") { - if let Some(host) = host_header { - // Use the Host header value, which contains the actual IP/domain the client used - host.to_string() - } else { - // Fallback to localhost if no Host header is present - bind_address.replacen("0.0.0.0", "localhost", 1) - } - } else { - bind_address.to_string() - } -} - -/// OAuth authorization server metadata endpoint -/// Returns server capabilities and endpoint URLs per RFC 8414 pub async fn oauth_authorization_server( - bind_address: &str, - headers: HeaderMap, + State(state): State>, ) -> impl IntoResponse { - let host_header = headers.get("host").and_then(|h| h.to_str().ok()); - let endpoint_address = get_endpoint_address(bind_address, host_header); - + let base = state.public_base_url(); + let endpoint = |path: &str| { + base.join(path) + .expect("validated public base URL") + .to_string() + }; let mut additional_fields = HashMap::new(); additional_fields.insert( "response_types_supported".into(), @@ -691,589 +859,572 @@ pub async fn oauth_authorization_server( "code_challenge_methods_supported".into(), Value::Array(vec![Value::String("S256".into())]), ); - let metadata = AuthorizationMetadata { - authorization_endpoint: format!("http://{endpoint_address}/authorize"), - token_endpoint: format!("http://{endpoint_address}/token"), - scopes_supported: Some(vec!["profile".to_string(), "email".to_string()]), - registration_endpoint: format!("http://{endpoint_address}/register"), - issuer: Some(format!("http://{endpoint_address}")), - jwks_uri: Some(format!("http://{endpoint_address}/jwks")), - additional_fields, - }; - debug!("metadata: {:?}", metadata); - (StatusCode::OK, Json(metadata)) + additional_fields.insert( + "grant_types_supported".into(), + Value::Array(vec![ + Value::String("authorization_code".into()), + Value::String("refresh_token".into()), + ]), + ); + additional_fields.insert( + "token_endpoint_auth_methods_supported".into(), + Value::Array(vec![ + Value::String("none".into()), + Value::String("client_secret_post".into()), + ]), + ); + let issuer = base.as_str().trim_end_matches('/').to_string(); + ( + StatusCode::OK, + Json(AuthorizationMetadata { + authorization_endpoint: endpoint("authorize"), + token_endpoint: endpoint("token"), + registration_endpoint: endpoint("register"), + issuer: Some(issuer), + jwks_uri: None, + scopes_supported: Some(SUPPORTED_SCOPES.iter().map(ToString::to_string).collect()), + additional_fields, + }), + ) } -/// Dynamic client registration endpoint -/// Allows clients to register themselves with the OAuth server pub async fn oauth_register( State(state): State>, - Json(req): Json, -) -> impl IntoResponse { - debug!("register request: {:?}", req); - if req.redirect_uris.is_empty() { - return ( + request: Request, +) -> Response { + let bytes = match read_limited_body(request).await { + Ok(bytes) => bytes, + Err(response) => return response, + }; + let request = match serde_json::from_slice::(&bytes) { + Ok(request) => request, + Err(_) => { + return oauth_json_error( + StatusCode::BAD_REQUEST, + "invalid_client_metadata", + "registration request is invalid", + ); + } + }; + if request.client_name.trim().is_empty() || request.client_name.len() > 100 { + return oauth_json_error( StatusCode::BAD_REQUEST, - Json(serde_json::json!({ - "error": "invalid_request", - "error_description": "at least one redirect uri is required" - })), - ) - .into_response(); + "invalid_client_metadata", + "client_name must contain between 1 and 100 characters", + ); } - - // generate client id and secret - let client_id = format!("client-{}", Uuid::new_v4()); - let client_secret = generate_random_string(32); - - let client = OAuthClientConfig { - client_id: client_id.clone(), - client_secret: Some(client_secret.clone()), - redirect_uri: req.redirect_uris[0].clone(), - scopes: vec![], + if request.redirect_uris.is_empty() + || request.redirect_uris.len() > 10 + || request + .redirect_uris + .iter() + .any(|uri| !validate_redirect_uri(uri)) + { + return oauth_json_error( + StatusCode::BAD_REQUEST, + "invalid_redirect_uri", + "redirect URIs must use HTTPS or an HTTP loopback address", + ); + } + if !request + .grant_types + .iter() + .any(|grant| grant == "authorization_code") + || request + .grant_types + .iter() + .any(|grant| !matches!(grant.as_str(), "authorization_code" | "refresh_token")) + || !request + .response_types + .iter() + .any(|response| response == "code") + || request + .response_types + .iter() + .any(|response| response != "code") + { + return oauth_json_error( + StatusCode::BAD_REQUEST, + "invalid_client_metadata", + "only authorization_code, optional refresh_token, and code response are supported", + ); + } + let Some(auth_method) = TokenEndpointAuthMethod::parse(&request.token_endpoint_auth_method) + else { + return oauth_json_error( + StatusCode::BAD_REQUEST, + "invalid_client_metadata", + "token_endpoint_auth_method must be none or client_secret_post", + ); }; - state - .clients - .write() - .await - .insert(client_id.clone(), client); - - // return client information - let response = ClientRegistrationResponse { - client_id, - client_secret: Some(client_secret), - client_name: req.client_name, - redirect_uris: req.redirect_uris, - additional_fields: HashMap::new(), - }; + let client_id = random_prefixed("client"); + let client_secret = (auth_method == TokenEndpointAuthMethod::ClientSecretPost) + .then(|| generate_random_string(48)); + let mut clients = state.clients.write().await; + if clients.len() >= MAX_CLIENTS { + return OAuthError::TemporarilyUnavailable.response(); + } + let mut unique_redirects = Vec::new(); + for redirect in &request.redirect_uris { + if !unique_redirects.contains(redirect) { + unique_redirects.push(redirect.clone()); + } + } + clients.insert( + client_id.clone(), + RegisteredClient { + client_id: client_id.clone(), + client_secret: client_secret.clone(), + redirect_uris: unique_redirects.clone(), + token_endpoint_auth_method: auth_method.clone(), + }, + ); + drop(clients); - (StatusCode::CREATED, Json(response)).into_response() + let mut additional_fields = HashMap::new(); + additional_fields.insert( + "token_endpoint_auth_method".to_string(), + Value::String(auth_method.as_str().to_string()), + ); + additional_fields.insert( + "grant_types".to_string(), + serde_json::json!(request.grant_types), + ); + additional_fields.insert( + "response_types".to_string(), + serde_json::json!(request.response_types), + ); + ( + StatusCode::CREATED, + Json(ClientRegistrationResponse { + client_id, + client_secret, + client_name: request.client_name, + redirect_uris: unique_redirects, + additional_fields, + }), + ) + .into_response() } #[cfg(test)] mod tests { use super::*; + use axum::body::to_bytes; - fn create_test_oauth_store() -> McpOAuthStore { - McpOAuthStore::new() - } + const VERIFIER: &str = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-._~"; - #[tokio::test] - async fn test_oauth_store_creation() { - let store = create_test_oauth_store(); + fn challenge() -> String { + URL_SAFE_NO_PAD.encode(Sha256::digest(VERIFIER.as_bytes())) + } - // Check that default client exists - let clients = store.clients.read().await; - assert!(clients.contains_key("mcp-client")); + fn authorize_query(client_id: &str, redirect_uri: &str) -> AuthorizeQuery { + AuthorizeQuery { + response_type: "code".to_string(), + client_id: client_id.to_string(), + redirect_uri: redirect_uri.to_string(), + scope: Some("email profile".to_string()), + state: Some("state with reserved & characters".to_string()), + code_challenge: Some(challenge()), + code_challenge_method: Some("S256".to_string()), + } + } - let default_client = clients.get("mcp-client").unwrap(); - assert_eq!(default_client.client_id, "mcp-client"); + async fn approve_transaction(store: &Arc, query: AuthorizeQuery) -> String { + let (transaction_id, transaction) = store + .create_authorization_transaction(&query) + .await + .unwrap(); + let body = serde_urlencoded::to_string(ApprovalForm { + transaction_id, + consent_nonce: transaction.consent_nonce, + approved: "true".to_string(), + approval_secret: store.approval_secret.as_ref().clone(), + }) + .unwrap(); + let response = oauth_approve( + State(store.clone()), + Request::builder().body(Body::from(body)).unwrap(), + ) + .await; + assert_eq!(response.status(), StatusCode::SEE_OTHER); + let location = response + .headers() + .get("location") + .unwrap() + .to_str() + .unwrap(); + let redirect = Url::parse(location).unwrap(); assert_eq!( - default_client.client_secret, - Some("mcp-client-secret".to_string()) + redirect + .query_pairs() + .find(|(key, _)| key == "state") + .unwrap() + .1, + "state with reserved & characters" ); - assert!(default_client.scopes.contains(&"profile".to_string())); - assert!(default_client.scopes.contains(&"email".to_string())); + redirect + .query_pairs() + .find(|(key, _)| key == "code") + .unwrap() + .1 + .into_owned() } - #[tokio::test] - async fn test_validate_client_success() { - let store = create_test_oauth_store(); - - let result = store - .validate_client("mcp-client", "http://localhost:8080/callback") - .await; - assert!(result.is_some()); - - let client = result.unwrap(); - assert_eq!(client.client_id, "mcp-client"); + async fn token_request(store: Arc, request: TokenRequest) -> Response { + oauth_token( + State(store), + Request::builder() + .body(Body::from(serde_urlencoded::to_string(request).unwrap())) + .unwrap(), + ) + .await } #[tokio::test] - async fn test_validate_client_invalid_client_id() { - let store = create_test_oauth_store(); - - let result = store - .validate_client("invalid-client", "http://localhost:8080/callback") - .await; - assert!(result.is_none()); + async fn public_client_completes_pkce_flow_and_code_is_single_use() { + let store = Arc::new(McpOAuthStore::new()); + let code = approve_transaction( + &store, + authorize_query("public-test-client", "http://127.0.0.1:8080/callback"), + ) + .await; + let request = TokenRequest { + grant_type: "authorization_code".to_string(), + code, + client_id: "public-test-client".to_string(), + client_secret: String::new(), + redirect_uri: "http://127.0.0.1:8080/callback".to_string(), + code_verifier: Some(VERIFIER.to_string()), + refresh_token: String::new(), + }; + let first = token_request(store.clone(), request.clone()).await; + assert_eq!(first.status(), StatusCode::OK); + let replay = token_request(store, request).await; + assert_eq!(replay.status(), StatusCode::BAD_REQUEST); } #[tokio::test] - async fn test_validate_client_invalid_redirect_uri() { - let store = create_test_oauth_store(); - - let result = store - .validate_client("mcp-client", "http://malicious.com/callback") - .await; - assert!(result.is_none()); + async fn incorrect_pkce_verifier_is_rejected_and_consumes_code() { + let store = Arc::new(McpOAuthStore::new()); + let code = approve_transaction( + &store, + authorize_query("public-test-client", "http://127.0.0.1:8080/callback"), + ) + .await; + let mut request = TokenRequest { + grant_type: "authorization_code".to_string(), + code, + client_id: "public-test-client".to_string(), + client_secret: String::new(), + redirect_uri: "http://127.0.0.1:8080/callback".to_string(), + code_verifier: Some(format!("{VERIFIER}x")), + refresh_token: String::new(), + }; + assert_eq!( + token_request(store.clone(), request.clone()).await.status(), + StatusCode::BAD_REQUEST + ); + request.code_verifier = Some(VERIFIER.to_string()); + assert_eq!( + token_request(store, request).await.status(), + StatusCode::BAD_REQUEST + ); } #[tokio::test] - async fn test_oauth_approve_rejects_direct_unauthenticated_post() { - let store = Arc::new(create_test_oauth_store()); - let form = ApprovalForm { - client_id: "mcp-client".to_string(), - redirect_uri: "http://localhost:8080/callback".to_string(), - scope: "profile".to_string(), - state: "state123".to_string(), - approved: "true".to_string(), - approval_secret: String::new(), + async fn confidential_client_requires_its_generated_auth_semantics() { + let store = Arc::new(McpOAuthStore::new()); + let code = approve_transaction( + &store, + authorize_query( + "confidential-test-client", + "https://client.example/callback", + ), + ) + .await; + let mut request = TokenRequest { + grant_type: "authorization_code".to_string(), + code, + client_id: "confidential-test-client".to_string(), + client_secret: "wrong".to_string(), + redirect_uri: "https://client.example/callback".to_string(), + code_verifier: Some(VERIFIER.to_string()), + refresh_token: String::new(), }; - - let response = oauth_approve(State(store.clone()), Form(form)) - .await - .into_response(); - - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - assert!(store.auth_sessions.read().await.is_empty()); + assert_eq!( + token_request(store.clone(), request.clone()).await.status(), + StatusCode::UNAUTHORIZED + ); + request.client_secret = "test-only-confidential-secret".to_string(); + assert_eq!(token_request(store, request).await.status(), StatusCode::OK); } #[tokio::test] - async fn test_oauth_token_rejects_invalid_client_secret() { - let store = Arc::new(create_test_oauth_store()); - let session_id = store - .create_auth_session( - "mcp-client".to_string(), + async fn refresh_token_rotates_and_old_value_cannot_be_replayed() { + let store = Arc::new(McpOAuthStore::new()); + let issued = store + .issue_tokens( + "public-test-client".to_string(), Some("profile".to_string()), - None, - "secret-check-session".to_string(), ) - .await; - let auth_token = AuthToken::new( - AccessToken::new("third-party-token".to_string()), - oauth2::basic::BasicTokenType::Bearer, - EmptyExtraTokenFields {}, - ); - store - .update_auth_session_token(&session_id, auth_token) .await .unwrap(); - let request_body = serde_urlencoded::to_string(TokenRequest { - grant_type: "authorization_code".to_string(), - code: format!("mcp-code-{session_id}"), - client_id: "mcp-client".to_string(), - client_secret: "wrong-secret".to_string(), - redirect_uri: "http://localhost:8080/callback".to_string(), + let refresh_token = issued.access.refresh_token.unwrap(); + let request = TokenRequest { + grant_type: "refresh_token".to_string(), + code: String::new(), + client_id: "public-test-client".to_string(), + client_secret: String::new(), + redirect_uri: String::new(), code_verifier: None, - refresh_token: String::new(), - }) - .unwrap(); - let request = Request::builder() - .header("content-type", "application/x-www-form-urlencoded") - .body(Body::from(request_body)) - .unwrap(); - - let response = oauth_token(State(store), request).await.into_response(); - - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + refresh_token, + }; + assert_eq!( + token_request(store.clone(), request.clone()).await.status(), + StatusCode::OK + ); + assert_eq!( + token_request(store, request).await.status(), + StatusCode::BAD_REQUEST + ); } #[tokio::test] - async fn test_oauth_approve_requires_valid_client_and_resource_owner_secret() { - let store = Arc::new(McpOAuthStore::with_approval_secret( - "resource-owner-secret-with-32-characters".to_string(), - )); - let valid_form = ApprovalForm { - client_id: "mcp-client".to_string(), - redirect_uri: "http://localhost:8080/callback".to_string(), - scope: "profile".to_string(), - state: "state123".to_string(), - approved: "true".to_string(), - approval_secret: "resource-owner-secret-with-32-characters".to_string(), - }; - - let response = oauth_approve(State(store.clone()), Form(valid_form)) - .await - .into_response(); - - assert_eq!(response.status(), StatusCode::SEE_OTHER); - assert_eq!(store.auth_sessions.read().await.len(), 1); - - let invalid_redirect_form = ApprovalForm { - client_id: "mcp-client".to_string(), - redirect_uri: "http://attacker.invalid/callback".to_string(), - scope: "profile".to_string(), - state: "state123".to_string(), - approved: "true".to_string(), - approval_secret: "resource-owner-secret-with-32-characters".to_string(), + async fn expired_access_token_is_removed() { + let store = McpOAuthStore::new(); + let token = McpAccessToken { + access_token: "expired-token".to_string(), + token_type: "bearer".to_string(), + expires_in: 0, + refresh_token: None, + scope: None, + client_id: "public-test-client".to_string(), + issued_at: Utc::now() - Duration::hours(2), + expires_at: Utc::now() - Duration::hours(1), }; - let response = oauth_approve(State(store.clone()), Form(invalid_redirect_form)) + store + .access_tokens + .write() .await - .into_response(); + .insert(token.access_token.clone(), token); - assert_eq!(response.status(), StatusCode::BAD_REQUEST); - assert_eq!(store.auth_sessions.read().await.len(), 1); + assert!(store.validate_token("expired-token").await.is_none()); + assert!(store.access_tokens.read().await.is_empty()); } #[tokio::test] - async fn test_oauth_token_consumes_authorization_code_once() { - let store = Arc::new(create_test_oauth_store()); - let session_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile".to_string()), - None, - "single-use-session".to_string(), - ) - .await; - let auth_token = AuthToken::new( - AccessToken::new("third-party-token".to_string()), - oauth2::basic::BasicTokenType::Bearer, - EmptyExtraTokenFields {}, - ); - store - .update_auth_session_token(&session_id, auth_token) + async fn authorization_code_and_transaction_expiry_are_enforced() { + let store = Arc::new(McpOAuthStore::new()); + let query = authorize_query("public-test-client", "http://127.0.0.1:8080/callback"); + let (transaction_id, transaction) = store + .create_authorization_transaction(&query) .await .unwrap(); - let request_body = serde_urlencoded::to_string(TokenRequest { - grant_type: "authorization_code".to_string(), - code: format!("mcp-code-{session_id}"), - client_id: "mcp-client".to_string(), - client_secret: "mcp-client-secret".to_string(), - redirect_uri: "http://localhost:8080/callback".to_string(), - code_verifier: None, - refresh_token: String::new(), - }) - .unwrap(); - - let first_request = Request::builder() - .header("content-type", "application/x-www-form-urlencoded") - .body(Body::from(request_body.clone())) - .unwrap(); - let first_response = oauth_token(State(store.clone()), first_request) + store + .auth_transactions + .write() .await - .into_response(); - assert_eq!(first_response.status(), StatusCode::OK); + .get_mut(&transaction_id) + .unwrap() + .expires_at = Utc::now() - Duration::seconds(1); + assert!( + store + .consume_authorization_transaction(&transaction_id, &transaction.consent_nonce) + .await + .is_err() + ); - let replay_request = Request::builder() - .header("content-type", "application/x-www-form-urlencoded") - .body(Body::from(request_body)) + let code = store + .create_authorization_code(AuthorizationTransaction { + expires_at: Utc::now() + Duration::minutes(1), + ..transaction + }) + .await .unwrap(); - let replay_response = oauth_token(State(store), replay_request) + store + .authorization_codes + .write() .await - .into_response(); - assert_eq!(replay_response.status(), StatusCode::BAD_REQUEST); - } - - #[tokio::test] - async fn test_create_auth_session() { - let store = create_test_oauth_store(); - - let session_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile email".to_string()), - Some("state123".to_string()), - "session123".to_string(), - ) - .await; - - assert_eq!(session_id, "session123"); - - // Verify session exists - let sessions = store.auth_sessions.read().await; - assert!(sessions.contains_key("session123")); - - let session = sessions.get("session123").unwrap(); - assert_eq!(session.client_id, "mcp-client"); - assert_eq!(session.scope, Some("profile email".to_string())); - assert!(session.auth_token.is_none()); - } - - #[tokio::test] - async fn test_update_auth_session_token() { - let store = create_test_oauth_store(); - - // Create session first - let session_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile".to_string()), - None, - "session456".to_string(), - ) - .await; - - // Create a mock token - let token = AuthToken::new( - AccessToken::new("access_token_123".to_string()), - oauth2::basic::BasicTokenType::Bearer, - EmptyExtraTokenFields {}, + .get_mut(&code) + .unwrap() + .expires_at = Utc::now() - Duration::seconds(1); + let request = TokenRequest { + grant_type: "authorization_code".to_string(), + code, + client_id: "public-test-client".to_string(), + client_secret: String::new(), + redirect_uri: "http://127.0.0.1:8080/callback".to_string(), + code_verifier: Some(VERIFIER.to_string()), + refresh_token: String::new(), + }; + assert_eq!( + token_request(store, request).await.status(), + StatusCode::BAD_REQUEST ); - - // Update session with token - let result = store.update_auth_session_token(&session_id, token).await; - assert!(result.is_ok()); - - // Verify token was added - let sessions = store.auth_sessions.read().await; - let session = sessions.get("session456").unwrap(); - assert!(session.auth_token.is_some()); } #[tokio::test] - async fn test_update_auth_session_token_invalid_session() { - let store = create_test_oauth_store(); - - let token = AuthToken::new( - AccessToken::new("access_token_123".to_string()), - oauth2::basic::BasicTokenType::Bearer, - EmptyExtraTokenFields {}, + async fn approval_requires_bound_nonce_and_is_single_use() { + let store = Arc::new(McpOAuthStore::new()); + let (transaction_id, transaction) = store + .create_authorization_transaction(&authorize_query( + "public-test-client", + "http://127.0.0.1:8080/callback", + )) + .await + .unwrap(); + let invalid_body = serde_urlencoded::to_string(ApprovalForm { + transaction_id: transaction_id.clone(), + consent_nonce: "attacker-controlled".to_string(), + approved: "true".to_string(), + approval_secret: store.approval_secret.as_ref().clone(), + }) + .unwrap(); + assert_eq!( + oauth_approve( + State(store.clone()), + Request::builder().body(Body::from(invalid_body)).unwrap() + ) + .await + .status(), + StatusCode::BAD_REQUEST ); - let result = store.update_auth_session_token("nonexistent", token).await; - assert!(result.is_err()); - assert_eq!(result.unwrap_err(), "Session not found"); + let valid_body = serde_urlencoded::to_string(ApprovalForm { + transaction_id, + consent_nonce: transaction.consent_nonce, + approved: "true".to_string(), + approval_secret: store.approval_secret.as_ref().clone(), + }) + .unwrap(); + let first = oauth_approve( + State(store.clone()), + Request::builder() + .body(Body::from(valid_body.clone())) + .unwrap(), + ) + .await; + assert_eq!(first.status(), StatusCode::SEE_OTHER); + let replay = oauth_approve( + State(store), + Request::builder().body(Body::from(valid_body)).unwrap(), + ) + .await; + assert_eq!(replay.status(), StatusCode::BAD_REQUEST); } #[tokio::test] - async fn test_create_mcp_token_success() { - let store = create_test_oauth_store(); - - // Create session and update with auth token - let session_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile".to_string()), - None, - "session789".to_string(), + async fn token_and_approval_endpoints_reject_oversized_bodies() { + let store = Arc::new(McpOAuthStore::new()); + let oversized = "x".repeat(MAX_OAUTH_BODY_BYTES + 1); + assert_eq!( + oauth_token( + State(store.clone()), + Request::builder() + .body(Body::from(oversized.clone())) + .unwrap() ) - .await; - - let auth_token = AuthToken::new( - AccessToken::new("third_party_token".to_string()), - oauth2::basic::BasicTokenType::Bearer, - EmptyExtraTokenFields {}, + .await + .status(), + StatusCode::PAYLOAD_TOO_LARGE ); - - store - .update_auth_session_token(&session_id, auth_token) + assert_eq!( + oauth_approve( + State(store), + Request::builder().body(Body::from(oversized)).unwrap() + ) .await - .unwrap(); - - // Create MCP token - let result = store.create_mcp_token(&session_id).await; - assert!(result.is_ok()); - - let mcp_token = result.unwrap(); - assert!(mcp_token.access_token.starts_with("mcp-token-")); - assert!( - mcp_token - .refresh_token - .as_ref() - .unwrap() - .starts_with("mcp-refresh-") + .status(), + StatusCode::PAYLOAD_TOO_LARGE ); - assert_eq!(mcp_token.token_type, "bearer"); - assert_eq!(mcp_token.expires_in, Some(3600)); - assert_eq!(mcp_token.scope, Some("profile".to_string())); - assert_eq!(mcp_token.client_id, "mcp-client"); - } - - #[tokio::test] - async fn test_create_mcp_token_no_session() { - let store = create_test_oauth_store(); - - let result = store.create_mcp_token("nonexistent").await; - assert!(result.is_err()); - assert_eq!(result.unwrap_err(), "Session not found"); } #[tokio::test] - async fn test_create_mcp_token_no_auth_token() { - let store = create_test_oauth_store(); + async fn authorize_validates_response_scope_redirect_and_pkce() { + let store = Arc::new(McpOAuthStore::new()); + let mut query = authorize_query("public-test-client", "http://127.0.0.1:8080/callback"); + query.response_type = "token".to_string(); + assert_eq!( + oauth_authorize(Query(query), State(store.clone())) + .await + .status(), + StatusCode::BAD_REQUEST + ); - // Create session without auth token - let session_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile".to_string()), - None, - "session_no_token".to_string(), - ) - .await; + let mut query = authorize_query("public-test-client", "http://127.0.0.1:8080/callback"); + query.scope = Some("admin".to_string()); + assert_eq!( + oauth_authorize(Query(query), State(store.clone())) + .await + .status(), + StatusCode::BAD_REQUEST + ); - let result = store.create_mcp_token(&session_id).await; - assert!(result.is_err()); + let mut query = authorize_query("public-test-client", "https://attacker.example/callback"); + query.code_challenge = None; assert_eq!( - result.unwrap_err(), - "No third-party token available for session" + oauth_authorize(Query(query), State(store)).await.status(), + StatusCode::BAD_REQUEST ); } #[tokio::test] - async fn test_validate_token_success() { - let store = create_test_oauth_store(); - - // Create a complete flow to get a valid token - let session_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile".to_string()), - None, - "token_test_session".to_string(), - ) - .await; - - let auth_token = AuthToken::new( - AccessToken::new("third_party_token".to_string()), - oauth2::basic::BasicTokenType::Bearer, - EmptyExtraTokenFields {}, - ); - - store - .update_auth_session_token(&session_id, auth_token) + async fn dynamic_registration_distinguishes_public_and_confidential_clients() { + let store = Arc::new(McpOAuthStore::new()); + let public = ClientRegistrationRequest { + client_name: "public".to_string(), + redirect_uris: vec!["http://127.0.0.1:9911/callback".to_string()], + grant_types: vec![ + "authorization_code".to_string(), + "refresh_token".to_string(), + ], + token_endpoint_auth_method: "none".to_string(), + response_types: vec!["code".to_string()], + }; + let response = oauth_register( + State(store.clone()), + Request::builder() + .body(Body::from(serde_json::to_vec(&public).unwrap())) + .unwrap(), + ) + .await; + assert_eq!(response.status(), StatusCode::CREATED); + let body = to_bytes(response.into_body(), MAX_OAUTH_BODY_BYTES) .await .unwrap(); - let mcp_token = store.create_mcp_token(&session_id).await.unwrap(); - - // Validate the token - let result = store.validate_token(&mcp_token.access_token).await; - assert!(result.is_some()); + let registered: ClientRegistrationResponse = serde_json::from_slice(&body).unwrap(); + assert!(registered.client_secret.is_none()); - let validated_token = result.unwrap(); - assert_eq!(validated_token.access_token, mcp_token.access_token); - assert_eq!(validated_token.client_id, "mcp-client"); - } - - #[tokio::test] - async fn test_validate_token_invalid() { - let store = create_test_oauth_store(); - - let result = store.validate_token("invalid_token").await; - assert!(result.is_none()); - } - - #[tokio::test] - async fn test_mcp_access_token_serialization() { - let auth_token = AuthToken::new( - AccessToken::new("test_token".to_string()), - oauth2::basic::BasicTokenType::Bearer, - EmptyExtraTokenFields {}, - ); - - let mcp_token = McpAccessToken { - access_token: "mcp-token-123".to_string(), - token_type: "bearer".to_string(), - expires_in: Some(3600), - refresh_token: Some("mcp-refresh-123".to_string()), - scope: Some("profile email".to_string()), - auth_token, - client_id: "test-client".to_string(), + let confidential = ClientRegistrationRequest { + token_endpoint_auth_method: "client_secret_post".to_string(), + ..public }; - - // Test that it can be serialized to JSON - let json_result = serde_json::to_string(&mcp_token); - assert!(json_result.is_ok()); - - let json_str = json_result.unwrap(); - assert!(json_str.contains("mcp-token-123")); - assert!(json_str.contains("bearer")); - assert!(json_str.contains("3600")); - } - - #[tokio::test] - async fn test_auth_session_creation_with_minimal_data() { - let store = create_test_oauth_store(); - - let session_id = store - .create_auth_session( - "test-client".to_string(), - None, // No scope - None, // No state - "minimal_session".to_string(), - ) - .await; - - assert_eq!(session_id, "minimal_session"); - - let sessions = store.auth_sessions.read().await; - let session = sessions.get("minimal_session").unwrap(); - assert_eq!(session.client_id, "test-client"); - assert!(session.scope.is_none()); - assert!(session._state.is_none()); - assert!(session.auth_token.is_none()); - } - - #[tokio::test] - async fn test_concurrent_access() { - let store = Arc::new(create_test_oauth_store()); - - // Test concurrent session creation - let mut handles = vec![]; - for i in 0..10 { - let store_clone = store.clone(); - let handle = tokio::spawn(async move { - store_clone - .create_auth_session( - "mcp-client".to_string(), - Some("profile".to_string()), - None, - format!("concurrent_session_{i}"), - ) - .await - }); - handles.push(handle); - } - - for handle in handles { - let session_id = handle.await.unwrap(); - assert!(session_id.starts_with("concurrent_session_")); - } - - // Verify all sessions were created - let sessions = store.auth_sessions.read().await; - assert_eq!(sessions.len(), 10); - } - #[test] - fn test_get_endpoint_address_with_zero_ip() { - let result = get_endpoint_address("0.0.0.0:8080", Some("192.168.1.100:8080")); - assert_eq!(result, "192.168.1.100:8080"); - } - - #[test] - fn test_get_endpoint_address_with_zero_ip_no_port() { - let result = get_endpoint_address("0.0.0.0", Some("192.168.1.100")); - assert_eq!(result, "192.168.1.100"); - } - - #[test] - fn test_get_endpoint_address_with_zero_ip_no_host_header() { - let result = get_endpoint_address("0.0.0.0:8080", None); - assert_eq!(result, "localhost:8080"); - } - - #[test] - fn test_get_endpoint_address_with_specific_ip() { - let result = get_endpoint_address("192.168.1.100:8080", Some("192.168.1.100:8080")); - assert_eq!(result, "192.168.1.100:8080"); - } - - #[test] - fn test_get_endpoint_address_with_localhost() { - let result = get_endpoint_address("localhost:8080", Some("localhost:8080")); - assert_eq!(result, "localhost:8080"); + let response = oauth_register( + State(store), + Request::builder() + .body(Body::from(serde_json::to_vec(&confidential).unwrap())) + .unwrap(), + ) + .await; + let body = to_bytes(response.into_body(), MAX_OAUTH_BODY_BYTES) + .await + .unwrap(); + let registered: ClientRegistrationResponse = serde_json::from_slice(&body).unwrap(); + assert_eq!(registered.client_secret.unwrap().len(), 48); } #[test] - fn test_get_endpoint_address_with_domain() { - let result = get_endpoint_address("example.com:8080", Some("example.com:8080")); - assert_eq!(result, "example.com:8080"); - } - - #[tokio::test] - async fn test_oauth_authorization_server_with_zero_ip() { - use axum::http::HeaderMap; - - let mut headers = HeaderMap::new(); - headers.insert("host", "192.168.1.100:8080".parse().unwrap()); - - let _response = oauth_authorization_server("0.0.0.0:8080", headers).await; - - // This is a basic test to ensure the function doesn't panic - // In a real test, you'd want to extract and verify the JSON response - // to ensure the URLs contain "192.168.1.100:8080" instead of "0.0.0.0:8080" + fn redirect_and_public_base_urls_have_safe_schemes_and_shapes() { + assert!(validate_redirect_uri("https://client.example/callback")); + assert!(validate_redirect_uri("http://127.0.0.1:8080/callback")); + assert!(!validate_redirect_uri("http://client.example/callback")); + assert!(!validate_redirect_uri( + "https://user:password@client.example/callback" + )); + assert!(validate_public_base_url("https://mcp.example/", true).is_ok()); + assert!(validate_public_base_url("http://mcp.example/", true).is_err()); + assert!(validate_public_base_url("https://mcp.example/path?query=1", true).is_err()); } } diff --git a/mcp-servers/mcp-bash-server/src/html/mcp_oauth_index.html b/mcp-servers/mcp-bash-server/src/html/mcp_oauth_index.html index 76afe635ef3..e37caeda4e0 100644 --- a/mcp-servers/mcp-bash-server/src/html/mcp_oauth_index.html +++ b/mcp-servers/mcp-bash-server/src/html/mcp_oauth_index.html @@ -31,7 +31,7 @@

MCP OAuth Server

-

This is an MCP server with OAuth 2.0 integration to a third-party authorization server.

+

This MCP server uses an OAuth 2.0 authorization-code flow with PKCE.

Available Endpoints:

@@ -41,10 +41,12 @@

Authorization Endpoint

Parameters:

  • response_type - Must be "code"
    • -
  • client_id - Client identifier (e.g., "mcp-client")
    • -
  • redirect_uri - URI to redirect after authorization
    • -
  • scope - Optional requested scope
    • -
  • state - Optional state value for CSRF prevention
    • +
  • client_id - Dynamically registered client identifier
    • +
  • redirect_uri - Exact registered redirect URI
    • +
  • scope - Optional supported scopes
    • +
  • state - Recommended client transaction state
    • +
  • code_challenge - PKCE S256 challenge
    • +
  • code_challenge_method - Must be "S256"
@@ -53,11 +55,13 @@

Token Endpoint

POST /token

Parameters:

    -
  • grant_type - Must be "authorization_code"
    • -
  • code - The authorization code
    • +
  • grant_type - "authorization_code" or "refresh_token"
    • +
  • code - One-time authorization code
  • client_id - Client identifier
    • -
  • client_secret - Client secret
    • +
  • client_secret - Required only for confidential clients
  • redirect_uri - Redirect URI used in authorization request
    • +
  • code_verifier - PKCE verifier for authorization-code exchange
    • +
  • refresh_token - Rotating token for refresh grants
@@ -69,13 +73,12 @@

MCP streamablehttp Endpoints

OAuth Flow:

    -
  1. MCP Client initiates OAuth flow with this MCP Server
    2. -
  2. MCP Server redirects to Third-Party OAuth Server
    3. -
  3. User authenticates with Third-Party Server
    4. -
  4. Third-Party Server redirects back to MCP Server with auth code
    5. -
  5. MCP Server exchanges the code for a third-party access token
    6. -
  6. MCP Server generates its own token bound to the third-party session
    7. -
  7. MCP Server completes the OAuth flow with the MCP Client
    8. +
  8. The MCP client discovers metadata and dynamically registers.
    9. +
  9. The client starts authorization with an S256 PKCE challenge.
    10. +
  10. The resource owner authenticates and approves the bound transaction.
    11. +
  11. The server returns a short-lived, one-time authorization code.
    12. +
  12. The client exchanges the code and PKCE verifier for expiring tokens.
    13. +
  13. The client sends the bearer access token to the MCP endpoint.
diff --git a/mcp-servers/mcp-bash-server/src/main.rs b/mcp-servers/mcp-bash-server/src/main.rs index c715003d04f..e3e5c11070f 100644 --- a/mcp-servers/mcp-bash-server/src/main.rs +++ b/mcp-servers/mcp-bash-server/src/main.rs @@ -1,42 +1,33 @@ /* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to You under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. */ -//! MCP Bash Server - A Model Context Protocol server for executing bash commands -//! -//! This server provides secure bash command execution capabilities through the MCP protocol. -//! It includes OAuth2 authentication, command validation, and cross-platform support. -//! -//! Features: -//! - Secure command execution with blacklist validation -//! - OAuth2 authentication for client authorization -//! - Cross-platform shell support (Linux, Windows, macOS) -//! - Built-in system information tools -//! - Configurable timeout and environment settings +//! MCP Bash Server. -use std::sync::OnceLock; use std::{net::SocketAddr, sync::Arc}; use anyhow::{Context, Result, bail}; use axum::{ Router, body::Body, - http::{HeaderMap, Request}, + http::Request, middleware::{self, Next}, - response::{Html, IntoResponse, Response}, + response::{Html, Response}, routing::{get, post}, }; use rmcp::transport::streamable_http_server::{ @@ -46,77 +37,39 @@ use tower_http::cors::{Any, CorsLayer}; use tracing::info; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; -// Import modules mod common; use common::bash_server::BashServer; use common::config; use common::oauth::{ McpOAuthStore, generate_random_string, oauth_approve, oauth_authorization_server, - oauth_authorize, oauth_register, oauth_token, validate_token_middleware, + oauth_authorize, oauth_register, oauth_token, validate_public_base_url, + validate_token_middleware, }; const INDEX_HTML: &str = include_str!("html/mcp_oauth_index.html"); -/// Global storage for server bind address, initialized once at startup -// Init once from environment variable BIND_ADDRESS -pub static BIND_ADDRESS: OnceLock = OnceLock::new(); - -/// Root path handler -/// Serves the main OAuth authorization index page async fn index() -> Html<&'static str> { Html(INDEX_HTML) } -/// Wrapper function for oauth_authorization_server to handle BIND_ADDRESS -async fn oauth_authorization_server_handler(headers: HeaderMap) -> impl IntoResponse { - let bind_address = BIND_ADDRESS - .get() - .expect("BIND_ADDRESS must be initialized in main()"); - oauth_authorization_server(bind_address, headers).await -} - -/// HTTP request logging middleware -/// Logs all incoming requests including method, URI, headers and response status +/// Log request metadata without reading form bodies or emitting credentials. async fn log_request(request: Request, next: Next) -> Response { let method = request.method().clone(); let uri = request.uri().clone(); let version = request.version(); - - // Log headers let headers = request.headers().clone(); let mut header_log = String::new(); for (key, value) in &headers { - let value_str = if key == "authorization" || key == "cookie" { + let value = if key == "authorization" || key == "cookie" { "" } else { value.to_str().unwrap_or("") }; - header_log.push_str(&format!("\n {key}: {value_str}")); + header_log.push_str(&format!("\n {key}: {value}")); } - - // Try to get request body for form submissions - let content_type = headers - .get("content-type") - .and_then(|v| v.to_str().ok()) - .unwrap_or(""); - - let request_info = if content_type.contains("application/x-www-form-urlencoded") - || content_type.contains("application/json") - { - format!("{method} {uri} {version:?}{header_log}\nContent-Type: {content_type}") - } else { - format!("{method} {uri} {version:?}{header_log}") - }; - - info!("REQUEST: {}", request_info); - - // Call the actual handler + info!("REQUEST: {method} {uri} {version:?}{header_log}"); let response = next.run(request).await; - - // Log response status - let status = response.status(); - info!("RESPONSE: {} for {} {}", status, method, uri); - + info!("RESPONSE: {} for {} {}", response.status(), method, uri); response } @@ -132,11 +85,25 @@ fn approval_secret_for_mode(is_dev: bool, configured: Option) -> Result< Ok(approval_secret) } -/// Main application entry point -/// Sets up logging, OAuth store, HTTP server, and starts the MCP bash server +fn public_base_url_for_mode( + is_dev: bool, + configured: Option, + bind_address: &str, +) -> Result { + if let Some(configured) = configured { + return validate_public_base_url(&configured, !is_dev) + .map_err(|message| anyhow::anyhow!(message)); + } + if !is_dev { + bail!("MCP_OAUTH_PUBLIC_BASE_URL must be set when the server runs in production mode"); + } + let local_address = bind_address.replacen("0.0.0.0", "127.0.0.1", 1); + validate_public_base_url(&format!("http://{local_address}"), false) + .map_err(|message| anyhow::anyhow!(message)) +} + #[tokio::main] async fn main() -> Result<()> { - // Initialize logging let logs = tracing_appender::rolling::daily("logs", "mcp.log"); let (non_blocking, _guard) = tracing_appender::non_blocking(logs); let log_setting = tracing_subscriber::fmt::layer().with_writer(non_blocking); @@ -149,7 +116,6 @@ async fn main() -> Result<()> { .with(log_setting) .init(); - // Read environment mode from config file, default to "production" let config = config::Config::read_config("config.toml")?; let env_mode = config .settings @@ -157,30 +123,29 @@ async fn main() -> Result<()> { .clone() .unwrap_or_else(|| "production".to_string()); let is_dev = env_mode == "development"; - - let approval_secret = - approval_secret_for_mode(is_dev, std::env::var("MCP_OAUTH_APPROVAL_SECRET").ok())?; - - // Create the OAuth store with an operator-controlled resource-owner secret. - let oauth_store = Arc::new(McpOAuthStore::with_approval_secret(approval_secret)); - let host = config.settings.host.clone(); let port = config.settings.port; let bind_address = format!("{host}:{port}"); - let addr = bind_address.parse::()?; - let _ = BIND_ADDRESS.set(bind_address); - // Create StreamableHttpServer + let approval_secret = + approval_secret_for_mode(is_dev, std::env::var("MCP_OAUTH_APPROVAL_SECRET").ok())?; + let public_base_url = public_base_url_for_mode( + is_dev, + std::env::var("MCP_OAUTH_PUBLIC_BASE_URL").ok(), + &bind_address, + )?; + let oauth_store = Arc::new(McpOAuthStore::with_settings( + approval_secret, + public_base_url, + )); + let service = StreamableHttpService::new( || Ok(BashServer::new()), LocalSessionManager::default().into(), Default::default(), ); - let server_router = Router::new().nest_service("/mcp", service); - - // Add OAuth authentication middleware only if not in development mode let protected_server_router = if is_dev { server_router } else { @@ -190,107 +155,52 @@ async fn main() -> Result<()> { )) }; - // Create CORS layer for the oauth authorization server endpoint let cors_layer = CorsLayer::new() .allow_origin(Any) .allow_methods(Any) .allow_headers(Any); - - // Create a sub-router for the oauth authorization server endpoint with CORS let oauth_server_router = Router::new() .route( "/.well-known/oauth-authorization-server", - get(oauth_authorization_server_handler).options(oauth_authorization_server_handler), + get(oauth_authorization_server).options(oauth_authorization_server), ) .route("/token", post(oauth_token).options(oauth_token)) .route("/register", post(oauth_register).options(oauth_register)) .layer(cors_layer) .with_state(oauth_store.clone()); - // Create HTTP router with request logging middleware let app = Router::new() .route("/", get(index)) .route("/authorize", get(oauth_authorize)) .route("/approve", post(oauth_approve)) - .merge(oauth_server_router) // Merge the CORS-enabled oauth server router + .merge(oauth_server_router) .merge(protected_server_router) - .with_state(oauth_store.clone()) + .with_state(oauth_store) .layer(middleware::from_fn(log_request)); - // Start HTTP server info!("MCP OAuth Server started on {}", addr); let listener = tokio::net::TcpListener::bind(addr).await?; let _ = axum::serve(listener, app) .with_graceful_shutdown(async { tokio::signal::ctrl_c().await.unwrap() }) .await; - Ok(()) } #[cfg(test)] mod tests { use super::*; - use axum::body::Body; - use axum::http::Method; - use axum::http::Request; #[tokio::test] - async fn test_index_handler() { + async fn index_handler_returns_oauth_page() { let response = index().await; - let html_content = response.0; - - // Verify it returns the expected HTML content - assert_eq!(html_content, INDEX_HTML); - assert!(html_content.contains("OAuth")); - } - - #[tokio::test] - async fn test_oauth_authorization_server_handler() { - use axum::http::HeaderMap; - - // Set up BIND_ADDRESS for testing - let _ = BIND_ADDRESS.set("localhost:8080".to_string()); - - let mut headers = HeaderMap::new(); - headers.insert("host", "localhost:8080".parse().unwrap()); - - let response = oauth_authorization_server_handler(headers).await; - - // Test that the handler returns a response - // We can't easily test the exact content without mocking, but we can verify it doesn't panic - let _response_body = response.into_response(); + assert_eq!(response.0, INDEX_HTML); + assert!(response.0.contains("OAuth")); } #[test] - fn test_bind_address_initialization() { - // Create a new OnceLock for testing to avoid conflicts - let test_bind_address: OnceLock = OnceLock::new(); - - // Test that we can set the value once - let result = test_bind_address.set("127.0.0.1:9090".to_string()); - assert!(result.is_ok()); - - // Test that we can get the value - let value = test_bind_address.get(); - assert!(value.is_some()); - assert_eq!(value.unwrap(), "127.0.0.1:9090"); - - // Test that we can't set it again - let result2 = test_bind_address.set("different:port".to_string()); - assert!(result2.is_err()); - } - - #[test] - fn test_index_html_constant() { - // Test that INDEX_HTML is not empty and contains expected content - assert!(INDEX_HTML.contains("html") || INDEX_HTML.contains("HTML")); - } - - #[test] - fn test_production_requires_strong_oauth_approval_secret() { + fn production_requires_strong_approval_secret() { assert!(approval_secret_for_mode(false, None).is_err()); assert!(approval_secret_for_mode(false, Some("too-short".to_string())).is_err()); - let configured = "resource-owner-secret-with-32-characters".to_string(); assert_eq!( approval_secret_for_mode(false, Some(configured.clone())).unwrap(), @@ -299,357 +209,35 @@ mod tests { } #[test] - fn test_development_generates_oauth_approval_secret() { - let generated = approval_secret_for_mode(true, None).unwrap(); - assert_eq!(generated.len(), 32); - } - - #[tokio::test] - async fn test_log_request_middleware_functionality() { - // Test basic properties of log_request function - // Since it requires complex setup with actual middleware, - // we focus on testing the types and structure - - let request = Request::builder() - .method(Method::GET) - .uri("/test") - .body(Body::empty()) - .unwrap(); - - // Verify request properties that log_request would process - assert_eq!(request.method(), Method::GET); - assert_eq!(request.uri().path(), "/test"); - assert!(request.headers().is_empty()); - } - - #[test] - fn test_module_imports() { - // Test that our modules are properly imported and accessible - let _server = BashServer::new(); - let _store = McpOAuthStore::new(); - - // Test config module - let config_result = config::Config::read_config("nonexistent.toml"); - assert!(config_result.is_err()); // Should fail gracefully - } - - #[test] - fn test_error_handling_types() { - // Test that Result type is properly used - let test_result: Result = Ok("test".to_string()); - assert!(test_result.is_ok()); - - let test_error: Result = Err(anyhow::anyhow!("test error")); - assert!(test_error.is_err()); - } - - #[test] - fn test_dependencies_availability() { - // Test that critical dependencies are available - use std::sync::Arc; - - let _arc_store = Arc::new(McpOAuthStore::new()); - - // Test that we can create basic types - let _socket_addr: Result = "127.0.0.1:8080".parse(); - } - - // ========== OAuth Mock Tests ========== - - #[tokio::test] - async fn test_oauth_store_functionality() { - use common::oauth::{AuthToken, McpOAuthStore}; - use oauth2::{AccessToken, EmptyExtraTokenFields}; - - let store = McpOAuthStore::new(); - - // Test client validation - let client = store - .validate_client("mcp-client", "http://localhost:8080/callback") - .await; - assert!(client.is_some()); - - let invalid_client = store - .validate_client("invalid-client", "http://localhost:8080/callback") - .await; - assert!(invalid_client.is_none()); - - // Test auth session creation - let session_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile email".to_string()), - Some("test-state".to_string()), - "test-session-123".to_string(), - ) - .await; - - assert_eq!(session_id, "test-session-123"); - - // Test token update and MCP token creation - let auth_token = AuthToken::new( - AccessToken::new("mock-third-party-token".to_string()), - oauth2::basic::BasicTokenType::Bearer, - EmptyExtraTokenFields {}, - ); - - let update_result = store - .update_auth_session_token(&session_id, auth_token) - .await; - assert!(update_result.is_ok()); - - let mcp_token = store.create_mcp_token(&session_id).await; - assert!(mcp_token.is_ok()); - - let token = mcp_token.unwrap(); - assert!(token.access_token.starts_with("mcp-token-")); - assert_eq!(token.client_id, "mcp-client"); - - // Test token validation - let validated = store.validate_token(&token.access_token).await; - assert!(validated.is_some()); - } - - #[tokio::test] - async fn test_oauth_authorization_flow_mock() { - use common::oauth::{AuthorizeQuery, McpOAuthStore}; - use std::sync::Arc; - - let store = Arc::new(McpOAuthStore::new()); - - // Mock authorization request - let auth_query = AuthorizeQuery { - response_type: "code".to_string(), - client_id: "mcp-client".to_string(), - redirect_uri: "http://localhost:8080/callback".to_string(), - scope: Some("profile email".to_string()), - state: Some("test-state-456".to_string()), - }; - - // Test that oauth_authorize function can be called - // Note: In a real test, we'd use test frameworks like tower::ServiceExt - // but here we're testing the basic functionality - let store_clone = store.clone(); - let sessions_before = store_clone.auth_sessions.read().await.len(); - - // Verify store is accessible and functional - assert_eq!(sessions_before, 0); - - // Test client validation within the flow - let client_validation = store - .validate_client(&auth_query.client_id, &auth_query.redirect_uri) - .await; - assert!(client_validation.is_some()); - } - - #[tokio::test] - async fn test_oauth_token_exchange_mock() { - use common::oauth::AuthToken; - use common::oauth::{McpOAuthStore, TokenRequest}; - use oauth2::{AccessToken, EmptyExtraTokenFields}; - use std::sync::Arc; - - let store = Arc::new(McpOAuthStore::new()); - - // Create a session and add auth token (simulating successful OAuth flow) - let session_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile".to_string()), - Some("test-state".to_string()), - "token-exchange-session".to_string(), + fn production_requires_explicit_https_public_base_url() { + assert!(public_base_url_for_mode(false, None, "0.0.0.0:4000").is_err()); + assert!( + public_base_url_for_mode( + false, + Some("http://mcp.example".to_string()), + "0.0.0.0:4000" ) - .await; - - let auth_token = AuthToken::new( - AccessToken::new("mock-external-token".to_string()), - oauth2::basic::BasicTokenType::Bearer, - EmptyExtraTokenFields {}, + .is_err() ); - - store - .update_auth_session_token(&session_id, auth_token) - .await - .unwrap(); - - // Mock token request (just for structure validation) - let _token_request = TokenRequest { - grant_type: "authorization_code".to_string(), - code: "mock-auth-code".to_string(), - client_id: "mcp-client".to_string(), - client_secret: "mcp-client-secret".to_string(), - redirect_uri: "http://localhost:8080/callback".to_string(), - code_verifier: None, - refresh_token: "".to_string(), - }; - - // Test token creation - let mcp_token_result = store.create_mcp_token(&session_id).await; - assert!(mcp_token_result.is_ok()); - - let mcp_token = mcp_token_result.unwrap(); - assert_eq!(mcp_token.token_type, "bearer"); - assert_eq!(mcp_token.expires_in, Some(3600)); - assert!(mcp_token.refresh_token.is_some()); - - // Verify token can be validated - let validation_result = store.validate_token(&mcp_token.access_token).await; - assert!(validation_result.is_some()); - } - - #[tokio::test] - async fn test_oauth_middleware_functionality() { - use common::oauth::AuthToken; - use common::oauth::McpOAuthStore; - use oauth2::{AccessToken, EmptyExtraTokenFields}; - use std::sync::Arc; - - let store = Arc::new(McpOAuthStore::new()); - - // Create a valid token for middleware testing - let session_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile".to_string()), - None, - "middleware-test-session".to_string(), + assert_eq!( + public_base_url_for_mode( + false, + Some("https://mcp.example".to_string()), + "0.0.0.0:4000" ) - .await; - - let auth_token = AuthToken::new( - AccessToken::new("middleware-test-token".to_string()), - oauth2::basic::BasicTokenType::Bearer, - EmptyExtraTokenFields {}, + .unwrap() + .as_str(), + "https://mcp.example/" ); - - store - .update_auth_session_token(&session_id, auth_token) - .await - .unwrap(); - let mcp_token = store.create_mcp_token(&session_id).await.unwrap(); - - // Test token validation (simulating middleware behavior) - let valid_token_check = store.validate_token(&mcp_token.access_token).await; - assert!(valid_token_check.is_some()); - - // Test invalid token - let invalid_token_check = store.validate_token("invalid-token-12345").await; - assert!(invalid_token_check.is_none()); - - // Test empty token - let empty_token_check = store.validate_token("").await; - assert!(empty_token_check.is_none()); } - #[tokio::test] - async fn test_oauth_error_handling() { - use common::oauth::McpOAuthStore; - use std::sync::Arc; - - let store = Arc::new(McpOAuthStore::new()); - - // Test creating MCP token without session - let no_session_result = store.create_mcp_token("nonexistent-session").await; - assert!(no_session_result.is_err()); - assert_eq!(no_session_result.unwrap_err(), "Session not found"); - - // Test creating MCP token without auth token in session - let session_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile".to_string()), - None, - "no-auth-token-session".to_string(), - ) - .await; - - let no_auth_token_result = store.create_mcp_token(&session_id).await; - assert!(no_auth_token_result.is_err()); + #[test] + fn development_uses_safe_loopback_metadata_address() { assert_eq!( - no_auth_token_result.unwrap_err(), - "No third-party token available for session" - ); - - // Test updating nonexistent session - let auth_token = oauth2::StandardTokenResponse::new( - oauth2::AccessToken::new("test-token".to_string()), - oauth2::basic::BasicTokenType::Bearer, - oauth2::EmptyExtraTokenFields {}, + public_base_url_for_mode(true, None, "0.0.0.0:4000") + .unwrap() + .as_str(), + "http://127.0.0.1:4000/" ); - - let update_nonexistent = store - .update_auth_session_token("nonexistent", auth_token) - .await; - assert!(update_nonexistent.is_err()); - assert_eq!(update_nonexistent.unwrap_err(), "Session not found"); - } - - #[tokio::test] - async fn test_oauth_security_validations() { - use common::oauth::McpOAuthStore; - use std::sync::Arc; - - let store = Arc::new(McpOAuthStore::new()); - - // Test invalid client ID - let invalid_client = store - .validate_client("malicious-client", "http://localhost:8080/callback") - .await; - assert!(invalid_client.is_none()); - - // Test invalid redirect URI (potential open redirect attack) - let malicious_redirect = store - .validate_client("mcp-client", "http://evil.com/steal-tokens") - .await; - assert!(malicious_redirect.is_none()); - - // Test valid client with valid redirect URI - let valid_client = store - .validate_client("mcp-client", "http://localhost:8080/callback") - .await; - assert!(valid_client.is_some()); - - // Test that tokens are properly random and unique - let session1_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile".to_string()), - None, - "security-test-1".to_string(), - ) - .await; - - let session2_id = store - .create_auth_session( - "mcp-client".to_string(), - Some("profile".to_string()), - None, - "security-test-2".to_string(), - ) - .await; - - // Add auth tokens to both sessions - for (i, session_id) in [&session1_id, &session2_id].iter().enumerate() { - let auth_token = oauth2::StandardTokenResponse::new( - oauth2::AccessToken::new(format!("security-token-{i}")), - oauth2::basic::BasicTokenType::Bearer, - oauth2::EmptyExtraTokenFields {}, - ); - store - .update_auth_session_token(session_id, auth_token) - .await - .unwrap(); - } - - let token1 = store.create_mcp_token(&session1_id).await.unwrap(); - let token2 = store.create_mcp_token(&session2_id).await.unwrap(); - - // Tokens should be different - assert_ne!(token1.access_token, token2.access_token); - assert_ne!(token1.refresh_token, token2.refresh_token); - - // Both should be valid - assert!(store.validate_token(&token1.access_token).await.is_some()); - assert!(store.validate_token(&token2.access_token).await.is_some()); } } diff --git a/mcp-servers/mcp-bash-server/templates/mcp_oauth_authorize.html b/mcp-servers/mcp-bash-server/templates/mcp_oauth_authorize.html index b51483b45b9..0944846d5dc 100644 --- a/mcp-servers/mcp-bash-server/templates/mcp_oauth_authorize.html +++ b/mcp-servers/mcp-bash-server/templates/mcp_oauth_authorize.html @@ -129,10 +129,8 @@

MCP OAuth

- - - - + +
From ac02a5db87752b754d8532737a842ac29f342571 Mon Sep 17 00:00:00 2001 From: Logic Date: Thu, 30 Jul 2026 21:03:20 +0800 Subject: [PATCH 3/5] make Rust workflow self-contained --- .github/workflows/mcp-bashserver-test.yml | 10 +- mcp-servers/mcp-bash-server/README.md | 9 +- .../mcp-bash-server/src/common/oauth.rs | 153 +++++++++++++++++- 3 files changed, 160 insertions(+), 12 deletions(-) diff --git a/.github/workflows/mcp-bashserver-test.yml b/.github/workflows/mcp-bashserver-test.yml index e18bada8e36..3f7f34949a1 100644 --- a/.github/workflows/mcp-bashserver-test.yml +++ b/.github/workflows/mcp-bashserver-test.yml @@ -44,10 +44,12 @@ jobs: uses: actions/checkout@v4 - name: Setup Rust toolchain - uses: dtolnay/rust-toolchain@4be7066ada62dd38de10e7b70166bc74ed198c30 # stable - with: - toolchain: ${{ env.RUST_VERSION }} - components: rustfmt, clippy + run: | + rustup toolchain install "$RUST_VERSION" \ + --profile minimal \ + --component rustfmt \ + --component clippy + rustup default "$RUST_VERSION" - name: Cache cargo registry uses: actions/cache@v4 diff --git a/mcp-servers/mcp-bash-server/README.md b/mcp-servers/mcp-bash-server/README.md index fd339514c83..408aedbbf06 100644 --- a/mcp-servers/mcp-bash-server/README.md +++ b/mcp-servers/mcp-bash-server/README.md @@ -37,8 +37,13 @@ Dynamic registration supports public clients (`token_endpoint_auth_method: none`) and confidential clients (`client_secret_post`). Authorization requests must use PKCE S256. Authorization transactions and codes are one-time and short-lived; access tokens expire after one hour, and refresh tokens expire -after one day and rotate on every use. OAuth form and JSON bodies are limited -to 16 KiB. +after one day and rotate on every use. Open client registration is limited to +16 successful registrations per minute. An unused registered client expires +after one hour; a successful authorization-code or refresh-token exchange +renews that idle period. Expired clients are removed before the 1,024-client +capacity check, so anonymous registration cannot fill the client store +permanently. A client that receives `invalid_client` after an idle period must +dynamically register again. OAuth form and JSON bodies are limited to 16 KiB. For information on how to use the modelcontextprotocol/inspector tool, refer to the [inspector documentation](https://github.com/modelcontextprotocol/inspector). diff --git a/mcp-servers/mcp-bash-server/src/common/oauth.rs b/mcp-servers/mcp-bash-server/src/common/oauth.rs index 1acf7535284..5294aac619c 100644 --- a/mcp-servers/mcp-bash-server/src/common/oauth.rs +++ b/mcp-servers/mcp-bash-server/src/common/oauth.rs @@ -23,10 +23,11 @@ //! it still enforces the protocol properties on which bearer-token safety //! depends: registered redirect URIs, PKCE S256, one-time authorization and //! consent transactions, expiring codes and tokens, refresh-token rotation, -//! bounded request bodies, and a configured public issuer URL. +//! bounded request bodies and registration state, and a configured public +//! issuer URL. use std::{ - collections::{HashMap, HashSet}, + collections::{HashMap, HashSet, VecDeque}, sync::Arc, }; @@ -35,7 +36,7 @@ use axum::{ Json, body::Body, extract::{Query, State}, - http::{Request, StatusCode}, + http::{HeaderValue, Request, StatusCode, header::RETRY_AFTER}, middleware::Next, response::{Html, IntoResponse, Redirect, Response}, }; @@ -49,7 +50,7 @@ use rmcp::transport::auth::{ use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; -use tokio::sync::RwLock; +use tokio::sync::{Mutex, RwLock}; use tracing::{debug, info, warn}; use url::Url; use uuid::Uuid; @@ -59,6 +60,9 @@ const AUTH_TRANSACTION_TTL_SECONDS: i64 = 5 * 60; const AUTHORIZATION_CODE_TTL_SECONDS: i64 = 2 * 60; const ACCESS_TOKEN_TTL_SECONDS: i64 = 60 * 60; const REFRESH_TOKEN_TTL_SECONDS: i64 = 24 * 60 * 60; +const REGISTERED_CLIENT_IDLE_TTL_SECONDS: i64 = 60 * 60; +const REGISTRATION_RATE_WINDOW_SECONDS: i64 = 60; +const MAX_REGISTRATIONS_PER_WINDOW: usize = 16; const MAX_CLIENTS: usize = 1_024; const MAX_AUTH_TRANSACTIONS: usize = 4_096; const MAX_AUTHORIZATION_CODES: usize = 4_096; @@ -95,6 +99,7 @@ struct RegisteredClient { client_secret: Option, redirect_uris: Vec, token_endpoint_auth_method: TokenEndpointAuthMethod, + expires_at: DateTime, } #[derive(Clone, Debug)] @@ -153,6 +158,7 @@ pub struct McpOAuthStore { authorization_codes: Arc>>, access_tokens: Arc>>, refresh_tokens: Arc>>, + registration_attempts: Arc>>>, approval_secret: Arc, public_base_url: Arc, } @@ -167,6 +173,7 @@ impl McpOAuthStore { authorization_codes: Arc::new(RwLock::new(HashMap::new())), access_tokens: Arc::new(RwLock::new(HashMap::new())), refresh_tokens: Arc::new(RwLock::new(HashMap::new())), + registration_attempts: Arc::new(Mutex::new(VecDeque::new())), approval_secret: Arc::new(approval_secret), public_base_url: Arc::new(public_base_url), } @@ -186,6 +193,7 @@ impl McpOAuthStore { client_secret: None, redirect_uris: vec!["http://127.0.0.1:8080/callback".to_string()], token_endpoint_auth_method: TokenEndpointAuthMethod::None, + expires_at: Utc::now() + Duration::seconds(REGISTERED_CLIENT_IDLE_TTL_SECONDS), }, ); clients.insert( @@ -195,6 +203,7 @@ impl McpOAuthStore { client_secret: Some("test-only-confidential-secret".to_string()), redirect_uris: vec!["https://client.example/callback".to_string()], token_endpoint_auth_method: TokenEndpointAuthMethod::ClientSecretPost, + expires_at: Utc::now() + Duration::seconds(REGISTERED_CLIENT_IDLE_TTL_SECONDS), }, ); store @@ -213,7 +222,9 @@ impl McpOAuthStore { client_id: &str, redirect_uri: &str, ) -> Option { - let clients = self.clients.read().await; + let now = Utc::now(); + let mut clients = self.clients.write().await; + clients.retain(|_, client| client.expires_at > now); let client = clients.get(client_id)?; client .redirect_uris @@ -227,7 +238,9 @@ impl McpOAuthStore { client_id: &str, client_secret: &str, ) -> Option { - let clients = self.clients.read().await; + let now = Utc::now(); + let mut clients = self.clients.write().await; + clients.retain(|_, client| client.expires_at > now); let client = clients.get(client_id)?; match client.token_endpoint_auth_method { TokenEndpointAuthMethod::None => client_secret.is_empty().then(|| client.clone()), @@ -239,6 +252,29 @@ impl McpOAuthStore { } } + async fn allow_registration(&self, now: DateTime) -> bool { + let window_start = now - Duration::seconds(REGISTRATION_RATE_WINDOW_SECONDS); + let mut attempts = self.registration_attempts.lock().await; + while attempts + .front() + .is_some_and(|attempt| *attempt <= window_start) + { + attempts.pop_front(); + } + if attempts.len() >= MAX_REGISTRATIONS_PER_WINDOW { + return false; + } + attempts.push_back(now); + true + } + + async fn refresh_client_expiry(&self, client_id: &str, now: DateTime) { + let mut clients = self.clients.write().await; + if let Some(client) = clients.get_mut(client_id) { + client.expires_at = now + Duration::seconds(REGISTERED_CLIENT_IDLE_TTL_SECONDS); + } + } + async fn create_authorization_transaction( &self, params: &AuthorizeQuery, @@ -377,6 +413,7 @@ impl McpOAuthStore { .ok_or(OAuthError::InvalidGrant)?; drop(codes); + self.refresh_client_expiry(&code.client_id, now).await; self.issue_tokens(code.client_id, code.scope).await } @@ -403,6 +440,7 @@ impl McpOAuthStore { .ok_or(OAuthError::InvalidGrant)?; drop(refresh_tokens); + self.refresh_client_expiry(&record.client_id, now).await; self.issue_tokens(record.client_id, record.scope).await } @@ -467,6 +505,7 @@ enum OAuthError { InvalidClient, InvalidGrant, InvalidTransaction, + RateLimited, TemporarilyUnavailable, } @@ -491,6 +530,17 @@ impl OAuthError { "invalid_request", "authorization transaction is invalid, expired, or already used", ), + Self::RateLimited => { + let mut response = oauth_json_error( + StatusCode::TOO_MANY_REQUESTS, + "temporarily_unavailable", + "client registration rate limit exceeded", + ); + response + .headers_mut() + .insert(RETRY_AFTER, HeaderValue::from_static("60")); + response + } Self::TemporarilyUnavailable => oauth_json_error( StatusCode::SERVICE_UNAVAILABLE, "temporarily_unavailable", @@ -961,7 +1011,12 @@ pub async fn oauth_register( let client_id = random_prefixed("client"); let client_secret = (auth_method == TokenEndpointAuthMethod::ClientSecretPost) .then(|| generate_random_string(48)); + let now = Utc::now(); + if !state.allow_registration(now).await { + return OAuthError::RateLimited.response(); + } let mut clients = state.clients.write().await; + clients.retain(|_, client| client.expires_at > now); if clients.len() >= MAX_CLIENTS { return OAuthError::TemporarilyUnavailable.response(); } @@ -978,6 +1033,7 @@ pub async fn oauth_register( client_secret: client_secret.clone(), redirect_uris: unique_redirects.clone(), token_endpoint_auth_method: auth_method.clone(), + expires_at: now + Duration::seconds(REGISTERED_CLIENT_IDLE_TTL_SECONDS), }, ); drop(clients); @@ -1415,6 +1471,91 @@ mod tests { assert_eq!(registered.client_secret.unwrap().len(), 48); } + #[tokio::test] + async fn dynamic_registration_is_rate_limited_before_client_capacity_is_exhausted() { + assert!( + MAX_REGISTRATIONS_PER_WINDOW + * ((REGISTERED_CLIENT_IDLE_TTL_SECONDS / REGISTRATION_RATE_WINDOW_SECONDS) + as usize) + < MAX_CLIENTS + ); + let store = Arc::new(McpOAuthStore::new()); + let registration = ClientRegistrationRequest { + client_name: "bounded-client".to_string(), + redirect_uris: vec!["http://127.0.0.1:9911/callback".to_string()], + grant_types: vec!["authorization_code".to_string()], + token_endpoint_auth_method: "none".to_string(), + response_types: vec!["code".to_string()], + }; + + for _ in 0..MAX_REGISTRATIONS_PER_WINDOW { + let response = oauth_register( + State(store.clone()), + Request::builder() + .body(Body::from(serde_json::to_vec(®istration).unwrap())) + .unwrap(), + ) + .await; + assert_eq!(response.status(), StatusCode::CREATED); + } + + let limited = oauth_register( + State(store), + Request::builder() + .body(Body::from(serde_json::to_vec(®istration).unwrap())) + .unwrap(), + ) + .await; + assert_eq!(limited.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!( + limited.headers().get(RETRY_AFTER), + Some(&HeaderValue::from_static("60")) + ); + } + + #[tokio::test] + async fn expired_dynamic_clients_are_pruned_before_the_capacity_check() { + let store = Arc::new(McpOAuthStore::new()); + let mut clients = store.clients.write().await; + let template = clients.get("public-test-client").unwrap().clone(); + clients.clear(); + for index in 0..MAX_CLIENTS { + clients.insert( + format!("expired-{index}"), + RegisteredClient { + client_id: format!("expired-{index}"), + expires_at: Utc::now() - Duration::seconds(1), + ..template.clone() + }, + ); + } + drop(clients); + + let registration = ClientRegistrationRequest { + client_name: "replacement-client".to_string(), + redirect_uris: vec!["http://127.0.0.1:9911/callback".to_string()], + grant_types: vec!["authorization_code".to_string()], + token_endpoint_auth_method: "none".to_string(), + response_types: vec!["code".to_string()], + }; + let response = oauth_register( + State(store.clone()), + Request::builder() + .body(Body::from(serde_json::to_vec(®istration).unwrap())) + .unwrap(), + ) + .await; + + assert_eq!(response.status(), StatusCode::CREATED); + let clients = store.clients.read().await; + assert_eq!(clients.len(), 1); + assert!( + clients + .values() + .all(|client| client.expires_at > Utc::now()) + ); + } + #[test] fn redirect_and_public_base_urls_have_safe_schemes_and_shapes() { assert!(validate_redirect_uri("https://client.example/callback")); From a6429a297be063d58a115bd975125a4cc3a9a30c Mon Sep 17 00:00:00 2001 From: Logic Date: Fri, 31 Jul 2026 08:49:21 +0800 Subject: [PATCH 4/5] maintenance: align OAuth client lifecycle --- mcp-servers/mcp-bash-server/Cargo.lock | 1 + mcp-servers/mcp-bash-server/Cargo.toml | 1 + mcp-servers/mcp-bash-server/README.md | 23 +- .../mcp-bash-server/src/common/oauth.rs | 351 ++++++++++++++++-- mcp-servers/mcp-bash-server/src/main.rs | 19 +- 5 files changed, 351 insertions(+), 44 deletions(-) diff --git a/mcp-servers/mcp-bash-server/Cargo.lock b/mcp-servers/mcp-bash-server/Cargo.lock index bbf0d6d796c..d8d638e586a 100644 --- a/mcp-servers/mcp-bash-server/Cargo.lock +++ b/mcp-servers/mcp-bash-server/Cargo.lock @@ -1065,6 +1065,7 @@ dependencies = [ "base64", "chrono", "hyper", + "ipnet", "rand 0.8.5", "regex", "rmcp", diff --git a/mcp-servers/mcp-bash-server/Cargo.toml b/mcp-servers/mcp-bash-server/Cargo.toml index 5863483a5e1..b5b1ed38bb8 100644 --- a/mcp-servers/mcp-bash-server/Cargo.toml +++ b/mcp-servers/mcp-bash-server/Cargo.toml @@ -46,6 +46,7 @@ regex = "1.11.1" shlex = "1.3" subtle = "2.6" url = "2.5" +ipnet = "2.11" [dev-dependencies] tokio-test = "0.4" diff --git a/mcp-servers/mcp-bash-server/README.md b/mcp-servers/mcp-bash-server/README.md index 408aedbbf06..2171878638c 100644 --- a/mcp-servers/mcp-bash-server/README.md +++ b/mcp-servers/mcp-bash-server/README.md @@ -37,13 +37,22 @@ Dynamic registration supports public clients (`token_endpoint_auth_method: none`) and confidential clients (`client_secret_post`). Authorization requests must use PKCE S256. Authorization transactions and codes are one-time and short-lived; access tokens expire after one hour, and refresh tokens expire -after one day and rotate on every use. Open client registration is limited to -16 successful registrations per minute. An unused registered client expires -after one hour; a successful authorization-code or refresh-token exchange -renews that idle period. Expired clients are removed before the 1,024-client -capacity check, so anonymous registration cannot fill the client store -permanently. A client that receives `invalid_client` after an idle period must -dynamically register again. OAuth form and JSON bodies are limited to 16 KiB. +after one day and rotate on every use. A client that holds a refresh token +remains registered for at least that token's full lifetime. Open client +registration is limited to 16 successful registrations per minute for each +TCP peer, so one source cannot consume every caller's admission window. The +rate-limit source table is bounded. Deployments behind a reverse proxy can set +`MCP_OAUTH_TRUSTED_PROXY_CIDRS` to a comma-separated list of the exact proxy +networks. Only connections from those networks may supply `X-Forwarded-For`; +untrusted peers cannot spoof the limiter identity, and `/0` networks are +rejected. The proxy must overwrite or safely append that header. An unused +registered client expires after one hour; expired clients are removed before +the 1,024-client capacity check. +When anonymous registrations fill the remaining capacity, the oldest client +that has never received a refresh token is reclaimed; clients with live +refresh credentials are not evicted. A client that receives `invalid_client` +after an unused registration expires must dynamically register again. OAuth +form and JSON bodies are limited to 16 KiB. For information on how to use the modelcontextprotocol/inspector tool, refer to the [inspector documentation](https://github.com/modelcontextprotocol/inspector). diff --git a/mcp-servers/mcp-bash-server/src/common/oauth.rs b/mcp-servers/mcp-bash-server/src/common/oauth.rs index 5294aac619c..f229bcd60cf 100644 --- a/mcp-servers/mcp-bash-server/src/common/oauth.rs +++ b/mcp-servers/mcp-bash-server/src/common/oauth.rs @@ -28,6 +28,7 @@ use std::{ collections::{HashMap, HashSet, VecDeque}, + net::{IpAddr, Ipv4Addr, SocketAddr}, sync::Arc, }; @@ -35,13 +36,14 @@ use askama::Template; use axum::{ Json, body::Body, - extract::{Query, State}, + extract::{ConnectInfo, Query, State}, http::{HeaderValue, Request, StatusCode, header::RETRY_AFTER}, middleware::Next, response::{Html, IntoResponse, Redirect, Response}, }; use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use chrono::{DateTime, Duration, Utc}; +use ipnet::IpNet; use rand::{Rng, distributions::Alphanumeric}; use rmcp::serde_json::{self, Value}; use rmcp::transport::auth::{ @@ -63,6 +65,8 @@ const REFRESH_TOKEN_TTL_SECONDS: i64 = 24 * 60 * 60; const REGISTERED_CLIENT_IDLE_TTL_SECONDS: i64 = 60 * 60; const REGISTRATION_RATE_WINDOW_SECONDS: i64 = 60; const MAX_REGISTRATIONS_PER_WINDOW: usize = 16; +const MAX_REGISTRATION_SOURCES: usize = 4_096; +const MAX_FORWARDED_FOR_HOPS: usize = 32; const MAX_CLIENTS: usize = 1_024; const MAX_AUTH_TRANSACTIONS: usize = 4_096; const MAX_AUTHORIZATION_CODES: usize = 4_096; @@ -99,7 +103,9 @@ struct RegisteredClient { client_secret: Option, redirect_uris: Vec, token_endpoint_auth_method: TokenEndpointAuthMethod, + registered_at: DateTime, expires_at: DateTime, + refresh_expires_at: Option>, } #[derive(Clone, Debug)] @@ -158,7 +164,8 @@ pub struct McpOAuthStore { authorization_codes: Arc>>, access_tokens: Arc>>, refresh_tokens: Arc>>, - registration_attempts: Arc>>>, + registration_attempts: Arc>>>>, + trusted_proxies: Arc>, approval_secret: Arc, public_base_url: Arc, } @@ -166,14 +173,19 @@ pub struct McpOAuthStore { impl McpOAuthStore { /// Create the production store. Clients are registered dynamically; there /// is intentionally no repository-known default confidential credential. - pub fn with_settings(approval_secret: String, public_base_url: Url) -> Self { + pub fn with_trusted_proxies( + approval_secret: String, + public_base_url: Url, + trusted_proxies: Vec, + ) -> Self { Self { clients: Arc::new(RwLock::new(HashMap::new())), auth_transactions: Arc::new(RwLock::new(HashMap::new())), authorization_codes: Arc::new(RwLock::new(HashMap::new())), access_tokens: Arc::new(RwLock::new(HashMap::new())), refresh_tokens: Arc::new(RwLock::new(HashMap::new())), - registration_attempts: Arc::new(Mutex::new(VecDeque::new())), + registration_attempts: Arc::new(Mutex::new(HashMap::new())), + trusted_proxies: Arc::new(trusted_proxies), approval_secret: Arc::new(approval_secret), public_base_url: Arc::new(public_base_url), } @@ -181,11 +193,13 @@ impl McpOAuthStore { #[cfg(test)] fn new() -> Self { - let mut store = Self::with_settings( + let mut store = Self::with_trusted_proxies( generate_random_string(32), Url::parse("http://127.0.0.1:4000/").unwrap(), + Vec::new(), ); let clients = Arc::get_mut(&mut store.clients).unwrap().get_mut(); + let now = Utc::now(); clients.insert( "public-test-client".to_string(), RegisteredClient { @@ -193,7 +207,9 @@ impl McpOAuthStore { client_secret: None, redirect_uris: vec!["http://127.0.0.1:8080/callback".to_string()], token_endpoint_auth_method: TokenEndpointAuthMethod::None, - expires_at: Utc::now() + Duration::seconds(REGISTERED_CLIENT_IDLE_TTL_SECONDS), + registered_at: now, + expires_at: now + Duration::seconds(REGISTERED_CLIENT_IDLE_TTL_SECONDS), + refresh_expires_at: None, }, ); clients.insert( @@ -203,7 +219,9 @@ impl McpOAuthStore { client_secret: Some("test-only-confidential-secret".to_string()), redirect_uris: vec!["https://client.example/callback".to_string()], token_endpoint_auth_method: TokenEndpointAuthMethod::ClientSecretPost, - expires_at: Utc::now() + Duration::seconds(REGISTERED_CLIENT_IDLE_TTL_SECONDS), + registered_at: now, + expires_at: now + Duration::seconds(REGISTERED_CLIENT_IDLE_TTL_SECONDS), + refresh_expires_at: None, }, ); store @@ -252,27 +270,78 @@ impl McpOAuthStore { } } - async fn allow_registration(&self, now: DateTime) -> bool { + async fn allow_registration(&self, source: IpAddr, now: DateTime) -> bool { let window_start = now - Duration::seconds(REGISTRATION_RATE_WINDOW_SECONDS); - let mut attempts = self.registration_attempts.lock().await; - while attempts - .front() - .is_some_and(|attempt| *attempt <= window_start) + let mut attempts_by_source = self.registration_attempts.lock().await; + attempts_by_source.retain(|_, attempts| { + while attempts + .front() + .is_some_and(|attempt| *attempt <= window_start) + { + attempts.pop_front(); + } + !attempts.is_empty() + }); + if !attempts_by_source.contains_key(&source) + && attempts_by_source.len() >= MAX_REGISTRATION_SOURCES { - attempts.pop_front(); + let oldest_source = attempts_by_source + .iter() + .min_by_key(|(_, attempts)| attempts.back().copied()) + .map(|(source, _)| *source); + if let Some(oldest_source) = oldest_source { + attempts_by_source.remove(&oldest_source); + } } - if attempts.len() >= MAX_REGISTRATIONS_PER_WINDOW { + let source_attempts = attempts_by_source.entry(source).or_default(); + if source_attempts.len() >= MAX_REGISTRATIONS_PER_WINDOW { return false; } - attempts.push_back(now); + source_attempts.push_back(now); true } - async fn refresh_client_expiry(&self, client_id: &str, now: DateTime) { - let mut clients = self.clients.write().await; - if let Some(client) = clients.get_mut(client_id) { - client.expires_at = now + Duration::seconds(REGISTERED_CLIENT_IDLE_TTL_SECONDS); + fn registration_source(&self, request: &Request) -> IpAddr { + let peer = request + .extensions() + .get::>() + .map(|connect_info| connect_info.0.ip()) + .unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED)); + if !self + .trusted_proxies + .iter() + .any(|trusted_proxy| trusted_proxy.contains(&peer)) + { + return peer; + } + let Some(forwarded_for) = request + .headers() + .get("x-forwarded-for") + .and_then(|value| value.to_str().ok()) + else { + return peer; + }; + let hops: Vec<&str> = forwarded_for.split(',').map(str::trim).collect(); + if hops.is_empty() || hops.len() > MAX_FORWARDED_FOR_HOPS { + return peer; } + let mut parsed_hops = Vec::with_capacity(hops.len()); + for hop in hops { + let Ok(hop) = hop.parse::() else { + return peer; + }; + parsed_hops.push(hop); + } + parsed_hops + .into_iter() + .rev() + .find(|hop| { + !self + .trusted_proxies + .iter() + .any(|trusted_proxy| trusted_proxy.contains(hop)) + }) + .unwrap_or(peer) } async fn create_authorization_transaction( @@ -413,7 +482,6 @@ impl McpOAuthStore { .ok_or(OAuthError::InvalidGrant)?; drop(codes); - self.refresh_client_expiry(&code.client_id, now).await; self.issue_tokens(code.client_id, code.scope).await } @@ -440,7 +508,6 @@ impl McpOAuthStore { .ok_or(OAuthError::InvalidGrant)?; drop(refresh_tokens); - self.refresh_client_expiry(&record.client_id, now).await; self.issue_tokens(record.client_id, record.scope).await } @@ -452,6 +519,7 @@ impl McpOAuthStore { let now = Utc::now(); let access_token_value = random_prefixed("mcp-token"); let refresh_token_value = random_prefixed("mcp-refresh"); + let refresh_expires_at = now + Duration::seconds(REFRESH_TOKEN_TTL_SECONDS); let access = McpAccessToken { access_token: access_token_value.clone(), token_type: "bearer".to_string(), @@ -463,6 +531,11 @@ impl McpOAuthStore { expires_at: now + Duration::seconds(ACCESS_TOKEN_TTL_SECONDS), }; + let mut clients = self.clients.write().await; + clients.retain(|_, client| client.expires_at > now); + let client = clients + .get_mut(&client_id) + .ok_or(OAuthError::InvalidClient)?; let mut access_tokens = self.access_tokens.write().await; access_tokens.retain(|_, token| token.expires_at > now); if access_tokens.len() >= MAX_ACCESS_TOKENS { @@ -479,9 +552,11 @@ impl McpOAuthStore { RefreshTokenRecord { client_id, scope, - expires_at: now + Duration::seconds(REFRESH_TOKEN_TTL_SECONDS), + expires_at: refresh_expires_at, }, ); + client.expires_at = refresh_expires_at; + client.refresh_expires_at = Some(refresh_expires_at); Ok(IssuedTokens { access }) } @@ -608,6 +683,30 @@ fn random_prefixed(prefix: &str) -> String { format!("{prefix}-{}-{}", Uuid::new_v4(), generate_random_string(24)) } +pub fn parse_trusted_proxy_cidrs(configured: Option<&str>) -> Result, String> { + let Some(configured) = configured + .map(str::trim) + .filter(|configured| !configured.is_empty()) + else { + return Ok(Vec::new()); + }; + configured + .split(',') + .map(|network| { + let network = network.trim().parse::().map_err(|_| { + "MCP_OAUTH_TRUSTED_PROXY_CIDRS contains an invalid CIDR".to_string() + })?; + if network.prefix_len() == 0 { + return Err( + "MCP_OAUTH_TRUSTED_PROXY_CIDRS must not trust an entire address family" + .to_string(), + ); + } + Ok(network) + }) + .collect() +} + fn validate_scope(scope: Option<&str>) -> Result, OAuthError> { let Some(scope) = scope.map(str::trim).filter(|scope| !scope.is_empty()) else { return Ok(None); @@ -942,6 +1041,7 @@ pub async fn oauth_register( State(state): State>, request: Request, ) -> Response { + let registration_source = state.registration_source(&request); let bytes = match read_limited_body(request).await { Ok(bytes) => bytes, Err(response) => return response, @@ -1012,13 +1112,26 @@ pub async fn oauth_register( let client_secret = (auth_method == TokenEndpointAuthMethod::ClientSecretPost) .then(|| generate_random_string(48)); let now = Utc::now(); - if !state.allow_registration(now).await { + if !state.allow_registration(registration_source, now).await { return OAuthError::RateLimited.response(); } let mut clients = state.clients.write().await; clients.retain(|_, client| client.expires_at > now); if clients.len() >= MAX_CLIENTS { - return OAuthError::TemporarilyUnavailable.response(); + let reclaimable_client = clients + .iter() + .filter(|(_, client)| { + client + .refresh_expires_at + .is_none_or(|expires_at| expires_at <= now) + }) + .min_by_key(|(_, client)| client.registered_at) + .map(|(client_id, _)| client_id.clone()); + if let Some(reclaimable_client) = reclaimable_client { + clients.remove(&reclaimable_client); + } else { + return OAuthError::TemporarilyUnavailable.response(); + } } let mut unique_redirects = Vec::new(); for redirect in &request.redirect_uris { @@ -1033,7 +1146,9 @@ pub async fn oauth_register( client_secret: client_secret.clone(), redirect_uris: unique_redirects.clone(), token_endpoint_auth_method: auth_method.clone(), + registered_at: now, expires_at: now + Duration::seconds(REGISTERED_CLIENT_IDLE_TTL_SECONDS), + refresh_expires_at: None, }, ); drop(clients); @@ -1472,13 +1587,7 @@ mod tests { } #[tokio::test] - async fn dynamic_registration_is_rate_limited_before_client_capacity_is_exhausted() { - assert!( - MAX_REGISTRATIONS_PER_WINDOW - * ((REGISTERED_CLIENT_IDLE_TTL_SECONDS / REGISTRATION_RATE_WINDOW_SECONDS) - as usize) - < MAX_CLIENTS - ); + async fn dynamic_registration_is_rate_limited_per_tcp_peer() { let store = Arc::new(McpOAuthStore::new()); let registration = ClientRegistrationRequest { client_name: "bounded-client".to_string(), @@ -1513,6 +1622,186 @@ mod tests { ); } + #[tokio::test] + async fn registration_rate_limit_is_isolated_by_tcp_peer() { + let store = Arc::new(McpOAuthStore::new()); + let registration = ClientRegistrationRequest { + client_name: "peer-bounded-client".to_string(), + redirect_uris: vec!["http://127.0.0.1:9911/callback".to_string()], + grant_types: vec!["authorization_code".to_string()], + token_endpoint_auth_method: "none".to_string(), + response_types: vec!["code".to_string()], + }; + let first_peer = std::net::SocketAddr::from(([192, 0, 2, 10], 41000)); + let second_peer = std::net::SocketAddr::from(([198, 51, 100, 20], 42000)); + + for _ in 0..MAX_REGISTRATIONS_PER_WINDOW { + let mut request = Request::builder() + .body(Body::from(serde_json::to_vec(®istration).unwrap())) + .unwrap(); + request + .extensions_mut() + .insert(axum::extract::ConnectInfo(first_peer)); + assert_eq!( + oauth_register(State(store.clone()), request).await.status(), + StatusCode::CREATED + ); + } + + let mut first_peer_request = Request::builder() + .body(Body::from(serde_json::to_vec(®istration).unwrap())) + .unwrap(); + first_peer_request + .extensions_mut() + .insert(axum::extract::ConnectInfo(first_peer)); + assert_eq!( + oauth_register(State(store.clone()), first_peer_request) + .await + .status(), + StatusCode::TOO_MANY_REQUESTS + ); + + let mut second_peer_request = Request::builder() + .body(Body::from(serde_json::to_vec(®istration).unwrap())) + .unwrap(); + second_peer_request + .extensions_mut() + .insert(axum::extract::ConnectInfo(second_peer)); + assert_eq!( + oauth_register(State(store), second_peer_request) + .await + .status(), + StatusCode::CREATED + ); + } + + #[tokio::test] + async fn registration_source_limiter_remains_bounded() { + let store = McpOAuthStore::new(); + let now = Utc::now(); + for source in 0..=MAX_REGISTRATION_SOURCES { + assert!( + store + .allow_registration(IpAddr::V6(std::net::Ipv6Addr::from(source as u128)), now,) + .await + ); + } + + assert_eq!( + store.registration_attempts.lock().await.len(), + MAX_REGISTRATION_SOURCES + ); + } + + #[test] + fn registration_source_uses_forwarding_only_from_configured_proxies() { + let trusted_proxies = parse_trusted_proxy_cidrs(Some("10.0.0.0/8,fd00::/8")).unwrap(); + let store = McpOAuthStore::with_trusted_proxies( + generate_random_string(32), + Url::parse("https://mcp.example/").unwrap(), + trusted_proxies, + ); + + let mut trusted_request = Request::builder() + .header("x-forwarded-for", "192.0.2.44, 10.1.0.8") + .body(Body::empty()) + .unwrap(); + trusted_request + .extensions_mut() + .insert(ConnectInfo("10.2.0.9:443".parse::().unwrap())); + assert_eq!( + store.registration_source(&trusted_request), + "192.0.2.44".parse::().unwrap() + ); + + let mut untrusted_request = Request::builder() + .header("x-forwarded-for", "192.0.2.99") + .body(Body::empty()) + .unwrap(); + untrusted_request.extensions_mut().insert(ConnectInfo( + "198.51.100.7:443".parse::().unwrap(), + )); + assert_eq!( + store.registration_source(&untrusted_request), + "198.51.100.7".parse::().unwrap() + ); + } + + #[test] + fn trusted_proxy_configuration_rejects_invalid_or_unbounded_networks() { + assert!(parse_trusted_proxy_cidrs(None).unwrap().is_empty()); + assert!(parse_trusted_proxy_cidrs(Some("127.0.0.1/32")).is_ok()); + assert!(parse_trusted_proxy_cidrs(Some("not-a-cidr")).is_err()); + assert!(parse_trusted_proxy_cidrs(Some("0.0.0.0/0")).is_err()); + assert!(parse_trusted_proxy_cidrs(Some("::/0")).is_err()); + } + + #[tokio::test] + async fn refresh_token_lifetime_keeps_its_registered_client_valid() { + let store = McpOAuthStore::new(); + let issued = store + .issue_tokens("public-test-client".to_string(), Some("email".to_string())) + .await + .unwrap(); + let refresh_token = issued.access.refresh_token.unwrap(); + let refresh_expires_at = store + .refresh_tokens + .read() + .await + .get(&refresh_token) + .unwrap() + .expires_at; + let client_expires_at = store + .clients + .read() + .await + .get("public-test-client") + .unwrap() + .expires_at; + + assert!( + client_expires_at >= refresh_expires_at, + "client registration must remain valid for the full refresh-token lifetime" + ); + } + + #[tokio::test] + async fn registration_capacity_reclaims_an_unactivated_client() { + let store = Arc::new(McpOAuthStore::new()); + let mut clients = store.clients.write().await; + let template = clients.get("public-test-client").unwrap().clone(); + clients.clear(); + for index in 0..MAX_CLIENTS { + clients.insert( + format!("unused-{index}"), + RegisteredClient { + client_id: format!("unused-{index}"), + expires_at: Utc::now() + Duration::seconds(REGISTERED_CLIENT_IDLE_TTL_SECONDS), + ..template.clone() + }, + ); + } + drop(clients); + + let registration = ClientRegistrationRequest { + client_name: "replacement-client".to_string(), + redirect_uris: vec!["http://127.0.0.1:9911/callback".to_string()], + grant_types: vec!["authorization_code".to_string()], + token_endpoint_auth_method: "none".to_string(), + response_types: vec!["code".to_string()], + }; + let response = oauth_register( + State(store.clone()), + Request::builder() + .body(Body::from(serde_json::to_vec(®istration).unwrap())) + .unwrap(), + ) + .await; + + assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(store.clients.read().await.len(), MAX_CLIENTS); + } + #[tokio::test] async fn expired_dynamic_clients_are_pruned_before_the_capacity_check() { let store = Arc::new(McpOAuthStore::new()); diff --git a/mcp-servers/mcp-bash-server/src/main.rs b/mcp-servers/mcp-bash-server/src/main.rs index e3e5c11070f..cd906439502 100644 --- a/mcp-servers/mcp-bash-server/src/main.rs +++ b/mcp-servers/mcp-bash-server/src/main.rs @@ -42,8 +42,8 @@ use common::bash_server::BashServer; use common::config; use common::oauth::{ McpOAuthStore, generate_random_string, oauth_approve, oauth_authorization_server, - oauth_authorize, oauth_register, oauth_token, validate_public_base_url, - validate_token_middleware, + oauth_authorize, oauth_register, oauth_token, parse_trusted_proxy_cidrs, + validate_public_base_url, validate_token_middleware, }; const INDEX_HTML: &str = include_str!("html/mcp_oauth_index.html"); @@ -135,9 +135,13 @@ async fn main() -> Result<()> { std::env::var("MCP_OAUTH_PUBLIC_BASE_URL").ok(), &bind_address, )?; - let oauth_store = Arc::new(McpOAuthStore::with_settings( + let trusted_proxy_config = std::env::var("MCP_OAUTH_TRUSTED_PROXY_CIDRS").ok(); + let trusted_proxies = parse_trusted_proxy_cidrs(trusted_proxy_config.as_deref()) + .map_err(|message| anyhow::anyhow!(message))?; + let oauth_store = Arc::new(McpOAuthStore::with_trusted_proxies( approval_secret, public_base_url, + trusted_proxies, )); let service = StreamableHttpService::new( @@ -180,9 +184,12 @@ async fn main() -> Result<()> { info!("MCP OAuth Server started on {}", addr); let listener = tokio::net::TcpListener::bind(addr).await?; - let _ = axum::serve(listener, app) - .with_graceful_shutdown(async { tokio::signal::ctrl_c().await.unwrap() }) - .await; + let _ = axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(async { tokio::signal::ctrl_c().await.unwrap() }) + .await; Ok(()) } From 9ee8a65a2759993f09cf8810a9c9108cd4cad2e2 Mon Sep 17 00:00:00 2001 From: Logic Date: Tue, 4 Aug 2026 20:34:29 +0800 Subject: [PATCH 5/5] Harden OAuth exchange state handling --- .../mcp-bash-server/src/common/oauth.rs | 426 +++++++++++++++++- 1 file changed, 403 insertions(+), 23 deletions(-) diff --git a/mcp-servers/mcp-bash-server/src/common/oauth.rs b/mcp-servers/mcp-bash-server/src/common/oauth.rs index f229bcd60cf..9a4b18be10b 100644 --- a/mcp-servers/mcp-bash-server/src/common/oauth.rs +++ b/mcp-servers/mcp-bash-server/src/common/oauth.rs @@ -58,6 +58,7 @@ use url::Url; use uuid::Uuid; pub const MAX_OAUTH_BODY_BYTES: usize = 16 * 1024; +const MAX_OAUTH_STATE_BYTES: usize = 2_048; const AUTH_TRANSACTION_TTL_SECONDS: i64 = 5 * 60; const AUTHORIZATION_CODE_TTL_SECONDS: i64 = 2 * 60; const ACCESS_TOKEN_TTL_SECONDS: i64 = 60 * 60; @@ -65,6 +66,7 @@ const REFRESH_TOKEN_TTL_SECONDS: i64 = 24 * 60 * 60; const REGISTERED_CLIENT_IDLE_TTL_SECONDS: i64 = 60 * 60; const REGISTRATION_RATE_WINDOW_SECONDS: i64 = 60; const MAX_REGISTRATIONS_PER_WINDOW: usize = 16; +const MAX_APPROVAL_FAILURES_PER_WINDOW: usize = 8; const MAX_REGISTRATION_SOURCES: usize = 4_096; const MAX_FORWARDED_FOR_HOPS: usize = 32; const MAX_CLIENTS: usize = 1_024; @@ -72,6 +74,7 @@ const MAX_AUTH_TRANSACTIONS: usize = 4_096; const MAX_AUTHORIZATION_CODES: usize = 4_096; const MAX_ACCESS_TOKENS: usize = 4_096; const MAX_REFRESH_TOKENS: usize = 4_096; +const MAX_USED_REFRESH_TOKENS: usize = 4_096; const SUPPORTED_SCOPES: [&str; 2] = ["profile", "email"]; #[derive(Clone, Debug, PartialEq, Eq)] @@ -132,9 +135,23 @@ struct AuthorizationCode { struct RefreshTokenRecord { client_id: String, scope: Option, + family_id: String, + family_expires_at: DateTime, expires_at: DateTime, } +#[derive(Clone, Debug)] +struct UsedRefreshToken { + family_id: String, + expires_at: DateTime, +} + +#[derive(Clone, Debug)] +struct RefreshRotation { + token: String, + record: RefreshTokenRecord, +} + /// Access-token state retained by the authorization server. #[derive(Clone, Debug, Serialize)] pub struct McpAccessToken { @@ -148,6 +165,8 @@ pub struct McpAccessToken { pub issued_at: DateTime, #[serde(skip)] pub expires_at: DateTime, + #[serde(skip)] + token_family_id: String, } #[derive(Clone, Debug)] @@ -164,7 +183,10 @@ pub struct McpOAuthStore { authorization_codes: Arc>>, access_tokens: Arc>>, refresh_tokens: Arc>>, + used_refresh_tokens: Arc>>, + token_issuance: Arc>, registration_attempts: Arc>>>>, + approval_failures: Arc>>>>, trusted_proxies: Arc>, approval_secret: Arc, public_base_url: Arc, @@ -184,7 +206,10 @@ impl McpOAuthStore { authorization_codes: Arc::new(RwLock::new(HashMap::new())), access_tokens: Arc::new(RwLock::new(HashMap::new())), refresh_tokens: Arc::new(RwLock::new(HashMap::new())), + used_refresh_tokens: Arc::new(RwLock::new(HashMap::new())), + token_issuance: Arc::new(Mutex::new(())), registration_attempts: Arc::new(Mutex::new(HashMap::new())), + approval_failures: Arc::new(Mutex::new(HashMap::new())), trusted_proxies: Arc::new(trusted_proxies), approval_secret: Arc::new(approval_secret), public_base_url: Arc::new(public_base_url), @@ -301,7 +326,47 @@ impl McpOAuthStore { true } - fn registration_source(&self, request: &Request) -> IpAddr { + async fn approval_locked_out(&self, source: IpAddr, now: DateTime) -> bool { + let window_start = now - Duration::seconds(REGISTRATION_RATE_WINDOW_SECONDS); + let mut failures_by_source = self.approval_failures.lock().await; + failures_by_source.retain(|_, failures| { + while failures + .front() + .is_some_and(|failure| *failure <= window_start) + { + failures.pop_front(); + } + !failures.is_empty() + }); + failures_by_source + .get(&source) + .is_some_and(|failures| failures.len() >= MAX_APPROVAL_FAILURES_PER_WINDOW) + } + + async fn record_approval_failure(&self, source: IpAddr, now: DateTime) { + let mut failures_by_source = self.approval_failures.lock().await; + if !failures_by_source.contains_key(&source) + && failures_by_source.len() >= MAX_REGISTRATION_SOURCES + { + let oldest_source = failures_by_source + .iter() + .min_by_key(|(_, failures)| failures.back().copied()) + .map(|(source, _)| *source); + if let Some(oldest_source) = oldest_source { + failures_by_source.remove(&oldest_source); + } + } + let failures = failures_by_source.entry(source).or_default(); + if failures.len() < MAX_APPROVAL_FAILURES_PER_WINDOW { + failures.push_back(now); + } + } + + async fn clear_approval_failures(&self, source: IpAddr) { + self.approval_failures.lock().await.remove(&source); + } + + fn request_source(&self, request: &Request) -> IpAddr { let peer = request .extensions() .get::>() @@ -358,6 +423,15 @@ impl McpOAuthStore { "code_challenge_method must be S256".to_string(), )); } + if params + .state + .as_ref() + .is_some_and(|state| state.len() > MAX_OAUTH_STATE_BYTES) + { + return Err(OAuthError::InvalidRequest( + "state exceeds the authorization endpoint limit".to_string(), + )); + } let code_challenge = params .code_challenge .as_deref() @@ -457,6 +531,7 @@ impl McpOAuthStore { .await .ok_or(OAuthError::InvalidClient)?; + let _issuance = self.token_issuance.lock().await; let now = Utc::now(); let mut codes = self.authorization_codes.write().await; codes.retain(|_, code| code.expires_at > now); @@ -477,12 +552,11 @@ impl McpOAuthStore { codes.remove(&request.code); return Err(OAuthError::InvalidGrant); } - let code = codes - .remove(&request.code) - .ok_or(OAuthError::InvalidGrant)?; - drop(codes); - - self.issue_tokens(code.client_id, code.scope).await + let issued = self + .issue_tokens_locked(code.client_id, code.scope, None) + .await?; + codes.remove(&request.code); + Ok(issued) } async fn exchange_refresh_token( @@ -493,33 +567,67 @@ impl McpOAuthStore { .await .ok_or(OAuthError::InvalidClient)?; + let _issuance = self.token_issuance.lock().await; let now = Utc::now(); let mut refresh_tokens = self.refresh_tokens.write().await; refresh_tokens.retain(|_, token| token.expires_at > now); - let record = refresh_tokens - .get(&request.refresh_token) - .cloned() - .ok_or(OAuthError::InvalidGrant)?; + let record = refresh_tokens.get(&request.refresh_token).cloned(); + drop(refresh_tokens); + let Some(record) = record else { + self.revoke_replayed_refresh_family(&request.refresh_token, now) + .await; + return Err(OAuthError::InvalidGrant); + }; if record.client_id != request.client_id { return Err(OAuthError::InvalidGrant); } - let record = refresh_tokens - .remove(&request.refresh_token) - .ok_or(OAuthError::InvalidGrant)?; - drop(refresh_tokens); - - self.issue_tokens(record.client_id, record.scope).await + self.issue_tokens_locked( + record.client_id.clone(), + record.scope.clone(), + Some(RefreshRotation { + token: request.refresh_token.clone(), + record, + }), + ) + .await } + #[cfg(test)] async fn issue_tokens( &self, client_id: String, scope: Option, + ) -> Result { + let _issuance = self.token_issuance.lock().await; + self.issue_tokens_locked(client_id, scope, None).await + } + + async fn issue_tokens_locked( + &self, + client_id: String, + scope: Option, + rotation: Option, ) -> Result { let now = Utc::now(); let access_token_value = random_prefixed("mcp-token"); let refresh_token_value = random_prefixed("mcp-refresh"); - let refresh_expires_at = now + Duration::seconds(REFRESH_TOKEN_TTL_SECONDS); + let (token_family_id, refresh_expires_at) = rotation + .as_ref() + .map(|rotation| { + ( + rotation.record.family_id.clone(), + rotation.record.family_expires_at, + ) + }) + .unwrap_or_else(|| { + ( + Uuid::new_v4().to_string(), + now + Duration::seconds(REFRESH_TOKEN_TTL_SECONDS), + ) + }); + if refresh_expires_at <= now { + return Err(OAuthError::InvalidGrant); + } let access = McpAccessToken { access_token: access_token_value.clone(), token_type: "bearer".to_string(), @@ -529,6 +637,7 @@ impl McpOAuthStore { client_id: client_id.clone(), issued_at: now, expires_at: now + Duration::seconds(ACCESS_TOKEN_TTL_SECONDS), + token_family_id: token_family_id.clone(), }; let mut clients = self.clients.write().await; @@ -543,15 +652,40 @@ impl McpOAuthStore { } let mut refresh_tokens = self.refresh_tokens.write().await; refresh_tokens.retain(|_, token| token.expires_at > now); - if refresh_tokens.len() >= MAX_REFRESH_TOKENS { + if refresh_tokens.len() >= MAX_REFRESH_TOKENS && rotation.is_none() { return Err(OAuthError::TemporarilyUnavailable); } + let mut used_refresh_tokens = self.used_refresh_tokens.write().await; + used_refresh_tokens.retain(|_, token| token.expires_at > now); + if rotation.is_some() && used_refresh_tokens.len() >= MAX_USED_REFRESH_TOKENS { + return Err(OAuthError::TemporarilyUnavailable); + } + if let Some(rotation) = rotation { + let active = refresh_tokens + .get(&rotation.token) + .ok_or(OAuthError::InvalidGrant)?; + if active.family_id != rotation.record.family_id { + return Err(OAuthError::InvalidGrant); + } + let removed = refresh_tokens + .remove(&rotation.token) + .ok_or(OAuthError::InvalidGrant)?; + used_refresh_tokens.insert( + rotation.token, + UsedRefreshToken { + family_id: removed.family_id, + expires_at: removed.family_expires_at, + }, + ); + } access_tokens.insert(access_token_value, access.clone()); refresh_tokens.insert( refresh_token_value, RefreshTokenRecord { client_id, scope, + family_id: token_family_id, + family_expires_at: refresh_expires_at, expires_at: refresh_expires_at, }, ); @@ -560,6 +694,23 @@ impl McpOAuthStore { Ok(IssuedTokens { access }) } + async fn revoke_replayed_refresh_family(&self, refresh_token: &str, now: DateTime) { + let mut used_refresh_tokens = self.used_refresh_tokens.write().await; + used_refresh_tokens.retain(|_, token| token.expires_at > now); + let Some(used) = used_refresh_tokens.get(refresh_token).cloned() else { + return; + }; + drop(used_refresh_tokens); + + let mut access_tokens = self.access_tokens.write().await; + access_tokens + .retain(|_, token| token.expires_at > now && token.token_family_id != used.family_id); + let mut refresh_tokens = self.refresh_tokens.write().await; + refresh_tokens + .retain(|_, token| token.expires_at > now && token.family_id != used.family_id); + warn!("Revoked an OAuth refresh-token family after replay detection"); + } + /// Validate a bearer token and atomically discard it after expiry. pub async fn validate_token(&self, token: &str) -> Option { let now = Utc::now(); @@ -581,6 +732,7 @@ enum OAuthError { InvalidGrant, InvalidTransaction, RateLimited, + ApprovalRateLimited, TemporarilyUnavailable, } @@ -616,6 +768,17 @@ impl OAuthError { .insert(RETRY_AFTER, HeaderValue::from_static("60")); response } + Self::ApprovalRateLimited => { + let mut response = oauth_json_error( + StatusCode::TOO_MANY_REQUESTS, + "temporarily_unavailable", + "approval authentication rate limit exceeded", + ); + response + .headers_mut() + .insert(RETRY_AFTER, HeaderValue::from_static("60")); + response + } Self::TemporarilyUnavailable => oauth_json_error( StatusCode::SERVICE_UNAVAILABLE, "temporarily_unavailable", @@ -858,6 +1021,11 @@ pub async fn oauth_approve( State(state): State>, request: Request, ) -> Response { + let approval_source = state.request_source(&request); + let now = Utc::now(); + if state.approval_locked_out(approval_source, now).await { + return OAuthError::ApprovalRateLimited.response(); + } let bytes = match read_limited_body(request).await { Ok(bytes) => bytes, Err(response) => return response, @@ -873,6 +1041,7 @@ pub async fn oauth_approve( } }; if !state.validate_approval_secret(&form.approval_secret) { + state.record_approval_failure(approval_source, now).await; warn!("Rejected OAuth approval with an invalid resource-owner credential"); return oauth_json_error( StatusCode::UNAUTHORIZED, @@ -880,6 +1049,7 @@ pub async fn oauth_approve( "approval authentication failed", ); } + state.clear_approval_failures(approval_source).await; let transaction = match state .consume_authorization_transaction(&form.transaction_id, &form.consent_nonce) .await @@ -1041,7 +1211,7 @@ pub async fn oauth_register( State(state): State>, request: Request, ) -> Response { - let registration_source = state.registration_source(&request); + let registration_source = state.request_source(&request); let bytes = match read_limited_body(request).await { Ok(bytes) => bytes, Err(response) => return response, @@ -1362,6 +1532,158 @@ mod tests { ); } + #[tokio::test] + async fn replayed_refresh_token_revokes_the_rotated_token_family() { + let store = McpOAuthStore::new(); + let issued = store + .issue_tokens( + "public-test-client".to_string(), + Some("profile".to_string()), + ) + .await + .unwrap(); + let old_refresh_token = issued.access.refresh_token.unwrap(); + let request = TokenRequest { + grant_type: "refresh_token".to_string(), + code: String::new(), + client_id: "public-test-client".to_string(), + client_secret: String::new(), + redirect_uri: String::new(), + code_verifier: None, + refresh_token: old_refresh_token.clone(), + }; + let rotated = store.exchange_refresh_token(&request).await.unwrap(); + let rotated_access_token = rotated.access.access_token.clone(); + let rotated_refresh_token = rotated.access.refresh_token.unwrap(); + + assert!(store.exchange_refresh_token(&request).await.is_err()); + assert!(store.validate_token(&rotated_access_token).await.is_none()); + assert!( + store + .exchange_refresh_token(&TokenRequest { + refresh_token: rotated_refresh_token, + ..request + }) + .await + .is_err() + ); + } + + #[tokio::test] + async fn token_capacity_failure_preserves_refresh_token_for_retry() { + let store = McpOAuthStore::new(); + let issued = store + .issue_tokens( + "public-test-client".to_string(), + Some("profile".to_string()), + ) + .await + .unwrap(); + let refresh_token = issued.access.refresh_token.clone().unwrap(); + let mut access_tokens = store.access_tokens.write().await; + for index in 1..MAX_ACCESS_TOKENS { + access_tokens.insert(format!("capacity-token-{index}"), issued.access.clone()); + } + drop(access_tokens); + let request = TokenRequest { + grant_type: "refresh_token".to_string(), + code: String::new(), + client_id: "public-test-client".to_string(), + client_secret: String::new(), + redirect_uri: String::new(), + code_verifier: None, + refresh_token: refresh_token.clone(), + }; + + assert!(store.exchange_refresh_token(&request).await.is_err()); + assert!( + store + .refresh_tokens + .read() + .await + .contains_key(&refresh_token) + ); + } + + #[tokio::test] + async fn token_capacity_failure_preserves_authorization_code_for_retry() { + let store = Arc::new(McpOAuthStore::new()); + let code = approve_transaction( + &store, + authorize_query("public-test-client", "http://127.0.0.1:8080/callback"), + ) + .await; + let issued = store + .issue_tokens( + "public-test-client".to_string(), + Some("profile".to_string()), + ) + .await + .unwrap(); + let mut access_tokens = store.access_tokens.write().await; + for index in 1..MAX_ACCESS_TOKENS { + access_tokens.insert( + format!("capacity-code-token-{index}"), + issued.access.clone(), + ); + } + drop(access_tokens); + let request = TokenRequest { + grant_type: "authorization_code".to_string(), + code: code.clone(), + client_id: "public-test-client".to_string(), + client_secret: String::new(), + redirect_uri: "http://127.0.0.1:8080/callback".to_string(), + code_verifier: Some(VERIFIER.to_string()), + refresh_token: String::new(), + }; + + assert!(store.exchange_authorization_code(&request).await.is_err()); + assert!(store.authorization_codes.read().await.contains_key(&code)); + } + + #[tokio::test] + async fn replay_cache_capacity_failure_preserves_refresh_token_for_retry() { + let store = McpOAuthStore::new(); + let issued = store + .issue_tokens( + "public-test-client".to_string(), + Some("profile".to_string()), + ) + .await + .unwrap(); + let refresh_token = issued.access.refresh_token.clone().unwrap(); + let mut used_refresh_tokens = store.used_refresh_tokens.write().await; + for index in 0..MAX_USED_REFRESH_TOKENS { + used_refresh_tokens.insert( + format!("used-refresh-{index}"), + UsedRefreshToken { + family_id: format!("used-family-{index}"), + expires_at: Utc::now() + Duration::hours(1), + }, + ); + } + drop(used_refresh_tokens); + let request = TokenRequest { + grant_type: "refresh_token".to_string(), + code: String::new(), + client_id: "public-test-client".to_string(), + client_secret: String::new(), + redirect_uri: String::new(), + code_verifier: None, + refresh_token: refresh_token.clone(), + }; + + assert!(store.exchange_refresh_token(&request).await.is_err()); + assert!( + store + .refresh_tokens + .read() + .await + .contains_key(&refresh_token) + ); + } + #[tokio::test] async fn expired_access_token_is_removed() { let store = McpOAuthStore::new(); @@ -1374,6 +1696,7 @@ mod tests { client_id: "public-test-client".to_string(), issued_at: Utc::now() - Duration::hours(2), expires_at: Utc::now() - Duration::hours(1), + token_family_id: "expired-family".to_string(), }; store .access_tokens @@ -1541,6 +1864,63 @@ mod tests { ); } + #[tokio::test] + async fn authorize_rejects_oversized_state_before_storing_a_transaction() { + let store = Arc::new(McpOAuthStore::new()); + let mut query = authorize_query("public-test-client", "http://127.0.0.1:8080/callback"); + query.state = Some("x".repeat(2_049)); + + assert_eq!( + oauth_authorize(Query(query), State(store.clone())) + .await + .status(), + StatusCode::BAD_REQUEST + ); + assert!(store.auth_transactions.read().await.is_empty()); + } + + #[tokio::test] + async fn repeated_invalid_approval_credentials_are_rate_limited() { + let store = Arc::new(McpOAuthStore::new()); + let (transaction_id, transaction) = store + .create_authorization_transaction(&authorize_query( + "public-test-client", + "http://127.0.0.1:8080/callback", + )) + .await + .unwrap(); + let invalid_body = serde_urlencoded::to_string(ApprovalForm { + transaction_id, + consent_nonce: transaction.consent_nonce, + approved: "true".to_string(), + approval_secret: "wrong-resource-owner-secret".to_string(), + }) + .unwrap(); + + for _ in 0..8 { + assert_eq!( + oauth_approve( + State(store.clone()), + Request::builder() + .body(Body::from(invalid_body.clone())) + .unwrap(), + ) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + } + assert_eq!( + oauth_approve( + State(store), + Request::builder().body(Body::from(invalid_body)).unwrap(), + ) + .await + .status(), + StatusCode::TOO_MANY_REQUESTS + ); + } + #[tokio::test] async fn dynamic_registration_distinguishes_public_and_confidential_clients() { let store = Arc::new(McpOAuthStore::new()); @@ -1694,7 +2074,7 @@ mod tests { } #[test] - fn registration_source_uses_forwarding_only_from_configured_proxies() { + fn request_source_uses_forwarding_only_from_configured_proxies() { let trusted_proxies = parse_trusted_proxy_cidrs(Some("10.0.0.0/8,fd00::/8")).unwrap(); let store = McpOAuthStore::with_trusted_proxies( generate_random_string(32), @@ -1710,7 +2090,7 @@ mod tests { .extensions_mut() .insert(ConnectInfo("10.2.0.9:443".parse::().unwrap())); assert_eq!( - store.registration_source(&trusted_request), + store.request_source(&trusted_request), "192.0.2.44".parse::().unwrap() ); @@ -1722,7 +2102,7 @@ mod tests { "198.51.100.7:443".parse::().unwrap(), )); assert_eq!( - store.registration_source(&untrusted_request), + store.request_source(&untrusted_request), "198.51.100.7".parse::().unwrap() ); }