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//! The configured forwarding header ([`TrustedProxies::header`]) is read
10//! only when that peer is one of the [`TrustedProxies`]: any client can
11//! write it. Without a peer address the IP limit is skipped.
12//!
13//! With [`AccountLimit::ByEmail`], a request whose JSON body names an
14//! `email` is also counted against that account, so spreading attempts on
15//! one account over many IP addresses does not help.
16//!
17//! Responses carry standard rate-limit headers:
18//! - `X-RateLimit-Limit` -- maximum requests allowed per window
19//! - `X-RateLimit-Remaining` -- requests left in the current window
20//! - `X-RateLimit-Reset` -- epoch timestamp when the window resets
21//!
22//! A request is counted under an API key only once its secret is verified,
23//! and the key's `rate_limit_override` replaces the server-wide default only
24//! when the key is owned by an admin. Secrets not yet verified are checked
25//! at most `KEY_CHECKS_PER_MINUTE` times per client and key each minute.
26
27mod account;
28mod client_ip;
29#[cfg(test)]
30mod tests;
31
32use std::net::{IpAddr, SocketAddr};
33use std::sync::Arc;
34use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
35use std::time::{SystemTime, UNIX_EPOCH};
36
37use axum::Json;
38use axum::body::{Body, to_bytes};
39use axum::extract::{ConnectInfo, Request};
40use axum::http::{HeaderMap, HeaderValue, StatusCode};
41use axum::middleware::Next;
42use axum::response::{IntoResponse, Response};
43use chrono::{DateTime, Utc};
44use dashmap::DashMap;
45use serde_json::json;
46use tokio::sync::Mutex;
47use tracing::warn;
48use uuid::Uuid;
49
50use ironflow_auth::cookies::AUTH_COOKIE_NAME;
51use ironflow_auth::extractor::{API_KEY_PREFIX, API_KEY_SUFFIX_LEN};
52use ironflow_auth::jwt::{AccessToken, JwtConfig, token_hash};
53use ironflow_auth::password;
54use ironflow_store::store::Store;
55
56use self::account::{ACCOUNT_BODY_LIMIT, account_key};
57use self::client_ip::client_ip;
58
59pub use self::account::AccountLimit;
60pub use self::client_ip::{
61    ClientIpHeader, InvalidClientIpHeader, InvalidTrustedProxy, TrustedProxies,
62};
63
64/// Counters kept before expired windows are swept. Account keys are chosen
65/// by the caller, so the map must not grow without bound.
66const SWEEP_THRESHOLD: usize = 10_000;
67
68/// Seconds a checked API key token keeps its verdict. Argon2 is too slow to
69/// run on every request, and a minute bounds how long a deactivated key or a
70/// demoted owner keeps its bucket.
71const VERIFIED_KEY_TTL_SECS: u64 = 60;
72
73/// Argon2 verifications allowed per client and per API key each minute for
74/// tokens not in the cache. A legitimate client needs one per key every
75/// `VERIFIED_KEY_TTL_SECS`; more is someone guessing secrets, refused before
76/// Argon2 so the limiter is not a CPU amplifier (a fresh secret on every
77/// request misses the cache every time).
78const KEY_CHECKS_PER_MINUTE: u32 = 10;
79
80/// Rate limit key: identity-based when authenticated, IP-based otherwise.
81///
82/// # Examples
83///
84/// ```
85/// use std::net::{IpAddr, Ipv4Addr};
86/// use uuid::Uuid;
87/// use ironflow_api::rate_limit::RateLimitKey;
88///
89/// let ip_key = RateLimitKey::Ip(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
90/// let user_key = RateLimitKey::User(Uuid::nil());
91/// let api_key = RateLimitKey::ApiKey(Uuid::nil());
92/// let account_key = RateLimitKey::Account {
93///     route: "/sign-in".to_string(),
94///     email: "alice@ironflow.dev".to_string(),
95/// };
96/// let key_check = RateLimitKey::KeyCheck {
97///     ip: Some(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1))),
98///     api_key: Uuid::nil(),
99/// };
100/// ```
101#[derive(Debug, Clone, PartialEq, Eq, Hash)]
102pub enum RateLimitKey {
103    /// Authenticated via API key.
104    ApiKey(Uuid),
105    /// Authenticated via JWT.
106    User(Uuid),
107    /// Unauthenticated, identified by client IP.
108    Ip(IpAddr),
109    /// The account a credential request targets on one route.
110    Account {
111        /// Path of the credential route, so sign-up and sign-in are counted
112        /// apart.
113        route: String,
114        /// Normalized email of the account.
115        email: String,
116    },
117    /// Argon2 verifications of an API key secret not yet in the cache, by
118    /// client IP, or by key alone when the server has no peer address.
119    ///
120    /// Keyed on both so an attacker cannot exhaust the budget of a victim key
121    /// from another address, and one NAT address using many keys is not
122    /// throttled.
123    KeyCheck {
124        /// Client IP, `None` when the server has no peer address.
125        ip: Option<IpAddr>,
126        /// ID of the API key whose secret is checked.
127        api_key: Uuid,
128    },
129}
130
131struct WindowEntry {
132    count: AtomicU32,
133    window_start: AtomicU64,
134}
135
136/// The verdict on one API key token, kept `VERIFIED_KEY_TTL_SECS`.
137struct CachedVerification {
138    /// Key ID and the override it is granted; `None` when the secret was
139    /// checked by Argon2 and is wrong, so a replayed bad token does not cost
140    /// one Argon2 per request.
141    identity: Option<(Uuid, Option<u32>)>,
142    valid_until: u64,
143}
144
145/// Shared rate limiter state with fixed-window counters.
146///
147/// # Examples
148///
149/// ```
150/// use ironflow_api::rate_limit::per_minute;
151///
152/// let limiter = per_minute(60);
153/// ```
154#[derive(Clone)]
155pub struct RateLimitState {
156    counters: Arc<DashMap<RateLimitKey, WindowEntry>>,
157    /// Checked API key tokens, by `token_hash`: the raw token is never kept.
158    verified_keys: Arc<DashMap<String, CachedVerification>>,
159    /// Verifications in flight, by `token_hash`: concurrent requests with the
160    /// same token wait for the first verdict instead of each spending the
161    /// verification budget.
162    pending_checks: Arc<DashMap<String, Arc<Mutex<()>>>>,
163    burst: u32,
164}
165
166/// Context needed by the rate limit middleware.
167///
168/// Combines the limiter state with auth dependencies for identity extraction.
169///
170/// # Examples
171///
172/// ```no_run
173/// use std::sync::Arc;
174/// use ironflow_api::rate_limit::{AccountLimit, RateLimitContext, TrustedProxies, per_minute};
175/// use ironflow_auth::jwt::JwtConfig;
176/// use ironflow_store::memory::InMemoryStore;
177/// use ironflow_store::store::Store;
178///
179/// let store: Arc<dyn Store> = Arc::new(InMemoryStore::new());
180/// let jwt_config = Arc::new(JwtConfig {
181///     secret: "secret".to_string(),
182///     access_token_ttl_secs: 900,
183///     refresh_token_ttl_secs: 604800,
184///     cookie_domain: None,
185///     cookie_secure: false,
186/// });
187/// let ctx = RateLimitContext {
188///     store,
189///     jwt_config,
190///     limiter: per_minute(60),
191///     trusted_proxies: TrustedProxies::default(),
192///     account_limit: AccountLimit::Off,
193/// };
194/// ```
195#[derive(Clone)]
196pub struct RateLimitContext {
197    /// Backing store for API key prefix lookups.
198    pub store: Arc<dyn Store>,
199    /// JWT config for token decoding.
200    pub jwt_config: Arc<JwtConfig>,
201    /// The rate limiter itself.
202    pub limiter: RateLimitState,
203    /// Proxies allowed to report the client IP through forwarding headers.
204    pub trusted_proxies: TrustedProxies,
205    /// Whether requests are also counted against the account they target.
206    pub account_limit: AccountLimit,
207}
208
209struct RateLimitResult {
210    limit: u32,
211    remaining: u32,
212    reset: u64,
213    allowed: bool,
214}
215
216/// Build a per-identity rate limiter allowing `requests_per_minute` requests
217/// per minute.
218///
219/// # Panics
220///
221/// Panics if `requests_per_minute` is 0.
222///
223/// # Examples
224///
225/// ```
226/// use ironflow_api::rate_limit::per_minute;
227///
228/// let auth_limiter = per_minute(10);
229/// let general_limiter = per_minute(60);
230/// ```
231pub fn per_minute(requests_per_minute: u32) -> RateLimitState {
232    assert!(requests_per_minute > 0, "burst must be > 0");
233    RateLimitState {
234        counters: Arc::new(DashMap::new()),
235        verified_keys: Arc::new(DashMap::new()),
236        pending_checks: Arc::new(DashMap::new()),
237        burst: requests_per_minute,
238    }
239}
240
241const WINDOW_SECS: u64 = 60;
242
243fn now_epoch() -> u64 {
244    SystemTime::now()
245        .duration_since(UNIX_EPOCH)
246        .expect("system clock before UNIX epoch")
247        .as_secs()
248}
249
250fn check_rate_limit(limiter: &RateLimitState, key: RateLimitKey, limit: u32) -> RateLimitResult {
251    let now = now_epoch();
252
253    if limiter.counters.len() >= SWEEP_THRESHOLD && !limiter.counters.contains_key(&key) {
254        limiter
255            .counters
256            .retain(|_, entry| entry.window_start.load(Ordering::Acquire) + WINDOW_SECS > now);
257    }
258
259    let entry = limiter.counters.entry(key).or_insert_with(|| WindowEntry {
260        count: AtomicU32::new(0),
261        window_start: AtomicU64::new(now),
262    });
263
264    let window_start = entry.window_start.load(Ordering::Acquire);
265    if now >= window_start + WINDOW_SECS
266        && entry
267            .window_start
268            .compare_exchange(window_start, now, Ordering::AcqRel, Ordering::Relaxed)
269            .is_ok()
270    {
271        entry.count.store(0, Ordering::Release);
272    }
273
274    let count = entry.count.fetch_add(1, Ordering::AcqRel) + 1;
275    let reset = entry.window_start.load(Ordering::Acquire) + WINDOW_SECS;
276
277    if count > limit {
278        entry.count.fetch_sub(1, Ordering::AcqRel);
279        RateLimitResult {
280            limit,
281            remaining: 0,
282            reset,
283            allowed: false,
284        }
285    } else {
286        RateLimitResult {
287            limit,
288            remaining: limit.saturating_sub(count),
289            reset,
290            allowed: true,
291        }
292    }
293}
294
295/// Axum middleware that enforces identity-aware rate limiting.
296///
297/// Extracts the caller's identity from the request:
298/// - API key (`irfl_*` Bearer token whose secret verifies) ->
299///   `RateLimitKey::ApiKey(id)`
300/// - JWT (Bearer token or auth cookie) -> `RateLimitKey::User(id)`
301/// - Fallback -> `RateLimitKey::Ip(addr)`, the TCP peer or, behind a
302///   [trusted proxy](TrustedProxies), the address it forwarded. Without a
303///   peer address no caller limit applies.
304///
305/// A verified API key owned by an admin uses its `rate_limit_override`
306/// instead of the default burst. Once a client has spent its budget of
307/// API key secret verifications for a key, its next unverified secret is
308/// answered with `429` before Argon2 runs. With [`AccountLimit::ByEmail`],
309/// the targeted account is counted too, per route, at the default burst; a
310/// body over 64 KiB is answered with `413` before reaching the handler.
311pub async fn rate_limit(mut req: Request, next: Next) -> Response {
312    let ctx = req.extensions_mut().remove::<RateLimitContext>();
313    let Some(ctx) = ctx else {
314        return next.run(req).await;
315    };
316    let peer = req
317        .extensions()
318        .get::<ConnectInfo<SocketAddr>>()
319        .map(|ConnectInfo(addr)| addr.ip());
320    let identified = match extract_rate_limit_key(req.headers(), peer, &ctx).await {
321        Ok(identified) => identified,
322        Err(result) => return too_many_requests(&result),
323    };
324
325    let caller = match identified {
326        Some((key, override_limit)) => {
327            let limit = override_limit.unwrap_or(ctx.limiter.burst);
328            if limit == 0 {
329                return next.run(req).await;
330            }
331            let result = check_rate_limit(&ctx.limiter, key, limit);
332            if !result.allowed {
333                return too_many_requests(&result);
334            }
335            Some(result)
336        }
337        None => None,
338    };
339
340    let (req, result) = match ctx.account_limit {
341        AccountLimit::Off => (req, caller),
342        AccountLimit::ByEmail => {
343            let (parts, body) = req.into_parts();
344            let Ok(bytes) = to_bytes(body, ACCOUNT_BODY_LIMIT).await else {
345                return payload_too_large();
346            };
347            // The layer only wraps matched static routes (`/sign-up`,
348            // `/sign-in`, `/password`), so a client cannot mint new buckets
349            // through the path.
350            let account = account_key(parts.uri.path(), &bytes);
351            let req = Request::from_parts(parts, Body::from(bytes));
352            match account {
353                None => (req, caller),
354                Some(key) => {
355                    let account = check_rate_limit(&ctx.limiter, key, ctx.limiter.burst);
356                    if !account.allowed {
357                        return too_many_requests(&account);
358                    }
359                    let tightest = match caller {
360                        Some(caller) if caller.remaining <= account.remaining => caller,
361                        _ => account,
362                    };
363                    (req, Some(tightest))
364                }
365            }
366        }
367    };
368
369    let mut resp = next.run(req).await;
370    if let Some(result) = &result {
371        insert_rate_limit_headers(resp.headers_mut(), result);
372    }
373    resp
374}
375
376fn too_many_requests(result: &RateLimitResult) -> Response {
377    let retry_after = result.reset.saturating_sub(now_epoch()).max(1);
378
379    let body = json!({
380        "error": {
381            "code": "RATE_LIMIT_EXCEEDED",
382            "message": "Too many requests, please try again later",
383            "retry_after_secs": retry_after,
384        }
385    });
386
387    let mut resp = (StatusCode::TOO_MANY_REQUESTS, Json(body)).into_response();
388    resp.headers_mut()
389        .insert("retry-after", HeaderValue::from(retry_after));
390    insert_rate_limit_headers(resp.headers_mut(), result);
391    resp
392}
393
394fn payload_too_large() -> Response {
395    let body = json!({
396        "error": {
397            "code": "PAYLOAD_TOO_LARGE",
398            "message": "Request body too large or unreadable",
399        }
400    });
401    (StatusCode::PAYLOAD_TOO_LARGE, Json(body)).into_response()
402}
403
404fn insert_rate_limit_headers(headers: &mut HeaderMap, result: &RateLimitResult) {
405    headers.insert("x-ratelimit-limit", HeaderValue::from(result.limit));
406    headers.insert("x-ratelimit-remaining", HeaderValue::from(result.remaining));
407    headers.insert("x-ratelimit-reset", HeaderValue::from(result.reset));
408}
409
410/// Extract the rate limit key and optional override from the request.
411///
412/// Attempts lightweight identity extraction:
413/// - API key prefix lookup plus an Argon2 verification of the secret, cached
414///   `VERIFIED_KEY_TTL_SECS` and budgeted by `KEY_CHECKS_PER_MINUTE`
415/// - JWT decode (local crypto, no DB)
416/// - Fallback to the client IP
417///
418/// `Ok(None)` when no identity matched and there is no client IP; `Err`
419/// when the API key verification budget is spent.
420async fn extract_rate_limit_key(
421    headers: &HeaderMap,
422    peer: Option<IpAddr>,
423    ctx: &RateLimitContext,
424) -> Result<Option<(RateLimitKey, Option<u32>)>, RateLimitResult> {
425    let ip = client_ip(peer, headers, &ctx.trusted_proxies);
426    let bearer = headers
427        .get("authorization")
428        .and_then(|v| v.to_str().ok())
429        .and_then(|v| v.strip_prefix("Bearer "))
430        .map(|s| s.to_string());
431
432    if let Some(ref token) = bearer {
433        if token.starts_with(API_KEY_PREFIX) {
434            if let Some(found) = try_api_key_identity(token, ip, ctx).await? {
435                return Ok(Some(found));
436            }
437        } else if let Some(key) = try_jwt_identity(token, ctx) {
438            return Ok(Some((key, None)));
439        }
440    }
441
442    // Try JWT from the dashboard auth cookie
443    if let Some(cookie_header) = headers.get("cookie").and_then(|v| v.to_str().ok()) {
444        for part in cookie_header.split(';') {
445            let cookie = part.trim().strip_prefix(AUTH_COOKIE_NAME);
446            if let Some(value) = cookie.and_then(|rest| rest.strip_prefix('='))
447                && let Some(key) = try_jwt_identity(value, ctx)
448            {
449                return Ok(Some((key, None)));
450            }
451        }
452    }
453
454    Ok(ip.map(|ip| (RateLimitKey::Ip(ip), None)))
455}
456
457/// Outcome of checking an API key token against the store.
458enum Verification {
459    /// The secret matches an active, unexpired key.
460    Valid {
461        id: Uuid,
462        /// The key's override, kept only when its owner is an admin.
463        rate_limit_override: Option<u32>,
464        expires_at: Option<DateTime<Utc>>,
465    },
466    /// No such key, or an inactive or expired key. Cheap to tell again, so
467    /// not cached: caching it would let a client grow the cache with fresh
468    /// random tokens.
469    Unknown,
470    /// The secret does not match the key, as checked by Argon2.
471    WrongSecret,
472    /// The store or Argon2 failed: not a verdict, so not cached.
473    Unavailable,
474    /// The verification budget of this client and key is spent.
475    Throttled(RateLimitResult),
476}
477
478/// The API key bucket of `token`, once its secret is verified.
479///
480/// `Ok(None)` when the token is not a usable key, so the request falls back
481/// to the cookie and IP buckets. `Err` when the verification budget is
482/// spent: the request must be answered `429`.
483async fn try_api_key_identity(
484    token: &str,
485    ip: Option<IpAddr>,
486    ctx: &RateLimitContext,
487) -> Result<Option<(RateLimitKey, Option<u32>)>, RateLimitResult> {
488    let now = now_epoch();
489    let hash = token_hash(token);
490    let as_key = |(id, limit): (Uuid, Option<u32>)| (RateLimitKey::ApiKey(id), limit);
491
492    if let Some(identity) = cached_verdict(ctx, &hash, now) {
493        return Ok(identity.map(as_key));
494    }
495
496    // The map guard is dropped at the end of the statement, before the await.
497    let lock = Arc::clone(
498        ctx.limiter
499            .pending_checks
500            .entry(hash.clone())
501            .or_default()
502            .value(),
503    );
504    let _pending = PendingCheck {
505        pending_checks: &ctx.limiter.pending_checks,
506        hash: &hash,
507        lock: Arc::clone(&lock),
508    };
509    let _turn = lock.lock().await;
510    // A concurrent request with the same token may have settled it meanwhile.
511    if let Some(identity) = cached_verdict(ctx, &hash, now) {
512        return Ok(identity.map(as_key));
513    }
514
515    let ttl_end = now + VERIFIED_KEY_TTL_SECS;
516    let (identity, valid_until) = match verify_api_key(token, ip, ctx).await {
517        Verification::Valid {
518            id,
519            rate_limit_override,
520            expires_at,
521        } => {
522            let expiry = expires_at.map_or(u64::MAX, |at| at.timestamp().max(0) as u64);
523            (Some((id, rate_limit_override)), ttl_end.min(expiry))
524        }
525        Verification::WrongSecret => (None, ttl_end),
526        Verification::Unknown | Verification::Unavailable => return Ok(None),
527        Verification::Throttled(result) => return Err(result),
528    };
529
530    let verified_keys = &ctx.limiter.verified_keys;
531    if verified_keys.len() >= SWEEP_THRESHOLD && !verified_keys.contains_key(&hash) {
532        verified_keys.retain(|_, cached| cached.valid_until > now);
533    }
534    let verdict = CachedVerification {
535        identity,
536        valid_until,
537    };
538    verified_keys.insert(hash.clone(), verdict);
539    Ok(identity.map(as_key))
540}
541
542/// The unexpired cached verdict on the token hashed to `hash`.
543fn cached_verdict(
544    ctx: &RateLimitContext,
545    hash: &str,
546    now: u64,
547) -> Option<Option<(Uuid, Option<u32>)>> {
548    // The map guard is dropped inside `map`, before any await or insert.
549    ctx.limiter
550        .verified_keys
551        .get(hash)
552        .filter(|cached| cached.valid_until > now)
553        .map(|cached| cached.identity)
554}
555
556/// Removes the in-flight entry of a verification when it ends, even when the
557/// request is cancelled mid-await.
558struct PendingCheck<'a> {
559    pending_checks: &'a DashMap<String, Arc<Mutex<()>>>,
560    hash: &'a str,
561    lock: Arc<Mutex<()>>,
562}
563
564impl Drop for PendingCheck<'_> {
565    fn drop(&mut self) {
566        // A later request may have registered a new lock once this one was
567        // removed: only remove our own.
568        self.pending_checks
569            .remove_if(self.hash, |_, lock| Arc::ptr_eq(lock, &self.lock));
570    }
571}
572
573/// Check `token` the way the `ApiKeyAuth` extractor does, without touching
574/// the key's `last_used_at`.
575async fn verify_api_key(token: &str, ip: Option<IpAddr>, ctx: &RateLimitContext) -> Verification {
576    let suffix_len = (token.len() - API_KEY_PREFIX.len()).min(API_KEY_SUFFIX_LEN);
577    let prefix = &token[..API_KEY_PREFIX.len() + suffix_len];
578
579    let api_key = match ctx.store.find_api_key_by_prefix(prefix).await {
580        Ok(Some(api_key)) => api_key,
581        Ok(None) => return Verification::Unknown,
582        Err(err) => {
583            warn!(error = %err, "rate limiter could not look up an API key");
584            return Verification::Unavailable;
585        }
586    };
587
588    if !api_key.is_active || api_key.expires_at.is_some_and(|at| at <= Utc::now()) {
589        return Verification::Unknown;
590    }
591
592    // Only reached on a cache miss for an existing, active, unexpired key, so
593    // each Argon2 below is paid for out of this budget.
594    let check = check_rate_limit(
595        &ctx.limiter,
596        RateLimitKey::KeyCheck {
597            ip,
598            api_key: api_key.id,
599        },
600        KEY_CHECKS_PER_MINUTE,
601    );
602    if !check.allowed {
603        return Verification::Throttled(check);
604    }
605
606    match password::verify(token, &api_key.key_hash) {
607        Ok(true) => {}
608        Ok(false) => return Verification::WrongSecret,
609        Err(err) => {
610            warn!(error = %err, "rate limiter could not verify an API key");
611            return Verification::Unavailable;
612        }
613    }
614
615    // Only admins may set an override; one a member set before that rule is
616    // ignored.
617    let rate_limit_override = match ctx.store.find_user_by_id(api_key.user_id).await {
618        Ok(Some(owner)) if owner.is_admin => api_key.rate_limit_override,
619        Ok(_) => None,
620        Err(err) => {
621            warn!(error = %err, "rate limiter could not look up an API key owner");
622            return Verification::Unavailable;
623        }
624    };
625
626    Verification::Valid {
627        id: api_key.id,
628        rate_limit_override,
629        expires_at: api_key.expires_at,
630    }
631}
632
633fn try_jwt_identity(token: &str, ctx: &RateLimitContext) -> Option<RateLimitKey> {
634    let claims = AccessToken::decode(token, &ctx.jwt_config).ok()?;
635    Some(RateLimitKey::User(claims.user_id))
636}
ironflow_api/rate_limit.rs

ironflow_api/
rate_limit.rs
