systemprompt_api/services/middleware/
rate_limit.rs

1//! Router extension traits for rate limiting and authenticated route groups.
2//!
3//! `RouterExt::with_auth` attaches authentication and authorization in one
4//! call: it requires an `AuthzPolicy`, so a route group cannot be mounted
5//! authenticated-but-unauthorized — omitting the policy is a compile error.
6//!
7//! Copyright (c) systemprompt.io — Business Source License 1.1.
8//! See <https://systemprompt.io> for licensing details.
9
10use crate::services::middleware::authz::{AuthzPolicy, authz_gate};
11use crate::services::middleware::client_addr::resolve_client_ip;
12use crate::services::middleware::context::{
13    A2AContextMiddleware, McpContextMiddleware, PublicContextMiddleware, UserOnlyContextMiddleware,
14};
15use axum::Router;
16use axum::extract::{ConnectInfo, Request};
17use axum::middleware::Next;
18use axum::response::{IntoResponse, Response};
19use governor::clock::DefaultClock;
20use governor::state::keyed::DefaultKeyedStateStore;
21use governor::{Quota, RateLimiter};
22use ipnet::IpNet;
23use std::future::Future;
24use std::net::SocketAddr;
25use std::num::NonZeroU32;
26use std::sync::Arc;
27use systemprompt_models::RequestContext;
28use systemprompt_models::api::{ApiError, ErrorCode};
29use systemprompt_models::auth::RateLimitTier;
30use systemprompt_models::config::RateLimitConfig;
31use tower_governor::key_extractor::SmartIpKeyExtractor;
32use tracing::warn;
33
34pub trait ContextLayer: Clone + Send + Sync + 'static {
35    fn handle(self, req: Request, next: Next) -> impl Future<Output = Response> + Send;
36}
37
38impl ContextLayer for PublicContextMiddleware {
39    async fn handle(self, req: Request, next: Next) -> Response {
40        Self::handle(&self, req, next).await
41    }
42}
43
44impl ContextLayer for UserOnlyContextMiddleware {
45    async fn handle(self, req: Request, next: Next) -> Response {
46        Self::handle(&self, req, next).await
47    }
48}
49
50impl ContextLayer for A2AContextMiddleware {
51    async fn handle(self, req: Request, next: Next) -> Response {
52        Self::handle(&self, req, next).await
53    }
54}
55
56impl ContextLayer for McpContextMiddleware {
57    async fn handle(self, req: Request, next: Next) -> Response {
58        Self::handle(&self, req, next).await
59    }
60}
61
62pub trait RouterExt<S> {
63    fn with_rate_limit(self, rate_config: &RateLimitConfig, per_second: u64) -> Self;
64
65    fn with_auth<L: ContextLayer>(self, auth: L, policy: AuthzPolicy) -> Self;
66}
67
68impl<S> RouterExt<S> for Router<S>
69where
70    S: Clone + Send + Sync + 'static,
71{
72    fn with_rate_limit(self, rate_config: &RateLimitConfig, per_second: u64) -> Self {
73        if rate_config.disabled {
74            return self;
75        }
76
77        let rate_limit_result = tower_governor::governor::GovernorConfigBuilder::default()
78            .per_second(per_second)
79            .burst_size((per_second * rate_config.burst_multiplier) as u32)
80            .key_extractor(SmartIpKeyExtractor)
81            .use_headers()
82            .finish();
83
84        if let Some(rate_limit) = rate_limit_result {
85            self.layer(tower_governor::GovernorLayer::new(rate_limit))
86        } else {
87            warn!("Failed to configure rate limiting - rate limiting disabled for this route");
88            self
89        }
90    }
91
92    fn with_auth<L: ContextLayer>(self, auth: L, policy: AuthzPolicy) -> Self {
93        self.layer(axum::middleware::from_fn(move |req, next| async move {
94            authz_gate(policy, req, next).await
95        }))
96        .layer(axum::middleware::from_fn(move |req, next| {
97            let auth = auth.clone();
98            async move { auth.handle(req, next).await }
99        }))
100    }
101}
102
103type KeyedRateLimiter = RateLimiter<String, DefaultKeyedStateStore<String>, DefaultClock>;
104
105#[derive(Clone, Debug)]
106pub struct TieredRateLimiter {
107    admin_limiter: Arc<KeyedRateLimiter>,
108    user_limiter: Arc<KeyedRateLimiter>,
109    a2a_limiter: Arc<KeyedRateLimiter>,
110    mcp_limiter: Arc<KeyedRateLimiter>,
111    service_limiter: Arc<KeyedRateLimiter>,
112    anon_limiter: Arc<KeyedRateLimiter>,
113    disabled: bool,
114    trusted_proxies: Arc<Vec<IpNet>>,
115}
116
117impl TieredRateLimiter {
118    pub fn new(config: &RateLimitConfig, base_per_second: u64) -> Self {
119        Self::with_trusted_proxies(config, base_per_second, Vec::new())
120    }
121
122    pub fn with_trusted_proxies(
123        config: &RateLimitConfig,
124        base_per_second: u64,
125        trusted_proxies: Vec<IpNet>,
126    ) -> Self {
127        let create_limiter = |tier: RateLimitTier| -> Arc<KeyedRateLimiter> {
128            let effective = config.effective_limit(base_per_second, tier);
129            let burst = effective.saturating_mul(config.burst_multiplier);
130            let effective_u32 = u32::try_from(effective).unwrap_or(u32::MAX).max(1);
131            let burst_u32 = u32::try_from(burst).unwrap_or(u32::MAX).max(1);
132            let quota =
133                Quota::per_second(NonZeroU32::new(effective_u32).unwrap_or(NonZeroU32::MIN))
134                    .allow_burst(NonZeroU32::new(burst_u32).unwrap_or(NonZeroU32::MIN));
135            Arc::new(RateLimiter::keyed(quota))
136        };
137
138        Self {
139            admin_limiter: create_limiter(RateLimitTier::Admin),
140            user_limiter: create_limiter(RateLimitTier::User),
141            a2a_limiter: create_limiter(RateLimitTier::A2a),
142            mcp_limiter: create_limiter(RateLimitTier::Mcp),
143            service_limiter: create_limiter(RateLimitTier::Service),
144            anon_limiter: create_limiter(RateLimitTier::Anon),
145            disabled: config.disabled,
146            trusted_proxies: Arc::new(trusted_proxies),
147        }
148    }
149
150    pub fn disabled() -> Self {
151        let quota = Quota::per_second(NonZeroU32::MAX);
152        let limiter = Arc::new(RateLimiter::keyed(quota));
153        Self {
154            admin_limiter: Arc::clone(&limiter),
155            user_limiter: Arc::clone(&limiter),
156            a2a_limiter: Arc::clone(&limiter),
157            mcp_limiter: Arc::clone(&limiter),
158            service_limiter: Arc::clone(&limiter),
159            anon_limiter: Arc::clone(&limiter),
160            disabled: true,
161            trusted_proxies: Arc::new(Vec::new()),
162        }
163    }
164
165    #[must_use]
166    pub fn trusted_proxies(&self) -> &[IpNet] {
167        &self.trusted_proxies
168    }
169
170    fn limiter_for_tier(&self, tier: RateLimitTier) -> &KeyedRateLimiter {
171        match tier {
172            RateLimitTier::Admin => &self.admin_limiter,
173            RateLimitTier::User => &self.user_limiter,
174            RateLimitTier::A2a => &self.a2a_limiter,
175            RateLimitTier::Mcp => &self.mcp_limiter,
176            RateLimitTier::Service => &self.service_limiter,
177            RateLimitTier::Anon => &self.anon_limiter,
178        }
179    }
180
181    pub fn check(&self, tier: RateLimitTier, key: &str) -> bool {
182        if self.disabled {
183            return true;
184        }
185        self.limiter_for_tier(tier)
186            .check_key(&key.to_owned())
187            .is_ok()
188    }
189}
190
191pub async fn tiered_rate_limit_middleware(
192    limiter: axum::extract::State<TieredRateLimiter>,
193    request: Request,
194    next: Next,
195) -> Response {
196    if limiter.disabled {
197        return next.run(request).await;
198    }
199
200    let (tier, key) = request.extensions().get::<RequestContext>().map_or_else(
201        || {
202            let connect_info = request.extensions().get::<ConnectInfo<SocketAddr>>();
203            let ip = resolve_client_ip(request.headers(), connect_info, limiter.trusted_proxies())
204                .map_or_else(|| "unknown".to_owned(), |a| a.to_string());
205            (RateLimitTier::Anon, ip)
206        },
207        |ctx| {
208            let tier = ctx.rate_limit_tier();
209            let key = ctx.user_id().to_string();
210            (tier, key)
211        },
212    );
213
214    if limiter.check(tier, &key) {
215        next.run(request).await
216    } else {
217        warn!(
218            tier = %tier.as_str(),
219            key = %key,
220            "Rate limit exceeded"
221        );
222        let api_error = ApiError::new(ErrorCode::RateLimited, "Rate limit exceeded");
223        let mut response = api_error.into_response();
224        response
225            .headers_mut()
226            .insert("Retry-After", http::HeaderValue::from_static("1"));
227        if let Ok(tier_value) = http::HeaderValue::from_str(tier.as_str()) {
228            response
229                .headers_mut()
230                .insert("X-Rate-Limit-Tier", tier_value);
231        }
232        response
233    }
234}
systemprompt_api/services/middleware/rate_limit.rs

systemprompt_api/services/middleware/
rate_limit.rs
