//! Abuse handling for the open UDP port.
//!
//! Layered in packet-touch order, cheapest check first, so that spending CPU on a
//! packet is always justified by the packet having survived everything cheaper:
//!
//! 1. length / version / type filter (in `otproto`, no state at all)
//! 2. per-IP token bucket
//! 3. unknown-token blocklist
//! 4. per-token token bucket
//! 5. AEAD verification — the first genuinely expensive step (~1 µs)
//! 6. AEAD-failure blocklist
//! 7. writer saturation → `THROTTLE`
//!
//! Two cross-cutting rules:
//!
//! * **Never reply to an unauthenticated packet.** Not with an error, not with a
//!   `NACK`. Anything else makes the port a reflector and a forgery oracle.
//! * **The rate limiter must not itself be a memory-exhaustion vector.** Every map
//!   here is capacity-capped and swept, because an attacker chooses the keys.

use std::net::IpAddr;
use std::num::NonZeroU32;
use std::time::{Duration, Instant};

use dashmap::DashMap;
use governor::clock::DefaultClock;
use governor::state::keyed::DefaultKeyedStateStore;
use governor::state::{InMemoryState, NotKeyed};
use governor::{Quota, RateLimiter};

/// Sustained packets per second from one IP, and how many may arrive at once.
/// Generous: a single IP can be a whole household behind NAT, all flushing queues
/// after a dead spot.
const IP_RATE: u32 = 25;
const IP_BURST: u32 = 75;

/// Per token. One report per second is already far above any profile; the burst
/// covers a queue flush after an offline stretch.
const TOKEN_RATE: u32 = 3;
const TOKEN_BURST: u32 = 40;

/// Revocation notices, the one reply sent without authenticating the request.
///
/// These numbers are the entire mitigation, so they are deliberately mean. The
/// key is the *destination* address, which for a spoofed datagram is the victim:
/// one packet a minute is a message, not a flood. A revoked device needs exactly
/// one to act, and it retries on its own schedule anyway.
///
/// The global ceiling bounds a distributed spoof, where each victim stays under
/// the per-address limit but the server is still made to emit broadly.
const NOTICE_PER_ADDR_PER_MINUTE: u32 = 1;
const NOTICE_GLOBAL_RATE: u32 = 10;
const NOTICE_GLOBAL_BURST: u32 = 20;

/// Strikes within [`STRIKE_WINDOW`] before an IP is banned for [`BAN_DURATION`].
const UNKNOWN_TOKEN_STRIKES: u32 = 50;
const AEAD_FAIL_STRIKES: u32 = 50;
const STRIKE_WINDOW: Duration = Duration::from_secs(60);
const BAN_DURATION: Duration = Duration::from_secs(600);

/// Hard ceiling on tracked IPs. Reached only under attack; when it is, the map is
/// swept rather than grown, because unbounded growth is the vulnerability.
const MAX_TRACKED_IPS: usize = 100_000;

type Keyed<K> = RateLimiter<K, DefaultKeyedStateStore<K>, DefaultClock>;
type Direct = RateLimiter<NotKeyed, InMemoryState, DefaultClock>;

fn quota(rate: u32, burst: u32) -> Quota {
    Quota::per_second(NonZeroU32::new(rate).expect("rate is a non-zero literal"))
        .allow_burst(NonZeroU32::new(burst).expect("burst is a non-zero literal"))
}

#[derive(Debug, Default)]
struct Strikes {
    count: u32,
    window_started: Option<Instant>,
    banned_until: Option<Instant>,
}

pub struct Limits {
    per_ip: Keyed<IpAddr>,
    per_token: Keyed<u64>,
    notice_per_addr: Keyed<IpAddr>,
    notice_global: Direct,
    unknown_token: DashMap<IpAddr, Strikes>,
    aead_fail: DashMap<IpAddr, Strikes>,
}

/// Why a packet was dropped. Used for aggregate counters only — never logged per
/// packet, because per-packet logging is itself a denial-of-service amplifier.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Drop {
    IpRate,
    TokenRate,
    Banned,
    /// The unauthenticated-notice budget, per destination or global.
    NoticeRate,
}

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

impl Limits {
    pub fn new() -> Self {
        Self {
            per_ip: RateLimiter::keyed(quota(IP_RATE, IP_BURST)),
            per_token: RateLimiter::keyed(quota(TOKEN_RATE, TOKEN_BURST)),
            notice_per_addr: RateLimiter::keyed(Quota::per_minute(
                NonZeroU32::new(NOTICE_PER_ADDR_PER_MINUTE).expect("non-zero literal"),
            )),
            notice_global: RateLimiter::direct(quota(NOTICE_GLOBAL_RATE, NOTICE_GLOBAL_BURST)),
            unknown_token: DashMap::new(),
            aead_fail: DashMap::new(),
        }
    }

    /// Step 2 and 3: is this source allowed to send anything right now?
    pub fn check_ip(&self, ip: IpAddr) -> Result<(), Drop> {
        if self.is_banned(ip) {
            return Err(Drop::Banned);
        }
        self.per_ip.check_key(&ip).map_err(|_| Drop::IpRate)
    }

    /// Per-credential budget, applied *after* the AEAD check.
    ///
    /// Deliberately not before it, even though that would be cheaper. Exceeding
    /// this budget produces a NACK, and a reply to an unverified datagram is a
    /// reflector: `token_id` is cleartext, so an attacker can name a real token,
    /// spoof a victim's source address, and have the server packet the victim.
    /// The per-IP budget is the cheap pre-crypto filter; this one is not.
    pub fn check_token(&self, token_id: u64) -> Result<(), Drop> {
        self.per_token
            .check_key(&token_id)
            .map_err(|_| Drop::TokenRate)
    }

    /// Called when a datagram names a token that does not exist.
    ///
    /// Scanning for valid token ids is the cheapest attack against this design —
    /// no key needed, just 2⁶⁴ guesses — so it is met with a ban rather than a
    /// reply. The port stays silent either way.
    pub fn note_unknown_token(&self, ip: IpAddr) {
        self.strike(&self.unknown_token, ip, UNKNOWN_TOKEN_STRIKES);
    }

    /// Called when a datagram named a real token but failed AEAD. Either
    /// corruption or forgery; both mean stop listening to this source.
    pub fn note_aead_failure(&self, ip: IpAddr) {
        self.strike(&self.aead_fail, ip, AEAD_FAIL_STRIKES);
    }

    fn is_banned(&self, ip: IpAddr) -> bool {
        let now = Instant::now();
        let banned = |m: &DashMap<IpAddr, Strikes>| {
            m.get(&ip)
                .and_then(|s| s.banned_until)
                .is_some_and(|t| t > now)
        };
        banned(&self.unknown_token) || banned(&self.aead_fail)
    }

    fn strike(&self, map: &DashMap<IpAddr, Strikes>, ip: IpAddr, threshold: u32) {
        // Sweep before inserting a new key, so the map can never exceed the cap.
        if map.len() >= MAX_TRACKED_IPS && !map.contains_key(&ip) {
            sweep(map);
            if map.len() >= MAX_TRACKED_IPS {
                // Still full after sweeping: we are under a wide attack. Dropping
                // the strike is the right failure mode — the per-IP rate limiter
                // is still holding, and growing this map is what the attacker
                // wants.
                return;
            }
        }

        let now = Instant::now();
        let mut entry = map.entry(ip).or_default();
        match entry.window_started {
            Some(started) if now.duration_since(started) <= STRIKE_WINDOW => entry.count += 1,
            _ => {
                entry.window_started = Some(now);
                entry.count = 1;
            }
        }
        if entry.count >= threshold {
            entry.banned_until = Some(now + BAN_DURATION);
            entry.count = 0;
            entry.window_started = None;
        }
    }

    /// May we send an unauthenticated revocation notice to `addr` right now?
    ///
    /// Both budgets are checked, and both must pass. This is called with the
    /// address the datagram *claimed* to come from, which is exactly the point:
    /// an attacker spoofing a victim gets the victim's budget, not their own.
    pub fn check_notice(&self, addr: IpAddr) -> Result<(), Drop> {
        if self.notice_global.check().is_err() {
            return Err(Drop::NoticeRate);
        }
        self.notice_per_addr
            .check_key(&addr)
            .map_err(|_| Drop::NoticeRate)
    }

    /// Periodic maintenance. Drops expired bans and stale rate-limiter state so
    /// idle keys do not accumulate.
    pub fn gc(&self) {
        self.per_ip.retain_recent();
        self.per_token.retain_recent();
        self.notice_per_addr.retain_recent();
        sweep(&self.unknown_token);
        sweep(&self.aead_fail);
    }

    /// For `/metrics`.
    pub fn tracked_ips(&self) -> usize {
        self.unknown_token.len() + self.aead_fail.len()
    }
}

fn sweep(map: &DashMap<IpAddr, Strikes>) {
    let now = Instant::now();
    map.retain(|_, s| {
        let ban_active = s.banned_until.is_some_and(|t| t > now);
        let window_active = s
            .window_started
            .is_some_and(|t| now.duration_since(t) <= STRIKE_WINDOW);
        ban_active || window_active
    });
}

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

    const IP: IpAddr = IpAddr::V4(std::net::Ipv4Addr::new(203, 0, 113, 5));

    #[test]
    fn a_burst_is_allowed_then_the_rate_bites() {
        let limits = Limits::new();
        for i in 0..IP_BURST {
            assert!(
                limits.check_ip(IP).is_ok(),
                "packet {i} of the burst was dropped"
            );
        }
        assert_eq!(limits.check_ip(IP), Err(Drop::IpRate));
    }

    #[test]
    fn per_token_limits_are_independent_of_each_other() {
        let limits = Limits::new();
        for _ in 0..TOKEN_BURST {
            assert!(limits.check_token(1).is_ok());
        }
        assert_eq!(limits.check_token(1), Err(Drop::TokenRate));
        // A different token is unaffected: one noisy phone must not silence a
        // whole household behind the same NAT.
        assert!(limits.check_token(2).is_ok());
    }

    #[test]
    fn token_scanning_earns_a_ban() {
        let limits = Limits::new();
        for _ in 0..UNKNOWN_TOKEN_STRIKES {
            limits.note_unknown_token(IP);
        }
        assert_eq!(limits.check_ip(IP), Err(Drop::Banned));
    }

    #[test]
    fn forged_packets_earn_a_ban() {
        let limits = Limits::new();
        for _ in 0..AEAD_FAIL_STRIKES {
            limits.note_aead_failure(IP);
        }
        assert_eq!(limits.check_ip(IP), Err(Drop::Banned));
    }

    #[test]
    fn a_ban_is_specific_to_the_offending_address() {
        let limits = Limits::new();
        for _ in 0..UNKNOWN_TOKEN_STRIKES {
            limits.note_unknown_token(IP);
        }
        let other = IpAddr::V4(std::net::Ipv4Addr::new(198, 51, 100, 9));
        assert!(limits.check_ip(other).is_ok());
    }

    #[test]
    fn strikes_below_the_threshold_do_not_ban() {
        let limits = Limits::new();
        for _ in 0..UNKNOWN_TOKEN_STRIKES - 1 {
            limits.note_unknown_token(IP);
        }
        assert!(limits.check_ip(IP).is_ok());
    }

    #[test]
    fn gc_drops_entries_with_nothing_left_to_remember() {
        let limits = Limits::new();
        limits.note_unknown_token(IP);
        assert_eq!(limits.tracked_ips(), 1);

        // Force the window to look expired rather than sleeping for a minute.
        limits
            .unknown_token
            .get_mut(&IP)
            .expect("entry")
            .window_started = Some(Instant::now() - STRIKE_WINDOW * 2);
        limits.gc();
        assert_eq!(limits.tracked_ips(), 0);
    }

    #[test]
    fn a_ban_survives_gc_until_it_expires() {
        let limits = Limits::new();
        for _ in 0..UNKNOWN_TOKEN_STRIKES {
            limits.note_unknown_token(IP);
        }
        limits.gc();
        assert_eq!(
            limits.check_ip(IP),
            Err(Drop::Banned),
            "gc must not lift an active ban"
        );
    }
}