ironflow_api/
rate_limit.rs

1//! Identity-aware rate limiting middleware.
2//!
3//! Requests are keyed by authenticated identity (API key ID or user ID) when
4//! available, falling back to the client IP address. Each key gets a
5//! fixed-window counter that resets every 60 seconds.
6//!
7//! The client IP is the address of the TCP peer, so the server must be
8//! served with `into_make_service_with_connect_info::<SocketAddr>()`.
9//! `X-Forwarded-For` and `X-Real-IP` are read only when that peer is one of
10//! the [`TrustedProxies`]: any client can write those headers.
11//!
12//! With [`AccountLimit::ByEmail`], a request whose JSON body names an
13//! `email` is also counted against that account, so spreading attempts on
14//! one account over many IP addresses does not help.
15//!
16//! Responses carry standard rate-limit headers:
17//! - `X-RateLimit-Limit` -- maximum requests allowed per window
18//! - `X-RateLimit-Remaining` -- requests left in the current window
19//! - `X-RateLimit-Reset` -- epoch timestamp when the window resets
20//!
21//! API keys with a `rate_limit_override` use that value instead of the
22//! server-wide default.
23
24mod account;
25mod client_ip;
26#[cfg(test)]
27mod tests;
28
29use std::net::{IpAddr, SocketAddr};
30use std::sync::Arc;
31use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
32use std::time::{SystemTime, UNIX_EPOCH};
33
34use axum::Json;
35use axum::body::{Body, to_bytes};
36use axum::extract::{ConnectInfo, Request};
37use axum::http::{HeaderMap, HeaderValue, StatusCode};
38use axum::middleware::Next;
39use axum::response::{IntoResponse, Response};
40use dashmap::DashMap;
41use serde_json::json;
42use uuid::Uuid;
43
44use ironflow_auth::extractor::{API_KEY_PREFIX, API_KEY_SUFFIX_LEN};
45use ironflow_auth::jwt::{AccessToken, JwtConfig};
46use ironflow_store::store::Store;
47
48use self::account::{ACCOUNT_BODY_LIMIT, account_key};
49use self::client_ip::client_ip;
50
51pub use self::account::AccountLimit;
52pub use self::client_ip::{InvalidTrustedProxy, TrustedProxies};
53
54/// Counters kept before expired windows are swept. Account keys are chosen
55/// by the caller, so the map must not grow without bound.
56const SWEEP_THRESHOLD: usize = 10_000;
57
58/// Rate limit key: identity-based when authenticated, IP-based otherwise.
59///
60/// # Examples
61///
62/// ```
63/// use std::net::{IpAddr, Ipv4Addr};
64/// use uuid::Uuid;
65/// use ironflow_api::rate_limit::RateLimitKey;
66///
67/// let ip_key = RateLimitKey::Ip(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
68/// let user_key = RateLimitKey::User(Uuid::nil());
69/// let api_key = RateLimitKey::ApiKey(Uuid::nil());
70/// let account_key = RateLimitKey::Account("alice@ironflow.dev".to_string());
71/// ```
72#[derive(Debug, Clone, PartialEq, Eq, Hash)]
73pub enum RateLimitKey {
74    /// Authenticated via API key.
75    ApiKey(Uuid),
76    /// Authenticated via JWT.
77    User(Uuid),
78    /// Unauthenticated, identified by client IP.
79    Ip(IpAddr),
80    /// The account a credential request targets, by normalized email.
81    Account(String),
82}
83
84struct WindowEntry {
85    count: AtomicU32,
86    window_start: AtomicU64,
87}
88
89/// Shared rate limiter state with fixed-window counters.
90///
91/// # Examples
92///
93/// ```
94/// use ironflow_api::rate_limit::per_minute;
95///
96/// let limiter = per_minute(60);
97/// ```
98#[derive(Clone)]
99pub struct RateLimitState {
100    counters: Arc<DashMap<RateLimitKey, WindowEntry>>,
101    burst: u32,
102}
103
104/// Context needed by the rate limit middleware.
105///
106/// Combines the limiter state with auth dependencies for identity extraction.
107///
108/// # Examples
109///
110/// ```no_run
111/// use std::sync::Arc;
112/// use ironflow_api::rate_limit::{AccountLimit, RateLimitContext, TrustedProxies, per_minute};
113/// use ironflow_auth::jwt::JwtConfig;
114/// use ironflow_store::memory::InMemoryStore;
115/// use ironflow_store::store::Store;
116///
117/// let store: Arc<dyn Store> = Arc::new(InMemoryStore::new());
118/// let jwt_config = Arc::new(JwtConfig {
119///     secret: "secret".to_string(),
120///     access_token_ttl_secs: 900,
121///     refresh_token_ttl_secs: 604800,
122///     cookie_domain: None,
123///     cookie_secure: false,
124/// });
125/// let ctx = RateLimitContext {
126///     store,
127///     jwt_config,
128///     limiter: per_minute(60),
129///     trusted_proxies: TrustedProxies::default(),
130///     account_limit: AccountLimit::Off,
131/// };
132/// ```
133#[derive(Clone)]
134pub struct RateLimitContext {
135    /// Backing store for API key prefix lookups.
136    pub store: Arc<dyn Store>,
137    /// JWT config for token decoding.
138    pub jwt_config: Arc<JwtConfig>,
139    /// The rate limiter itself.
140    pub limiter: RateLimitState,
141    /// Proxies allowed to report the client IP through forwarding headers.
142    pub trusted_proxies: TrustedProxies,
143    /// Whether requests are also counted against the account they target.
144    pub account_limit: AccountLimit,
145}
146
147struct RateLimitResult {
148    limit: u32,
149    remaining: u32,
150    reset: u64,
151    allowed: bool,
152}
153
154/// Build a per-identity rate limiter allowing `requests_per_minute` requests
155/// per minute.
156///
157/// # Panics
158///
159/// Panics if `requests_per_minute` is 0.
160///
161/// # Examples
162///
163/// ```
164/// use ironflow_api::rate_limit::per_minute;
165///
166/// let auth_limiter = per_minute(10);
167/// let general_limiter = per_minute(60);
168/// ```
169pub fn per_minute(requests_per_minute: u32) -> RateLimitState {
170    assert!(requests_per_minute > 0, "burst must be > 0");
171    RateLimitState {
172        counters: Arc::new(DashMap::new()),
173        burst: requests_per_minute,
174    }
175}
176
177const WINDOW_SECS: u64 = 60;
178
179fn now_epoch() -> u64 {
180    SystemTime::now()
181        .duration_since(UNIX_EPOCH)
182        .expect("system clock before UNIX epoch")
183        .as_secs()
184}
185
186fn check_rate_limit(limiter: &RateLimitState, key: RateLimitKey, limit: u32) -> RateLimitResult {
187    let now = now_epoch();
188
189    if limiter.counters.len() >= SWEEP_THRESHOLD && !limiter.counters.contains_key(&key) {
190        limiter
191            .counters
192            .retain(|_, entry| entry.window_start.load(Ordering::Acquire) + WINDOW_SECS > now);
193    }
194
195    let entry = limiter.counters.entry(key).or_insert_with(|| WindowEntry {
196        count: AtomicU32::new(0),
197        window_start: AtomicU64::new(now),
198    });
199
200    let window_start = entry.window_start.load(Ordering::Acquire);
201    if now >= window_start + WINDOW_SECS
202        && entry
203            .window_start
204            .compare_exchange(window_start, now, Ordering::AcqRel, Ordering::Relaxed)
205            .is_ok()
206    {
207        entry.count.store(0, Ordering::Release);
208    }
209
210    let count = entry.count.fetch_add(1, Ordering::AcqRel) + 1;
211    let reset = entry.window_start.load(Ordering::Acquire) + WINDOW_SECS;
212
213    if count > limit {
214        entry.count.fetch_sub(1, Ordering::AcqRel);
215        RateLimitResult {
216            limit,
217            remaining: 0,
218            reset,
219            allowed: false,
220        }
221    } else {
222        RateLimitResult {
223            limit,
224            remaining: limit.saturating_sub(count),
225            reset,
226            allowed: true,
227        }
228    }
229}
230
231/// Axum middleware that enforces identity-aware rate limiting.
232///
233/// Extracts the caller's identity from the request:
234/// - API key (`irfl_*` Bearer token) -> `RateLimitKey::ApiKey(id)`
235/// - JWT (Bearer token or cookie) -> `RateLimitKey::User(id)`
236/// - Fallback -> `RateLimitKey::Ip(addr)`, the TCP peer or, behind a
237///   [trusted proxy](TrustedProxies), the address it forwarded
238///
239/// API keys with a `rate_limit_override` use that value instead of
240/// the default burst. With [`AccountLimit::ByEmail`], the targeted account
241/// is counted too, at the default burst; a body over 64 KiB is answered
242/// with `413` before reaching the handler.
243pub async fn rate_limit(mut req: Request, next: Next) -> Response {
244    let ctx = req.extensions_mut().remove::<RateLimitContext>();
245    let Some(ctx) = ctx else {
246        return next.run(req).await;
247    };
248    let peer = req
249        .extensions()
250        .get::<ConnectInfo<SocketAddr>>()
251        .map(|ConnectInfo(addr)| addr.ip());
252    let (key, override_limit) = extract_rate_limit_key(req.headers(), peer, &ctx).await;
253    let limit = override_limit.unwrap_or(ctx.limiter.burst);
254
255    if limit == 0 {
256        return next.run(req).await;
257    }
258
259    let caller = check_rate_limit(&ctx.limiter, key, limit);
260    if !caller.allowed {
261        return too_many_requests(&caller);
262    }
263
264    let (req, result) = match ctx.account_limit {
265        AccountLimit::Off => (req, caller),
266        AccountLimit::ByEmail => {
267            let (parts, body) = req.into_parts();
268            let Ok(bytes) = to_bytes(body, ACCOUNT_BODY_LIMIT).await else {
269                return payload_too_large();
270            };
271            let account = account_key(&bytes);
272            let req = Request::from_parts(parts, Body::from(bytes));
273            match account {
274                None => (req, caller),
275                Some(key) => {
276                    let account = check_rate_limit(&ctx.limiter, key, ctx.limiter.burst);
277                    if !account.allowed {
278                        return too_many_requests(&account);
279                    }
280                    let tightest = if account.remaining < caller.remaining {
281                        account
282                    } else {
283                        caller
284                    };
285                    (req, tightest)
286                }
287            }
288        }
289    };
290
291    let mut resp = next.run(req).await;
292    insert_rate_limit_headers(resp.headers_mut(), &result);
293    resp
294}
295
296fn too_many_requests(result: &RateLimitResult) -> Response {
297    let retry_after = result.reset.saturating_sub(now_epoch()).max(1);
298
299    let body = json!({
300        "error": {
301            "code": "RATE_LIMIT_EXCEEDED",
302            "message": "Too many requests, please try again later",
303            "retry_after_secs": retry_after,
304        }
305    });
306
307    let mut resp = (StatusCode::TOO_MANY_REQUESTS, Json(body)).into_response();
308    resp.headers_mut()
309        .insert("retry-after", HeaderValue::from(retry_after));
310    insert_rate_limit_headers(resp.headers_mut(), result);
311    resp
312}
313
314fn payload_too_large() -> Response {
315    let body = json!({
316        "error": {
317            "code": "PAYLOAD_TOO_LARGE",
318            "message": "Request body too large or unreadable",
319        }
320    });
321    (StatusCode::PAYLOAD_TOO_LARGE, Json(body)).into_response()
322}
323
324fn insert_rate_limit_headers(headers: &mut HeaderMap, result: &RateLimitResult) {
325    headers.insert("x-ratelimit-limit", HeaderValue::from(result.limit));
326    headers.insert("x-ratelimit-remaining", HeaderValue::from(result.remaining));
327    headers.insert("x-ratelimit-reset", HeaderValue::from(result.reset));
328}
329
330/// Extract the rate limit key and optional override from the request.
331///
332/// Attempts lightweight identity extraction without full auth verification:
333/// - API key prefix lookup (cheap index scan, no argon2)
334/// - JWT decode (local crypto, no DB)
335/// - Fallback to the client IP
336async fn extract_rate_limit_key(
337    headers: &HeaderMap,
338    peer: Option<IpAddr>,
339    ctx: &RateLimitContext,
340) -> (RateLimitKey, Option<u32>) {
341    let bearer = headers
342        .get("authorization")
343        .and_then(|v| v.to_str().ok())
344        .and_then(|v| v.strip_prefix("Bearer "))
345        .map(|s| s.to_string());
346
347    if let Some(ref token) = bearer {
348        if token.starts_with(API_KEY_PREFIX) {
349            if let Some((key, override_limit)) = try_api_key_identity(token, ctx).await {
350                return (key, override_limit);
351            }
352        } else if let Some(key) = try_jwt_identity(token, ctx) {
353            return (key, None);
354        }
355    }
356
357    // Try JWT from cookie
358    if let Some(cookie_header) = headers.get("cookie").and_then(|v| v.to_str().ok()) {
359        for part in cookie_header.split(';') {
360            let part = part.trim();
361            if let Some(value) = part.strip_prefix("ironflow_session=")
362                && let Some(key) = try_jwt_identity(value, ctx)
363            {
364                return (key, None);
365            }
366        }
367    }
368
369    let ip = client_ip(peer, headers, &ctx.trusted_proxies);
370    (RateLimitKey::Ip(ip), None)
371}
372
373async fn try_api_key_identity(
374    token: &str,
375    ctx: &RateLimitContext,
376) -> Option<(RateLimitKey, Option<u32>)> {
377    let suffix_len = (token.len() - API_KEY_PREFIX.len()).min(API_KEY_SUFFIX_LEN);
378    let prefix = &token[..API_KEY_PREFIX.len() + suffix_len];
379
380    let api_key = ctx.store.find_api_key_by_prefix(prefix).await.ok()??;
381
382    let override_limit = api_key.rate_limit_override;
383    Some((RateLimitKey::ApiKey(api_key.id), override_limit))
384}
385
386fn try_jwt_identity(token: &str, ctx: &RateLimitContext) -> Option<RateLimitKey> {
387    let claims = AccessToken::decode(token, &ctx.jwt_config).ok()?;
388    Some(RateLimitKey::User(claims.user_id))
389}
ironflow_api/rate_limit.rs

ironflow_api/
rate_limit.rs
