Oregami
Repositories/oxedyne/fe2o3

oxedyne/fe2o3/fe2o3_net/src/guard/addr.rs

35.8 KiB, 123 runs

created by r1870400018:10862, which is this file's identity for as long as the history lasts, whatever it is later renamed to

download · who wrote it · its history

1//! Generic per-IP rate-limiting and blacklisting guard.
2//!
3//! `AddressGuard` tracks a per-IP state machine -- `Monitor` then `Throttle` then
4//! `Blacklist`, plus a manual `Whitelist` escape hatch -- driven by a sliding window of
5//! request timestamps. Each caller-visible primitive returns an `Outcome<GuardDecision>`
6//! describing whether the request should be allowed through, throttled, or blocked.
7//!
8//! The guard is transport-agnostic; HTTPS, SMTP, SHIELD's UDP wire protocol or any other
9//! caller can plug it into its accept path. Protocol-specific extensions (for example
10//! SHIELD's handshake-sequence check) are layered on top via the low-level
11//! [`AddressGuard::update_log`] helper, which exposes the per-address log under the same
12//! shard lock acquired by the rate-limit check.
13//!
14//! [Written with AI entirely](https://need2know.ai/entirely-ai/code)\
15//! Anthropic Claude
16
17use oxedyne_fe2o3_core::{
18 prelude::*,
19 map::MapMut,
20};
21use oxedyne_fe2o3_data::ring::RingTimer;
22use oxedyne_fe2o3_hash::map::ShardMap;
23use oxedyne_fe2o3_iop_hash::api::{
24 Hasher,
25 HashForm,
26};
27
28use std::{
29 clone::Clone,
30 fmt::Debug,
31 net::{
32 IpAddr,
33 Ipv4Addr,
34 SocketAddr,
35 },
36 sync::{
37 Arc,
38 RwLock,
39 atomic::{
40 AtomicUsize,
41 Ordering,
42 },
43 },
44 time::{
45 Duration,
46 SystemTime,
47 },
48};
49
50/// Per-address state in the guard state machine.
51#[derive(Clone, Debug)]
52pub enum AddressState<const N: usize> {
53 Monitor(RingTimer<N>), // watching the rate only, nothing dropped
54 Throttle { // a request closer than `tint_min` to the last is dropped
55 reqs: RingTimer<N>,
56 tint_min: Duration,
57 start: SystemTime,
58 sunset: Duration, // cooldown, after which the address returns to Monitor
59 },
60 Blacklist { // blocked outright
61 since: SystemTime,
62 reason: BlacklistReason,
63 },
64 Whitelist, // always allowed through
65}
66
67impl<const N: usize> Default for AddressState<N> {
68 fn default() -> Self {
69 Self::Monitor(RingTimer::default())
70 }
71}
72
73impl<const N: usize> AddressState<N> {
74 /// One of "monitor", "throttle", "blacklist" or "whitelist".
75 pub fn label(&self) -> &'static str {
76 match self {
77 Self::Monitor(_) => "monitor",
78 Self::Throttle{..} => "throttle",
79 Self::Blacklist{..} => "blacklist",
80 Self::Whitelist => "whitelist",
81 }
82 }
83}
84
85/// Reason an address is in the `Blacklist` state.
86#[derive(Clone, Copy, Debug, Eq, PartialEq)]
87pub enum BlacklistReason {
88 AutoRateLimit, // the state machine moved it there after sustained abuse
89 Manual, // administrator action
90}
91
92/// Per-address log: state, counters, and caller-supplied extension data `D`.
93#[derive(Clone, Debug)]
94pub struct AddressLog<
95 const N: usize,
96 D: Clone + Debug + Default,
97> {
98 pub ip: Option<IpAddr>, // held so a snapshot need not reverse a shard key
99 pub state: AddressState<N>,
100 pub throttle_cnt: u16, // throttling episodes so far
101 pub first_seen: SystemTime,
102 pub last_seen: SystemTime,
103 pub total_reqs: u64,
104 // Live connections from this address right now, held behind an `Arc` so a
105 // `ConnPermit` handed to a spawned task can decrement it on drop without
106 // re-acquiring the shard lock, and so an eviction of the log cannot strand
107 // an in-flight permit. This is a *concurrency* gauge, orthogonal to the
108 // *rate* window above.
109 pub conns: Arc<AtomicUsize>,
110 // When this address last crossed into an offending state (throttled or auto-
111 // blacklisted), so a decay can relax it after a quiet spell. `None` once it
112 // has decayed or if it has never offended. Distinct from `last_seen`, which a
113 // still-active address refreshes on every request: the whole point is to
114 // relax an address that is *alive but no longer offending*.
115 pub offended_at: Option<SystemTime>,
116 pub data: D, // caller-supplied extension payload
117}
118
119impl<
120 const N: usize,
121 D: Clone + Debug + Default,
122>
123 Default for AddressLog<N, D>
124{
125 fn default() -> Self {
126 let now = SystemTime::now();
127 Self {
128 ip: None,
129 state: AddressState::default(),
130 throttle_cnt: 0,
131 first_seen: now,
132 last_seen: now,
133 total_reqs: 0,
134 conns: Arc::new(AtomicUsize::new(0)),
135 offended_at: None,
136 data: D::default(),
137 }
138 }
139}
140
141/// Decision returned by the guard check APIs.
142#[derive(Clone, Copy, Debug, Eq, PartialEq)]
143pub enum GuardDecision {
144 Allow,
145 Throttled, // inside an active throttle window
146 Blocked(BlacklistReason),
147}
148
149impl GuardDecision {
150 pub fn should_drop(&self) -> bool {
151 !matches!(self, Self::Allow)
152 }
153}
154
155/// Aggregate tallies across all known addresses, suitable for a dashboard chip row.
156#[derive(Clone, Copy, Debug, Default)]
157pub struct GuardCounts {
158 pub monitor: usize,
159 pub throttle: usize,
160 pub blacklist: usize,
161 pub whitelist: usize,
162 pub total: usize, // distinct addresses observed
163 pub total_reqs: u64, // across every address
164}
165
166/// One row in a guard snapshot table.
167#[derive(Clone, Debug)]
168pub struct GuardEntry {
169 pub ip: IpAddr, // a row appears only where the log stored one
170 pub state: &'static str, // "monitor", "throttle", "blacklist" or "whitelist"
171 pub throttle_cnt: u16,
172 pub total_reqs: u64,
173 pub first_seen: SystemTime,
174 pub last_seen: SystemTime,
175}
176
177/// Snapshot of the guard: counts plus up to `max` per-address entries.
178#[derive(Clone, Debug, Default)]
179pub struct GuardSnapshot {
180 pub counts: GuardCounts,
181 pub entries: Vec<GuardEntry>, // capped by the caller
182}
183
184/// An RAII grant of one concurrent connection from an address.
185///
186/// Returned by [`AddressGuard::acquire`] and held for the lifetime of the
187/// connection it admits, typically moved into the task that serves it. The two
188/// counters it holds -- the per-address one and the guard-wide one -- both fall
189/// in `Drop`, so a task that panics is not counted against its address for ever
190/// after, and neither counter needs the shard lock to be released.
191#[derive(Debug)]
192pub struct ConnPermit {
193 per_ip: Arc<AtomicUsize>,
194 total: Arc<AtomicUsize>,
195}
196
197impl Drop for ConnPermit {
198 fn drop(&mut self) {
199 self.per_ip.fetch_sub(1, Ordering::AcqRel);
200 self.total.fetch_sub(1, Ordering::AcqRel);
201 }
202}
203
204/// Generic per-address guard.
205#[derive(Debug)]
206pub struct AddressGuard<
207 const C: usize, // ShardMap capacity.
208 M: MapMut<HashForm, AddressLog<N, D>> + Clone + Debug,
209 H: Hasher + Send + Sync + 'static,
210 const S: usize, // Hasher salt length.
211 const N: usize, // Request timer ring length.
212 D: Clone + Debug + Default,
213> {
214 // Per-address logs, keyed by a hash of the IP octets.
215 pub amap: ShardMap<C, S, AddressLog<N, D>, M, H>,
216 pub arps_max: u64, // maximum average requests per second in Monitor
217 pub tint_min: Duration, // minimum interval between requests in Throttle
218 pub tsunset_base: Duration, // base throttle cooldown
219 pub tsunset_spread: Duration, // jitter ceiling on it, to spread the expiries
220 pub blist_cnt: u16, // throttling episodes before blacklisting
221 // Maximum concurrent connections from any one address; 0 disables the cap,
222 // which is the inert default. A concurrency bound orthogonal to the rate
223 // window, closing the slow-hold shape a rate limiter alone cannot see.
224 pub conn_max: usize,
225 // Live connections across every address right now, a guard-wide gauge kept
226 // in step by `acquire` and `ConnPermit::drop`. Read for a health body.
227 pub live_total: Arc<AtomicUsize>,
228 // How long an alive-but-quiet record keeps its throttle history before it
229 // decays: an address past this long with no fresh offence has its throttle
230 // count reset and an auto-blacklist or throttle relaxed to Monitor. 0
231 // disables the decay, which is the inert default -- but a running process
232 // then never forgets a one-off burst, so a shared NAT address that was heavy
233 // once stays near the blacklist threshold until a restart.
234 pub decay_after: Duration,
235}
236
237impl<
238 const C: usize,
239 M: MapMut<HashForm, AddressLog<N, D>> + Clone + Debug,
240 H: Hasher + Send + Sync + 'static,
241 const S: usize,
242 const N: usize,
243 D: Clone + Debug + Default,
244>
245 AddressGuard<C, M, H, S, N, D>
246{
247 fn ip_bytes(addr: &IpAddr) -> Vec<u8> {
248 match addr {
249 IpAddr::V4(a) => a.octets().to_vec(),
250 IpAddr::V6(a) => a.octets().to_vec(),
251 }
252 }
253
254 /// The jitter is coarse and deterministic, taken from the system clock: its
255 /// purpose is to stop cooldowns expiring together, not to be unguessable.
256 fn sunset(&self) -> Duration {
257 if self.tsunset_spread.is_zero() {
258 return self.tsunset_base;
259 }
260 let spread = self.tsunset_spread.as_nanos() as u64;
261 let now = match SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) {
262 Ok(d) => d.as_nanos() as u64,
263 Err(_) => 0,
264 };
265 self.tsunset_base + Duration::from_nanos(now % spread)
266 }
267
268 /// An address not seen before is inserted in `Monitor` and allowed through.
269 pub fn check(&self, addr: &IpAddr) -> Outcome<GuardDecision> {
270 let (decision, _) = res!(self.update_log(addr, |_log, _new| Ok(())));
271 Ok(decision)
272 }
273
274 /// Claim one concurrent connection slot for an address.
275 ///
276 /// Returns `Some(permit)` when the address is under its concurrency cap, and
277 /// `None` when it is at or over it -- the caller then drops the connection.
278 /// The permit decrements both the per-address and the guard-wide live count
279 /// when it is dropped, so it must be held for exactly as long as the
280 /// connection lives, which is what moving it into the serving task achieves.
281 ///
282 /// Orthogonal to [`Self::check`]: `check` is the *rate* limit run before the
283 /// handshake, `acquire` is the *concurrency* limit that bounds how many
284 /// connections one address may hold open at once. A distributed flood is
285 /// caught by the rate window; a slow-hold from a single source is caught
286 /// here. A `conn_max` of 0 never refuses, only counts.
287 pub fn acquire(&self, addr: &IpAddr) -> Outcome<Option<ConnPermit>> {
288 let key = self.amap.key(&Self::ip_bytes(addr));
289 let locked_map = res!(self.amap.get_shard_using_hash(&key));
290 let counter = {
291 let mut unlocked_map = lock_write!(locked_map);
292 match unlocked_map.get(&key).map(|log| log.conns.clone()) {
293 Some(c) => c,
294 None => {
295 let mut log = AddressLog::<N, D>::default();
296 log.ip = Some(*addr);
297 let c = log.conns.clone();
298 unlocked_map.insert(key, log);
299 c
300 },
301 }
302 };
303 // The increment is speculative and rolled back on refusal, so a rejected
304 // connection leaves the count exactly where it was.
305 let now = counter.fetch_add(1, Ordering::AcqRel) + 1;
306 if self.conn_max != 0 && now > self.conn_max {
307 counter.fetch_sub(1, Ordering::AcqRel);
308 return Ok(None);
309 }
310 self.live_total.fetch_add(1, Ordering::AcqRel);
311 Ok(Some(ConnPermit {
312 per_ip: counter,
313 total: self.live_total.clone(),
314 }))
315 }
316
317 /// Live connections across every address right now.
318 pub fn live_conns(&self) -> usize {
319 self.live_total.load(Ordering::Acquire)
320 }
321
322 /// Prove in-process that the guard's state machine is armed, without any
323 /// external probe that a live guard would itself blacklist.
324 ///
325 /// Blacklists a documentation address (TEST-NET-1, `192.0.2.1`, RFC 5737),
326 /// confirms it is then blocked, releases it, and confirms it is allowed
327 /// again. A `true` result is what a health body reports as `guard_selftest`.
328 /// Leaves one monitor entry for that reserved address behind, which is inert.
329 pub fn self_test(&self) -> bool {
330 let ip: IpAddr = Ipv4Addr::new(192, 0, 2, 1).into();
331 if self.blacklist(&ip).is_err() {
332 return false;
333 }
334 let blocked = matches!(self.check(&ip), Ok(d) if d.should_drop());
335 let _ = self.unblock(&ip);
336 let allowed = matches!(self.check(&ip), Ok(GuardDecision::Allow));
337 blocked && allowed
338 }
339
340 /// `extra` runs while the shard write lock is still held, so a caller can
341 /// compose its own check on top of the generic rate limit without a second
342 /// acquisition. Its boolean argument says whether this call created the log.
343 pub fn update_log<F, T>(
344 &self,
345 addr: &IpAddr,
346 extra: F,
347 )
348 -> Outcome<(GuardDecision, T)>
349 where
350 F: FnOnce(&mut AddressLog<N, D>, bool) -> Outcome<T>,
351 {
352 let key = self.amap.key(&Self::ip_bytes(addr));
353 let locked_map = res!(self.amap.get_shard_using_hash(&key));
354 let mut unlocked_map = lock_write!(locked_map);
355 let now = SystemTime::now();
356
357 // Fast path for new entries: insert a fresh Monitor log, record the first request,
358 // run the caller's extra hook and return.
359 if unlocked_map.get(&key).is_none() {
360 let mut log = AddressLog::<N, D>::default();
361 log.ip = Some(*addr);
362 log.first_seen = now;
363 log.last_seen = now;
364 log.total_reqs = 1;
365 if let AddressState::Monitor(ref mut reqs) = log.state {
366 reqs.update();
367 }
368 let extra_val = res!(extra(&mut log, true));
369 unlocked_map.insert(key, log);
370 return Ok((GuardDecision::Allow, extra_val));
371 }
372
373 // Existing entry: update, run state machine, run extra hook, all under one lock.
374 let sunset = self.sunset();
375 let decision = {
376 let log = match unlocked_map.get_mut(&key) {
377 Some(l) => l,
378 None => return Err(err!(
379 "Address log for {} vanished between contains and get_mut.", addr;
380 Bug, Missing)),
381 };
382 log.last_seen = now;
383 log.total_reqs = log.total_reqs.saturating_add(1);
384 self.evaluate(log, sunset)
385 };
386 let extra_val = match unlocked_map.get_mut(&key) {
387 Some(l) => res!(extra(l, false)),
388 None => return Err(err!(
389 "Address log for {} vanished mid-update.", addr;
390 Bug, Missing)),
391 };
392 Ok((decision, extra_val))
393 }
394
395 /// One Monitor -> Throttle -> Blacklist step, for an entry already present.
396 fn evaluate(
397 &self,
398 log: &mut AddressLog<N, D>,
399 sunset: Duration,
400 )
401 -> GuardDecision
402 {
403 // Decay first: an address that has served the decay window with no fresh
404 // offence has its throttle history forgiven and an auto-blacklist or
405 // throttle relaxed to Monitor, so a one-off burst does not hold a shared
406 // NAT address near the blacklist threshold for the life of the process.
407 // A manual blacklist and a whitelist are operator decisions and never
408 // decay. This runs on a live request, so it relaxes an alive-but-quiet
409 // record, which the idle sweep (that evicts a *dead* record) does not.
410 if !self.decay_after.is_zero() {
411 let quiet_enough = log.offended_at
412 .and_then(|t| t.elapsed().ok())
413 .map(|since| since >= self.decay_after)
414 .unwrap_or(false);
415 if quiet_enough {
416 let auto_black = matches!(log.state,
417 AddressState::Blacklist { reason: BlacklistReason::AutoRateLimit, .. });
418 let throttled = matches!(log.state, AddressState::Throttle { .. });
419 if auto_black || throttled {
420 log.state = AddressState::Monitor(RingTimer::default());
421 }
422 log.throttle_cnt = 0;
423 log.offended_at = None;
424 }
425 }
426
427 // Sunset expired throttled addresses back to Monitor before this step.
428 if let AddressState::Throttle{ start, sunset: cool, .. } = &log.state {
429 if let Ok(elapsed) = start.elapsed() {
430 if elapsed >= *cool {
431 log.state = AddressState::Monitor(RingTimer::default());
432 }
433 }
434 }
435
436 match &mut log.state {
437 AddressState::Monitor(ref mut reqs) => {
438 reqs.update();
439 if reqs.avg_rps() > self.arps_max {
440 let next_cnt = log.throttle_cnt.saturating_add(1);
441 if next_cnt >= self.blist_cnt {
442 log.state = AddressState::Blacklist {
443 since: SystemTime::now(),
444 reason: BlacklistReason::AutoRateLimit,
445 };
446 log.throttle_cnt = next_cnt;
447 log.offended_at = Some(SystemTime::now());
448 return GuardDecision::Blocked(BlacklistReason::AutoRateLimit);
449 }
450 log.state = AddressState::Throttle {
451 reqs: RingTimer::default(),
452 tint_min: self.tint_min,
453 start: SystemTime::now(),
454 sunset,
455 };
456 log.throttle_cnt = next_cnt;
457 log.offended_at = Some(SystemTime::now());
458 return GuardDecision::Throttled;
459 }
460 GuardDecision::Allow
461 },
462 AddressState::Throttle{ ref mut reqs, tint_min, .. } => {
463 reqs.update();
464 if reqs.last_duration() < *tint_min {
465 return GuardDecision::Throttled;
466 }
467 GuardDecision::Allow
468 },
469 AddressState::Blacklist{ reason, .. } => {
470 GuardDecision::Blocked(*reason)
471 },
472 AddressState::Whitelist => GuardDecision::Allow,
473 }
474 }
475
476 /// For a caller that wants to drive the map itself rather than go through
477 /// `update_log`.
478 pub fn get_locked_map(
479 &self,
480 addr: &SocketAddr,
481 )
482 -> Outcome<(HashForm, &RwLock<M>)>
483 {
484 let ip_addr = addr.ip();
485 let key = self.amap.key(&Self::ip_bytes(&ip_addr));
486 let locked_map = res!(self.amap.get_shard_using_hash(&key));
487 Ok((key, locked_map))
488 }
489
490 /// Creates the log where the address has not been seen before.
491 pub fn whitelist(&self, addr: &IpAddr) -> Outcome<()> {
492 let key = self.amap.key(&Self::ip_bytes(addr));
493 let locked_map = res!(self.amap.get_shard_using_hash(&key));
494 let mut unlocked_map = lock_write!(locked_map);
495 match unlocked_map.get_mut(&key) {
496 Some(log) => log.state = AddressState::Whitelist,
497 None => {
498 let mut log = AddressLog::<N, D>::default();
499 log.ip = Some(*addr);
500 log.state = AddressState::Whitelist;
501 unlocked_map.insert(key, log);
502 }
503 }
504 Ok(())
505 }
506
507 /// The reason recorded is `Manual`, and the log is created where the address
508 /// has not been seen before.
509 pub fn blacklist(&self, addr: &IpAddr) -> Outcome<()> {
510 let key = self.amap.key(&Self::ip_bytes(addr));
511 let locked_map = res!(self.amap.get_shard_using_hash(&key));
512 let mut unlocked_map = lock_write!(locked_map);
513 let bl = AddressState::Blacklist {
514 since: SystemTime::now(),
515 reason: BlacklistReason::Manual,
516 };
517 match unlocked_map.get_mut(&key) {
518 Some(log) => log.state = bl,
519 None => {
520 let mut log = AddressLog::<N, D>::default();
521 log.ip = Some(*addr);
522 log.state = bl;
523 unlocked_map.insert(key, log);
524 }
525 }
526 Ok(())
527 }
528
529 /// Back to `Monitor`, with the throttle count zeroed, so the address starts
530 /// again from nothing.
531 pub fn unblock(&self, addr: &IpAddr) -> Outcome<()> {
532 let key = self.amap.key(&Self::ip_bytes(addr));
533 let locked_map = res!(self.amap.get_shard_using_hash(&key));
534 let mut unlocked_map = lock_write!(locked_map);
535 if let Some(log) = unlocked_map.get_mut(&key) {
536 log.state = AddressState::default();
537 log.throttle_cnt = 0;
538 }
539 Ok(())
540 }
541
542 /// Reads without mutating, unlike `check`. `None` where the address has never
543 /// been observed.
544 pub fn peek(&self, addr: &IpAddr) -> Outcome<Option<&'static str>> {
545 let key = self.amap.key(&Self::ip_bytes(addr));
546 let locked_map = res!(self.amap.get_shard_using_hash(&key));
547 let unlocked_map = lock_read!(locked_map);
548 Ok(unlocked_map.get(&key).map(|l| l.state.label()))
549 }
550
551 /// Walks every address, so the cost is linear in the number of them.
552 pub fn counts(&self) -> Outcome<GuardCounts> {
553 let mut c = GuardCounts::default();
554 for i in 0..self.amap.n {
555 if let Some(locked_map) = self.amap.shards[i].as_ref() {
556 let unlocked = lock_read!(locked_map);
557 for (_k, log) in unlocked.iter() {
558 c.total += 1;
559 c.total_reqs = c.total_reqs.saturating_add(log.total_reqs);
560 match log.state {
561 AddressState::Monitor(_) => c.monitor += 1,
562 AddressState::Throttle{..} => c.throttle += 1,
563 AddressState::Blacklist{..} => c.blacklist += 1,
564 AddressState::Whitelist => c.whitelist += 1,
565 }
566 }
567 }
568 }
569 Ok(c)
570 }
571
572 /// Counts plus up to `max` per-address entries, in no particular order: a
573 /// caller wanting a stable view sorts the entries itself.
574 pub fn snapshot(&self, max: usize) -> Outcome<GuardSnapshot> {
575 let mut snap = GuardSnapshot::default();
576 for i in 0..self.amap.n {
577 if let Some(locked_map) = self.amap.shards[i].as_ref() {
578 let unlocked = lock_read!(locked_map);
579 for (_k, log) in unlocked.iter() {
580 snap.counts.total += 1;
581 snap.counts.total_reqs = snap.counts.total_reqs.saturating_add(log.total_reqs);
582 match log.state {
583 AddressState::Monitor(_) => snap.counts.monitor += 1,
584 AddressState::Throttle{..} => snap.counts.throttle += 1,
585 AddressState::Blacklist{..} => snap.counts.blacklist += 1,
586 AddressState::Whitelist => snap.counts.whitelist += 1,
587 }
588 if snap.entries.len() < max {
589 if let Some(ip) = log.ip {
590 snap.entries.push(GuardEntry {
591 ip,
592 state: log.state.label(),
593 throttle_cnt: log.throttle_cnt,
594 total_reqs: log.total_reqs,
595 first_seen: log.first_seen,
596 last_seen: log.last_seen,
597 });
598 }
599 }
600 }
601 }
602 }
603 Ok(snap)
604 }
605
606 /// Evict idle `Monitor` records, returning how many were dropped.
607 ///
608 /// A distributed flood is a great many addresses each making one request
609 /// below the rate floor: every one mints a `Monitor` log that is never
610 /// throttled and so, without this, is never reclaimed -- the memory-
611 /// exhaustion shape the rate window cannot itself close. This drops a record
612 /// only when all three hold: it is in `Monitor` (a transient rate-watching
613 /// state, not an operator `Whitelist` nor an active `Throttle`/`Blacklist`
614 /// carrying a security decision that must outlive it until its own sunset);
615 /// it has no live connection (`conns == 0`, so eviction can never reset a
616 /// per-IP concurrency cap out from under a connection still holding a
617 /// `ConnPermit`); and it has not been seen within `idle`. A re-created record
618 /// starts clean, which for a `Monitor` address reaches the identical
619 /// decision, so nothing is lost by forgetting it.
620 pub fn sweep_idle(&self, idle: Duration) -> Outcome<usize> {
621 let now = SystemTime::now();
622 let mut evicted = 0usize;
623 for i in 0..self.amap.n {
624 if let Some(locked_map) = self.amap.shards[i].as_ref() {
625 let mut unlocked = lock_write!(locked_map);
626 unlocked.retain(|_k, log| {
627 let stale = now.duration_since(log.last_seen)
628 .map(|age| age > idle)
629 .unwrap_or(false);
630 let quiet = log.conns.load(Ordering::Acquire) == 0;
631 let transient = matches!(log.state, AddressState::Monitor(_));
632 let drop_it = stale && quiet && transient;
633 if drop_it {
634 evicted += 1;
635 }
636 !drop_it
637 });
638 }
639 }
640 Ok(evicted)
641 }
642}
643
644#[cfg(test)]
645mod tests {
646 use super::*;
647
648 use oxedyne_fe2o3_hash::{
649 hash::HashScheme,
650 map::ShardMap,
651 };
652 use oxedyne_fe2o3_iop_hash::api::HashForm;
653
654 use std::{
655 collections::BTreeMap,
656 net::Ipv4Addr,
657 };
658
659 const N: usize = 16; // Short ring so tests can saturate it quickly.
660
661 type TestGuard = AddressGuard<
662 4, // C: shards
663 BTreeMap<HashForm, AddressLog<N, ()>>, // M: inner map
664 HashScheme, // H: hasher
665 8, // S: salt len
666 N, // N: ring length
667 (), // D: no extension data
668 >;
669
670 fn make_guard(rps_max: u64, blist_cnt: u16) -> TestGuard {
671 let salt = [1u8; 8];
672 AddressGuard {
673 amap: ShardMap::<4, 8, AddressLog<N, ()>, BTreeMap<HashForm, AddressLog<N, ()>>, HashScheme>::new(
674 4,
675 salt,
676 BTreeMap::new(),
677 HashScheme::try_from("Seahash").expect("seahash scheme"),
678 ).expect("shard map"),
679 arps_max: rps_max,
680 tint_min: Duration::from_millis(10),
681 tsunset_base: Duration::from_millis(50),
682 tsunset_spread: Duration::ZERO,
683 blist_cnt,
684 conn_max: 0,
685 live_total: Arc::new(AtomicUsize::new(0)),
686 decay_after: Duration::ZERO,
687 }
688 }
689
690 fn make_conn_guard(conn_max: usize) -> TestGuard {
691 let mut g = make_guard(1_000_000, 1_000);
692 g.conn_max = conn_max;
693 g
694 }
695
696 #[test]
697 fn check_allows_first_request() {
698 let guard = make_guard(100, 5);
699 let addr: IpAddr = Ipv4Addr::new(10, 0, 0, 1).into();
700 let d = guard.check(&addr).expect("check");
701 assert_eq!(d, GuardDecision::Allow);
702 let counts = guard.counts().expect("counts");
703 assert_eq!(counts.total, 1);
704 assert_eq!(counts.monitor, 1);
705 }
706
707 #[test]
708 fn manual_blacklist_blocks_subsequent_checks() {
709 let guard = make_guard(100, 5);
710 let addr: IpAddr = Ipv4Addr::new(10, 0, 0, 2).into();
711 guard.blacklist(&addr).expect("blacklist");
712 let d = guard.check(&addr).expect("check");
713 match d {
714 GuardDecision::Blocked(BlacklistReason::Manual) => (),
715 other => panic!("expected Blocked(Manual), got {:?}", other),
716 }
717 }
718
719 #[test]
720 fn unblock_restores_monitor() {
721 let guard = make_guard(100, 5);
722 let addr: IpAddr = Ipv4Addr::new(10, 0, 0, 3).into();
723 guard.blacklist(&addr).expect("blacklist");
724 guard.unblock(&addr).expect("unblock");
725 let d = guard.check(&addr).expect("check");
726 assert_eq!(d, GuardDecision::Allow);
727 }
728
729 #[test]
730 fn whitelist_overrides_rate_limit() {
731 let guard = make_guard(1, 5); // very tight
732 let addr: IpAddr = Ipv4Addr::new(10, 0, 0, 4).into();
733 guard.whitelist(&addr).expect("whitelist");
734 // Drive many rapid checks; all should be Allow.
735 for _ in 0..64 {
736 assert_eq!(guard.check(&addr).expect("check"), GuardDecision::Allow);
737 }
738 }
739
740 #[test]
741 fn concurrency_admits_up_to_the_cap_and_refuses_beyond() {
742 let guard = make_conn_guard(3);
743 let addr: IpAddr = Ipv4Addr::new(10, 0, 1, 1).into();
744 // Three admitted, held live.
745 let p1 = guard.acquire(&addr).expect("acquire 1");
746 let p2 = guard.acquire(&addr).expect("acquire 2");
747 let p3 = guard.acquire(&addr).expect("acquire 3");
748 assert!(p1.is_some() && p2.is_some() && p3.is_some(),
749 "the first N connections must be admitted");
750 assert_eq!(guard.live_conns(), 3);
751 // The N+1th is refused while the first three are still held.
752 let p4 = guard.acquire(&addr).expect("acquire 4");
753 assert!(p4.is_none(), "the N+1th connection must be refused");
754 assert_eq!(guard.live_conns(), 3, "a refusal must not count");
755 }
756
757 #[test]
758 fn concurrency_releases_on_drop() {
759 let guard = make_conn_guard(2);
760 let addr: IpAddr = Ipv4Addr::new(10, 0, 1, 2).into();
761 let p1 = guard.acquire(&addr).expect("acquire 1");
762 {
763 let p2 = guard.acquire(&addr).expect("acquire 2");
764 assert!(p2.is_some());
765 // At the cap with p1 and p2 held: a third is refused, and the refusal
766 // does not count.
767 assert!(guard.acquire(&addr).expect("acquire 3").is_none(),
768 "at the cap, a third is refused");
769 assert_eq!(guard.live_conns(), 2);
770 }
771 // p2 dropped at the end of the block: a slot is free again.
772 assert_eq!(guard.live_conns(), 1, "dropping a permit must free its slot");
773 let p4 = guard.acquire(&addr).expect("acquire 4");
774 assert!(p4.is_some(), "a released slot must be reusable");
775 // p1 and p4 held.
776 assert_eq!(guard.live_conns(), 2);
777 drop(p1);
778 // Only p4 remains.
779 assert_eq!(guard.live_conns(), 1);
780 }
781
782 #[test]
783 fn concurrency_is_isolated_per_address() {
784 let guard = make_conn_guard(1);
785 let a: IpAddr = Ipv4Addr::new(10, 0, 1, 3).into();
786 let b: IpAddr = Ipv4Addr::new(10, 0, 1, 4).into();
787 let _pa = guard.acquire(&a).expect("acquire a").expect("a admitted");
788 // a is at its cap, but b is unaffected.
789 assert!(guard.acquire(&a).expect("acquire a2").is_none(),
790 "a's second connection is refused");
791 assert!(guard.acquire(&b).expect("acquire b").is_some(),
792 "a second address must be unaffected by the first's cap");
793 }
794
795 #[test]
796 fn a_zero_cap_never_refuses_but_still_counts() {
797 let guard = make_conn_guard(0);
798 let addr: IpAddr = Ipv4Addr::new(10, 0, 1, 5).into();
799 let mut held = Vec::new();
800 for _ in 0..100 {
801 held.push(guard.acquire(&addr).expect("acquire").expect("admitted"));
802 }
803 assert_eq!(guard.live_conns(), 100);
804 held.clear();
805 assert_eq!(guard.live_conns(), 0);
806 }
807
808 #[test]
809 fn self_test_reports_armed() {
810 let guard = make_guard(100, 5);
811 assert!(guard.self_test(), "a constructed guard must report armed");
812 }
813
814 #[test]
815 fn throttle_history_decays_after_a_quiet_window() {
816 // A high rate ceiling so the decayed record does not immediately re-offend
817 // on the very check that relaxes it; the decay is what is under test.
818 let mut guard = make_guard(1_000_000, 3);
819 guard.decay_after = Duration::from_millis(20);
820 let addr: IpAddr = Ipv4Addr::new(10, 0, 4, 1).into();
821
822 // Plant a record auto-blacklisted well in the past, with no fresh offence:
823 // an address that was heavy once and has since gone quiet.
824 let old = SystemTime::now()
825 .checked_sub(Duration::from_millis(200))
826 .expect("time before now");
827 guard.update_log(&addr, |log, _new| {
828 log.state = AddressState::Blacklist {
829 since: old,
830 reason: BlacklistReason::AutoRateLimit,
831 };
832 log.throttle_cnt = 3;
833 log.offended_at = Some(old);
834 Ok(())
835 }).expect("plant blacklisted record");
836
837 // The next check sees the quiet window has elapsed: the auto-blacklist is
838 // relaxed and the throttle history forgiven, so the address is allowed.
839 let decision = guard.check(&addr).expect("check");
840 assert_eq!(decision, GuardDecision::Allow,
841 "a quiet auto-blacklisted address must decay back to allowed");
842
843 // A manual blacklist is an operator decision and must NOT decay.
844 let banned: IpAddr = Ipv4Addr::new(10, 0, 4, 2).into();
845 guard.update_log(&banned, |log, _new| {
846 log.state = AddressState::Blacklist {
847 since: old,
848 reason: BlacklistReason::Manual,
849 };
850 log.offended_at = Some(old);
851 Ok(())
852 }).expect("plant manual blacklist");
853 assert!(matches!(guard.check(&banned), Ok(d) if d.should_drop()),
854 "a manual blacklist must survive the decay window");
855 }
856
857 #[test]
858 fn sweep_evicts_idle_monitor_records_but_spares_live_and_blacklisted() {
859 let guard = make_conn_guard(4);
860 let idle: IpAddr = Ipv4Addr::new(10, 0, 2, 1).into();
861 let live: IpAddr = Ipv4Addr::new(10, 0, 2, 2).into();
862 let banned: IpAddr = Ipv4Addr::new(10, 0, 2, 3).into();
863
864 // An idle Monitor record: acquired then released, so conns == 0.
865 drop(guard.acquire(&idle).expect("acquire idle"));
866 // A live record still holding a connection.
867 let _held = guard.acquire(&live).expect("acquire live").expect("admitted");
868 // A blacklisted record, which carries a security decision.
869 guard.blacklist(&banned).expect("blacklist");
870
871 // Everything is younger than an hour, so an hour-idle sweep evicts nothing.
872 assert_eq!(guard.sweep_idle(Duration::from_secs(3600)).expect("sweep"), 0,
873 "a sweep must spare records seen within the idle window");
874
875 // A zero idle window makes every past-seen record stale: the idle Monitor
876 // one goes, the live one is spared (conns > 0), the blacklisted one is
877 // spared (not Monitor).
878 let evicted = guard.sweep_idle(Duration::ZERO).expect("sweep");
879 assert_eq!(evicted, 1, "only the idle Monitor record must be evicted");
880 // The live cap is intact: its slot was never reset by an eviction.
881 assert_eq!(guard.live_conns(), 1);
882 // The blacklist survived the sweep.
883 assert!(matches!(guard.check(&banned), Ok(d) if d.should_drop()),
884 "a blacklisted address must not be forgotten by an idle sweep");
885 }
886
887 #[test]
888 fn snapshot_reports_per_address_entries() {
889 let guard = make_guard(100, 5);
890 let a: IpAddr = Ipv4Addr::new(10, 0, 0, 5).into();
891 let b: IpAddr = Ipv4Addr::new(10, 0, 0, 6).into();
892 guard.check(&a).expect("a");
893 guard.check(&b).expect("b");
894 let snap = guard.snapshot(16).expect("snapshot");
895 assert_eq!(snap.counts.total, 2);
896 assert_eq!(snap.entries.len(), 2);
897 assert!(snap.entries.iter().any(|e| e.ip == a));
898 assert!(snap.entries.iter().any(|e| e.ip == b));
899 }
900}