chanora/crates/chanora_protocol/src/poke_limiter.rs

//! Per-connection poke strength classification.

use std::collections::HashMap;
use std::time::{Duration, Instant};

/// Notification strength assigned to an inbound poke.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum PokeStrength {
    /// Poke should be surfaced at full strength.
    Strong,
    /// Poke is rate-limited but below overflow severity.
    Suppressed,
    /// Poke remains suppressed after repeated suppressed pokes.
    SuppressedOverflow,
}

/// Per-connection poke limiter.
#[derive(Debug)]
pub struct PokeLimiter {
    window: Duration,
    entries: HashMap<u64, PokeEntry>,
}

#[derive(Debug)]
struct PokeEntry {
    tokens: u8,
    last_refill: Instant,
    suppressed_in_window: u8,
}

impl PokeLimiter {
    const CAPACITY: u8 = 2;
    const OVERFLOW_THRESHOLD: u8 = 3;

    /// Create a limiter using the default five-minute refill interval.
    pub fn new() -> Self {
        Self {
            window: Duration::from_secs(5 * 60),
            entries: HashMap::new(),
        }
    }

    /// Record a poke at the current instant.
    pub fn record(&mut self, sender_id: u64, own_client_id: Option<u64>) -> PokeStrength {
        self.record_at(sender_id, own_client_id, Instant::now())
    }

    /// Record a poke at an injected instant.
    pub fn record_at(
        &mut self,
        sender_id: u64,
        own_client_id: Option<u64>,
        now: Instant,
    ) -> PokeStrength {
        if own_client_id == Some(sender_id) {
            return PokeStrength::Suppressed;
        }

        let entry = self.entries.entry(sender_id).or_insert(PokeEntry {
            tokens: Self::CAPACITY,
            last_refill: now,
            suppressed_in_window: 0,
        });

        if now.duration_since(entry.last_refill) >= self.window {
            entry.tokens = Self::CAPACITY;
            entry.last_refill = now;
            entry.suppressed_in_window = 0;
        }

        if entry.tokens > 0 {
            entry.tokens -= 1;
            return PokeStrength::Strong;
        }

        entry.suppressed_in_window = entry.suppressed_in_window.saturating_add(1);
        if entry.suppressed_in_window >= Self::OVERFLOW_THRESHOLD {
            PokeStrength::SuppressedOverflow
        } else {
            PokeStrength::Suppressed
        }
    }
}

impl Default for PokeLimiter {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn first_two_pokes_are_strong_third_is_suppressed() {
        let mut limiter = PokeLimiter::new();
        let now = Instant::now();

        assert_eq!(limiter.record_at(7, Some(1), now), PokeStrength::Strong);
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(1)),
            PokeStrength::Strong
        );
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(2)),
            PokeStrength::Suppressed
        );
    }

    #[test]
    fn self_poke_is_suppressed_without_consuming_token() {
        let mut limiter = PokeLimiter::new();
        let now = Instant::now();

        assert_eq!(limiter.record_at(7, Some(7), now), PokeStrength::Suppressed);
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(1)),
            PokeStrength::Strong
        );
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(2)),
            PokeStrength::Strong
        );
    }

    #[test]
    fn overflow_after_three_suppressed_pokes_in_five_minutes() {
        let mut limiter = PokeLimiter::new();
        let now = Instant::now();

        assert_eq!(limiter.record_at(7, Some(1), now), PokeStrength::Strong);
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(1)),
            PokeStrength::Strong
        );
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(2)),
            PokeStrength::Suppressed
        );
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(3)),
            PokeStrength::Suppressed
        );
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(4)),
            PokeStrength::SuppressedOverflow
        );
    }

    #[test]
    fn refill_after_interval_uses_injected_time_without_sleeping() {
        let mut limiter = PokeLimiter::new();
        let now = Instant::now();

        assert_eq!(limiter.record_at(7, Some(1), now), PokeStrength::Strong);
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(1)),
            PokeStrength::Strong
        );
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(2)),
            PokeStrength::Suppressed
        );
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(5 * 60)),
            PokeStrength::Strong
        );
        assert_eq!(
            limiter.record_at(7, Some(1), now + Duration::from_secs(5 * 60 + 1)),
            PokeStrength::Strong
        );
    }
}