113 lines
3.5 KiB
Rust
113 lines
3.5 KiB
Rust
use std::sync::Arc;
|
|
|
|
use axum::{
|
|
http::{HeaderMap, HeaderValue, StatusCode, header::RETRY_AFTER},
|
|
response::{IntoResponse, Response},
|
|
};
|
|
use crank_core::PlatformApiKeyKind;
|
|
use crank_runtime::RateLimitCheckError;
|
|
use serde_json::{Value, json};
|
|
|
|
use crate::{
|
|
access::{bearer_token, hash_access_secret},
|
|
app::{AgentRoutePath, AppState},
|
|
jsonrpc::{correlated_error_data, request_id},
|
|
transport::{ResponseMode, session_id_from_headers, transport_response},
|
|
};
|
|
|
|
pub(super) async fn enforce_transport_rate_limit(
|
|
state: &Arc<AppState>,
|
|
path: &AgentRoutePath,
|
|
headers: &HeaderMap,
|
|
) -> Result<(), RateLimitCheckError> {
|
|
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,
|
|
error: RateLimitCheckError,
|
|
) -> Response {
|
|
let RateLimitCheckError::Rejected(rejection) = error else {
|
|
return transport_response(
|
|
StatusCode::SERVICE_UNAVAILABLE,
|
|
json!({
|
|
"jsonrpc": "2.0",
|
|
"id": request_id(message),
|
|
"error": {
|
|
"code": -32603,
|
|
"message": "rate limit service unavailable",
|
|
"data": correlated_error_data(json!({ "code": "rate_limit_unavailable" }))
|
|
}
|
|
}),
|
|
response_mode,
|
|
None,
|
|
Some(protocol_version),
|
|
);
|
|
};
|
|
let payload = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": request_id(message),
|
|
"error": {
|
|
"code": -32029,
|
|
"message": "request rate limit exceeded",
|
|
"data": correlated_error_data(json!({
|
|
"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(error: RateLimitCheckError) -> Response {
|
|
match error {
|
|
RateLimitCheckError::Rejected(rejection) => {
|
|
let mut response = StatusCode::TOO_MANY_REQUESTS.into_response();
|
|
attach_retry_after_header(&mut response, rejection.retry_after_ms);
|
|
response
|
|
}
|
|
RateLimitCheckError::StoreUnavailable => StatusCode::SERVICE_UNAVAILABLE.into_response(),
|
|
}
|
|
}
|
|
|
|
fn rate_limit_key(path: &AgentRoutePath, headers: &HeaderMap) -> String {
|
|
let access_secret = bearer_token(headers);
|
|
if let Some(secret) = access_secret
|
|
&& secret.starts_with(PlatformApiKeyKind::Approval.secret_marker())
|
|
{
|
|
return format!("api_key:{}", hash_access_secret(secret));
|
|
}
|
|
|
|
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) = access_secret {
|
|
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);
|
|
}
|
|
}
|