mold_server/
rate_limit.rs

1use axum::{
2    extract::{ConnectInfo, Request},
3    http::{HeaderValue, StatusCode},
4    middleware::Next,
5    response::{IntoResponse, Response},
6    Json,
7};
8use governor::{
9    clock::DefaultClock,
10    state::{InMemoryState, NotKeyed},
11    Quota, RateLimiter,
12};
13use serde::Serialize;
14use std::collections::HashMap;
15use std::net::{IpAddr, SocketAddr};
16use std::num::NonZeroU32;
17use std::sync::{Arc, Mutex};
18use std::time::Duration;
19use tracing::warn;
20
21type IpRateLimiter = RateLimiter<NotKeyed, InMemoryState, DefaultClock>;
22
23/// Maximum number of per-IP limiter entries before eviction.
24/// At ~200 bytes per entry, 10,000 entries ≈ 2 MB — bounded and safe.
25pub(crate) const MAX_LIMITER_ENTRIES: usize = 10_000;
26
27/// Per-IP rate limiter state. Each IP gets its own limiter instance.
28/// Maps are bounded to [`MAX_LIMITER_ENTRIES`] — when full, the entire map
29/// is cleared to reclaim memory (simple but effective against OOM from
30/// IP-spoofing attacks; legitimate clients just get a fresh bucket).
31pub struct RateLimitState {
32    pub generation_quota: Quota,
33    pub read_quota: Quota,
34    pub(crate) generation_limiters: Mutex<HashMap<IpAddr, Arc<IpRateLimiter>>>,
35    pub(crate) read_limiters: Mutex<HashMap<IpAddr, Arc<IpRateLimiter>>>,
36}
37
38impl RateLimitState {
39    pub(crate) fn new(generation_quota: Quota, read_quota: Quota) -> Self {
40        Self {
41            generation_quota,
42            read_quota,
43            generation_limiters: Mutex::new(HashMap::new()),
44            read_limiters: Mutex::new(HashMap::new()),
45        }
46    }
47
48    pub(crate) fn get_generation_limiter(&self, ip: IpAddr) -> Arc<IpRateLimiter> {
49        let mut map = self.generation_limiters.lock().unwrap();
50        if map.len() >= MAX_LIMITER_ENTRIES {
51            warn!(
52                entries = map.len(),
53                "generation rate limiter map exceeded cap, evicting all entries"
54            );
55            map.clear();
56        }
57        map.entry(ip)
58            .or_insert_with(|| Arc::new(RateLimiter::direct(self.generation_quota)))
59            .clone()
60    }
61
62    fn get_read_limiter(&self, ip: IpAddr) -> Arc<IpRateLimiter> {
63        let mut map = self.read_limiters.lock().unwrap();
64        if map.len() >= MAX_LIMITER_ENTRIES {
65            warn!(
66                entries = map.len(),
67                "read rate limiter map exceeded cap, evicting all entries"
68            );
69            map.clear();
70        }
71        map.entry(ip)
72            .or_insert_with(|| Arc::new(RateLimiter::direct(self.read_quota)))
73            .clone()
74    }
75}
76
77/// Shared rate limit state — `None` means rate limiting is disabled.
78pub type RateLimitConfig = Option<Arc<RateLimitState>>;
79
80#[derive(Debug, Serialize)]
81struct RateLimitError {
82    error: String,
83    code: String,
84}
85
86/// Parse the `MOLD_RATE_LIMIT` env var and build rate limiter state.
87///
88/// Format: `N/period` where period is `sec`, `min`, or `hour`.
89/// Examples: `10/min`, `5/sec`, `100/hour`
90///
91/// Returns `None` when unset (rate limiting disabled).
92pub fn load_rate_limit_config() -> anyhow::Result<RateLimitConfig> {
93    let raw = match std::env::var("MOLD_RATE_LIMIT") {
94        Ok(v) if !v.is_empty() => v,
95        _ => return Ok(None),
96    };
97
98    let (count, period) = parse_rate_spec(&raw)?;
99
100    let burst = match std::env::var("MOLD_RATE_LIMIT_BURST") {
101        Ok(v) if !v.is_empty() => v
102            .parse::<u32>()
103            .map_err(|_| anyhow::anyhow!("MOLD_RATE_LIMIT_BURST must be a positive integer"))?,
104        _ => (count * 2).min(100), // Default: 2x rate, capped at 100
105    };
106
107    let generation_quota = build_quota(count, burst, period)?;
108
109    // Read endpoints get 10x the generation limit.
110    let read_count = (count * 10).min(1000);
111    let read_burst = (burst * 10).min(1000);
112    let read_quota = build_quota(read_count, read_burst, period)?;
113
114    tracing::info!(
115        rate = %raw,
116        burst,
117        read_multiplier = 10,
118        "rate limiting enabled"
119    );
120
121    Ok(Some(Arc::new(RateLimitState::new(
122        generation_quota,
123        read_quota,
124    ))))
125}
126
127fn parse_rate_spec(spec: &str) -> anyhow::Result<(u32, Duration)> {
128    let parts: Vec<&str> = spec.splitn(2, '/').collect();
129    if parts.len() != 2 {
130        anyhow::bail!("MOLD_RATE_LIMIT must be in the format N/period (e.g., 10/min)");
131    }
132
133    let count: u32 = parts[0]
134        .trim()
135        .parse()
136        .map_err(|_| anyhow::anyhow!("MOLD_RATE_LIMIT count must be a positive integer"))?;
137    if count == 0 {
138        anyhow::bail!("MOLD_RATE_LIMIT count must be greater than 0");
139    }
140
141    let period = match parts[1].trim() {
142        "sec" | "second" | "s" => Duration::from_secs(1),
143        "min" | "minute" | "m" => Duration::from_secs(60),
144        "hour" | "hr" | "h" => Duration::from_secs(3600),
145        other => anyhow::bail!("unknown MOLD_RATE_LIMIT period: {other} (use sec, min, or hour)"),
146    };
147
148    Ok((count, period))
149}
150
151fn build_quota(count: u32, burst: u32, period: Duration) -> anyhow::Result<Quota> {
152    let replenish_interval = period / count;
153    let burst_nz = NonZeroU32::new(burst)
154        .ok_or_else(|| anyhow::anyhow!("rate limit burst must be greater than 0"))?;
155    Ok(Quota::with_period(replenish_interval)
156        .ok_or_else(|| anyhow::anyhow!("invalid rate limit period"))?
157        .allow_burst(burst_nz))
158}
159
160/// Route tier for rate limiting.
161#[derive(Clone, Copy, Debug, PartialEq, Eq)]
162pub enum RouteTier {
163    /// Expensive operations: generate, expand, model load/pull/unload.
164    Generation,
165    /// Cheap reads: list models, status, gallery.
166    Read,
167}
168
169/// Paths and their rate limit tiers.
170pub fn classify_route(path: &str, method: &axum::http::Method) -> Option<RouteTier> {
171    // Health/docs are not rate limited.
172    if path == "/health" || path == "/api/docs" || path == "/api/openapi.json" {
173        return None;
174    }
175
176    match (method.as_str(), path) {
177        (
178            "POST",
179            "/api/generate"
180            | "/api/generate/stream"
181            | "/api/expand"
182            | "/api/upscale"
183            | "/api/upscale/stream",
184        ) => Some(RouteTier::Generation),
185        ("POST", "/api/models/load" | "/api/models/pull") => Some(RouteTier::Generation),
186        ("DELETE", "/api/models/unload") => Some(RouteTier::Generation),
187        ("DELETE", _) if path.starts_with("/api/gallery/") => Some(RouteTier::Generation),
188        ("GET", _) => Some(RouteTier::Read),
189        _ => Some(RouteTier::Read), // Default unknown routes to read tier
190    }
191}
192
193/// Axum middleware that enforces per-IP rate limiting.
194pub async fn rate_limit_middleware(
195    connect_info: ConnectInfo<SocketAddr>,
196    request: Request,
197    next: Next,
198) -> Response {
199    let rl_config = request.extensions().get::<RateLimitConfig>().cloned();
200
201    let state = match rl_config.as_ref().and_then(|c| c.as_ref()) {
202        Some(s) => s,
203        None => return next.run(request).await, // Rate limiting disabled
204    };
205
206    let ip = connect_info.0.ip();
207    let tier = classify_route(request.uri().path(), request.method());
208
209    let tier = match tier {
210        Some(t) => t,
211        None => return next.run(request).await, // Exempt path
212    };
213
214    let limiter = match tier {
215        RouteTier::Generation => state.get_generation_limiter(ip),
216        RouteTier::Read => state.get_read_limiter(ip),
217    };
218
219    match limiter.check() {
220        Ok(_) => next.run(request).await,
221        Err(not_until) => {
222            let retry_after = not_until.wait_time_from(governor::clock::Clock::now(
223                &governor::clock::DefaultClock::default(),
224            ));
225            let retry_secs = retry_after.as_secs().max(1);
226            warn!(
227                ip = %ip,
228                retry_after_secs = retry_secs,
229                tier = ?match tier { RouteTier::Generation => "generation", RouteTier::Read => "read" },
230                "rate limit exceeded"
231            );
232            rate_limited_response(retry_secs)
233        }
234    }
235}
236
237fn rate_limited_response(retry_after_secs: u64) -> Response {
238    let body = RateLimitError {
239        error: "rate limit exceeded".to_string(),
240        code: "RATE_LIMITED".to_string(),
241    };
242    let mut response = (StatusCode::TOO_MANY_REQUESTS, Json(body)).into_response();
243    if let Ok(val) = HeaderValue::from_str(&retry_after_secs.to_string()) {
244        response.headers_mut().insert("retry-after", val);
245    }
246    response
247}
248
249/// Injects the `RateLimitConfig` as a request extension.
250pub async fn inject_rate_limit_state(
251    axum::extract::State(rl): axum::extract::State<RateLimitConfig>,
252    mut request: Request,
253    next: Next,
254) -> Response {
255    request.extensions_mut().insert(rl);
256    next.run(request).await
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262
263    #[test]
264    fn parse_rate_per_minute() {
265        let (count, period) = parse_rate_spec("10/min").unwrap();
266        assert_eq!(count, 10);
267        assert_eq!(period, Duration::from_secs(60));
268    }
269
270    #[test]
271    fn parse_rate_per_second() {
272        let (count, period) = parse_rate_spec("5/sec").unwrap();
273        assert_eq!(count, 5);
274        assert_eq!(period, Duration::from_secs(1));
275    }
276
277    #[test]
278    fn parse_rate_per_hour() {
279        let (count, period) = parse_rate_spec("100/hour").unwrap();
280        assert_eq!(count, 100);
281        assert_eq!(period, Duration::from_secs(3600));
282    }
283
284    #[test]
285    fn parse_rate_short_aliases() {
286        assert!(parse_rate_spec("1/s").is_ok());
287        assert!(parse_rate_spec("1/m").is_ok());
288        assert!(parse_rate_spec("1/h").is_ok());
289    }
290
291    #[test]
292    fn parse_rate_invalid_format() {
293        assert!(parse_rate_spec("10").is_err());
294        assert!(parse_rate_spec("10/xyz").is_err());
295        assert!(parse_rate_spec("0/min").is_err());
296        assert!(parse_rate_spec("abc/min").is_err());
297    }
298
299    #[test]
300    fn classify_generation_routes() {
301        use axum::http::Method;
302        assert_eq!(
303            classify_route("/api/generate", &Method::POST),
304            Some(RouteTier::Generation)
305        );
306        assert_eq!(
307            classify_route("/api/generate/stream", &Method::POST),
308            Some(RouteTier::Generation)
309        );
310        assert_eq!(
311            classify_route("/api/expand", &Method::POST),
312            Some(RouteTier::Generation)
313        );
314        assert_eq!(
315            classify_route("/api/models/load", &Method::POST),
316            Some(RouteTier::Generation)
317        );
318        assert_eq!(
319            classify_route("/api/models/pull", &Method::POST),
320            Some(RouteTier::Generation)
321        );
322        assert_eq!(
323            classify_route("/api/models/unload", &Method::DELETE),
324            Some(RouteTier::Generation)
325        );
326        assert_eq!(
327            classify_route("/api/upscale", &Method::POST),
328            Some(RouteTier::Generation)
329        );
330        assert_eq!(
331            classify_route("/api/upscale/stream", &Method::POST),
332            Some(RouteTier::Generation)
333        );
334    }
335
336    #[test]
337    fn classify_read_routes() {
338        use axum::http::Method;
339        assert_eq!(
340            classify_route("/api/models", &Method::GET),
341            Some(RouteTier::Read)
342        );
343        assert_eq!(
344            classify_route("/api/status", &Method::GET),
345            Some(RouteTier::Read)
346        );
347        assert_eq!(
348            classify_route("/api/gallery", &Method::GET),
349            Some(RouteTier::Read)
350        );
351    }
352
353    #[test]
354    fn classify_exempt_routes() {
355        use axum::http::Method;
356        assert_eq!(classify_route("/health", &Method::GET), None);
357        assert_eq!(classify_route("/api/docs", &Method::GET), None);
358        assert_eq!(classify_route("/api/openapi.json", &Method::GET), None);
359    }
360
361    #[test]
362    fn build_quota_valid() {
363        let q = build_quota(10, 20, Duration::from_secs(60));
364        assert!(q.is_ok());
365    }
366
367    #[test]
368    fn rate_limiter_rejects_after_burst() {
369        let quota = build_quota(1, 2, Duration::from_secs(60)).unwrap();
370        let limiter = RateLimiter::direct(quota);
371        assert!(limiter.check().is_ok()); // 1st
372        assert!(limiter.check().is_ok()); // 2nd (burst)
373        assert!(limiter.check().is_err()); // 3rd → rejected
374    }
375}
mold_server/rate_limit.rs

mold_server/
rate_limit.rs
