//! Which client packets reach the ET server: in-game packets always, connectionless ones
//! only when a client or a master would send them, queries rate limited per source IP.
use crate::OOB;
use std::collections::HashMap;
use std::net::IpAddr;
use std::time::Duration;
use tokio::time::Instant;

/// Connectionless commands clients and masters send a server. Anything else is dropped,
/// rcon above all: only etadm uses it, from the server itself, and a stream of rcon
/// packets would block it (2.60b runs one rcon per 500 ms, whoever sends it).
const ALLOWED: [&str; 4] = ["getstatus", "getinfo", "getchallenge", "connect"];
/// PunkBuster's connectionless packets, matched as is by the server (no tokenizing).
const PUNKBUSTER: &[u8] = b"PB_";
/// Queries answered with a much bigger packet: from a spoofed source they would make the
/// server flood a victim (UDP amplification). Limited per source IP.
const QUERIES: [&str; 2] = ["getstatus", "getinfo"];
/// A burst (a browser refresh, a master's challenge), then one query a second.
const QUERY_BURST: f64 = 5.0;
const QUERY_RATE: f64 = 1.0;
/// Source IPs the limiter remembers: past this, queries from new IPs are dropped, so a
/// flood of spoofed sources can't grow memory.
const MAX_QUERY_SOURCES: usize = 65_536;
/// Idle IPs are forgotten after this (their bucket is full again long before).
const QUERY_FORGET: Duration = Duration::from_secs(60);

pub struct Filter {
    queries: HashMap<IpAddr, Bucket>,
    forgotten: Instant,
}

/// Token bucket of one source IP's queries.
struct Bucket {
    tokens: f64,
    at: Instant,
}

impl Filter {
    pub fn new(now: Instant) -> Self {
        Self {
            queries: HashMap::new(),
            forgotten: now,
        }
    }

    /// Why the packet is dropped, None to forward it. In-game packets always pass.
    pub fn refuse(&mut self, packet: &[u8], ip: IpAddr, now: Instant) -> Option<&'static str> {
        let body = packet.strip_prefix(OOB)?;
        if body.starts_with(PUNKBUSTER) {
            return None;
        }

        let name = first_token(body);
        if !ALLOWED
            .iter()
            .any(|c| c.as_bytes().eq_ignore_ascii_case(name))
        {
            return Some("connectionless command not allowed");
        }

        if QUERIES
            .iter()
            .any(|q| q.as_bytes().eq_ignore_ascii_case(name))
            && !self.query(ip, now)
        {
            return Some("query rate");
        }

        None
    }

    /// Takes a token from ip's bucket.
    fn query(&mut self, ip: IpAddr, now: Instant) -> bool {
        if now.duration_since(self.forgotten) >= QUERY_FORGET {
            self.queries
                .retain(|_, b| now.duration_since(b.at) < QUERY_FORGET);
            self.forgotten = now;
        }

        if !self.queries.contains_key(&ip) && self.queries.len() >= MAX_QUERY_SOURCES {
            return false;
        }

        let bucket = self.queries.entry(ip).or_insert(Bucket {
            tokens: QUERY_BURST,
            at: now,
        });
        let refill = now.duration_since(bucket.at).as_secs_f64() * QUERY_RATE;
        bucket.tokens = (bucket.tokens + refill).min(QUERY_BURST);
        bucket.at = now;

        if bucket.tokens < 1.0 {
            return false;
        }

        bucket.tokens -= 1.0;
        true
    }
}

/// First word of a connectionless packet as the server's tokenizer reads it: leading
/// bytes up to space skipped (control characters included). Only compared against
/// ALLOWED: a word the tokenizer would read differently (comments, quotes) matches
/// nothing and is dropped.
fn first_token(body: &[u8]) -> &[u8] {
    let start = body.iter().position(|b| *b > b' ').unwrap_or(body.len());
    let body = &body[start..];
    let end = body.iter().position(|b| *b <= b' ').unwrap_or(body.len());

    &body[..end]
}

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

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

    fn oob(body: &[u8]) -> Vec<u8> {
        [OOB, body].concat()
    }

    #[test]
    fn forwards_game_traffic() {
        let mut f = Filter::new(Instant::now());
        let now = Instant::now();

        for body in [
            &b"getchallenge"[..],
            b"connect \x12\x34",
            b"getinfo xxx",
            b"PB_whatever",
        ] {
            assert_eq!(f.refuse(&oob(body), IP, now), None, "{body:?}");
        }

        // In-game packets (no OOB prefix).
        assert_eq!(f.refuse(b"\x01\x02\x03", IP, now), None);
    }

    #[test]
    fn drops_rcon_however_written() {
        let mut f = Filter::new(Instant::now());
        let now = Instant::now();

        for body in [
            &b"rcon pass status"[..],
            b"RCON pass status",
            b"  \t rcon pass status",
            b"\x01rcon pass status",
            b"/**/rcon pass status",
            b"\"rcon\" pass status",
            b"",
        ] {
            assert!(f.refuse(&oob(body), IP, now).is_some(), "{body:?}");
        }
    }

    #[test]
    fn limits_queries_per_ip() {
        let start = Instant::now();
        let mut f = Filter::new(start);

        let passed = (0..20)
            .filter(|_| f.refuse(&oob(b"getstatus"), IP, start).is_none())
            .count();
        assert_eq!(passed, QUERY_BURST as usize);

        // Another IP has its own bucket; connect isn't a query.
        let other = IpAddr::V4(std::net::Ipv4Addr::new(198, 51, 100, 1));
        assert_eq!(f.refuse(&oob(b"getinfo"), other, start), None);
        assert_eq!(f.refuse(&oob(b"connect x"), IP, start), None);

        // One token a second.
        let later = start + Duration::from_secs(1);
        assert_eq!(f.refuse(&oob(b"getinfo"), IP, later), None);
        assert!(f.refuse(&oob(b"getinfo"), IP, later).is_some());

        // Forgotten after QUERY_FORGET: full again.
        let much_later = start + QUERY_FORGET + Duration::from_secs(1);
        let passed = (0..20)
            .filter(|_| f.refuse(&oob(b"getinfo"), IP, much_later).is_none())
            .count();
        assert_eq!(passed, QUERY_BURST as usize);
    }
}