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::Response;
19use ipnet::IpNet;
20use std::future::Future;
21use std::net::SocketAddr;
22use std::sync::Arc;
23use systemprompt_extension::LoaderError;
24use systemprompt_models::auth::UserType;
25use systemprompt_models::{Config, RequestContext};
26
27#[derive(Clone, Debug)]
28pub struct IdentityOrTrustedIpKey {
29    trusted_proxies: Arc<Vec<IpNet>>,
30}
31
32impl IdentityOrTrustedIpKey {
33    const fn new(trusted_proxies: Arc<Vec<IpNet>>) -> Self {
34        Self { trusted_proxies }
35    }
36}
37
38impl tower_governor::key_extractor::KeyExtractor for IdentityOrTrustedIpKey {
39    type Key = String;
40
41    fn extract<T>(&self, req: &Request<T>) -> Result<Self::Key, tower_governor::GovernorError> {
42        // Why: an anonymous context's user id is a hash of the User-Agent and
43        // Accept-Language headers, so a caller who rotates either would mint a fresh
44        // bucket per request. Only a signature-verified identity is safe to key on.
45        if let Some(ctx) = req.extensions().get::<RequestContext>()
46            && ctx.auth.user_type != UserType::Anon
47        {
48            return Ok(format!("u:{}", ctx.user_id()));
49        }
50
51        resolve_client_ip(
52            req.headers(),
53            req.extensions().get::<ConnectInfo<SocketAddr>>(),
54            &self.trusted_proxies,
55        )
56        .map(|ip| format!("ip:{ip}"))
57        .ok_or(tower_governor::GovernorError::UnableToExtractKey)
58    }
59}
60
61pub trait ContextLayer: Clone + Send + Sync + 'static {
62    fn handle(self, req: Request, next: Next) -> impl Future<Output = Response> + Send;
63}
64
65impl ContextLayer for PublicContextMiddleware {
66    async fn handle(self, req: Request, next: Next) -> Response {
67        Self::handle(&self, req, next).await
68    }
69}
70
71impl ContextLayer for UserOnlyContextMiddleware {
72    async fn handle(self, req: Request, next: Next) -> Response {
73        Self::handle(&self, req, next).await
74    }
75}
76
77impl ContextLayer for A2AContextMiddleware {
78    async fn handle(self, req: Request, next: Next) -> Response {
79        Self::handle(&self, req, next).await
80    }
81}
82
83impl ContextLayer for McpContextMiddleware {
84    async fn handle(self, req: Request, next: Next) -> Response {
85        Self::handle(&self, req, next).await
86    }
87}
88
89pub trait RouterExt<S>: Sized {
90    fn with_rate_limit(self, config: &Config, per_second: u64) -> Result<Self, LoaderError>;
91
92    fn with_auth<L: ContextLayer>(self, auth: L, policy: AuthzPolicy) -> Self;
93}
94
95impl<S> RouterExt<S> for Router<S>
96where
97    S: Clone + Send + Sync + 'static,
98{
99    fn with_rate_limit(self, config: &Config, per_second: u64) -> Result<Self, LoaderError> {
100        let rate_config = &config.rate_limits;
101        if rate_config.disabled {
102            return Ok(self);
103        }
104
105        // Why: a truncating `as u32` turns any product that is a multiple of 2^32 into
106        // a zero burst, which `finish()` reports only by returning `None` —
107        // silently leaving the route unlimited. Saturate and clamp so the quota
108        // is always representable.
109        let burst = per_second.saturating_mul(rate_config.burst_multiplier);
110        let burst_u32 = u32::try_from(burst).unwrap_or(u32::MAX).max(1);
111        let per_second_clamped = per_second.max(1);
112
113        let rate_limit = tower_governor::governor::GovernorConfigBuilder::default()
114            .per_second(per_second_clamped)
115            .burst_size(burst_u32)
116            .key_extractor(IdentityOrTrustedIpKey::new(Arc::new(
117                config.trusted_proxies.clone(),
118            )))
119            .use_headers()
120            .finish()
121            .ok_or_else(|| LoaderError::InitializationFailed {
122                extension: "rate_limit".to_owned(),
123                message: format!(
124                    "rate limit rejected for {per_second_clamped}/s with burst {burst_u32}"
125                ),
126            })?;
127
128        Ok(self.layer(tower_governor::GovernorLayer::new(rate_limit)))
129    }
130
131    fn with_auth<L: ContextLayer>(self, auth: L, policy: AuthzPolicy) -> Self {
132        self.layer(axum::middleware::from_fn(move |req, next| async move {
133            authz_gate(policy, req, next).await
134        }))
135        .layer(axum::middleware::from_fn(move |req, next| {
136            let auth = auth.clone();
137            async move { auth.handle(req, next).await }
138        }))
139    }
140}
systemprompt_api/services/middleware/rate_limit.rs

systemprompt_api/services/middleware/
rate_limit.rs
