alice.pds.demo.boltless.dev
This repository has no description
19 kB
585 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
122const CONCENTRATED_REFUSALS: u32 = 1_024;
123
124#[derive(Default)]
125struct RefusalMajority {
126 peer: Option<IpAddr>,
127 votes: u32,
128 reported: bool,
129}
130
131impl RefusalMajority {
132 fn observe(&mut self, peer: IpAddr) -> Option<IpAddr> {
133 match (self.peer == Some(peer), self.votes) {
134 (true, _) => self.votes = self.votes.saturating_add(1),
135 (false, 0) => {
136 self.peer = Some(peer);
137 self.votes = 1;
138 }
139 (false, _) => self.votes -= 1,
140 }
141 let crossed = self.votes >= CONCENTRATED_REFUSALS && !self.reported;
142 self.reported |= crossed;
143 crossed.then_some(peer)
144 }
145}
146
147struct Inner {
148 peers: HashMap<Option<IpAddr>, PeerState>,
149 global_inflight: usize,
150 last_sweep: UnixMicros,
151 refusal_majority: RefusalMajority,
152}
153
154impl Inner {
155 fn has_room_for(
156 &mut self,
157 peer: Option<IpAddr>,
158 rate: Option<RateLimit>,
159 now: UnixMicros,
160 ) -> bool {
161 if rate.is_none() || self.peers.contains_key(&peer) || self.peers.len() < MAX_TRACKED_PEERS
162 {
163 return true;
164 }
165 if now.get().saturating_sub(self.last_sweep.get()) >= SWEEP_INTERVAL_MICROS {
166 self.last_sweep = now;
167 self.peers
168 .retain(|_, state| state.worth_tracking(rate, now));
169 }
170 self.peers.len() < MAX_TRACKED_PEERS
171 }
172}
173
174pub struct PreAuthLimiter {
175 config: LimitConfig,
176 inner: Mutex<Inner>,
177}
178
179#[derive(Debug, Clone, Copy, PartialEq, Eq)]
180pub enum Refusal {
181 RateLimited,
182 Saturated,
183}
184
185impl Default for PreAuthLimiter {
186 fn default() -> Self {
187 Self::with_config(LimitConfig::default())
188 }
189}
190
191impl PreAuthLimiter {
192 pub fn with_config(config: LimitConfig) -> Self {
193 Self {
194 config,
195 inner: Mutex::new(Inner {
196 peers: HashMap::new(),
197 global_inflight: 0,
198 last_sweep: UnixMicros::new(0),
199 refusal_majority: RefusalMajority::default(),
200 }),
201 }
202 }
203
204 fn lock(&self) -> std::sync::MutexGuard<'_, Inner> {
205 self.inner
206 .lock()
207 .unwrap_or_else(|poisoned| poisoned.into_inner())
208 }
209
210 pub fn admit(
211 self: &Arc<Self>,
212 peer: Option<IpAddr>,
213 now: UnixMicros,
214 ) -> Result<AdmitGuard, Refusal> {
215 let mut inner = self.lock();
216 let rate = self.config.rate;
217 let over_global = self
218 .config
219 .global_inflight
220 .is_some_and(|limit| inner.global_inflight >= limit.get());
221
222 let per_peer = self.config.per_peer_inflight;
223 let decision = match inner.has_room_for(peer, rate, now) {
224 false => Err(Refusal::Saturated),
225 true => {
226 let state = inner.peers.entry(peer).or_insert_with(|| PeerState {
227 bucket: rate.map(|rate| Bucket::new(rate, now)),
228 inflight: 0,
229 });
230 let ready = match (rate, state.bucket.as_mut()) {
231 (Some(rate), Some(bucket)) => bucket.replenish(rate, now),
232 _ => true,
233 };
234 let over_peer = per_peer.is_some_and(|limit| state.inflight >= limit.get());
235 match (ready, over_peer || over_global) {
236 (false, _) => Err(Refusal::RateLimited),
237 (_, true) => Err(Refusal::Saturated),
238 (true, false) => {
239 if let Some(bucket) = state.bucket.as_mut() {
240 bucket.tokens -= 1;
241 }
242 state.inflight += 1;
243 Ok(())
244 }
245 }
246 }
247 };
248 match decision {
249 Err(refusal) => {
250 if inner.peers.get(&peer).is_some_and(PeerState::forgettable) {
251 inner.peers.remove(&peer);
252 }
253 let concentrated = peer.and_then(|peer| inner.refusal_majority.observe(peer));
254 drop(inner);
255 if let Some(peer) = concentrated {
256 tracing::warn!(
257 %peer,
258 "one address has taken {CONCENTRATED_REFUSALS} more of the pre-authentication limiter's refusals than every other address combined. If this address is a proxy, set xrpc.trusted_proxy_header to the header it forwards the client address in and add the address to xrpc.trusted_proxies, since every client behind a proxy will share its one rate-limit bucket. This warning reports the first such address only."
259 );
260 }
261 Err(refusal)
262 }
263 Ok(()) => {
264 inner.global_inflight += 1;
265 Ok(AdmitGuard {
266 limiter: Arc::clone(self),
267 peer,
268 })
269 }
270 }
271 }
272
273 fn leave(&self, peer: Option<IpAddr>) {
274 let mut inner = self.lock();
275 inner.global_inflight = inner.global_inflight.saturating_sub(1);
276 let Some(state) = inner.peers.get_mut(&peer) else {
277 return;
278 };
279 state.inflight = state.inflight.saturating_sub(1);
280 if state.forgettable() {
281 inner.peers.remove(&peer);
282 }
283 }
284
285 fn repay(&self, peer: Option<IpAddr>) {
286 let Some(rate) = self.config.rate else {
287 return;
288 };
289 let mut inner = self.lock();
290 if let Some(bucket) = inner
291 .peers
292 .get_mut(&peer)
293 .and_then(|state| state.bucket.as_mut())
294 {
295 bucket.tokens = bucket.tokens.saturating_add(1).min(rate.burst.get());
296 }
297 }
298}
299
300pub struct AdmitGuard {
301 limiter: Arc<PreAuthLimiter>,
302 peer: Option<IpAddr>,
303}
304
305impl AdmitGuard {
306 pub fn refund(self) {
307 self.limiter.repay(self.peer);
308 }
309}
310
311impl Drop for AdmitGuard {
312 fn drop(&mut self) {
313 self.limiter.leave(self.peer);
314 }
315}
316
317#[cfg(test)]
318mod tests {
319 use super::*;
320 use std::net::Ipv4Addr;
321
322 fn limiter(config: LimitConfig) -> Arc<PreAuthLimiter> {
323 Arc::new(PreAuthLimiter::with_config(config))
324 }
325
326 fn rate(burst: u32, refill_micros: u64) -> Option<RateLimit> {
327 Some(RateLimit {
328 burst: Burst::new(burst),
329 refill: RefillMicros::new(refill_micros),
330 })
331 }
332
333 fn peer(last: u8) -> Option<IpAddr> {
334 Some(ip(last))
335 }
336
337 fn rotating(index: u64) -> Option<IpAddr> {
338 let octets = (index as u32).to_be_bytes();
339 Some(IpAddr::V4(Ipv4Addr::new(
340 1, octets[1], octets[2], octets[3],
341 )))
342 }
343
344 fn at(micros: u64) -> UnixMicros {
345 UnixMicros::new(micros)
346 }
347
348 fn tracked(limiter: &Arc<PreAuthLimiter>) -> usize {
349 limiter.lock().peers.len()
350 }
351
352 fn ip(last: u8) -> IpAddr {
353 IpAddr::V4(Ipv4Addr::new(127, 0, 0, last))
354 }
355
356 #[test]
357 fn only_an_address_taking_most_of_the_refusals_is_reported_and_only_once() {
358 let reported = |rounds, peer: fn(u32) -> u8| {
359 let mut majority = RefusalMajority::default();
360 (0..rounds)
361 .filter_map(|round| majority.observe(ip(peer(round))))
362 .collect::<Vec<IpAddr>>()
363 };
364 [
365 (CONCENTRATED_REFUSALS * 3, (|round| (round % 4 == 0) as u8) as fn(u32) -> u8, vec![ip(0)],
366 "the warning reports that address once, since three refusals in every four come from it"),
367 (CONCENTRATED_REFUSALS * 8, |round| (round % 251) as u8, vec![],
368 "a knot under scattered load doesn't have a proxy to point the operator at"),
369 (CONCENTRATED_REFUSALS - 1, |_| 1, vec![],
370 "a burst under the threshold is ordinary rate limiting and doesn't point at a proxy"),
371 (CONCENTRATED_REFUSALS, |_| 1, vec![ip(1)],
372 "the threshold itself is where an address earns the warning"),
373 ]
374 .into_iter()
375 .for_each(|(rounds, peer, expected, why)| {
376 assert_eq!(reported(rounds, peer), expected, "{why}");
377 });
378 }
379
380 #[test]
381 fn a_rate_budget_refuses_a_flood_refills_over_time_and_stays_per_peer() {
382 let limiter = limiter(LimitConfig {
383 rate: rate(2, 1_000),
384 per_peer_inflight: Some(PerPeerInflight::new(100)),
385 global_inflight: Some(GlobalInflight::new(100)),
386 });
387 assert!(limiter.admit(peer(1), at(0)).is_ok());
388 assert!(limiter.admit(peer(1), at(0)).is_ok());
389 assert_eq!(
390 limiter.admit(peer(1), at(0)).err(),
391 Some(Refusal::RateLimited),
392 "the limiter refuses a third request inside the same instant even when nothing is in flight"
393 );
394 assert!(
395 limiter.admit(peer(2), at(0)).is_ok(),
396 "a second peer keeps its own rate budget"
397 );
398 assert_eq!(
399 limiter.admit(peer(1), at(600)).err(),
400 Some(Refusal::RateLimited)
401 );
402 assert!(
403 limiter.admit(peer(1), at(1_000)).is_ok(),
404 "a refusal partway through the interval mustn't reset the clock the refill measures from"
405 );
406 }
407
408 #[test]
409 fn inflight_shedding_frees_on_drop_and_spends_no_rate_budget() {
410 let limiter = limiter(LimitConfig {
411 rate: rate(4, 1_000_000),
412 per_peer_inflight: Some(PerPeerInflight::new(1)),
413 global_inflight: Some(GlobalInflight::new(2)),
414 });
415 let held = limiter
416 .admit(peer(1), at(0))
417 .expect("the limiter admits the first request");
418 let other = limiter
419 .admit(peer(2), at(0))
420 .expect("a second peer fills the global budget");
421 assert_eq!(
422 limiter.admit(peer(1), at(0)).err(),
423 Some(Refusal::Saturated),
424 "the limiter sheds a second concurrent request from one peer"
425 );
426 assert_eq!(
427 limiter.admit(peer(3), at(0)).err(),
428 Some(Refusal::Saturated),
429 "the limiter sheds a third peer once the global in-flight limit is reached"
430 );
431 drop(other);
432 drop(
433 limiter
434 .admit(peer(3), at(0))
435 .expect("freeing a global slot admits the peer that was shed"),
436 );
437 drop(held);
438 (0..50).for_each(|_| {
439 limiter
440 .admit(peer(1), at(0))
441 .expect("a refunded admission must leave the full budget available")
442 .refund();
443 });
444 (0..3).for_each(|_| {
445 limiter
446 .admit(peer(1), at(0))
447 .expect("a shed request mustn't spend the rate budget it never used");
448 });
449 assert_eq!(
450 limiter.admit(peer(1), at(0)).err(),
451 Some(Refusal::RateLimited),
452 "a dropped guard without a refund keeps its token spent"
453 );
454 }
455
456 #[test]
457 fn a_budget_without_a_rate_never_rate_limits_and_keeps_no_idle_state() {
458 let overflowing = MAX_TRACKED_PEERS as u64 + 1_000;
459
460 let per_peer = limiter(LimitConfig::per_peer_only(PerPeerInflight::new(2)));
461 (0..1_000).for_each(|_| {
462 per_peer
463 .admit(peer(1), at(0))
464 .expect("a budget with no rate has nothing for a sequential flood to exhaust");
465 });
466 let concurrent: Vec<AdmitGuard> = (0..2)
467 .map(|_| {
468 per_peer
469 .admit(peer(1), at(0))
470 .expect("the limiter admits both concurrent operations from one peer")
471 })
472 .collect();
473 assert_eq!(
474 per_peer.admit(peer(1), at(0)).err(),
475 Some(Refusal::Saturated),
476 "only the per-peer count refuses a request in this budget"
477 );
478 drop(concurrent);
479 let held: Vec<AdmitGuard> = (0..overflowing)
480 .map(|index| {
481 per_peer
482 .admit(rotating(index), at(0))
483 .expect("a budget that keeps no idle state has no table to overflow")
484 })
485 .collect();
486 assert_eq!(tracked(&per_peer), held.len());
487 drop(held);
488 assert_eq!(
489 tracked(&per_peer),
490 0,
491 "with no tokens to remember, an idle peer leaves no entry, \
492 so address rotation mustn't fill the table and start shedding newcomers"
493 );
494
495 let global = limiter(LimitConfig {
496 rate: None,
497 per_peer_inflight: None,
498 global_inflight: Some(GlobalInflight::new(1)),
499 });
500 let _saturating = global
501 .admit(peer(1), at(0))
502 .expect("the limiter admits the first peer");
503 (0..overflowing).for_each(|index| {
504 assert_eq!(
505 global.admit(rotating(index), at(0)).err(),
506 Some(Refusal::Saturated)
507 );
508 });
509 assert_eq!(
510 tracked(&global),
511 1,
512 "a refusal returns no guard, so an entry it left behind would never be freed, \
513 and the sweep that bounds the table only reclaims idle rate state"
514 );
515
516 let unmetered = limiter(LimitConfig::unmetered());
517 let guards: Vec<AdmitGuard> = (0..512)
518 .map(|_| {
519 unmetered
520 .admit(peer(1), at(0))
521 .expect("an unmetered budget admits every request from every peer")
522 })
523 .collect();
524 drop(guards);
525 assert_eq!(tracked(&unmetered), 0);
526 }
527
528 #[test]
529 fn a_rate_budget_bounds_its_peer_map_and_sweeps_at_most_once_per_interval() {
530 let limiter = limiter(LimitConfig {
531 rate: rate(1, 1_000),
532 per_peer_inflight: Some(PerPeerInflight::new(8)),
533 global_inflight: Some(GlobalInflight::new(64)),
534 });
535 (0..MAX_TRACKED_PEERS as u64).for_each(|index| {
536 let _ = limiter.admit(rotating(index), at(0));
537 });
538 assert_eq!(
539 limiter
540 .admit(peer(201), at(SWEEP_INTERVAL_MICROS - 1))
541 .err(),
542 Some(Refusal::Saturated),
543 "inside the sweep interval a full map sheds unseen peers without rescanning"
544 );
545 assert!(
546 limiter.admit(peer(202), at(SWEEP_INTERVAL_MICROS)).is_ok(),
547 "once the interval elapses the sweep evicts replenished entries and admits the newcomer"
548 );
549 (0..(MAX_TRACKED_PEERS as u64 + 50_000)).for_each(|index| {
550 let _ = limiter.admit(
551 rotating(MAX_TRACKED_PEERS as u64 + index),
552 at(SWEEP_INTERVAL_MICROS + index),
553 );
554 });
555 let tracked = tracked(&limiter);
556 assert!(
557 tracked <= MAX_TRACKED_PEERS,
558 "a flood of distinct source addresses mustn't grow the peer map past its limit, saw {tracked}"
559 );
560 }
561
562 #[test]
563 fn the_majority_counter_sees_a_refusal_from_a_full_peer_table() {
564 let limiter = limiter(LimitConfig {
565 rate: rate(1, 1_000),
566 per_peer_inflight: Some(PerPeerInflight::new(8)),
567 global_inflight: Some(GlobalInflight::new(64)),
568 });
569 (0..MAX_TRACKED_PEERS as u64).for_each(|index| {
570 let _ = limiter.admit(rotating(index), at(0));
571 });
572 assert_eq!(
573 limiter
574 .admit(peer(201), at(SWEEP_INTERVAL_MICROS - 1))
575 .err(),
576 Some(Refusal::Saturated)
577 );
578 let inner = limiter.lock();
579 assert_eq!(
580 (inner.refusal_majority.peer, inner.refusal_majority.votes),
581 (peer(201), 1),
582 "that refusal has to reach the counter like every other refusal, since a peer shed by a full table is what the warning most needs to report"
583 );
584 }
585}