1use axum::extract::{Request, State};
2use axum::http::StatusCode;
3use axum::middleware::Next;
4use axum::response::{IntoResponse, Response};
5
6use crate::http::state::SharedState;
7
8pub async fn rate_limit_middleware(
9 State(state): State<SharedState>,
10 request: Request,
11 next: Next,
12) -> Response {
13 let path = request.uri().path();
15 let is_health = path == "/health" || path == "/stats" || path == "/metrics";
16 let is_registration = path == "/agents" && request.method() == "POST";
17
18 if is_health || is_registration {
19 return next.run(request).await;
20 }
21
22 let agent_id = request
23 .headers()
24 .get("x-agent-id")
25 .and_then(|v| v.to_str().ok())
26 .unwrap_or("");
27
28 if agent_id.is_empty() || !state.agent_registry.is_active(agent_id).unwrap_or(false) {
29 return StatusCode::UNAUTHORIZED.into_response();
30 }
31
32 let decision = {
33 let engine = state.engine.lock();
34 state
35 .rate_limiter
36 .check_rate_limit(engine.graph(), agent_id)
37 };
38
39 if !decision.allowed {
40 return StatusCode::TOO_MANY_REQUESTS.into_response();
41 }
42
43 next.run(request).await
44}