From 35841359b39602ae7c6e2c6452af1cc7b16e1d86 Mon Sep 17 00:00:00 2001 From: github-ops Date: Sun, 21 Jun 2026 09:48:30 +0000 Subject: [PATCH] Extract MCP rate limiting module --- crates/crank-community-mcp/src/app.rs | 90 +++----------------- crates/crank-community-mcp/src/lib.rs | 1 + crates/crank-community-mcp/src/rate_limit.rs | 90 ++++++++++++++++++++ 3 files changed, 101 insertions(+), 80 deletions(-) create mode 100644 crates/crank-community-mcp/src/rate_limit.rs diff --git a/crates/crank-community-mcp/src/app.rs b/crates/crank-community-mcp/src/app.rs index a9a5387..434f698 100644 --- a/crates/crank-community-mcp/src/app.rs +++ b/crates/crank-community-mcp/src/app.rs @@ -8,7 +8,7 @@ use std::{ use axum::{ Json, Router, extract::{Path, State}, - http::{HeaderMap, HeaderValue, StatusCode, header::RETRY_AFTER}, + http::{HeaderMap, StatusCode}, response::{IntoResponse, Response, sse::Event}, routing::get, }; @@ -18,8 +18,8 @@ use crank_core::{ }; use crank_registry::{CreateInvocationLogRequest, PostgresRegistry, PublishedAgentTool}; use crank_runtime::{ - RateLimitRejection, RequestRateLimiter, ResolvedAuth, RuntimeError, RuntimeExecutionRequest, - RuntimeExecutor, RuntimeOperation, RuntimeRequestContext, SecretCrypto, + RequestRateLimiter, ResolvedAuth, RuntimeError, RuntimeExecutionRequest, RuntimeExecutor, + RuntimeOperation, RuntimeRequestContext, SecretCrypto, }; use futures_util::stream; use serde::{Deserialize, Serialize}; @@ -29,8 +29,8 @@ use tracing::info; use crate::{ access::{ - bearer_token, credential_allows_security_level, hash_access_secret, require_machine_access, - serialize_machine_access_mode, serialize_security_level, + credential_allows_security_level, require_machine_access, serialize_machine_access_mode, + serialize_security_level, }, auth::{SharedMachineCredentialVerifier, VerifiedMachineCredential}, catalog::PublishedToolCatalog, @@ -39,6 +39,10 @@ use crate::{ jsonrpc_result, method_name, negotiated_protocol_version, params, request_id, }, manifest::tool_definitions, + rate_limit::{ + enforce_post_rate_limit, enforce_transport_rate_limit, rate_limited_jsonrpc_response, + rate_limited_status_response, + }, session::{SessionState, SharedSessionStore}, tool_error::{ ToolErrorContract, generic_tool_error_contract, runtime_error_code, @@ -59,7 +63,7 @@ pub(super) struct AppState { pub(super) registry: PostgresRegistry, catalog: PublishedToolCatalog, runtime: RuntimeExecutor, - api_rate_limiter: RequestRateLimiter, + pub(super) api_rate_limiter: RequestRateLimiter, secret_crypto: SecretCrypto, sessions: SharedSessionStore, pub(super) credential_verifier: SharedMachineCredentialVerifier, @@ -810,80 +814,6 @@ async fn require_initialized_session( Ok(session) } -async fn enforce_post_rate_limit( - state: &Arc, - path: &AgentRoutePath, - headers: &HeaderMap, -) -> Result<(), RateLimitRejection> { - enforce_transport_rate_limit(state, path, headers).await -} - -async fn enforce_transport_rate_limit( - state: &Arc, - path: &AgentRoutePath, - headers: &HeaderMap, -) -> Result<(), RateLimitRejection> { - let key = rate_limit_key(path, headers); - state.api_rate_limiter.check(&key).await -} - -fn rate_limit_key(path: &AgentRoutePath, headers: &HeaderMap) -> String { - if let Ok(Some(session_id)) = session_id_from_headers(headers) { - return format!( - "session:{}:{}:{}", - path.workspace_slug, path.agent_slug, session_id - ); - } - - if let Some(secret) = bearer_token(headers) { - return format!("api_key:{}", hash_access_secret(secret)); - } - - format!("workspace:{}:anonymous", path.workspace_slug) -} - -fn rate_limited_jsonrpc_response( - message: &Value, - response_mode: ResponseMode, - protocol_version: &str, - rejection: RateLimitRejection, -) -> Response { - let payload = json!({ - "jsonrpc": "2.0", - "id": request_id(message), - "error": { - "code": -32029, - "message": "request rate limit exceeded", - "data": { - "code": "request_rate_limited", - "retry_after_ms": rejection.retry_after_ms, - } - } - }); - - let mut response = transport_response( - StatusCode::TOO_MANY_REQUESTS, - payload, - response_mode, - None, - Some(protocol_version), - ); - let retry_after_seconds = rejection.retry_after_ms.div_ceil(1000); - if let Ok(value) = HeaderValue::from_str(&retry_after_seconds.to_string()) { - response.headers_mut().insert(RETRY_AFTER, value); - } - response -} - -fn rate_limited_status_response(rejection: RateLimitRejection) -> Response { - let mut response = StatusCode::TOO_MANY_REQUESTS.into_response(); - let retry_after_seconds = rejection.retry_after_ms.div_ceil(1000); - if let Ok(value) = HeaderValue::from_str(&retry_after_seconds.to_string()) { - response.headers_mut().insert(RETRY_AFTER, value); - } - response -} - fn internal_jsonrpc_error(message: &Value, error: impl std::fmt::Display) -> Response { transport_response( StatusCode::INTERNAL_SERVER_ERROR, diff --git a/crates/crank-community-mcp/src/lib.rs b/crates/crank-community-mcp/src/lib.rs index 1bc33e8..db2dc09 100644 --- a/crates/crank-community-mcp/src/lib.rs +++ b/crates/crank-community-mcp/src/lib.rs @@ -4,6 +4,7 @@ pub mod auth; pub mod catalog; pub mod jsonrpc; pub mod manifest; +mod rate_limit; pub mod session; pub mod tool_error; mod transport; diff --git a/crates/crank-community-mcp/src/rate_limit.rs b/crates/crank-community-mcp/src/rate_limit.rs new file mode 100644 index 0000000..63f9978 --- /dev/null +++ b/crates/crank-community-mcp/src/rate_limit.rs @@ -0,0 +1,90 @@ +use std::sync::Arc; + +use axum::{ + http::{HeaderMap, HeaderValue, StatusCode, header::RETRY_AFTER}, + response::{IntoResponse, Response}, +}; +use crank_runtime::RateLimitRejection; +use serde_json::{Value, json}; + +use crate::{ + access::{bearer_token, hash_access_secret}, + app::{AgentRoutePath, AppState}, + jsonrpc::request_id, + transport::{ResponseMode, session_id_from_headers, transport_response}, +}; + +pub(super) async fn enforce_post_rate_limit( + state: &Arc, + path: &AgentRoutePath, + headers: &HeaderMap, +) -> Result<(), RateLimitRejection> { + enforce_transport_rate_limit(state, path, headers).await +} + +pub(super) async fn enforce_transport_rate_limit( + state: &Arc, + path: &AgentRoutePath, + headers: &HeaderMap, +) -> Result<(), RateLimitRejection> { + let key = rate_limit_key(path, headers); + state.api_rate_limiter.check(&key).await +} + +pub(super) fn rate_limited_jsonrpc_response( + message: &Value, + response_mode: ResponseMode, + protocol_version: &str, + rejection: RateLimitRejection, +) -> Response { + let payload = json!({ + "jsonrpc": "2.0", + "id": request_id(message), + "error": { + "code": -32029, + "message": "request rate limit exceeded", + "data": { + "code": "request_rate_limited", + "retry_after_ms": rejection.retry_after_ms, + } + } + }); + + let mut response = transport_response( + StatusCode::TOO_MANY_REQUESTS, + payload, + response_mode, + None, + Some(protocol_version), + ); + attach_retry_after_header(&mut response, rejection.retry_after_ms); + response +} + +pub(super) fn rate_limited_status_response(rejection: RateLimitRejection) -> Response { + let mut response = StatusCode::TOO_MANY_REQUESTS.into_response(); + attach_retry_after_header(&mut response, rejection.retry_after_ms); + response +} + +fn rate_limit_key(path: &AgentRoutePath, headers: &HeaderMap) -> String { + if let Ok(Some(session_id)) = session_id_from_headers(headers) { + return format!( + "session:{}:{}:{}", + path.workspace_slug, path.agent_slug, session_id + ); + } + + if let Some(secret) = bearer_token(headers) { + return format!("api_key:{}", hash_access_secret(secret)); + } + + format!("workspace:{}:anonymous", path.workspace_slug) +} + +fn attach_retry_after_header(response: &mut Response, retry_after_ms: u64) { + let retry_after_seconds = retry_after_ms.div_ceil(1000); + if let Ok(value) = HeaderValue::from_str(&retry_after_seconds.to_string()) { + response.headers_mut().insert(RETRY_AFTER, value); + } +}