0.34.6_repo_cleanup

src/rate_limit/mod.rs

@@ -0,0 +1,419 @@
+use anyhow::{ensure, Result};
+use serde::{Deserialize, Serialize};
+use std::{
+    collections::{HashMap, HashSet},
+    net::IpAddr,
+    sync::{Mutex, MutexGuard},
+    time::{Duration, Instant},
+};
+
+const MAX_BUCKETS: usize = 65_536;
+const MAX_COUNTRY_RULES: usize = 64;
+const MAX_USER_AGENT_RULES: usize = 64;
+const MAX_USER_AGENT_BYTES: usize = 4096;
+
+fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
+    mutex.lock().unwrap_or_else(|e| e.into_inner())
+}
+fn default_requests() -> u32 {
+    120
+}
+fn default_window() -> u64 {
+    60
+}
+fn default_burst() -> u32 {
+    30
+}
+
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum CountryMode {
+    #[default]
+    Off,
+    Block,
+    Allow,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct CountryAccess {
+    pub mode: CountryMode,
+    pub countries: Vec<String>,
+}
+impl Default for CountryAccess {
+    fn default() -> Self {
+        Self {
+            mode: CountryMode::Off,
+            countries: Vec::new(),
+        }
+    }
+}
+impl CountryAccess {
+    pub fn validate(&mut self) -> Result<()> {
+        ensure!(
+            self.countries.len() <= 128,
+            "country access list allows at most 128 countries"
+        );
+        let mut countries = HashSet::new();
+        for country in &mut self.countries {
+            *country = normalize_country(country)?;
+            ensure!(
+                countries.insert(country.clone()),
+                "duplicate country access rule"
+            );
+        }
+        ensure!(
+            self.mode == CountryMode::Off || !self.countries.is_empty(),
+            "country allow/block mode requires at least one country"
+        );
+        Ok(())
+    }
+    pub fn needs_geoip(&self) -> bool {
+        self.mode != CountryMode::Off
+    }
+    pub fn allows(&self, country: Option<&str>) -> bool {
+        match self.mode {
+            CountryMode::Off => true,
+            CountryMode::Block => country.is_none_or(|code| {
+                !self
+                    .countries
+                    .iter()
+                    .any(|item| item.eq_ignore_ascii_case(code))
+            }),
+            CountryMode::Allow => country.is_some_and(|code| {
+                self.countries
+                    .iter()
+                    .any(|item| item.eq_ignore_ascii_case(code))
+            }),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct CountryLimit {
+    pub country: String,
+    pub requests: u32,
+    pub window_seconds: u64,
+    pub burst: u32,
+}
+impl Default for CountryLimit {
+    fn default() -> Self {
+        Self {
+            country: String::new(),
+            requests: default_requests(),
+            window_seconds: default_window(),
+            burst: default_burst(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct UserAgentLimit {
+    pub pattern: String,
+    pub requests: u32,
+    pub window_seconds: u64,
+    pub burst: u32,
+}
+impl Default for UserAgentLimit {
+    fn default() -> Self {
+        Self {
+            pattern: String::new(),
+            requests: default_requests(),
+            window_seconds: default_window(),
+            burst: default_burst(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct Settings {
+    pub enabled: bool,
+    pub requests: u32,
+    pub window_seconds: u64,
+    pub burst: u32,
+    pub country_access: CountryAccess,
+    pub country_limits: Vec<CountryLimit>,
+    pub user_agent_limits: Vec<UserAgentLimit>,
+}
+impl Default for Settings {
+    fn default() -> Self {
+        Self {
+            enabled: false,
+            requests: default_requests(),
+            window_seconds: default_window(),
+            burst: default_burst(),
+            country_access: CountryAccess::default(),
+            country_limits: Vec::new(),
+            user_agent_limits: Vec::new(),
+        }
+    }
+}
+
+fn normalize_country(value: &str) -> Result<String> {
+    let value = value.trim().to_ascii_uppercase();
+    ensure!(
+        value.len() == 2 && value.bytes().all(|b| b.is_ascii_uppercase()),
+        "country must use a two-letter ISO 3166-1 alpha-2 code"
+    );
+    Ok(value)
+}
+fn validate_limit(requests: u32, window_seconds: u64, burst: u32) -> Result<()> {
+    ensure!(
+        (1..=1_000_000).contains(&requests),
+        "rate limit requests must be 1..1000000"
+    );
+    ensure!(
+        (1..=3600).contains(&window_seconds),
+        "rate limit window_seconds must be 1..3600"
+    );
+    ensure!(burst <= 1_000_000, "rate limit burst must be 0..1000000");
+    ensure!(
+        requests.saturating_add(burst) <= 2_000_000,
+        "rate limit requests + burst is too large"
+    );
+    Ok(())
+}
+impl Settings {
+    pub fn validate(&mut self) -> Result<()> {
+        validate_limit(self.requests, self.window_seconds, self.burst)?;
+        self.country_access.validate()?;
+        ensure!(
+            self.country_limits.len() <= MAX_COUNTRY_RULES,
+            "at most 64 country rate-limit overrides"
+        );
+        let mut overrides = HashSet::new();
+        for rule in &mut self.country_limits {
+            rule.country = normalize_country(&rule.country)?;
+            ensure!(
+                overrides.insert(rule.country.clone()),
+                "duplicate country rate-limit override"
+            );
+            validate_limit(rule.requests, rule.window_seconds, rule.burst)?;
+        }
+        ensure!(
+            self.user_agent_limits.len() <= MAX_USER_AGENT_RULES,
+            "at most 64 User-Agent rate-limit overrides"
+        );
+        let mut patterns = HashSet::new();
+        for rule in &mut self.user_agent_limits {
+            rule.pattern = crate::user_agent::validate_pattern(&rule.pattern)?;
+            ensure!(
+                patterns.insert(rule.pattern.clone()),
+                "duplicate User-Agent rate-limit pattern"
+            );
+            validate_limit(rule.requests, rule.window_seconds, rule.burst)?;
+        }
+        Ok(())
+    }
+    pub fn needs_geoip(&self) -> bool {
+        self.country_access.mode != CountryMode::Off || !self.country_limits.is_empty()
+    }
+    pub fn country_allowed(&self, country: Option<&str>) -> bool {
+        self.country_access.allows(country)
+    }
+    fn effective(
+        &self,
+        country: Option<&str>,
+        user_agent: Option<&str>,
+    ) -> Option<(Limit, Option<usize>)> {
+        if let Some(agent) = user_agent.filter(|ua| ua.len() <= MAX_USER_AGENT_BYTES) {
+            let candidate = agent.to_ascii_lowercase();
+            if let Some((index, rule)) = self
+                .user_agent_limits
+                .iter()
+                .enumerate()
+                .find(|(_, rule)| crate::user_agent::wildcard_match(&candidate, &rule.pattern))
+            {
+                return Some((
+                    Limit {
+                        requests: rule.requests,
+                        window_seconds: rule.window_seconds,
+                        burst: rule.burst,
+                    },
+                    Some(index),
+                ));
+            }
+        }
+        if let Some(country) = country {
+            if let Some(rule) = self
+                .country_limits
+                .iter()
+                .find(|rule| rule.country.eq_ignore_ascii_case(country))
+            {
+                return Some((
+                    Limit {
+                        requests: rule.requests,
+                        window_seconds: rule.window_seconds,
+                        burst: rule.burst,
+                    },
+                    None,
+                ));
+            }
+        }
+        self.enabled.then_some((
+            Limit {
+                requests: self.requests,
+                window_seconds: self.window_seconds,
+                burst: self.burst,
+            },
+            None,
+        ))
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+struct Limit {
+    requests: u32,
+    window_seconds: u64,
+    burst: u32,
+}
+#[derive(Debug)]
+struct Bucket {
+    started: Instant,
+    count: u32,
+    window_seconds: u64,
+}
+#[derive(Debug, Default)]
+struct LimiterState {
+    revision: Option<u64>,
+    buckets: HashMap<String, HashMap<IpAddr, Bucket>>,
+    bucket_count: usize,
+    last_sweep: Option<Instant>,
+}
+impl LimiterState {
+    fn clear(&mut self) {
+        self.buckets.clear();
+        self.bucket_count = 0;
+        self.last_sweep = None;
+    }
+    fn sweep_expired(&mut self, now: Instant, force: bool) {
+        if !force
+            && self
+                .last_sweep
+                .is_some_and(|last| now.duration_since(last) < Duration::from_secs(5))
+        {
+            return;
+        }
+        self.last_sweep = Some(now);
+        let mut removed = 0usize;
+        self.buckets.retain(|_, clients| {
+            let before = clients.len();
+            clients.retain(|_, bucket| {
+                now.duration_since(bucket.started) < Duration::from_secs(bucket.window_seconds)
+            });
+            let removed_here = before.saturating_sub(clients.len());
+            removed = removed.saturating_add(removed_here);
+            if removed_here > 0
+                && !clients.is_empty()
+                && clients.capacity() > clients.len().saturating_mul(2).max(64)
+            {
+                clients.shrink_to_fit();
+            }
+            !clients.is_empty()
+        });
+        self.bucket_count = self.bucket_count.saturating_sub(removed);
+        if removed > 0 && self.buckets.capacity() > self.buckets.len().saturating_mul(2).max(16) {
+            self.buckets.shrink_to_fit();
+        }
+    }
+}
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct Limited {
+    pub retry_after_seconds: u64,
+}
+#[derive(Debug, Default)]
+pub struct Limiter {
+    state: Mutex<LimiterState>,
+}
+impl Limiter {
+    pub fn new() -> Self {
+        Self::default()
+    }
+    pub fn reset(&self) {
+        lock(&self.state).clear();
+    }
+    pub fn maintenance(&self) {
+        lock(&self.state).sweep_expired(Instant::now(), true);
+    }
+    pub fn check(
+        &self,
+        revision: u64,
+        vhost: &str,
+        client: IpAddr,
+        country: Option<&str>,
+        settings: &Settings,
+    ) -> std::result::Result<(), Limited> {
+        self.check_with_user_agent(revision, vhost, client, country, None, settings)
+    }
+    pub fn check_with_user_agent(
+        &self,
+        revision: u64,
+        vhost: &str,
+        client: IpAddr,
+        country: Option<&str>,
+        user_agent: Option<&str>,
+        settings: &Settings,
+    ) -> std::result::Result<(), Limited> {
+        let Some((limit, ua_index)) = settings.effective(country, user_agent) else {
+            return Ok(());
+        };
+        let bucket_key = match ua_index {
+            Some(index) => format!("{vhost}\u{1f}ua:{index}"),
+            None => vhost.to_owned(),
+        };
+        let now = Instant::now();
+        let mut state = lock(&self.state);
+        if state.revision != Some(revision) {
+            state.clear();
+            state.revision = Some(revision);
+        }
+        state.sweep_expired(now, false);
+        let new_client = !state
+            .buckets
+            .get(&bucket_key)
+            .is_some_and(|clients| clients.contains_key(&client));
+        if new_client && state.bucket_count >= MAX_BUCKETS {
+            state.sweep_expired(now, true);
+            if state.bucket_count >= MAX_BUCKETS {
+                return Err(Limited {
+                    retry_after_seconds: 1,
+                });
+            }
+        }
+        if new_client {
+            state.bucket_count = state.bucket_count.saturating_add(1);
+        }
+        let bucket = state
+            .buckets
+            .entry(bucket_key)
+            .or_default()
+            .entry(client)
+            .or_insert(Bucket {
+                started: now,
+                count: 0,
+                window_seconds: limit.window_seconds,
+            });
+        if bucket.window_seconds != limit.window_seconds
+            || now.duration_since(bucket.started) >= Duration::from_secs(limit.window_seconds)
+        {
+            bucket.started = now;
+            bucket.count = 0;
+            bucket.window_seconds = limit.window_seconds;
+        }
+        let capacity = limit.requests.saturating_add(limit.burst);
+        if bucket.count >= capacity {
+            let elapsed = now.duration_since(bucket.started).as_secs();
+            return Err(Limited {
+                retry_after_seconds: limit.window_seconds.saturating_sub(elapsed).max(1),
+            });
+        }
+        bucket.count = bucket.count.saturating_add(1);
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+include!("tests.rs");

src/rate_limit/tests.rs

@@ -129,7 +129,6 @@ mod tests {
             assert!(limiter.check_with_user_agent(7, "a", ip, None, Some("Browser"), &settings).is_ok());
         }
         assert!(limiter.check_with_user_agent(7, "a", ip, None, Some("Browser"), &settings).is_err());
-        // The same IP and UA on a different vhost gets a separate quota.
         assert!(limiter.check_with_user_agent(7, "b", ip, None, Some("BadBot"), &settings).is_ok());
     }
     #[test]