mcp: throttle transport read requests
This commit is contained in:
@@ -179,6 +179,10 @@ async fn mcp_get(
|
||||
return status.into_response();
|
||||
}
|
||||
|
||||
if let Err(rejection) = enforce_transport_rate_limit(&state, &path, &headers) {
|
||||
return rate_limited_status_response(rejection);
|
||||
}
|
||||
|
||||
let session_id = match session_id_from_headers(&headers) {
|
||||
Ok(Some(session_id)) => session_id,
|
||||
Ok(None) => return StatusCode::BAD_REQUEST.into_response(),
|
||||
@@ -223,6 +227,10 @@ async fn mcp_delete(
|
||||
return status.into_response();
|
||||
}
|
||||
|
||||
if let Err(rejection) = enforce_transport_rate_limit(&state, &path, &headers) {
|
||||
return rate_limited_status_response(rejection);
|
||||
}
|
||||
|
||||
match session_id_from_headers(&headers) {
|
||||
Ok(Some(session_id)) => match state.sessions.get(&session_id).await {
|
||||
Ok(Some(session))
|
||||
@@ -1832,6 +1840,14 @@ fn enforce_post_rate_limit(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<(), RateLimitRejection> {
|
||||
enforce_transport_rate_limit(state, path, headers)
|
||||
}
|
||||
|
||||
fn enforce_transport_rate_limit(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<(), RateLimitRejection> {
|
||||
let key = rate_limit_key(path, headers);
|
||||
state.api_rate_limiter.check(&key)
|
||||
@@ -1885,6 +1901,15 @@ fn rate_limited_jsonrpc_response(
|
||||
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 allows_scope(scopes: &[PlatformApiKeyScope], required_scope: PlatformApiKeyScope) -> bool {
|
||||
match required_scope {
|
||||
PlatformApiKeyScope::Read => scopes.iter().any(|scope| {
|
||||
|
||||
Reference in New Issue
Block a user