github-ops
91 lines
2.6 KiB
Rust
91 lines
2.6 KiB
Rust
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<AppState>,
|
|
path: &AgentRoutePath,
|
|
headers: &HeaderMap,
|
|
) -> Result<(), RateLimitRejection> {
|
|
enforce_transport_rate_limit(state, path, headers).await
|
|
}
|
|
|
|
pub(super) async 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).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);
|
|
}
|
|
}
|