alice.pds.demo.boltless.dev
This repository has no description
16 kB
497 lines
did:plc:3fwecdnvtcscjnrx2p4n7a…
1use std::collections::HashMap;
2use std::net::IpAddr;
3use std::sync::{Arc, Mutex};
4
5use knot_types::UnixMicros;
6
7const MAX_TRACKED_PEERS: usize = 100_000;
8
9const SWEEP_INTERVAL_MICROS: u64 = 1_000_000;
10
11knot_types::scalar_newtype! {
12 pub struct Burst(u32);
13 pub struct RefillMicros(u64);
14 pub struct PerPeerInflight(usize);
15 pub struct GlobalInflight(usize);
16}
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub struct RateLimit {
20 pub burst: Burst,
21 pub refill: RefillMicros,
22}
23
24impl RateLimit {
25 const fn interval(self) -> u64 {
26 match self.refill.get() {
27 0 => 1,
28 micros => micros,
29 }
30 }
31}
32
33#[derive(Debug, Clone, Copy)]
34pub struct LimitConfig {
35 pub rate: Option<RateLimit>,
36 pub per_peer_inflight: Option<PerPeerInflight>,
37 pub global_inflight: Option<GlobalInflight>,
38}
39
40impl Default for LimitConfig {
41 fn default() -> Self {
42 Self {
43 rate: Some(RateLimit {
44 burst: Burst::new(20),
45 refill: RefillMicros::new(100_000),
46 }),
47 per_peer_inflight: Some(PerPeerInflight::new(8)),
48 global_inflight: Some(GlobalInflight::new(64)),
49 }
50 }
51}
52
53impl LimitConfig {
54 pub const fn per_peer_only(per_peer_inflight: PerPeerInflight) -> Self {
55 Self {
56 rate: None,
57 per_peer_inflight: Some(per_peer_inflight),
58 global_inflight: None,
59 }
60 }
61
62 pub const fn unmetered() -> Self {
63 Self {
64 rate: None,
65 per_peer_inflight: None,
66 global_inflight: None,
67 }
68 }
69}
70
71struct Bucket {
72 tokens: u32,
73 last_refill: UnixMicros,
74}
75
76impl Bucket {
77 fn new(rate: RateLimit, now: UnixMicros) -> Self {
78 Self {
79 tokens: rate.burst.get(),
80 last_refill: now,
81 }
82 }
83
84 fn tokens_at(&self, rate: RateLimit, now: UnixMicros) -> u32 {
85 let elapsed = now.get().saturating_sub(self.last_refill.get());
86 let gained = (elapsed / rate.interval()).min(u64::from(rate.burst.get())) as u32;
87 self.tokens.saturating_add(gained).min(rate.burst.get())
88 }
89
90 fn replenish(&mut self, rate: RateLimit, now: UnixMicros) -> bool {
91 if now.get().saturating_sub(self.last_refill.get()) >= rate.interval() {
92 self.tokens = self.tokens_at(rate, now);
93 self.last_refill = now;
94 }
95 self.tokens > 0
96 }
97
98 fn full(&self, rate: RateLimit, now: UnixMicros) -> bool {
99 self.tokens_at(rate, now) >= rate.burst.get()
100 }
101}
102
103struct PeerState {
104 bucket: Option<Bucket>,
105 inflight: usize,
106}
107
108impl PeerState {
109 fn forgettable(&self) -> bool {
110 self.inflight == 0 && self.bucket.is_none()
111 }
112
113 fn worth_tracking(&self, rate: Option<RateLimit>, now: UnixMicros) -> bool {
114 match (self.inflight, &self.bucket, rate) {
115 (0, Some(bucket), Some(rate)) => !bucket.full(rate, now),
116 (0, _, _) => false,
117 _ => true,
118 }
119 }
120}
121
122struct Inner {
123 peers: HashMap<Option<IpAddr>, PeerState>,
124 global_inflight: usize,
125 last_sweep: UnixMicros,
126}
127
128pub struct PreAuthLimiter {
129 config: LimitConfig,
130 inner: Mutex<Inner>,
131}
132
133#[derive(Debug, Clone, Copy, PartialEq, Eq)]
134pub enum Refusal {
135 RateLimited,
136 Saturated,
137}
138
139impl Default for PreAuthLimiter {
140 fn default() -> Self {
141 Self::with_config(LimitConfig::default())
142 }
143}
144
145impl PreAuthLimiter {
146 pub fn with_config(config: LimitConfig) -> Self {
147 Self {
148 config,
149 inner: Mutex::new(Inner {
150 peers: HashMap::new(),
151 global_inflight: 0,
152 last_sweep: UnixMicros::new(0),
153 }),
154 }
155 }
156
157 fn lock(&self) -> std::sync::MutexGuard<'_, Inner> {
158 self.inner
159 .lock()
160 .unwrap_or_else(|poisoned| poisoned.into_inner())
161 }
162
163 pub fn admit(
164 self: &Arc<Self>,
165 peer: Option<IpAddr>,
166 now: UnixMicros,
167 ) -> Result<AdmitGuard, Refusal> {
168 let mut inner = self.lock();
169 let rate = self.config.rate;
170 let over_global = self
171 .config
172 .global_inflight
173 .is_some_and(|limit| inner.global_inflight >= limit.get());
174
175 // Only a rate budget outlives the work it admitted, so only a rate
176 // budget can pile up entries for peers that have gone away. Without one
177 // the last operation removes the entry and the table is already bounded
178 // by live concurrency, so capping it would shed callers over a table
179 // that can't grow.
180 if rate.is_some()
181 && !inner.peers.contains_key(&peer)
182 && inner.peers.len() >= MAX_TRACKED_PEERS
183 {
184 let sweep_due =
185 now.get().saturating_sub(inner.last_sweep.get()) >= SWEEP_INTERVAL_MICROS;
186 if sweep_due {
187 inner.last_sweep = now;
188 inner
189 .peers
190 .retain(|_, state| state.worth_tracking(rate, now));
191 }
192 if inner.peers.len() >= MAX_TRACKED_PEERS {
193 return Err(Refusal::Saturated);
194 }
195 }
196
197 let per_peer = self.config.per_peer_inflight;
198 let decision = {
199 let state = inner.peers.entry(peer).or_insert_with(|| PeerState {
200 bucket: rate.map(|rate| Bucket::new(rate, now)),
201 inflight: 0,
202 });
203 let ready = match (rate, state.bucket.as_mut()) {
204 (Some(rate), Some(bucket)) => bucket.replenish(rate, now),
205 _ => true,
206 };
207 let over_peer = per_peer.is_some_and(|limit| state.inflight >= limit.get());
208 match (ready, over_peer || over_global) {
209 (false, _) => Err(Refusal::RateLimited),
210 (_, true) => Err(Refusal::Saturated),
211 (true, false) => {
212 if let Some(bucket) = state.bucket.as_mut() {
213 bucket.tokens -= 1;
214 }
215 state.inflight += 1;
216 Ok(())
217 }
218 }
219 };
220 match decision {
221 Err(refusal) => {
222 if inner.peers.get(&peer).is_some_and(PeerState::forgettable) {
223 inner.peers.remove(&peer);
224 }
225 Err(refusal)
226 }
227 Ok(()) => {
228 inner.global_inflight += 1;
229 Ok(AdmitGuard {
230 limiter: Arc::clone(self),
231 peer,
232 })
233 }
234 }
235 }
236
237 fn leave(&self, peer: Option<IpAddr>) {
238 let mut inner = self.lock();
239 inner.global_inflight = inner.global_inflight.saturating_sub(1);
240 let Some(state) = inner.peers.get_mut(&peer) else {
241 return;
242 };
243 state.inflight = state.inflight.saturating_sub(1);
244 if state.forgettable() {
245 inner.peers.remove(&peer);
246 }
247 }
248
249 fn repay(&self, peer: Option<IpAddr>) {
250 let Some(rate) = self.config.rate else {
251 return;
252 };
253 let mut inner = self.lock();
254 if let Some(bucket) = inner
255 .peers
256 .get_mut(&peer)
257 .and_then(|state| state.bucket.as_mut())
258 {
259 bucket.tokens = bucket.tokens.saturating_add(1).min(rate.burst.get());
260 }
261 }
262}
263
264pub struct AdmitGuard {
265 limiter: Arc<PreAuthLimiter>,
266 peer: Option<IpAddr>,
267}
268
269impl AdmitGuard {
270 pub fn refund(self) {
271 self.limiter.repay(self.peer);
272 }
273}
274
275impl Drop for AdmitGuard {
276 fn drop(&mut self) {
277 self.limiter.leave(self.peer);
278 }
279}
280
281#[cfg(test)]
282mod tests {
283 use super::*;
284 use std::net::Ipv4Addr;
285
286 fn limiter(config: LimitConfig) -> Arc<PreAuthLimiter> {
287 Arc::new(PreAuthLimiter::with_config(config))
288 }
289
290 fn rate(burst: u32, refill_micros: u64) -> Option<RateLimit> {
291 Some(RateLimit {
292 burst: Burst::new(burst),
293 refill: RefillMicros::new(refill_micros),
294 })
295 }
296
297 fn peer(last: u8) -> Option<IpAddr> {
298 Some(IpAddr::V4(Ipv4Addr::new(127, 0, 0, last)))
299 }
300
301 fn rotating(index: u64) -> Option<IpAddr> {
302 let octets = (index as u32).to_be_bytes();
303 Some(IpAddr::V4(Ipv4Addr::new(
304 1, octets[1], octets[2], octets[3],
305 )))
306 }
307
308 fn at(micros: u64) -> UnixMicros {
309 UnixMicros::new(micros)
310 }
311
312 fn tracked(limiter: &Arc<PreAuthLimiter>) -> usize {
313 limiter.lock().peers.len()
314 }
315
316 #[test]
317 fn a_rate_budget_refuses_a_flood_refills_over_time_and_stays_per_peer() {
318 let limiter = limiter(LimitConfig {
319 rate: rate(2, 1_000),
320 per_peer_inflight: Some(PerPeerInflight::new(100)),
321 global_inflight: Some(GlobalInflight::new(100)),
322 });
323 assert!(limiter.admit(peer(1), at(0)).is_ok());
324 assert!(limiter.admit(peer(1), at(0)).is_ok());
325 assert_eq!(
326 limiter.admit(peer(1), at(0)).err(),
327 Some(Refusal::RateLimited),
328 "the limiter refuses a third request inside the same instant even when nothing is in flight"
329 );
330 assert!(
331 limiter.admit(peer(2), at(0)).is_ok(),
332 "a second peer keeps its own rate budget"
333 );
334 assert_eq!(
335 limiter.admit(peer(1), at(600)).err(),
336 Some(Refusal::RateLimited)
337 );
338 assert!(
339 limiter.admit(peer(1), at(1_000)).is_ok(),
340 "a refusal partway through the interval mustn't reset the clock the refill measures from"
341 );
342 }
343
344 #[test]
345 fn inflight_shedding_frees_on_drop_and_spends_no_rate_budget() {
346 let limiter = limiter(LimitConfig {
347 rate: rate(4, 1_000_000),
348 per_peer_inflight: Some(PerPeerInflight::new(1)),
349 global_inflight: Some(GlobalInflight::new(2)),
350 });
351 let held = limiter
352 .admit(peer(1), at(0))
353 .expect("the limiter admits the first request");
354 let other = limiter
355 .admit(peer(2), at(0))
356 .expect("a second peer fills the global budget");
357 assert_eq!(
358 limiter.admit(peer(1), at(0)).err(),
359 Some(Refusal::Saturated),
360 "the limiter sheds a second concurrent request from one peer"
361 );
362 assert_eq!(
363 limiter.admit(peer(3), at(0)).err(),
364 Some(Refusal::Saturated),
365 "the limiter sheds a third peer once the global in-flight limit is reached"
366 );
367 drop(other);
368 drop(
369 limiter
370 .admit(peer(3), at(0))
371 .expect("freeing a global slot admits the peer that was shed"),
372 );
373 drop(held);
374 (0..50).for_each(|_| {
375 limiter
376 .admit(peer(1), at(0))
377 .expect("a refunded admission must leave the full budget available")
378 .refund();
379 });
380 (0..3).for_each(|_| {
381 limiter
382 .admit(peer(1), at(0))
383 .expect("a shed request mustn't spend the rate budget it never used");
384 });
385 assert_eq!(
386 limiter.admit(peer(1), at(0)).err(),
387 Some(Refusal::RateLimited),
388 "a dropped guard without a refund keeps its token spent"
389 );
390 }
391
392 #[test]
393 fn a_budget_without_a_rate_never_rate_limits_and_keeps_no_idle_state() {
394 let overflowing = MAX_TRACKED_PEERS as u64 + 1_000;
395
396 let per_peer = limiter(LimitConfig::per_peer_only(PerPeerInflight::new(2)));
397 (0..1_000).for_each(|_| {
398 per_peer
399 .admit(peer(1), at(0))
400 .expect("a budget with no rate has nothing for a sequential flood to exhaust");
401 });
402 let concurrent: Vec<AdmitGuard> = (0..2)
403 .map(|_| {
404 per_peer
405 .admit(peer(1), at(0))
406 .expect("the limiter admits both concurrent operations from one peer")
407 })
408 .collect();
409 assert_eq!(
410 per_peer.admit(peer(1), at(0)).err(),
411 Some(Refusal::Saturated),
412 "only the per-peer count refuses a request in this budget"
413 );
414 drop(concurrent);
415 let held: Vec<AdmitGuard> = (0..overflowing)
416 .map(|index| {
417 per_peer
418 .admit(rotating(index), at(0))
419 .expect("a budget that keeps no idle state has no table to overflow")
420 })
421 .collect();
422 assert_eq!(tracked(&per_peer), held.len());
423 drop(held);
424 assert_eq!(
425 tracked(&per_peer),
426 0,
427 "with no tokens to remember, an idle peer leaves no entry, \
428 so address rotation mustn't fill the table and start shedding newcomers"
429 );
430
431 let global = limiter(LimitConfig {
432 rate: None,
433 per_peer_inflight: None,
434 global_inflight: Some(GlobalInflight::new(1)),
435 });
436 let _saturating = global
437 .admit(peer(1), at(0))
438 .expect("the limiter admits the first peer");
439 (0..overflowing).for_each(|index| {
440 assert_eq!(
441 global.admit(rotating(index), at(0)).err(),
442 Some(Refusal::Saturated)
443 );
444 });
445 assert_eq!(
446 tracked(&global),
447 1,
448 "a refusal returns no guard, so an entry it left behind would never be freed, \
449 and the sweep that bounds the table only reclaims idle rate state"
450 );
451
452 let unmetered = limiter(LimitConfig::unmetered());
453 let guards: Vec<AdmitGuard> = (0..512)
454 .map(|_| {
455 unmetered
456 .admit(peer(1), at(0))
457 .expect("an unmetered budget admits every request from every peer")
458 })
459 .collect();
460 drop(guards);
461 assert_eq!(tracked(&unmetered), 0);
462 }
463
464 #[test]
465 fn a_rate_budget_bounds_its_peer_map_and_sweeps_at_most_once_per_interval() {
466 let limiter = limiter(LimitConfig {
467 rate: rate(1, 1_000),
468 per_peer_inflight: Some(PerPeerInflight::new(8)),
469 global_inflight: Some(GlobalInflight::new(64)),
470 });
471 (0..MAX_TRACKED_PEERS as u64).for_each(|index| {
472 let _ = limiter.admit(rotating(index), at(0));
473 });
474 assert_eq!(
475 limiter
476 .admit(peer(201), at(SWEEP_INTERVAL_MICROS - 1))
477 .err(),
478 Some(Refusal::Saturated),
479 "inside the sweep interval a full map sheds unseen peers without rescanning"
480 );
481 assert!(
482 limiter.admit(peer(202), at(SWEEP_INTERVAL_MICROS)).is_ok(),
483 "once the interval elapses the sweep evicts replenished entries and admits the newcomer"
484 );
485 (0..(MAX_TRACKED_PEERS as u64 + 50_000)).for_each(|index| {
486 let _ = limiter.admit(
487 rotating(MAX_TRACKED_PEERS as u64 + index),
488 at(SWEEP_INTERVAL_MICROS + index),
489 );
490 });
491 let tracked = tracked(&limiter);
492 assert!(
493 tracked <= MAX_TRACKED_PEERS,
494 "a flood of distinct source addresses mustn't grow the peer map past its limit, saw {tracked}"
495 );
496 }
497}