1use std::collections::HashMap;
29use std::net::IpAddr;
30use std::sync::Mutex;
31use std::time::{SystemTime, UNIX_EPOCH};
32
33use anyhow::{Context, Result};
34use tracing::warn;
35
36use crate::config::{parse_rate, RateLimitCfg};
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum StoreMode {
41 Local,
43 Memory,
45 Redis,
47}
48
49impl StoreMode {
50 pub fn parse(s: &str) -> Result<StoreMode> {
51 match s.trim().to_ascii_lowercase().as_str() {
52 "local" | "governor" | "" => Ok(StoreMode::Local),
53 "memory" | "in-memory" => Ok(StoreMode::Memory),
54 "redis" => Ok(StoreMode::Redis),
55 other => {
56 anyhow::bail!("invalid ratelimit.store {other:?} (expected local|memory|redis)")
57 }
58 }
59 }
60
61 pub fn is_distributed(self) -> bool {
64 matches!(self, StoreMode::Memory | StoreMode::Redis)
65 }
66}
67
68#[derive(Debug, Clone, Copy)]
72pub struct Gcra {
73 emission_interval: u64,
74 tolerance: u64,
75}
76
77impl Gcra {
78 pub fn from_parts(count: u64, period: std::time::Duration, burst: u64) -> Result<Gcra> {
85 anyhow::ensure!(count > 0, "rate count must be > 0");
86 anyhow::ensure!(burst > 0, "burst must be > 0");
87 let period_us = period.as_micros() as u64;
88 let emission_interval = period_us / count;
89 anyhow::ensure!(emission_interval > 0, "rate is too high to represent");
90 Ok(Gcra {
91 emission_interval,
92 tolerance: emission_interval.saturating_mul(burst),
93 })
94 }
95
96 pub fn next_admit_at(&self, stored_tat: Option<u64>, now: u64) -> u64 {
103 let tat = stored_tat.unwrap_or(now).max(now);
104 tat.saturating_add(self.emission_interval)
105 .saturating_sub(self.tolerance)
106 }
107
108 pub fn remaining(&self, stored_tat: Option<u64>, now: u64) -> u64 {
115 let burst = self.tolerance / self.emission_interval.max(1);
116 let tat = stored_tat.unwrap_or(now).max(now);
117 let spent = (tat - now) / self.emission_interval.max(1);
118 burst.saturating_sub(spent)
119 }
120
121 pub fn from_rate(rate: &str, burst: u32) -> Result<Gcra> {
124 let (count, period) = parse_rate(rate)?;
125 anyhow::ensure!(count > 0, "rate count must be > 0 (got {rate:?})");
126 anyhow::ensure!(burst > 0, "burst must be > 0 (rate {rate:?})");
127 let period_us = period.as_micros() as u64;
128 let emission_interval = period_us / count as u64;
129 anyhow::ensure!(
130 emission_interval > 0,
131 "rate too high for a usable sub-microsecond interval: {rate:?}"
132 );
133 let tolerance = emission_interval.saturating_mul(burst as u64);
134 Ok(Gcra {
135 emission_interval,
136 tolerance,
137 })
138 }
139}
140
141pub(crate) fn gcra_admit(stored_tat: Option<u64>, now: u64, g: &Gcra) -> Option<u64> {
147 let tat = stored_tat.unwrap_or(now).max(now);
149 let new_tat = tat + g.emission_interval;
150 let allow_at = new_tat.saturating_sub(g.tolerance);
151 if now < allow_at {
152 None
153 } else {
154 Some(new_tat)
155 }
156}
157
158enum Store {
161 Memory(MemoryStore),
162 Redis(Box<RedisStore>),
163}
164
165impl Store {
166 async fn admit(&self, key: &str, g: &Gcra, now: u64) -> Result<bool> {
168 match self {
169 Store::Memory(s) => Ok(s.admit(key, g, now)),
170 Store::Redis(s) => s.admit(key, g, now).await,
171 }
172 }
173}
174
175#[derive(Default)]
179struct MemoryStore {
180 tats: Mutex<HashMap<String, u64>>,
181}
182
183impl MemoryStore {
184 fn admit(&self, key: &str, g: &Gcra, now: u64) -> bool {
185 let mut map = self.tats.lock().expect("limiter store mutex poisoned");
186 match gcra_admit(map.get(key).copied(), now, g) {
187 Some(new_tat) => {
188 map.insert(key.to_string(), new_tat);
189 true
190 }
191 None => false,
192 }
193 }
194}
195
196const GCRA_LUA: &str = r#"
201local tat = redis.call('GET', KEYS[1])
202local now = tonumber(ARGV[1])
203local interval = tonumber(ARGV[2])
204local tolerance = tonumber(ARGV[3])
205if tat == false then
206 tat = now
207else
208 tat = tonumber(tat)
209 if tat < now then tat = now end
210end
211local new_tat = tat + interval
212local allow_at = new_tat - tolerance
213if now < allow_at then
214 return 0
215end
216local ttl_ms = math.ceil((new_tat - now) / 1000)
217if ttl_ms < 1 then ttl_ms = 1 end
218redis.call('SET', KEYS[1], new_tat, 'PX', ttl_ms)
219return 1
220"#;
221
222struct RedisStore {
226 client: redis::Client,
227 conn: tokio::sync::OnceCell<redis::aio::ConnectionManager>,
228 script: redis::Script,
229}
230
231impl RedisStore {
232 fn new(url: &str) -> Result<RedisStore> {
233 anyhow::ensure!(
234 !url.trim().is_empty(),
235 "ratelimit.redis_url is required when ratelimit.store = \"redis\""
236 );
237 let client = redis::Client::open(url)
238 .with_context(|| format!("opening redis client for {url:?} (ratelimit.redis_url)"))?;
239 Ok(RedisStore {
240 client,
241 conn: tokio::sync::OnceCell::new(),
242 script: redis::Script::new(GCRA_LUA),
243 })
244 }
245
246 async fn admit(&self, key: &str, g: &Gcra, now: u64) -> Result<bool> {
247 let manager = self
248 .conn
249 .get_or_try_init(|| redis::aio::ConnectionManager::new(self.client.clone()))
250 .await
251 .context("connecting to redis rate-limit store")?;
252 let mut conn = manager.clone();
253 let admitted: i64 = self
254 .script
255 .key(key)
256 .arg(now)
257 .arg(g.emission_interval)
258 .arg(g.tolerance)
259 .invoke_async(&mut conn)
260 .await
261 .context("evaluating redis GCRA script")?;
262 Ok(admitted == 1)
263 }
264}
265
266#[derive(Debug, Clone, Copy, PartialEq, Eq)]
268pub enum Admit {
269 Allowed,
271 Limited(&'static str),
273 Error,
275}
276
277struct RouteGcra {
279 prefix: String,
280 gcra: Gcra,
281}
282
283pub struct DistributedLimiter {
287 store: Store,
288 key_prefix: String,
289 fail_open: bool,
290 global: Gcra,
291 routes: Vec<RouteGcra>,
292 per_key: Option<Gcra>,
293}
294
295impl DistributedLimiter {
296 pub fn build(rl: &RateLimitCfg, mode: StoreMode) -> Result<DistributedLimiter> {
300 let store = match mode {
301 StoreMode::Memory => Store::Memory(MemoryStore::default()),
302 StoreMode::Redis => Store::Redis(Box::new(RedisStore::new(&rl.redis_url)?)),
303 StoreMode::Local => {
304 anyhow::bail!("DistributedLimiter::build called for the local store")
305 }
306 };
307
308 let global = Gcra::from_rate(&rl.rate, rl.burst)?;
309 let mut routes = Vec::new();
310 for route in &rl.routes {
311 anyhow::ensure!(
312 !route.path.is_empty(),
313 "ratelimit.routes[].path must not be empty"
314 );
315 routes.push(RouteGcra {
316 prefix: route.path.clone(),
317 gcra: Gcra::from_rate(&route.rate, route.burst)?,
318 });
319 }
320 let per_key = if rl.per_key.enabled {
321 Some(Gcra::from_rate(&rl.per_key.rate, rl.per_key.burst)?)
322 } else {
323 None
324 };
325
326 Ok(DistributedLimiter {
327 store,
328 key_prefix: rl.redis_prefix.clone(),
329 fail_open: rl.fail_open,
330 global,
331 routes,
332 per_key,
333 })
334 }
335
336 pub async fn check_ip_route(&self, ip: IpAddr, path: &str) -> Admit {
339 let now = now_micros();
340 if let Some(route) = self
341 .routes
342 .iter()
343 .filter(|r| path.starts_with(&r.prefix))
344 .max_by_key(|r| r.prefix.len())
345 {
346 let key = format!("{}:route:{}:{}", self.key_prefix, route.prefix, ip);
347 self.admit(&key, &route.gcra, now, "route").await
348 } else {
349 let key = format!("{}:ip:{}", self.key_prefix, ip);
350 self.admit(&key, &self.global, now, "ip").await
351 }
352 }
353
354 pub async fn check_key(&self, principal: &str) -> Admit {
357 match &self.per_key {
358 Some(gcra) => {
359 let now = now_micros();
360 let key = format!("{}:key:{}", self.key_prefix, principal);
361 self.admit(&key, gcra, now, "key").await
362 }
363 None => Admit::Allowed,
364 }
365 }
366
367 async fn admit(&self, key: &str, g: &Gcra, now: u64, scope: &'static str) -> Admit {
368 match self.store.admit(key, g, now).await {
369 Ok(true) => Admit::Allowed,
370 Ok(false) => Admit::Limited(scope),
371 Err(e) => {
372 if self.fail_open {
373 warn!(error = %format!("{e:#}"), scope, "rate-limit store error; failing open (allowing request)");
374 Admit::Allowed
375 } else {
376 warn!(error = %format!("{e:#}"), scope, "rate-limit store error; failing closed (503)");
377 Admit::Error
378 }
379 }
380 }
381 }
382}
383
384fn now_micros() -> u64 {
386 SystemTime::now()
387 .duration_since(UNIX_EPOCH)
388 .map(|d| u64::try_from(d.as_micros()).unwrap_or(u64::MAX))
390 .unwrap_or(0)
391}
392
393#[cfg(test)]
394mod tests {
395 use super::*;
396 use crate::config::{PerKeyRateLimit, RouteRateLimit};
397
398 fn gcra(rate: &str, burst: u32) -> Gcra {
399 Gcra::from_rate(rate, burst).unwrap()
400 }
401
402 #[test]
403 fn store_mode_parses_and_classifies() {
404 assert_eq!(StoreMode::parse("local").unwrap(), StoreMode::Local);
405 assert_eq!(StoreMode::parse("").unwrap(), StoreMode::Local);
406 assert_eq!(StoreMode::parse("REDIS").unwrap(), StoreMode::Redis);
407 assert_eq!(StoreMode::parse(" memory ").unwrap(), StoreMode::Memory);
408 assert!(StoreMode::parse("dynamo").is_err());
409 assert!(!StoreMode::parse("local").unwrap().is_distributed());
410 assert!(StoreMode::parse("redis").unwrap().is_distributed());
411 assert!(StoreMode::parse("memory").unwrap().is_distributed());
412 }
413
414 #[test]
415 fn gcra_from_rate_rejects_degenerate_input() {
416 assert!(Gcra::from_rate("0/sec", 5).is_err()); assert!(Gcra::from_rate("10/sec", 0).is_err()); assert!(Gcra::from_rate("nonsense", 5).is_err());
419 }
420
421 #[test]
422 fn gcra_admit_allows_burst_then_rejects_at_same_instant() {
423 let g = gcra("1/sec", 3);
425 let now = 1_000_000_000;
426 let mut tat = None;
427 for _ in 0..3 {
428 let next = gcra_admit(tat, now, &g);
429 assert!(next.is_some(), "within-burst request should be admitted");
430 tat = next;
431 }
432 assert!(
433 gcra_admit(tat, now, &g).is_none(),
434 "the request past the burst must be rejected"
435 );
436 }
437
438 #[test]
439 fn gcra_admit_recovers_after_emission_interval() {
440 let g = gcra("1/sec", 1);
442 let t0 = 5_000_000_000;
443 let tat = gcra_admit(None, t0, &g).expect("first admitted");
444 assert!(
445 gcra_admit(Some(tat), t0, &g).is_none(),
446 "immediate second rejected"
447 );
448 assert!(
450 gcra_admit(Some(tat), t0 + 1_000_000, &g).is_some(),
451 "request after the interval admitted"
452 );
453 }
454
455 #[test]
456 fn gcra_admit_does_not_advance_tat_on_rejection() {
457 let g = gcra("1/min", 1);
458 let now = 2_000_000_000;
459 let tat = gcra_admit(None, now, &g).unwrap();
460 assert!(gcra_admit(Some(tat), now, &g).is_none());
463 assert!(gcra_admit(Some(tat), now, &g).is_none());
464 }
465
466 #[tokio::test]
467 async fn memory_store_enforces_global_limit() {
468 let rl = RateLimitCfg {
469 enabled: true,
470 rate: "1/min".into(),
471 burst: 1,
472 store: "memory".into(),
473 ..Default::default()
474 };
475 let limiter = DistributedLimiter::build(&rl, StoreMode::Memory).unwrap();
476 let ip: IpAddr = "203.0.113.7".parse().unwrap();
477
478 assert_eq!(limiter.check_ip_route(ip, "/").await, Admit::Allowed);
480 assert_eq!(limiter.check_ip_route(ip, "/").await, Admit::Limited("ip"));
481 let ip2: IpAddr = "203.0.113.8".parse().unwrap();
483 assert_eq!(limiter.check_ip_route(ip2, "/").await, Admit::Allowed);
484 }
485
486 #[tokio::test]
487 async fn memory_store_applies_per_route_override() {
488 let rl = RateLimitCfg {
489 enabled: true,
490 rate: "1000/min".into(), burst: 1000,
492 routes: vec![RouteRateLimit {
493 path: "/api/".into(),
494 rate: "1/min".into(),
495 burst: 1,
496 }],
497 store: "memory".into(),
498 ..Default::default()
499 };
500 let limiter = DistributedLimiter::build(&rl, StoreMode::Memory).unwrap();
501 let ip: IpAddr = "198.51.100.4".parse().unwrap();
502
503 assert_eq!(limiter.check_ip_route(ip, "/api/x").await, Admit::Allowed);
505 assert_eq!(
506 limiter.check_ip_route(ip, "/api/x").await,
507 Admit::Limited("route")
508 );
509 assert_eq!(limiter.check_ip_route(ip, "/public").await, Admit::Allowed);
510 }
511
512 #[tokio::test]
513 async fn memory_store_per_key_limit() {
514 let rl = RateLimitCfg {
515 enabled: true,
516 rate: "1000/min".into(),
517 burst: 1000,
518 per_key: PerKeyRateLimit {
519 enabled: true,
520 rate: "1/min".into(),
521 burst: 1,
522 },
523 store: "memory".into(),
524 ..Default::default()
525 };
526 let limiter = DistributedLimiter::build(&rl, StoreMode::Memory).unwrap();
527
528 assert_eq!(limiter.check_key("apikey:abc").await, Admit::Allowed);
529 assert_eq!(limiter.check_key("apikey:abc").await, Admit::Limited("key"));
530 assert_eq!(limiter.check_key("apikey:def").await, Admit::Allowed);
532 }
533
534 #[tokio::test]
535 async fn per_key_disabled_always_allows() {
536 let rl = RateLimitCfg {
537 enabled: true,
538 store: "memory".into(),
539 ..Default::default()
540 };
541 let limiter = DistributedLimiter::build(&rl, StoreMode::Memory).unwrap();
542 assert_eq!(limiter.check_key("whoever").await, Admit::Allowed);
543 }
544
545 #[test]
546 fn redis_store_requires_a_url() {
547 let rl = RateLimitCfg {
548 enabled: true,
549 store: "redis".into(),
550 redis_url: "".into(),
551 ..Default::default()
552 };
553 assert!(DistributedLimiter::build(&rl, StoreMode::Redis).is_err());
554 let bad = RateLimitCfg {
556 enabled: true,
557 store: "redis".into(),
558 redis_url: "not-a-redis-url".into(),
559 ..Default::default()
560 };
561 assert!(DistributedLimiter::build(&bad, StoreMode::Redis).is_err());
562 }
563
564 fn redis_url() -> String {
578 std::env::var("EDGEGUARD_TEST_REDIS_URL")
579 .unwrap_or_else(|_| "redis://127.0.0.1:6379".into())
580 }
581
582 #[tokio::test]
583 #[ignore = "requires a live Redis (EDGEGUARD_TEST_REDIS_URL, default redis://127.0.0.1:6379)"]
584 async fn redis_store_enforces_global_limit_live() {
585 let rl = RateLimitCfg {
586 enabled: true,
587 rate: "1/min".into(),
588 burst: 3,
589 store: "redis".into(),
590 redis_url: redis_url(),
591 redis_prefix: format!("egtest:global:{}:{}", std::process::id(), now_micros()),
592 ..Default::default()
593 };
594 let limiter = DistributedLimiter::build(&rl, StoreMode::Redis).unwrap();
595 let ip: IpAddr = "203.0.113.20".parse().unwrap();
596
597 match limiter.check_ip_route(ip, "/").await {
600 Admit::Error => {
601 eprintln!("skipping redis_store_enforces_global_limit_live: Redis unreachable");
602 return;
603 }
604 Admit::Allowed => {}
605 other => panic!("unexpected first admit: {other:?}"),
606 }
607 assert_eq!(limiter.check_ip_route(ip, "/").await, Admit::Allowed);
609 assert_eq!(limiter.check_ip_route(ip, "/").await, Admit::Allowed);
610 assert_eq!(limiter.check_ip_route(ip, "/").await, Admit::Limited("ip"));
611 let ip2: IpAddr = "203.0.113.21".parse().unwrap();
613 assert_eq!(limiter.check_ip_route(ip2, "/").await, Admit::Allowed);
614 }
615
616 #[tokio::test]
617 #[ignore = "requires a live Redis (EDGEGUARD_TEST_REDIS_URL, default redis://127.0.0.1:6379)"]
618 async fn redis_store_per_key_limit_live() {
619 let rl = RateLimitCfg {
620 enabled: true,
621 rate: "1000/min".into(), burst: 1000,
623 per_key: PerKeyRateLimit {
624 enabled: true,
625 rate: "1/min".into(),
626 burst: 1,
627 },
628 store: "redis".into(),
629 redis_url: redis_url(),
630 redis_prefix: format!("egtest:key:{}:{}", std::process::id(), now_micros()),
631 ..Default::default()
632 };
633 let limiter = DistributedLimiter::build(&rl, StoreMode::Redis).unwrap();
634
635 match limiter.check_key("apikey:abc").await {
636 Admit::Error => {
637 eprintln!("skipping redis_store_per_key_limit_live: Redis unreachable");
638 return;
639 }
640 Admit::Allowed => {}
641 other => panic!("unexpected first admit: {other:?}"),
642 }
643 assert_eq!(limiter.check_key("apikey:abc").await, Admit::Limited("key"));
645 assert_eq!(limiter.check_key("apikey:def").await, Admit::Allowed);
647 }
648}