zuka
zuka/src/http/ratelimit.rs

worklyn / zukapublic

Agent-first git hosting. One Rust binary: git over HTTP and SSH, a REST API, MCP, CI, and multi-tenant isolation.

Get a copy: git clone https://zuka.worklyn.com/worklyn/zuka.git
branchbc985d6bbd321be340e3750b3da4734ee1553cbd ▾
main All branches →
zuka/src/http/ratelimit.rs
RSratelimit.rs5.9 KBDownload
1// Per-credential rate limiting.
2//
3// This matters more here than on a typical service because the clients are agents,
4// which retry in loops — a misbehaving agent and an attack are indistinguishable
5// from the server's side, and both are answered the same way.
6//
7// A token bucket rather than a fixed window: a fixed window lets a caller spend its
8// whole allowance in the last instant of one window and again in the first instant
9// of the next, which is twice the intended rate at exactly the wrong moment.
10//
11// The git transport is deliberately *not* rate limited. A clone is one request that
12// may run for minutes, and it is already bounded by the concurrency semaphore; a
13// per-minute count is the wrong instrument for it.
14
15use std::collections::HashMap;
16use std::sync::Mutex;
17use std::time::{Duration, Instant};
18
19/// One caller's allowance.
20#[derive(Debug, Clone, Copy)]
21struct Bucket {
22 /// Whole requests still available, as a float so refill is smooth.
23 tokens: f64,
24 last: Instant,
25}
26
27pub struct RateLimiter {
28 /// Requests per minute. Zero disables the limiter entirely.
29 per_minute: u32,
30 buckets: Mutex<HashMap<String, Bucket>>,
31}
32
33/// How long an idle bucket is kept before it is dropped.
34///
35/// Without eviction the map grows once per distinct caller and never shrinks, which
36/// is a slow memory leak keyed by attacker-controlled input.
37const IDLE_EVICTION: Duration = Duration::from_secs(600);
38
39/// Evict only when the map is big enough for it to matter.
40const EVICTION_THRESHOLD: usize = 1024;
41
42impl RateLimiter {
43 pub fn new(per_minute: u32) -> Self {
44 RateLimiter {
45 per_minute,
46 buckets: Mutex::new(HashMap::new()),
47 }
48 }
49
50 pub fn enabled(&self) -> bool {
51 self.per_minute > 0
52 }
53
54 /// Take one request's worth of allowance.
55 ///
56 /// Returns `Err(retry_after_secs)` when the caller is out.
57 pub fn check(&self, key: &str) -> Result<(), u32> {
58 self.check_at(key, Instant::now())
59 }
60
61 fn check_at(&self, key: &str, now: Instant) -> Result<(), u32> {
62 if !self.enabled() {
63 return Ok(());
64 }
65
66 let capacity = self.per_minute as f64;
67 let per_second = capacity / 60.0;
68
69 let mut buckets = self.buckets.lock().expect("rate limiter lock");
70
71 if buckets.len() > EVICTION_THRESHOLD {
72 buckets.retain(|_, b| now.duration_since(b.last) < IDLE_EVICTION);
73 }
74
75 let bucket = buckets.entry(key.to_string()).or_insert(Bucket {
76 tokens: capacity,
77 last: now,
78 });
79
80 let elapsed = now.duration_since(bucket.last).as_secs_f64();
81 bucket.tokens = (bucket.tokens + elapsed * per_second).min(capacity);
82 bucket.last = now;
83
84 if bucket.tokens >= 1.0 {
85 bucket.tokens -= 1.0;
86 return Ok(());
87 }
88
89 // Seconds until one whole token is available again, rounded up so a client
90 // that obeys the header does not come back too early.
91 let wait = ((1.0 - bucket.tokens) / per_second).ceil().max(1.0);
92 Err(wait as u32)
93 }
94}
95
96#[cfg(test)]
97mod tests {
98 use super::*;
99
100 #[test]
101 fn a_zero_limit_disables_the_limiter() {
102 let limiter = RateLimiter::new(0);
103 assert!(!limiter.enabled());
104 for _ in 0..10_000 {
105 limiter.check("anyone").unwrap();
106 }
107 }
108
109 #[test]
110 fn a_caller_may_spend_its_allowance_and_is_then_refused() {
111 let limiter = RateLimiter::new(60);
112 let now = Instant::now();
113
114 for i in 0..60 {
115 limiter
116 .check_at("alice", now)
117 .unwrap_or_else(|_| panic!("request {i} should have been allowed"));
118 }
119 let wait = limiter.check_at("alice", now).unwrap_err();
120 assert!(wait >= 1, "a refusal must say when to come back");
121 }
122
123 #[test]
124 fn allowance_refills_over_time() {
125 let limiter = RateLimiter::new(60);
126 let start = Instant::now();
127
128 for _ in 0..60 {
129 limiter.check_at("alice", start).unwrap();
130 }
131 assert!(limiter.check_at("alice", start).is_err());
132
133 // One token per second at 60/minute.
134 limiter
135 .check_at("alice", start + Duration::from_secs(1))
136 .expect("a second later, one request should be available");
137 }
138
139 #[test]
140 fn refill_is_capped_at_the_bucket_size() {
141 let limiter = RateLimiter::new(60);
142 let start = Instant::now();
143 limiter.check_at("alice", start).unwrap();
144
145 // An hour of idling must not bank an hour of requests.
146 let later = start + Duration::from_secs(3600);
147 for _ in 0..60 {
148 limiter.check_at("alice", later).unwrap();
149 }
150 assert!(
151 limiter.check_at("alice", later).is_err(),
152 "idle time must not accumulate beyond one bucket"
153 );
154 }
155
156 #[test]
157 fn callers_are_limited_independently() {
158 let limiter = RateLimiter::new(10);
159 let now = Instant::now();
160
161 for _ in 0..10 {
162 limiter.check_at("alice", now).unwrap();
163 }
164 assert!(limiter.check_at("alice", now).is_err());
165 limiter
166 .check_at("bob", now)
167 .expect("one caller's spending must not affect another");
168 }
169
170 #[test]
171 fn idle_buckets_are_evicted_so_the_map_does_not_grow_without_bound() {
172 let limiter = RateLimiter::new(600);
173 let start = Instant::now();
174
175 for i in 0..EVICTION_THRESHOLD + 10 {
176 limiter.check_at(&format!("caller-{i}"), start).unwrap();
177 }
178 assert!(limiter.buckets.lock().unwrap().len() > EVICTION_THRESHOLD);
179
180 // Long after they all went idle, one more request triggers the sweep.
181 limiter
182 .check_at("fresh", start + IDLE_EVICTION + Duration::from_secs(1))
183 .unwrap();
184 assert!(
185 limiter.buckets.lock().unwrap().len() < EVICTION_THRESHOLD,
186 "idle buckets must not accumulate"
187 );
188 }
189}