Skip to main content

ipfrs_network/
flood_protection.rs

1//! Multi-layer flood/DoS protection for P2P networks.
2//!
3//! Provides three integrated protection layers:
4//! 1. **Message deduplication** – FNV-1a–hashed sliding-window cache rejects
5//!    replayed messages within a configurable time window.
6//! 2. **Per-peer rate limiting** – Token-bucket throttles each peer independently.
7//! 3. **Global rate limiting** – A single global token bucket caps aggregate
8//!    throughput regardless of how many peers contribute traffic.
9//!
10//! Peers that violate limits accumulate a violation counter; once they reach
11//! [`FloodConfig::ban_threshold`] they are temporarily banned for
12//! [`FloodConfig::ban_duration_ms`] milliseconds.
13//!
14//! ## Example
15//!
16//! ```rust
17//! use ipfrs_network::flood_protection::{
18//!     FloodConfig, FloodProtection, MessageId, fnv1a_message_id, CheckResult,
19//! };
20//!
21//! let config = FloodConfig::default();
22//! let mut fp = FloodProtection::new(config);
23//! let msg_id = fnv1a_message_id(b"hello");
24//! let now_ms: u64 = 1_000_000;
25//!
26//! match fp.check("peer-1", msg_id, now_ms) {
27//!     CheckResult::Allow => {
28//!         fp.record_allowed("peer-1", msg_id, now_ms);
29//!     }
30//!     other => eprintln!("blocked: {:?}", other),
31//! }
32//! ```
33
34use std::collections::{HashMap, VecDeque};
35
36// ── FNV-1a constants ───────────────────────────────────────────────────────────
37
38const FNV_OFFSET_BASIS_64: u64 = 14_695_981_039_346_656_037;
39const FNV_PRIME_64: u64 = 1_099_511_628_211;
40
41// ═══════════════════════════════════════════════════════════════════════════════
42// Public primitive types
43// ═══════════════════════════════════════════════════════════════════════════════
44
45/// A 64-bit message identifier derived via FNV-1a from message bytes.
46#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
47pub struct MessageId(pub u64);
48
49/// Compute a [`MessageId`] from raw message bytes using FNV-1a hashing.
50///
51/// ```rust
52/// use ipfrs_network::flood_protection::{fnv1a_message_id, MessageId};
53/// let id = fnv1a_message_id(b"test message");
54/// assert_ne!(id.0, 0);
55/// ```
56#[inline]
57pub fn fnv1a_message_id(data: &[u8]) -> MessageId {
58    let mut h: u64 = FNV_OFFSET_BASIS_64;
59    for &b in data {
60        h ^= b as u64;
61        h = h.wrapping_mul(FNV_PRIME_64);
62    }
63    MessageId(h)
64}
65
66// ═══════════════════════════════════════════════════════════════════════════════
67// Configuration
68// ═══════════════════════════════════════════════════════════════════════════════
69
70/// Configuration for [`FloodProtection`].
71#[derive(Debug, Clone)]
72pub struct FloodConfig {
73    /// Global maximum requests per second (token bucket).
74    pub global_rps: u64,
75    /// Per-peer maximum requests per second (token bucket).
76    pub per_peer_rps: u64,
77    /// Sliding window for message deduplication, in milliseconds.
78    pub dedup_window_ms: u64,
79    /// Maximum number of entries in the deduplication cache.
80    pub dedup_capacity: usize,
81    /// Number of violations before a peer is banned.
82    pub ban_threshold: u32,
83    /// How long a ban lasts, in milliseconds.
84    pub ban_duration_ms: u64,
85    /// Token-bucket burst capacity = `rps * burst_multiplier`.
86    pub burst_multiplier: f64,
87}
88
89impl Default for FloodConfig {
90    fn default() -> Self {
91        Self {
92            global_rps: 1_000,
93            per_peer_rps: 100,
94            dedup_window_ms: 30_000,
95            dedup_capacity: 10_000,
96            ban_threshold: 10,
97            ban_duration_ms: 300_000,
98            burst_multiplier: 2.0,
99        }
100    }
101}
102
103// ═══════════════════════════════════════════════════════════════════════════════
104// Result / violation types
105// ═══════════════════════════════════════════════════════════════════════════════
106
107/// Outcome of a [`FloodProtection::check`] call.
108#[derive(Debug, Clone, PartialEq, Eq)]
109pub enum CheckResult {
110    /// The message is allowed to pass through.
111    Allow,
112    /// The originating peer is currently banned.
113    Banned {
114        /// Timestamp (ms) when the ban expires.
115        until: u64,
116    },
117    /// The global token bucket is exhausted.
118    GlobalRateLimited,
119    /// The per-peer token bucket for this peer is exhausted.
120    PeerRateLimited,
121    /// This message was seen recently and is a duplicate.
122    Duplicate,
123}
124
125/// Categories of violations tracked by the flood-protection layer.
126#[derive(Debug, Clone, PartialEq, Eq, Hash)]
127pub enum ViolationType {
128    /// Peer exceeded its per-peer rate limit.
129    RateLimitExceeded,
130    /// Peer sent a message that was already seen.
131    DuplicateMessage,
132    /// Peer sent a message while banned.
133    BannedPeer,
134}
135
136impl ViolationType {
137    fn as_str(&self) -> &'static str {
138        match self {
139            Self::RateLimitExceeded => "RateLimitExceeded",
140            Self::DuplicateMessage => "DuplicateMessage",
141            Self::BannedPeer => "BannedPeer",
142        }
143    }
144}
145
146/// A single violation event.
147#[derive(Debug, Clone)]
148pub struct ViolationRecord {
149    /// Peer that triggered the violation.
150    pub peer_id: String,
151    /// Kind of violation.
152    pub violation_type: ViolationType,
153    /// Wall-clock timestamp of the violation (ms since epoch / reference point).
154    pub timestamp: u64,
155    /// Optional message that triggered the violation.
156    pub message_id: Option<MessageId>,
157}
158
159// ═══════════════════════════════════════════════════════════════════════════════
160// Per-peer state
161// ═══════════════════════════════════════════════════════════════════════════════
162
163/// Runtime state maintained for each peer.
164#[derive(Debug, Clone)]
165pub struct PeerState {
166    /// Peer identifier.
167    pub peer_id: String,
168    /// Current token count for the per-peer token bucket.
169    pub tokens: f64,
170    /// Timestamp (ms) of the last token-refill operation.
171    pub last_refill: u64,
172    /// Cumulative violation count since peer was first seen (or last unbanned).
173    pub violation_count: u32,
174    /// If `Some(t)`, the peer is banned until timestamp `t`.
175    pub banned_until: Option<u64>,
176}
177
178impl PeerState {
179    fn new(peer_id: &str, initial_tokens: f64, now: u64) -> Self {
180        Self {
181            peer_id: peer_id.to_string(),
182            tokens: initial_tokens,
183            last_refill: now,
184            violation_count: 0,
185            banned_until: None,
186        }
187    }
188}
189
190// ═══════════════════════════════════════════════════════════════════════════════
191// Aggregate statistics
192// ═══════════════════════════════════════════════════════════════════════════════
193
194/// Snapshot of flood-protection statistics.
195#[derive(Debug, Clone)]
196pub struct FloodStats {
197    /// Number of peers with tracked state.
198    pub total_peers: usize,
199    /// Number of currently-banned peers.
200    pub banned_peers: usize,
201    /// Number of entries in the deduplication cache.
202    pub dedup_cache_size: usize,
203    /// Total messages allowed since creation.
204    pub total_allowed: u64,
205    /// Total messages blocked since creation.
206    pub total_blocked: u64,
207    /// Current global token count.
208    pub global_tokens: f64,
209}
210
211// ═══════════════════════════════════════════════════════════════════════════════
212// Main struct
213// ═══════════════════════════════════════════════════════════════════════════════
214
215/// Multi-layer flood/DoS protection for P2P networks.
216///
217/// See the [module-level documentation](self) for a detailed description of each
218/// protection layer and usage examples.
219pub struct FloodProtection {
220    /// Immutable configuration.
221    pub config: FloodConfig,
222    /// Global token-bucket state.
223    pub global_tokens: f64,
224    /// Timestamp of last global token refill.
225    pub global_last_refill: u64,
226    /// Per-peer token-bucket and ban state.
227    pub peers: HashMap<String, PeerState>,
228    /// Sliding-window deduplication cache: `(message_id, insertion_timestamp_ms)`.
229    pub seen_messages: VecDeque<(MessageId, u64)>,
230    /// Ring buffer of recent violations (capped at 1 000).
231    pub violations: VecDeque<ViolationRecord>,
232    /// Monotonically increasing count of allowed messages.
233    pub total_allowed: u64,
234    /// Monotonically increasing count of blocked messages.
235    pub total_blocked: u64,
236}
237
238// ── Internal helpers ───────────────────────────────────────────────────────────
239
240impl FloodProtection {
241    /// Initial token count for a new bucket = `rps * burst_multiplier`.
242    #[inline]
243    fn initial_tokens(rps: u64, burst: f64) -> f64 {
244        rps as f64 * burst
245    }
246
247    /// Refill `tokens` according to elapsed time and return the updated value.
248    ///
249    /// ```text
250    /// new_tokens = min(old + elapsed_ms / 1000 * rps, rps * burst)
251    /// ```
252    #[inline]
253    fn refill(tokens: f64, last_refill: u64, now: u64, rps: u64, burst: f64) -> (f64, u64) {
254        let elapsed_ms = now.saturating_sub(last_refill);
255        let replenished = tokens + elapsed_ms as f64 / 1_000.0 * rps as f64;
256        let cap = rps as f64 * burst;
257        (replenished.min(cap), now)
258    }
259
260    /// Check whether a `MessageId` is already in the dedup cache.
261    fn is_duplicate(&self, msg_id: MessageId) -> bool {
262        self.seen_messages.iter().any(|(id, _)| *id == msg_id)
263    }
264}
265
266// ── Public API ─────────────────────────────────────────────────────────────────
267
268impl FloodProtection {
269    /// Create a new [`FloodProtection`] instance with the given configuration.
270    pub fn new(config: FloodConfig) -> Self {
271        let burst = config.burst_multiplier;
272        let g_rps = config.global_rps;
273        Self {
274            global_tokens: Self::initial_tokens(g_rps, burst),
275            global_last_refill: 0,
276            peers: HashMap::new(),
277            seen_messages: VecDeque::new(),
278            violations: VecDeque::new(),
279            total_allowed: 0,
280            total_blocked: 0,
281            config,
282        }
283    }
284
285    // ── Core check ──────────────────────────────────────────────────────────
286
287    /// Run all protection checks in priority order.
288    ///
289    /// The checks are applied in this order:
290    /// 1. Peer ban check.
291    /// 2. Global rate limit.
292    /// 3. Per-peer rate limit.
293    /// 4. Message deduplication.
294    ///
295    /// This method is **read-only**: it does not consume tokens or register the
296    /// message.  Call [`record_allowed`](Self::record_allowed) after receiving
297    /// [`CheckResult::Allow`] to commit the state change.
298    ///
299    /// Violations triggered here are recorded automatically.
300    pub fn check(&mut self, peer_id: &str, message_id: MessageId, now: u64) -> CheckResult {
301        // ── 1. Ban check ─────────────────────────────────────────────────────
302        if self.is_banned(peer_id, now) {
303            let until = self
304                .peers
305                .get(peer_id)
306                .and_then(|p| p.banned_until)
307                .unwrap_or(now);
308            self.record_violation(peer_id, ViolationType::BannedPeer, Some(message_id), now);
309            self.total_blocked += 1;
310            return CheckResult::Banned { until };
311        }
312
313        // ── 2. Global rate limit ─────────────────────────────────────────────
314        let burst = self.config.burst_multiplier;
315        let (new_global, new_global_ts) = Self::refill(
316            self.global_tokens,
317            self.global_last_refill,
318            now,
319            self.config.global_rps,
320            burst,
321        );
322        self.global_tokens = new_global;
323        self.global_last_refill = new_global_ts;
324
325        if self.global_tokens < 1.0 {
326            self.total_blocked += 1;
327            return CheckResult::GlobalRateLimited;
328        }
329
330        // ── 3. Per-peer rate limit ───────────────────────────────────────────
331        let peer_rps = self.config.per_peer_rps;
332        let peer_burst = self.config.burst_multiplier;
333        let initial_tokens = Self::initial_tokens(peer_rps, peer_burst);
334
335        let peer = self
336            .peers
337            .entry(peer_id.to_string())
338            .or_insert_with(|| PeerState::new(peer_id, initial_tokens, now));
339
340        let (new_peer_tokens, new_peer_ts) =
341            Self::refill(peer.tokens, peer.last_refill, now, peer_rps, peer_burst);
342        peer.tokens = new_peer_tokens;
343        peer.last_refill = new_peer_ts;
344
345        if peer.tokens < 1.0 {
346            self.record_violation(
347                peer_id,
348                ViolationType::RateLimitExceeded,
349                Some(message_id),
350                now,
351            );
352            self.total_blocked += 1;
353            return CheckResult::PeerRateLimited;
354        }
355
356        // ── 4. Deduplication ─────────────────────────────────────────────────
357        if self.is_duplicate(message_id) {
358            self.record_violation(
359                peer_id,
360                ViolationType::DuplicateMessage,
361                Some(message_id),
362                now,
363            );
364            self.total_blocked += 1;
365            return CheckResult::Duplicate;
366        }
367
368        CheckResult::Allow
369    }
370
371    // ── State mutation after allow ───────────────────────────────────────────
372
373    /// Consume one token from the global and per-peer buckets and record the
374    /// message in the deduplication cache.
375    ///
376    /// This **must** be called after [`check`](Self::check) returns
377    /// [`CheckResult::Allow`] and before the next call to `check` for the same
378    /// peer, otherwise the token counts will be stale.
379    pub fn record_allowed(&mut self, peer_id: &str, message_id: MessageId, now: u64) {
380        // Consume global token.
381        self.global_tokens -= 1.0;
382
383        // Consume per-peer token (create entry if somehow missing).
384        let peer_rps = self.config.per_peer_rps;
385        let burst = self.config.burst_multiplier;
386        let initial_tokens = Self::initial_tokens(peer_rps, burst);
387
388        let peer = self
389            .peers
390            .entry(peer_id.to_string())
391            .or_insert_with(|| PeerState::new(peer_id, initial_tokens, now));
392        peer.tokens -= 1.0;
393
394        // Register in dedup cache, evicting overflow if at capacity.
395        if self.seen_messages.len() >= self.config.dedup_capacity {
396            self.seen_messages.pop_front();
397        }
398        self.seen_messages.push_back((message_id, now));
399
400        self.total_allowed += 1;
401    }
402
403    // ── Violation tracking ───────────────────────────────────────────────────
404
405    /// Record a violation for `peer_id` and ban them if they hit the threshold.
406    ///
407    /// The violations ring-buffer is capped at 1 000 entries; oldest entries are
408    /// silently dropped when the buffer is full.
409    pub fn record_violation(
410        &mut self,
411        peer_id: &str,
412        vtype: ViolationType,
413        msg_id: Option<MessageId>,
414        now: u64,
415    ) {
416        // Push violation record, respecting the 1 000-entry cap.
417        if self.violations.len() >= 1_000 {
418            self.violations.pop_front();
419        }
420        self.violations.push_back(ViolationRecord {
421            peer_id: peer_id.to_string(),
422            violation_type: vtype,
423            timestamp: now,
424            message_id: msg_id,
425        });
426
427        // Bump the violation counter on the peer and potentially ban them.
428        let ban_threshold = self.config.ban_threshold;
429        let ban_duration = self.config.ban_duration_ms;
430        let peer_rps = self.config.per_peer_rps;
431        let burst = self.config.burst_multiplier;
432        let initial_tokens = Self::initial_tokens(peer_rps, burst);
433
434        let peer = self
435            .peers
436            .entry(peer_id.to_string())
437            .or_insert_with(|| PeerState::new(peer_id, initial_tokens, now));
438
439        peer.violation_count += 1;
440        if peer.violation_count >= ban_threshold {
441            peer.banned_until = Some(now + ban_duration);
442        }
443    }
444
445    // ── Ban management ───────────────────────────────────────────────────────
446
447    /// Return `true` if `peer_id` is currently banned.
448    ///
449    /// If a ban has expired at `now` the peer's `banned_until` field is cleared
450    /// automatically (lazy unban).
451    pub fn is_banned(&mut self, peer_id: &str, now: u64) -> bool {
452        if let Some(peer) = self.peers.get_mut(peer_id) {
453            if let Some(until) = peer.banned_until {
454                if now >= until {
455                    // Ban has expired – remove it.
456                    peer.banned_until = None;
457                    return false;
458                }
459                return true;
460            }
461        }
462        false
463    }
464
465    /// Unconditionally lift the ban on `peer_id`.
466    pub fn unban(&mut self, peer_id: &str) {
467        if let Some(peer) = self.peers.get_mut(peer_id) {
468            peer.banned_until = None;
469        }
470    }
471
472    // ── Eviction ─────────────────────────────────────────────────────────────
473
474    /// Remove deduplication entries older than [`FloodConfig::dedup_window_ms`].
475    ///
476    /// Returns the number of entries removed.
477    pub fn evict_expired_dedup(&mut self, now: u64) -> usize {
478        let window = self.config.dedup_window_ms;
479        let cutoff = now.saturating_sub(window);
480        let before = self.seen_messages.len();
481        self.seen_messages.retain(|(_, ts)| *ts > cutoff);
482        before - self.seen_messages.len()
483    }
484
485    /// Remove unbanned peers whose `last_refill` is older than `max_age_ms`.
486    ///
487    /// Banned peers are never evicted regardless of age.
488    /// Returns the number of peer entries removed.
489    pub fn evict_stale_peers(&mut self, max_age_ms: u64, now: u64) -> usize {
490        let cutoff = now.saturating_sub(max_age_ms);
491        let before = self.peers.len();
492        self.peers.retain(|_, p| {
493            // Keep if banned (regardless of staleness).
494            if p.banned_until.is_some() {
495                return true;
496            }
497            // Keep if recently active.
498            p.last_refill > cutoff
499        });
500        before - self.peers.len()
501    }
502
503    // ── Reporting ────────────────────────────────────────────────────────────
504
505    /// Summarise violations recorded since `since` (ms timestamp).
506    ///
507    /// Returns a map from violation-type name to occurrence count.
508    pub fn violation_summary(&self, since: u64) -> HashMap<String, usize> {
509        let mut summary: HashMap<String, usize> = HashMap::new();
510        for v in &self.violations {
511            if v.timestamp >= since {
512                *summary
513                    .entry(v.violation_type.as_str().to_string())
514                    .or_insert(0) += 1;
515            }
516        }
517        summary
518    }
519
520    /// Return a point-in-time statistics snapshot.
521    pub fn stats(&self, now: u64) -> FloodStats {
522        let banned_peers = self
523            .peers
524            .values()
525            .filter(|p| p.banned_until.map(|u| u > now).unwrap_or(false))
526            .count();
527
528        FloodStats {
529            total_peers: self.peers.len(),
530            banned_peers,
531            dedup_cache_size: self.seen_messages.len(),
532            total_allowed: self.total_allowed,
533            total_blocked: self.total_blocked,
534            global_tokens: self.global_tokens,
535        }
536    }
537}
538
539// ══════════════════════════════════════════════════════════════════════════════
540// Tests
541// ══════════════════════════════════════════════════════════════════════════════
542
543#[cfg(test)]
544mod tests {
545    use std::collections::HashSet;
546
547    use crate::flood_protection::{
548        fnv1a_message_id, CheckResult, FloodConfig, FloodProtection, MessageId, ViolationType,
549    };
550
551    // ── helpers ────────────────────────────────────────────────────────────
552
553    fn default_fp() -> FloodProtection {
554        FloodProtection::new(FloodConfig::default())
555    }
556
557    fn tight_config() -> FloodConfig {
558        FloodConfig {
559            global_rps: 10,
560            per_peer_rps: 5,
561            dedup_window_ms: 5_000,
562            dedup_capacity: 100,
563            ban_threshold: 3,
564            ban_duration_ms: 60_000,
565            burst_multiplier: 1.0, // no burst headroom
566        }
567    }
568
569    // ── fnv1a_message_id ───────────────────────────────────────────────────
570
571    #[test]
572    fn test_fnv1a_empty_slice() {
573        // FNV-1a of empty input equals the offset basis.
574        let id = fnv1a_message_id(&[]);
575        assert_eq!(id.0, 14_695_981_039_346_656_037_u64);
576    }
577
578    #[test]
579    fn test_fnv1a_deterministic() {
580        assert_eq!(fnv1a_message_id(b"hello"), fnv1a_message_id(b"hello"));
581    }
582
583    #[test]
584    fn test_fnv1a_different_inputs() {
585        assert_ne!(fnv1a_message_id(b"foo"), fnv1a_message_id(b"bar"));
586    }
587
588    #[test]
589    fn test_fnv1a_single_byte() {
590        let id = fnv1a_message_id(&[0x42]);
591        assert_ne!(id.0, 0);
592    }
593
594    #[test]
595    fn test_fnv1a_all_zeros_vs_empty() {
596        assert_ne!(fnv1a_message_id(&[0u8]), fnv1a_message_id(&[]));
597    }
598
599    #[test]
600    fn test_fnv1a_uniqueness_across_messages() {
601        let msgs: Vec<&[u8]> = vec![b"msg1", b"msg2", b"msg3", b"hello world", b"ipfrs"];
602        let ids: HashSet<u64> = msgs.iter().map(|m| fnv1a_message_id(m).0).collect();
603        assert_eq!(ids.len(), msgs.len(), "all hashes should be unique");
604    }
605
606    // ── MessageId ─────────────────────────────────────────────────────────
607
608    #[test]
609    fn test_message_id_equality() {
610        let a = MessageId(42);
611        let b = MessageId(42);
612        assert_eq!(a, b);
613    }
614
615    #[test]
616    fn test_message_id_inequality() {
617        assert_ne!(MessageId(1), MessageId(2));
618    }
619
620    #[test]
621    fn test_message_id_copy() {
622        let a = MessageId(99);
623        let b = a; // copy
624        assert_eq!(a, b);
625    }
626
627    // ── FloodConfig defaults ───────────────────────────────────────────────
628
629    #[test]
630    fn test_default_config() {
631        let cfg = FloodConfig::default();
632        assert_eq!(cfg.global_rps, 1_000);
633        assert_eq!(cfg.per_peer_rps, 100);
634        assert_eq!(cfg.dedup_window_ms, 30_000);
635        assert_eq!(cfg.dedup_capacity, 10_000);
636        assert_eq!(cfg.ban_threshold, 10);
637        assert_eq!(cfg.ban_duration_ms, 300_000);
638        assert!((cfg.burst_multiplier - 2.0).abs() < f64::EPSILON);
639    }
640
641    // ── Basic allow path ──────────────────────────────────────────────────
642
643    #[test]
644    fn test_fresh_message_allowed() {
645        let mut fp = default_fp();
646        let id = fnv1a_message_id(b"new message");
647        let result = fp.check("peer-a", id, 1_000_000);
648        assert_eq!(result, CheckResult::Allow);
649    }
650
651    #[test]
652    fn test_record_allowed_increments_counter() {
653        let mut fp = default_fp();
654        let id = fnv1a_message_id(b"m1");
655        fp.check("peer-a", id, 1_000);
656        fp.record_allowed("peer-a", id, 1_000);
657        assert_eq!(fp.total_allowed, 1);
658    }
659
660    #[test]
661    fn test_multiple_distinct_messages_allowed() {
662        let mut fp = default_fp();
663        let now = 1_000_000_u64;
664        for i in 0_u8..5 {
665            let id = fnv1a_message_id(&[i]);
666            assert_eq!(
667                fp.check("peer-a", id, now + u64::from(i) * 10),
668                CheckResult::Allow
669            );
670            fp.record_allowed("peer-a", id, now + u64::from(i) * 10);
671        }
672        assert_eq!(fp.total_allowed, 5);
673    }
674
675    // ── Deduplication ─────────────────────────────────────────────────────
676
677    #[test]
678    fn test_duplicate_message_blocked() {
679        let mut fp = FloodProtection::new(tight_config());
680        let id = fnv1a_message_id(b"dup");
681        let now = 1_000_u64;
682        assert_eq!(fp.check("peer-a", id, now), CheckResult::Allow);
683        fp.record_allowed("peer-a", id, now);
684        // Same message again – must be rejected.
685        let result = fp.check("peer-a", id, now + 1);
686        assert_eq!(result, CheckResult::Duplicate);
687    }
688
689    #[test]
690    fn test_duplicate_from_different_peer_also_blocked() {
691        let mut fp = FloodProtection::new(tight_config());
692        let id = fnv1a_message_id(b"shared");
693        fp.check("peer-a", id, 1_000);
694        fp.record_allowed("peer-a", id, 1_000);
695        // Different peer sending the same message.
696        let result = fp.check("peer-b", id, 1_100);
697        assert_eq!(result, CheckResult::Duplicate);
698    }
699
700    #[test]
701    fn test_evict_expired_dedup_removes_old_entries() {
702        let config = FloodConfig {
703            dedup_window_ms: 1_000,
704            global_rps: 10_000,
705            per_peer_rps: 10_000,
706            burst_multiplier: 2.0,
707            ..FloodConfig::default()
708        };
709        let mut fp = FloodProtection::new(config);
710        let id = fnv1a_message_id(b"old");
711        fp.check("peer-a", id, 0);
712        fp.record_allowed("peer-a", id, 0);
713        assert_eq!(fp.seen_messages.len(), 1);
714
715        // Time advances past the window.
716        let removed = fp.evict_expired_dedup(2_000);
717        assert_eq!(removed, 1);
718        assert_eq!(fp.seen_messages.len(), 0);
719    }
720
721    #[test]
722    fn test_evict_expired_dedup_keeps_fresh_entries() {
723        let config = FloodConfig {
724            dedup_window_ms: 60_000,
725            global_rps: 10_000,
726            per_peer_rps: 10_000,
727            burst_multiplier: 2.0,
728            ..FloodConfig::default()
729        };
730        let mut fp = FloodProtection::new(config);
731        let id = fnv1a_message_id(b"fresh");
732        fp.check("peer-a", id, 50_000);
733        fp.record_allowed("peer-a", id, 50_000);
734        let removed = fp.evict_expired_dedup(51_000); // only 1 s elapsed
735        assert_eq!(removed, 0);
736    }
737
738    #[test]
739    fn test_dedup_capacity_evicts_oldest() {
740        let config = FloodConfig {
741            dedup_capacity: 3,
742            global_rps: 100_000,
743            per_peer_rps: 100_000,
744            burst_multiplier: 10.0,
745            ..FloodConfig::default()
746        };
747        let mut fp = FloodProtection::new(config);
748        for i in 0_u8..4 {
749            let id = fnv1a_message_id(&[i]);
750            assert_eq!(
751                fp.check("peer-a", id, 1_000 + u64::from(i)),
752                CheckResult::Allow
753            );
754            fp.record_allowed("peer-a", id, 1_000 + u64::from(i));
755        }
756        // Cache should still be capped at 3.
757        assert_eq!(fp.seen_messages.len(), 3);
758        // The oldest entry (byte 0) was evicted; sending it again should be allowed.
759        let evicted_id = fnv1a_message_id(&[0_u8]);
760        assert_eq!(fp.check("peer-a", evicted_id, 2_000), CheckResult::Allow);
761    }
762
763    // ── Per-peer rate limiting ─────────────────────────────────────────────
764
765    #[test]
766    fn test_per_peer_rate_limit_exhausted() {
767        let config = FloodConfig {
768            per_peer_rps: 2,
769            global_rps: 1_000,
770            burst_multiplier: 1.0, // tokens start at 2
771            dedup_capacity: 10_000,
772            dedup_window_ms: 30_000,
773            ..FloodConfig::default()
774        };
775        let mut fp = FloodProtection::new(config);
776        let now = 0_u64;
777
778        // First two messages – consume all tokens.
779        let id1 = fnv1a_message_id(b"a");
780        assert_eq!(fp.check("p", id1, now), CheckResult::Allow);
781        fp.record_allowed("p", id1, now);
782
783        let id2 = fnv1a_message_id(b"b");
784        assert_eq!(fp.check("p", id2, now), CheckResult::Allow);
785        fp.record_allowed("p", id2, now);
786
787        // Third message – peer bucket empty.
788        let id3 = fnv1a_message_id(b"c");
789        let result = fp.check("p", id3, now);
790        assert_eq!(result, CheckResult::PeerRateLimited);
791    }
792
793    #[test]
794    fn test_per_peer_token_refill_over_time() {
795        let config = FloodConfig {
796            per_peer_rps: 1,
797            global_rps: 10_000,
798            burst_multiplier: 1.0,
799            dedup_capacity: 10_000,
800            dedup_window_ms: 30_000,
801            ..FloodConfig::default()
802        };
803        let mut fp = FloodProtection::new(config);
804
805        let id1 = fnv1a_message_id(b"x");
806        assert_eq!(fp.check("p", id1, 0), CheckResult::Allow);
807        fp.record_allowed("p", id1, 0);
808
809        // Immediately after – bucket empty.
810        let id2 = fnv1a_message_id(b"y");
811        assert_eq!(fp.check("p", id2, 0), CheckResult::PeerRateLimited);
812
813        // After 1 second the token is replenished.
814        let id3 = fnv1a_message_id(b"z");
815        assert_eq!(fp.check("p", id3, 1_000), CheckResult::Allow);
816    }
817
818    #[test]
819    fn test_different_peers_have_independent_buckets() {
820        let config = FloodConfig {
821            per_peer_rps: 1,
822            global_rps: 10_000,
823            burst_multiplier: 1.0,
824            dedup_capacity: 10_000,
825            dedup_window_ms: 30_000,
826            ..FloodConfig::default()
827        };
828        let mut fp = FloodProtection::new(config);
829
830        let id_a = fnv1a_message_id(b"pa");
831        let id_b = fnv1a_message_id(b"pb");
832
833        assert_eq!(fp.check("peer-a", id_a, 0), CheckResult::Allow);
834        fp.record_allowed("peer-a", id_a, 0);
835        // peer-a is exhausted but peer-b should still be fine.
836        assert_eq!(fp.check("peer-b", id_b, 0), CheckResult::Allow);
837    }
838
839    // ── Global rate limiting ───────────────────────────────────────────────
840
841    #[test]
842    fn test_global_rate_limit_exhausted() {
843        let config = FloodConfig {
844            global_rps: 2,
845            per_peer_rps: 1_000,
846            burst_multiplier: 1.0,
847            dedup_capacity: 10_000,
848            dedup_window_ms: 30_000,
849            ..FloodConfig::default()
850        };
851        let mut fp = FloodProtection::new(config);
852        let now = 0_u64;
853
854        let id1 = fnv1a_message_id(b"g1");
855        assert_eq!(fp.check("p1", id1, now), CheckResult::Allow);
856        fp.record_allowed("p1", id1, now);
857
858        let id2 = fnv1a_message_id(b"g2");
859        assert_eq!(fp.check("p2", id2, now), CheckResult::Allow);
860        fp.record_allowed("p2", id2, now);
861
862        // Global bucket exhausted.
863        let id3 = fnv1a_message_id(b"g3");
864        assert_eq!(fp.check("p3", id3, now), CheckResult::GlobalRateLimited);
865    }
866
867    #[test]
868    fn test_global_token_refill() {
869        let config = FloodConfig {
870            global_rps: 1,
871            per_peer_rps: 10_000,
872            burst_multiplier: 1.0,
873            dedup_capacity: 10_000,
874            dedup_window_ms: 30_000,
875            ..FloodConfig::default()
876        };
877        let mut fp = FloodProtection::new(config);
878
879        let id1 = fnv1a_message_id(b"r1");
880        assert_eq!(fp.check("p", id1, 0), CheckResult::Allow);
881        fp.record_allowed("p", id1, 0);
882
883        // No refill yet – blocked.
884        let id2 = fnv1a_message_id(b"r2");
885        assert_eq!(fp.check("p", id2, 0), CheckResult::GlobalRateLimited);
886
887        // After 1 second – token refilled.
888        let id3 = fnv1a_message_id(b"r3");
889        assert_eq!(fp.check("p", id3, 1_000), CheckResult::Allow);
890    }
891
892    // ── Ban management ────────────────────────────────────────────────────
893
894    #[test]
895    fn test_violation_ban_after_threshold() {
896        let config = FloodConfig {
897            ban_threshold: 2,
898            ban_duration_ms: 10_000,
899            global_rps: 1_000_000,
900            per_peer_rps: 1_000_000,
901            burst_multiplier: 10.0,
902            dedup_capacity: 10_000,
903            dedup_window_ms: 30_000,
904        };
905        let mut fp = FloodProtection::new(config);
906        let now = 1_000_u64;
907
908        fp.record_violation("bad-peer", ViolationType::RateLimitExceeded, None, now);
909        assert!(!fp.is_banned("bad-peer", now));
910        fp.record_violation("bad-peer", ViolationType::RateLimitExceeded, None, now);
911        assert!(fp.is_banned("bad-peer", now));
912    }
913
914    #[test]
915    fn test_ban_expires_after_duration() {
916        let config = FloodConfig {
917            ban_threshold: 1,
918            ban_duration_ms: 5_000,
919            ..FloodConfig::default()
920        };
921        let mut fp = FloodProtection::new(config);
922        let now = 1_000_u64;
923
924        fp.record_violation("bad-peer", ViolationType::BannedPeer, None, now);
925        assert!(fp.is_banned("bad-peer", now + 1_000)); // still banned
926        assert!(!fp.is_banned("bad-peer", now + 6_000)); // expired
927    }
928
929    #[test]
930    fn test_unban_clears_ban() {
931        let config = FloodConfig {
932            ban_threshold: 1,
933            ban_duration_ms: 100_000,
934            ..FloodConfig::default()
935        };
936        let mut fp = FloodProtection::new(config);
937        fp.record_violation("bad-peer", ViolationType::RateLimitExceeded, None, 0);
938        assert!(fp.is_banned("bad-peer", 100));
939        fp.unban("bad-peer");
940        assert!(!fp.is_banned("bad-peer", 100));
941    }
942
943    #[test]
944    fn test_check_returns_banned_for_banned_peer() {
945        let config = FloodConfig {
946            ban_threshold: 1,
947            ban_duration_ms: 100_000,
948            global_rps: 100_000,
949            per_peer_rps: 100_000,
950            burst_multiplier: 2.0,
951            dedup_capacity: 10_000,
952            dedup_window_ms: 30_000,
953        };
954        let mut fp = FloodProtection::new(config);
955        fp.record_violation("evil", ViolationType::RateLimitExceeded, None, 0);
956        let id = fnv1a_message_id(b"anything");
957        match fp.check("evil", id, 1_000) {
958            CheckResult::Banned { .. } => {}
959            other => panic!("expected Banned, got {other:?}"),
960        }
961    }
962
963    #[test]
964    fn test_unban_unknown_peer_is_noop() {
965        let mut fp = default_fp();
966        // Should not panic.
967        fp.unban("unknown-peer");
968    }
969
970    // ── Stale peer eviction ───────────────────────────────────────────────
971
972    #[test]
973    fn test_evict_stale_peers_removes_inactive() {
974        let mut fp = default_fp();
975        let id = fnv1a_message_id(b"ev");
976        fp.check("old-peer", id, 0);
977        fp.record_allowed("old-peer", id, 0);
978        assert_eq!(fp.peers.len(), 1);
979
980        let removed = fp.evict_stale_peers(1_000, 10_000);
981        assert_eq!(removed, 1);
982        assert!(fp.peers.is_empty());
983    }
984
985    #[test]
986    fn test_evict_stale_peers_keeps_banned() {
987        let config = FloodConfig {
988            ban_threshold: 1,
989            ban_duration_ms: 1_000_000,
990            ..FloodConfig::default()
991        };
992        let mut fp = FloodProtection::new(config);
993        fp.record_violation("banned-peer", ViolationType::RateLimitExceeded, None, 0);
994
995        // Even though last_refill is 0 and 10 000 ms have elapsed, banned peers stay.
996        let removed = fp.evict_stale_peers(1_000, 10_000);
997        assert_eq!(removed, 0);
998        assert_eq!(fp.peers.len(), 1);
999    }
1000
1001    #[test]
1002    fn test_evict_stale_peers_keeps_recently_active() {
1003        let mut fp = default_fp();
1004        let id = fnv1a_message_id(b"recent");
1005        fp.check("active-peer", id, 9_500);
1006        fp.record_allowed("active-peer", id, 9_500);
1007
1008        // Only 500 ms have elapsed – peer is still fresh.
1009        let removed = fp.evict_stale_peers(1_000, 10_000);
1010        assert_eq!(removed, 0);
1011    }
1012
1013    // ── Violation summary ─────────────────────────────────────────────────
1014
1015    #[test]
1016    fn test_violation_summary_counts_by_type() {
1017        let mut fp = default_fp();
1018        fp.record_violation("p", ViolationType::RateLimitExceeded, None, 1_000);
1019        fp.record_violation("p", ViolationType::DuplicateMessage, None, 2_000);
1020        fp.record_violation("p", ViolationType::RateLimitExceeded, None, 3_000);
1021
1022        let summary = fp.violation_summary(0);
1023        assert_eq!(summary.get("RateLimitExceeded").copied().unwrap_or(0), 2);
1024        assert_eq!(summary.get("DuplicateMessage").copied().unwrap_or(0), 1);
1025    }
1026
1027    #[test]
1028    fn test_violation_summary_respects_since_filter() {
1029        let mut fp = default_fp();
1030        fp.record_violation("p", ViolationType::RateLimitExceeded, None, 500);
1031        fp.record_violation("p", ViolationType::RateLimitExceeded, None, 1_500);
1032
1033        // Only the second violation is >= 1_000.
1034        let summary = fp.violation_summary(1_000);
1035        assert_eq!(summary.get("RateLimitExceeded").copied().unwrap_or(0), 1);
1036    }
1037
1038    #[test]
1039    fn test_violation_summary_empty_when_none() {
1040        let fp = default_fp();
1041        assert!(fp.violation_summary(0).is_empty());
1042    }
1043
1044    #[test]
1045    fn test_violation_buffer_capped_at_1000() {
1046        let mut fp = default_fp();
1047        for i in 0_u64..1_100 {
1048            fp.record_violation("p", ViolationType::RateLimitExceeded, None, i);
1049        }
1050        assert_eq!(fp.violations.len(), 1_000);
1051    }
1052
1053    // ── Stats ─────────────────────────────────────────────────────────────
1054
1055    #[test]
1056    fn test_stats_initial_state() {
1057        let fp = FloodProtection::new(FloodConfig {
1058            global_rps: 100,
1059            burst_multiplier: 2.0,
1060            ..FloodConfig::default()
1061        });
1062        let s = fp.stats(0);
1063        assert_eq!(s.total_peers, 0);
1064        assert_eq!(s.banned_peers, 0);
1065        assert_eq!(s.dedup_cache_size, 0);
1066        assert_eq!(s.total_allowed, 0);
1067        assert_eq!(s.total_blocked, 0);
1068        assert!((s.global_tokens - 200.0).abs() < 1e-6);
1069    }
1070
1071    #[test]
1072    fn test_stats_reflects_activity() {
1073        let mut fp = FloodProtection::new(FloodConfig {
1074            global_rps: 1_000,
1075            per_peer_rps: 1_000,
1076            burst_multiplier: 2.0,
1077            ..FloodConfig::default()
1078        });
1079        let id = fnv1a_message_id(b"stat-test");
1080        fp.check("p", id, 0);
1081        fp.record_allowed("p", id, 0);
1082
1083        let s = fp.stats(0);
1084        assert_eq!(s.total_peers, 1);
1085        assert_eq!(s.total_allowed, 1);
1086        assert_eq!(s.dedup_cache_size, 1);
1087    }
1088
1089    #[test]
1090    fn test_stats_banned_peers_counted() {
1091        let config = FloodConfig {
1092            ban_threshold: 1,
1093            ban_duration_ms: 100_000,
1094            ..FloodConfig::default()
1095        };
1096        let mut fp = FloodProtection::new(config);
1097        fp.record_violation("bad", ViolationType::RateLimitExceeded, None, 0);
1098        let s = fp.stats(500);
1099        assert_eq!(s.banned_peers, 1);
1100    }
1101
1102    // ── Interaction / integration checks ─────────────────────────────────
1103
1104    #[test]
1105    fn test_block_counts_increment_total_blocked() {
1106        let config = FloodConfig {
1107            global_rps: 1,
1108            burst_multiplier: 1.0,
1109            per_peer_rps: 10_000,
1110            dedup_capacity: 10_000,
1111            dedup_window_ms: 30_000,
1112            ..FloodConfig::default()
1113        };
1114        let mut fp = FloodProtection::new(config);
1115
1116        // Use up the one global token.
1117        let id1 = fnv1a_message_id(b"first");
1118        assert_eq!(fp.check("p", id1, 0), CheckResult::Allow);
1119        fp.record_allowed("p", id1, 0);
1120
1121        let id2 = fnv1a_message_id(b"second");
1122        fp.check("p", id2, 0); // GlobalRateLimited → blocked
1123        assert_eq!(fp.total_blocked, 1);
1124    }
1125
1126    #[test]
1127    fn test_check_ordering_ban_before_global() {
1128        // If a peer is banned the ban response must come before the global-limit
1129        // check (which might also be triggered).
1130        let config = FloodConfig {
1131            global_rps: 1,
1132            burst_multiplier: 1.0,
1133            per_peer_rps: 100_000,
1134            ban_threshold: 1,
1135            ban_duration_ms: 50_000,
1136            dedup_capacity: 10_000,
1137            dedup_window_ms: 30_000,
1138        };
1139        let mut fp = FloodProtection::new(config);
1140
1141        // Exhaust global bucket.
1142        let id1 = fnv1a_message_id(b"c1");
1143        fp.check("p", id1, 0);
1144        fp.record_allowed("p", id1, 0);
1145
1146        // Ban the peer.
1147        fp.record_violation("p", ViolationType::RateLimitExceeded, None, 0);
1148
1149        // Both global-limit and ban conditions hold; ban check fires first.
1150        let id2 = fnv1a_message_id(b"c2");
1151        match fp.check("p", id2, 0) {
1152            CheckResult::Banned { .. } => {}
1153            other => panic!("expected Banned, got {other:?}"),
1154        }
1155    }
1156
1157    #[test]
1158    fn test_burst_multiplier_affects_initial_tokens() {
1159        let config = FloodConfig {
1160            global_rps: 10,
1161            burst_multiplier: 3.0,
1162            per_peer_rps: 10,
1163            dedup_capacity: 10_000,
1164            dedup_window_ms: 30_000,
1165            ..FloodConfig::default()
1166        };
1167        let fp = FloodProtection::new(config);
1168        // Initial global tokens = rps * burst = 10 * 3 = 30.
1169        assert!((fp.global_tokens - 30.0).abs() < 1e-6);
1170    }
1171
1172    #[test]
1173    fn test_violation_with_message_id_stored() {
1174        let mut fp = default_fp();
1175        let msg = fnv1a_message_id(b"viol-msg");
1176        fp.record_violation("p", ViolationType::DuplicateMessage, Some(msg), 100);
1177        let v = fp.violations.back().expect("violation recorded");
1178        assert_eq!(v.message_id, Some(msg));
1179    }
1180}