RustProxy/src/rate_limit/tests.rs

mod tests {
    use super::*;
    #[test]
    fn access_modes_are_fail_closed_for_allow_lists() {
        let mut settings = Settings::default();
        settings.country_access.mode = CountryMode::Allow;
        settings.country_access.countries = vec!["pl".into()];
        settings.validate().unwrap();
        assert!(settings.country_allowed(Some("PL")));
        assert!(!settings.country_allowed(Some("DE")));
        assert!(!settings.country_allowed(None));
        settings.country_access.mode = CountryMode::Block;
        assert!(!settings.country_allowed(Some("PL")));
        assert!(settings.country_allowed(None));
    }
    #[test]
    fn country_limit_overrides_base_limit() {
        let mut settings = Settings {
            enabled: true,
            requests: 10,
            burst: 0,
            country_limits: vec![CountryLimit {
                country: "PL".into(),
                requests: 1,
                window_seconds: 60,
                burst: 0,
            }],
            ..Default::default()
        };
        settings.validate().unwrap();
        let limiter = Limiter::new();
        let ip = "192.0.2.1".parse().unwrap();
        assert!(limiter.check(1, "site", ip, Some("PL"), &settings).is_ok());
        assert!(limiter.check(1, "site", ip, Some("PL"), &settings).is_err());
    }

    #[test]
    fn expired_client_buckets_are_reclaimed_without_waiting_for_capacity_pressure() {
        let settings = Settings {
            enabled: true,
            requests: 10,
            window_seconds: 1,
            burst: 0,
            ..Default::default()
        };
        let limiter = Limiter::new();
        let first = "192.0.2.1".parse().unwrap();
        let second = "192.0.2.2".parse().unwrap();
        assert!(limiter.check(7, "site", first, None, &settings).is_ok());
        {
            let mut state = lock(&limiter.state);
            let bucket = state.buckets.get_mut("site").unwrap().get_mut(&first).unwrap();
            bucket.started = Instant::now()
                .checked_sub(Duration::from_secs(2))
                .unwrap();
            state.last_sweep = None;
        }
        assert!(limiter.check(7, "site", second, None, &settings).is_ok());
        let state = lock(&limiter.state);
        assert_eq!(state.bucket_count, 1);
        assert!(!state.buckets.get("site").unwrap().contains_key(&first));
        assert!(state.buckets.get("site").unwrap().contains_key(&second));
    }

    #[test]
    fn maintenance_reclaims_expired_clients_without_new_traffic() {
        let settings = Settings {
            enabled: true,
            requests: 10,
            window_seconds: 1,
            burst: 0,
            ..Default::default()
        };
        let limiter = Limiter::new();
        let ip = "192.0.2.10".parse().unwrap();
        assert!(limiter.check(9, "idle-site", ip, None, &settings).is_ok());
        {
            let mut state = lock(&limiter.state);
            let bucket = state
                .buckets
                .get_mut("idle-site")
                .unwrap()
                .get_mut(&ip)
                .unwrap();
            bucket.started = Instant::now()
                .checked_sub(Duration::from_secs(2))
                .unwrap();
        }
        limiter.maintenance();
        let state = lock(&limiter.state);
        assert_eq!(state.bucket_count, 0);
        assert!(state.buckets.is_empty());
    }

    #[test]
    fn vhost_name_is_stored_once_for_multiple_clients() {
        let settings = Settings {
            enabled: true,
            requests: 10,
            ..Default::default()
        };
        let limiter = Limiter::new();
        for ip in ["192.0.2.1", "192.0.2.2"] {
            assert!(limiter
                .check(1, "shared-vhost", ip.parse().unwrap(), None, &settings)
                .is_ok());
        }
        let state = lock(&limiter.state);
        assert_eq!(state.buckets.len(), 1);
        assert_eq!(state.bucket_count, 2);
    }
    #[test]
    fn user_agent_override_is_isolated_from_other_traffic() {
        let mut settings = Settings {
            enabled: true,
            requests: 3,
            burst: 0,
            user_agent_limits: vec![UserAgentLimit {
                pattern: "*BadBot*".into(), requests: 1, window_seconds: 60, burst: 0,
            }],
            ..Default::default()
        };
        settings.validate().unwrap();
        let limiter = Limiter::new();
        let ip = "198.51.100.12".parse().unwrap();
        assert!(limiter.check_with_user_agent(7, "a", ip, None, Some("badbot/1.0"), &settings).is_ok());
        assert!(limiter.check_with_user_agent(7, "a", ip, None, Some("BADBOT/2"), &settings).is_err());
        for _ in 0..3 {
            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());
        assert!(limiter.check_with_user_agent(7, "b", ip, None, Some("BadBot"), &settings).is_ok());
    }
    #[test]
    fn user_agent_overrides_work_without_base_rate_or_geoip() {
        let mut settings = Settings {
            enabled: false,
            user_agent_limits: vec![UserAgentLimit { pattern: "scanner".into(), requests: 1, window_seconds: 60, burst: 0 }],
            ..Default::default()
        };
        settings.validate().unwrap();
        let limiter = Limiter::new();
        let ip = "192.0.2.100".parse().unwrap();
        assert!(limiter.check_with_user_agent(1,"site",ip,None,Some("ScanNER"),&settings).is_ok());
        assert!(limiter.check_with_user_agent(1,"site",ip,None,Some("scanner"),&settings).is_err());
        assert!(limiter.check_with_user_agent(1,"site",ip,None,None,&settings).is_ok());
    }
    #[test]
    fn user_agent_override_validation_is_bounded_and_unique() {
        let mut settings = Settings { user_agent_limits: vec![
            UserAgentLimit { pattern: "TEST*".into(), ..Default::default() },
            UserAgentLimit { pattern: "test*".into(), ..Default::default() },
        ], ..Default::default() };
        assert!(settings.validate().is_err());
        settings.user_agent_limits[1].pattern = "other".into();
        assert!(settings.validate().is_ok());
    }

}