goalkeeper 0.3.0

//! Per-IP connection and bandwidth limiter.

use crate::rate_limiter::{RateLimiterProps, RateLimiterState, Units};
use fxhash::FxHashMap;
use log::warn;
use std::net::IpAddr;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

/// Provides a mutable reference to some [IpLimiter].
pub trait ProvideIpLimiter: Clone {
    /// Visit a mutable refererence to the [IpLimiter].
    fn provide_ip_limiter<R>(&self, f: impl FnOnce(&mut IpLimiter) -> R) -> R;

    /// Visit the [IpStats] for each tracked IP.
    ///
    /// The `visitor` should never block.
    fn stats(&self, mut visitor: impl FnMut(IpAddr, IpStats)) {
        self.provide_ip_limiter(|this| {
            for (ip, v) in &this.usage {
                visitor(*ip, v.stats);
            }
        });
    }

    /// Total number of outstanding [ConnectionPermit]s.
    fn total_connections(&self) -> u32 {
        self.provide_ip_limiter(|this| this.total_connections)
    }

    /// Set the bandwidth limits for the [IpLimiter].
    ///
    /// The `bytes_per_second` should be less than 1,000,000,000.
    ///
    /// Default: 500,000 bytes per second, 1,000,000 bytes burst.
    fn set_bandwidth_limits(&self, bytes_per_second: Units, bytes_burst: Units) {
        self.provide_ip_limiter(|this| {
            this.connection_rate_limit =
                RateLimiterProps::new_throughput(bytes_per_second, bytes_burst);
        })
    }

    /// Set the properties of the custom rate limit corresponding to each IP.
    fn set_custom_limits(&self, props: RateLimiterProps) {
        self.provide_ip_limiter(|this| {
            this.custom_rate_limit = props;
        })
    }

    /// Set the 90th percentile number of [ConnectionPermit]s required for
    /// an [ActiveSession] for the [IpLimiter]. This should always
    /// be less than or equal to the value set by
    /// [Self::set_connections_per_active_p99].
    ///
    /// For example:
    /// - if you have an HTTP/1-only server, set this to 4-6.
    /// - if you have an HTTP/2 server with HTTP/1 WebSockets, set this to 2.
    /// - if you have an HTTP/2 server (possibly with HTTP/2 WebSockets), set this to 1.
    ///
    /// Default: 1
    fn set_connections_per_active_p90(&self, connections_per_active_p90: u32) {
        self.provide_ip_limiter(|this| {
            this.connections_per_active_p90 = connections_per_active_p90;
            this.connections_per_active_p99 = this
                .connections_per_active_p99
                .max(connections_per_active_p90);
        })
    }

    /// Set the 99th+ percentile number of [ConnectionPermit]s required for
    /// an [ActiveSession] for the [IpLimiter]. This should always
    /// be greater than or equal to the value set by
    /// [Self::set_connections_per_active_p90].
    ///
    /// For example, if your server accepts HTTP/1 clients, set this to 6-12.
    ///
    /// Default: 6
    fn set_connections_per_active_p99(&self, connections_per_active_p99: u32) {
        self.provide_ip_limiter(|this| {
            this.connections_per_active_p99 = connections_per_active_p99;
            this.connections_per_active_p90 = this
                .connections_per_active_p90
                .min(connections_per_active_p99);
        })
    }

    /// Set the soft maximum number of [ConnectionPermit]s, across all IP's,
    /// before fewer new [ConnectionPermit]s are afforded to each IP.
    ///
    /// Default: 400
    fn set_total_connections_soft_limit(&self, total_connections_soft_limit: u32) {
        self.provide_ip_limiter(|this| {
            this.total_connections_soft_limit = total_connections_soft_limit;
        })
    }

    /// Set the hard limit on the number of [ConnectionPermit]s, across all IP's,
    /// before all new [ConnectionPermit]s are denied.
    ///
    /// Default 1000.
    fn set_total_connections_hard_limit(&self, total_connections_hard_limit: u32) {
        self.provide_ip_limiter(|this| {
            this.total_connections_hard_limit = total_connections_hard_limit;
        })
    }

    /// Call when processing a message of `bytes` bytes from `ip` at `now`.
    ///
    /// If this returns `true`, block the `usage` of bandwidth.
    fn should_limit_bandwidth(
        &self,
        ip: IpAddr,
        bytes: Units,
        label: &'static str,
        now: Instant,
    ) -> bool {
        self.provide_ip_limiter(|this| this.should_limit_bandwidth_inner(ip, bytes, label, now))
    }

    /// Call to rate limit some custom (perhaps expensive) action per IP.
    ///
    /// If this returns `true`, block the `usage`.
    fn should_limit_custom(&self, ip: IpAddr, usage: Units, now: Instant) -> bool {
        self.provide_ip_limiter(|this| {
            let entry = this
                .usage
                .entry(ip)
                .or_insert_with(|| Usage::new(now, &mut this.new_ip_counter));
            entry
                .custom_rate_limit
                .should_limit_rate_with_now_and_usage(&this.custom_rate_limit, now, usage)
        })
    }

    /// Set how long a DDoS incident will be be remembered.
    ///
    /// Default: 10m
    fn set_ddos_memory(&self, ddos_memory: Duration) {
        self.provide_ip_limiter(|this| {
            this.ddos_memory = ddos_memory;
        })
    }

    /// Gets the value set by [`Self::set_compute_pressure`].
    fn compute_pressure(&self) -> bool {
        self.provide_ip_limiter(|this| this.compute_pressure)
    }

    /// Call with `true` any time the host's CPU and/or RAM are nearly exausted,
    /// then call with `false` when they are back to normal.
    ///
    /// This can trigger additional connection limiting before the connection
    /// count hits the value set by [Self::set_total_connections_soft_limit].
    ///
    /// Default: `false`
    fn set_compute_pressure(&self, compute_pressure: bool) {
        self.provide_ip_limiter(|this| {
            this.compute_pressure = compute_pressure;
        });
    }
}

/// A singleton implementation of [ProvideIpLimiter], analogous to the
/// [std::alloc::System] allocator.
#[derive(Copy, Clone, Default, Debug)]
pub struct SystemIpLimiter;

impl ProvideIpLimiter for SystemIpLimiter {
    fn provide_ip_limiter<R>(&self, f: impl FnOnce(&mut IpLimiter) -> R) -> R {
        static SINGLETON: Mutex<Option<IpLimiter>> = Mutex::new(None);
        let mut opt = SINGLETON.lock().unwrap();
        let this = opt.get_or_insert_with(Default::default);
        // Yikes..
        f(this)
    }
}

/// A non-global implementation of [ProvideIpLimiter]. If you only have one instance,
/// prefer [SystemIpLimiter].
#[derive(Clone, Default, Debug)]
pub struct ArcIpLimiter(Arc<Mutex<IpLimiter>>);

impl ProvideIpLimiter for ArcIpLimiter {
    fn provide_ip_limiter<R>(&self, f: impl FnOnce(&mut IpLimiter) -> R) -> R {
        let mut this = self.0.lock().unwrap();
        // Yikes..
        f(&mut this)
    }
}

/// Limits connections and bandwidth based on client IP addresses.
#[derive(Debug)]
pub struct IpLimiter {
    usage: FxHashMap<IpAddr, Usage>,
    connection_rate_limit: RateLimiterProps,
    custom_rate_limit: RateLimiterProps,
    next_prune: Instant,
    warning_limiter: RateLimiterState,
    pending: WarningSet,
    connections_per_active_p90: u32,
    connections_per_active_p99: u32,
    total_connections: u32,
    total_connections_soft_limit: u32,
    total_connections_hard_limit: u32,
    last_soft_limit: Option<Instant>,
    ddos_memory: Duration,
    compute_pressure: bool,
    new_ip_counter: u32,
}

#[derive(Debug)]
struct WarningSet {
    small: Vec<Warning>,
    small_full: bool,
    bandwidth: FxHashMap<IpAddr, u32>,
    connections: FxHashMap<IpAddr, u32>,
    granted_rare_exemptions: u32,
}

#[derive(Debug)]
struct Warning {
    ip: IpAddr,
    kind: WarningKind,
    label: &'static str,
}

#[derive(Debug)]
enum WarningKind {
    Bandwidth {
        amount: u32,
    },
    ConnectionCount {
        connections: u32,
        active_sessions: u32,
        limit: u32,
        granted_rare_exemption: bool,
    },
}

/// Summary statistics for a given IP. They may be forgotten
/// if there are no connections for an extended period.
#[derive(Copy, Clone, Debug)]
#[non_exhaustive]
pub struct IpStats {
    /// First connection.
    pub first: Instant,
    /// Number of outstanding [ConnectionPermit]s.
    pub connections: u32,
    /// Number of outstanding [ActiveSession]s.
    pub active_sessions: u32,
    // Max concurrent [ActiveSession]s.
    // pub max_active_sessions: u32,
    /// Last time this IP hit a hard limit.
    pub last_limit: Option<Instant>,
}

impl IpStats {
    fn limited(&mut self, now: Instant) {
        self.last_limit = Some(now);
    }
}

impl Default for IpLimiter {
    fn default() -> Self {
        Self::new(500000, 1000000)
    }
}

impl IpLimiter {
    pub(crate) fn should_limit_bandwidth_inner(
        &mut self,
        ip: IpAddr,
        bytes: Units,
        label: &'static str,
        now: Instant,
    ) -> bool {
        let should_rate_limit = self.should_limit_bandwidth_inner_inner(ip, bytes, now);
        if should_rate_limit {
            self.warn(
                Warning {
                    label,
                    ip,
                    kind: WarningKind::Bandwidth { amount: bytes },
                },
                now,
            );
        }
        should_rate_limit
    }
}

#[derive(Debug)]
struct Usage {
    connection_rate_limit: RateLimiterState,
    custom_rate_limit: RateLimiterState,
    stats: IpStats,
    /// IP is in the top ~1% (eligible for p99 limits).
    granted_rare_exemption: bool,
}

impl Usage {
    fn new(now: Instant, new_ip_counter: &mut u32) -> Self {
        *new_ip_counter = new_ip_counter.saturating_add(1);
        Self {
            connection_rate_limit: RateLimiterState {
                until: now,
                burst_used: 0,
            },
            custom_rate_limit: RateLimiterState {
                until: now,
                burst_used: 0,
            },
            stats: IpStats {
                first: now,
                connections: 0,
                active_sessions: 0,
                last_limit: None,
            },
            granted_rare_exemption: false,
        }
    }
}

const WARNING_LIMIT: RateLimiterProps = RateLimiterProps::const_new(Duration::from_secs(1), 0);

/// A RAII guard representing a permissible, long-lived connection (such as a TCP stream).
#[derive(Debug)]
pub struct ConnectionPermit<P: ProvideIpLimiter = SystemIpLimiter>(IpAddr, P);

impl ConnectionPermit<SystemIpLimiter> {
    /// Check if a new connection is permissible. If this returns `Some`, accept the
    /// connection and keep the [ConnectionPermit] for its lifetime. If this returns
    /// `None`, reject the connection.
    ///
    /// The `label` should be something like `"TCP connection"`.
    pub fn new(ip: IpAddr, label: &'static str) -> Option<Self> {
        Self::new_with(ip, label, SystemIpLimiter)
    }
}

impl<P: ProvideIpLimiter> ConnectionPermit<P> {
    /// Like [Self::new] but with any [ProvideIpLimiter] implementation.
    pub fn new_with(ip: IpAddr, label: &'static str, provide: P) -> Option<Self> {
        let now = Instant::now();
        provide
            .provide_ip_limiter(|limiter| {
                let entry = limiter
                    .usage
                    .entry(ip)
                    .or_insert_with(|| Usage::new(now, &mut limiter.new_ip_counter));

                if limiter.total_connections >= limiter.total_connections_hard_limit {
                    let warning = Warning {
                        ip,
                        label,
                        kind: WarningKind::ConnectionCount {
                            connections: entry.stats.connections,
                            active_sessions: entry.stats.active_sessions,
                            limit: entry.stats.connections,
                            granted_rare_exemption: false,
                        },
                    };
                    limiter.warn(warning, now);
                    return None;
                }

                // Represents overhead of starting a new connection.
                let amount = 10000;
                let should_rate_limit = entry
                    .connection_rate_limit
                    .should_limit_rate_with_now_and_usage(
                        &limiter.connection_rate_limit,
                        now,
                        amount,
                    );

                if should_rate_limit {
                    entry.stats.limited(now);
                    limiter.warn(
                        Warning {
                            label,
                            ip,
                            kind: WarningKind::Bandwidth { amount },
                        },
                        now,
                    );
                    return None;
                }
                let old = now.duration_since(entry.stats.first) > Duration::from_secs(60);
                let soft_limit_reached = limiter.compute_pressure
                    || limiter.total_connections >= limiter.total_connections_soft_limit;
                if soft_limit_reached {
                    limiter.last_soft_limit = Some(now);
                }
                let recent_global_soft_limit = limiter
                    .last_soft_limit
                    .filter(|&last| now.duration_since(last) < limiter.ddos_memory)
                    .is_some();
                let recent_local_limit = entry
                    .stats
                    .last_limit
                    .filter(|&last| now.duration_since(last) < limiter.ddos_memory)
                    .is_some();
                let strict_limit = (!old || recent_local_limit) && recent_global_soft_limit;
                let mut just_granted_rare_exemption = false;
                let granted_rare_exemption = if entry.granted_rare_exemption {
                    true
                } else if strict_limit && !recent_local_limit && limiter.new_ip_counter >= 100 {
                    limiter.new_ip_counter =
                        (limiter.new_ip_counter - 100).min(limiter.new_ip_counter / 2);
                    entry.granted_rare_exemption = true;
                    just_granted_rare_exemption = true;
                    true
                } else {
                    false
                };
                let limit = (entry.stats.active_sessions + 1 + (!strict_limit) as u32)
                    .saturating_mul(if strict_limit && !granted_rare_exemption {
                        limiter.connections_per_active_p90
                    } else {
                        limiter.connections_per_active_p99
                    });
                if entry.stats.connections >= limit {
                    entry.stats.limited(now);
                    let warning = Warning {
                        ip,
                        label,
                        kind: WarningKind::ConnectionCount {
                            connections: entry.stats.connections,
                            active_sessions: entry.stats.active_sessions,
                            limit,
                            granted_rare_exemption: false,
                        },
                    };
                    limiter.warn(warning, now);
                    None
                } else {
                    entry.stats.connections += 1;
                    limiter.total_connections += 1;

                    if just_granted_rare_exemption {
                        let warning = Warning {
                            ip,
                            label,
                            kind: WarningKind::ConnectionCount {
                                connections: entry.stats.connections,
                                active_sessions: entry.stats.active_sessions,
                                limit,
                                granted_rare_exemption: true,
                            },
                        };
                        limiter.warn(warning, now);
                    }

                    Some(ip)
                }
            })
            .map(|ip| Self(ip, provide))
    }
}

impl<P: ProvideIpLimiter> Drop for ConnectionPermit<P> {
    fn drop(&mut self) {
        self.1.provide_ip_limiter(|limiter| {
            if let Some(usage) = limiter.usage.get_mut(&self.0) {
                debug_assert!(usage.stats.connections > 0);
                usage.stats.connections = usage.stats.connections.saturating_sub(1);
            } else {
                debug_assert!(false);
            }
            // Fail open by subtracting from the total even if `get_mut` returned `None`.
            debug_assert!(limiter.total_connections > 0);
            limiter.total_connections = limiter.total_connections.saturating_sub(1);
        })
    }
}

/// A RAII guard representing a connection with meaningful activity taking place,
/// such as an authenticated WebSocket on top of a TCP stream.
#[derive(Debug)]
pub struct ActiveSession<P: ProvideIpLimiter = SystemIpLimiter>(IpAddr, P);

impl ActiveSession<SystemIpLimiter> {
    /// Keep the [ActiveSession] as long as meaningful activity is taking place.
    pub fn new(addr: IpAddr) -> Self {
        Self::new_with(addr, SystemIpLimiter)
    }
}

impl<P: ProvideIpLimiter> ActiveSession<P> {
    /// Like [Self::new] but with any [ProvideIpLimiter] implementation.
    pub fn new_with(addr: IpAddr, provide: P) -> Self {
        provide.provide_ip_limiter(|limiter| {
            limiter
                .usage
                .entry(addr)
                .or_insert_with(|| Usage::new(Instant::now(), &mut limiter.new_ip_counter))
                .stats
                .active_sessions += 1;
        });
        Self(addr, provide)
    }
}

impl<P: ProvideIpLimiter> Drop for ActiveSession<P> {
    fn drop(&mut self) {
        self.1.provide_ip_limiter(|limiter| {
            if let Some(usage) = limiter.usage.get_mut(&self.0) {
                debug_assert!(usage.stats.active_sessions > 0);
                usage.stats.active_sessions = usage.stats.active_sessions.saturating_sub(1);
            } else {
                debug_assert!(false);
            }
        })
    }
}

impl IpLimiter {
    /// Uses [`Units`] to represent bytes, to limit bandwidth.
    pub(crate) fn new(bytes_per_second: Units, bytes_burst: Units) -> Self {
        Self {
            usage: FxHashMap::default(),
            connection_rate_limit: RateLimiterProps::new_throughput(bytes_per_second, bytes_burst),
            custom_rate_limit: RateLimiterProps::no_limit(),
            next_prune: Instant::now(),
            warning_limiter: Default::default(),
            pending: WarningSet {
                small: Vec::with_capacity(5),
                small_full: false,
                bandwidth: Default::default(),
                connections: Default::default(),
                granted_rare_exemptions: 0,
            },
            connections_per_active_p90: 1,
            connections_per_active_p99: 6,
            total_connections: 0,
            total_connections_soft_limit: 400,
            total_connections_hard_limit: 1000,
            last_soft_limit: None,
            ddos_memory: Duration::from_secs(5 * 60),
            compute_pressure: false,
            new_ip_counter: 200,
        }
    }

    fn warn(&mut self, warning: Warning, now: Instant) {
        if !self.pending.small_full && self.pending.small.len() < 5 {
            self.pending.small.push(warning);
        } else {
            self.pending.small_full = true;
            for warning in self.pending.small.drain(..).chain(std::iter::once(warning)) {
                match warning.kind {
                    WarningKind::Bandwidth { amount } => {
                        let entry = self.pending.bandwidth.entry(warning.ip).or_default();
                        *entry = entry.saturating_add(amount);
                    }
                    WarningKind::ConnectionCount {
                        granted_rare_exemption,
                        ..
                    } => {
                        if granted_rare_exemption {
                            self.pending.granted_rare_exemptions =
                                self.pending.granted_rare_exemptions.saturating_add(1);
                        } else {
                            let entry = self.pending.connections.entry(warning.ip).or_default();
                            *entry = entry.saturating_add(1);
                        }
                    }
                }
            }
        }

        if self
            .warning_limiter
            .should_limit_rate_with_now(&WARNING_LIMIT, now)
        {
            return;
        }

        if !self.pending.small_full {
            for Warning { ip, label, kind } in self.pending.small.drain(..) {
                match kind {
                    WarningKind::Bandwidth { amount } => {
                        warn!("{ip} exceeded bw limit with {label} ({amount}B)");
                    }
                    WarningKind::ConnectionCount {
                        connections,
                        active_sessions,
                        limit,
                        granted_rare_exemption,
                    } => {
                        let event = if granted_rare_exemption {
                            "granted special"
                        } else {
                            "hit"
                        };
                        warn!("{ip} {event} conn limit {limit} with {label} ({active_sessions} act, {connections})");
                    }
                }
            }
            return;
        }

        if let Some(sample) = self.pending.bandwidth.keys().next() {
            let mut bytes = self
                .pending
                .bandwidth
                .values()
                .copied()
                .map(|v| v as u64)
                .sum::<u64>();
            let mut unit = "B";
            if bytes >= 1000 {
                bytes /= 1000;
                unit = "KB";
            }
            if bytes >= 1000 {
                bytes /= 1000;
                unit = "MB";
            }
            if bytes >= 1000 {
                bytes /= 1000;
                unit = "GB";
            }
            warn!(
                "{} IP's, e.g. {sample}, hit bw limit with {bytes}{unit}",
                self.pending.bandwidth.len()
            );
            self.pending.bandwidth.clear();
        }
        if let Some(sample) = self.pending.connections.keys().next() {
            let attempts = self.pending.connections.values().copied().sum::<u32>();
            warn!(
                "{} IP's, e.g. {sample}, hit conn limit with {attempts} attempts ({} exempt. granted)",
                self.pending.connections.len(),
                self.pending.granted_rare_exemptions,
            );
            self.pending.connections.clear();
            self.pending.granted_rare_exemptions = 0;
        } else if self.pending.granted_rare_exemptions > 0 {
            warn!("{} exempt. granted", self.pending.granted_rare_exemptions);
            self.pending.granted_rare_exemptions = 0;
        }

        self.pending.small_full = false;
    }

    /// Marks usage as being performed by the ip address.
    /// Returns true if the action should be blocked (rate limited).
    pub(crate) fn should_limit_bandwidth_inner_inner(
        &mut self,
        ip: IpAddr,
        bytes: Units,
        now: Instant,
    ) -> bool {
        let entry = self
            .usage
            .entry(ip)
            .or_insert_with(|| Usage::new(now, &mut self.new_ip_counter));
        let should_limit_rate = entry
            .connection_rate_limit
            .should_limit_rate_with_now_and_usage(&self.connection_rate_limit, now, bytes);

        if should_limit_rate {
            entry.stats.limited(now);
        }

        self.maybe_prune(now);

        should_limit_rate
    }

    /// Clean up old items. Called automatically; it is not necessary to call manually.
    fn maybe_prune(&mut self, now: Instant) {
        if now < self.next_prune {
            return;
        }
        self.next_prune = now + Duration::from_secs(5).max(self.ddos_memory / 2);
        self.prune(now);
    }

    fn prune(&mut self, now: Instant) {
        let forget = now + self.ddos_memory;
        self.usage.retain(|_, usage: &mut Usage| {
            usage.connection_rate_limit.until > forget
                || usage.custom_rate_limit.until > forget
                || usage.stats.active_sessions > 0
                || usage.stats.connections > 0
        })
    }

    /// Returns number of IP addresses being tracked.
    #[allow(unused)]
    pub(crate) fn len(&self) -> usize {
        self.usage.len()
    }

    /// Returns `true` if any IP addresses are being tracked.
    #[allow(unused)]
    pub(crate) fn is_empty(&self) -> bool {
        self.usage.is_empty()
    }
}

#[cfg(test)]
mod test {
    use super::IpLimiter;
    use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
    use std::time::{Duration, Instant};

    #[test]
    pub fn ip_rate_limiter() {
        let ip_one = IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4));
        let ip_two = IpAddr::V6(Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8));
        let mut limiter = IpLimiter::new(10, 3);

        limiter.ddos_memory = Duration::ZERO;

        assert_eq!(limiter.len(), 0);
        assert!(!limiter.should_limit_bandwidth_inner_inner(ip_one, 1, Instant::now()));
        assert_eq!(limiter.len(), 1);
        assert!(!limiter.should_limit_bandwidth_inner_inner(ip_one, 1, Instant::now()));
        assert_eq!(limiter.len(), 1);

        limiter.prune(Instant::now());
        assert_eq!(limiter.len(), 1);

        assert!(!limiter.should_limit_bandwidth_inner_inner(ip_one, 1, Instant::now()));
        assert_eq!(limiter.len(), 1);

        limiter.prune(Instant::now());
        assert_eq!(limiter.len(), 1);

        assert!(limiter.should_limit_bandwidth_inner_inner(ip_one, 1, Instant::now()));
        assert_eq!(limiter.len(), 1);

        std::thread::sleep(Duration::from_millis(250));

        assert!(!limiter.should_limit_bandwidth_inner_inner(ip_two, 1, Instant::now()));
        assert_eq!(limiter.len(), 2);
        assert!(!limiter.should_limit_bandwidth_inner_inner(ip_two, 1, Instant::now()));
        assert_eq!(limiter.len(), 2);

        limiter.prune(Instant::now());
        assert_eq!(limiter.len(), 2);

        std::thread::sleep(Duration::from_millis(100));

        limiter.prune(Instant::now());
        assert_eq!(limiter.len(), 1);

        std::thread::sleep(Duration::from_millis(500));

        limiter.prune(Instant::now());
        assert_eq!(limiter.len(), 0);

        assert!(!limiter.should_limit_bandwidth_inner_inner(ip_one, 1, Instant::now()));
        assert_eq!(limiter.len(), 1);
    }
}