IdeA/crates/infrastructure/src/pair_attempt_limiter.rs

//! In-memory pairing-attempt limiter.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use application::{PairAttemptDecision, PairAttemptLimiter, RateLimitKey};
use async_trait::async_trait;
use domain::ports::{Clock, StoreError};

const WINDOW_MS: i64 = 60_000;
const PER_ORIGIN_LIMIT: usize = 5;
const GLOBAL_LIMIT: usize = 30;

/// Process-local limiter for failed pairing attempts.
pub struct InMemoryPairAttemptLimiter {
    clock: Arc<dyn Clock>,
    inner: Mutex<LimiterState>,
}

#[derive(Default)]
struct LimiterState {
    by_origin: HashMap<String, Vec<i64>>,
    global: Vec<i64>,
}

impl InMemoryPairAttemptLimiter {
    /// Creates a limiter using the supplied clock.
    #[must_use]
    pub fn new(clock: Arc<dyn Clock>) -> Self {
        Self {
            clock,
            inner: Mutex::new(LimiterState::default()),
        }
    }
}

#[async_trait]
impl PairAttemptLimiter for InMemoryPairAttemptLimiter {
    async fn check(&self, key: RateLimitKey) -> Result<PairAttemptDecision, StoreError> {
        let now = self.clock.now_millis().max(0);
        let cutoff = now.saturating_sub(WINDOW_MS);
        let mut inner = self.inner.lock().expect("pair limiter mutex poisoned");
        prune(&mut inner.global, cutoff);
        let origin = origin_key(&key);
        let origin_bucket = inner.by_origin.entry(origin).or_default();
        prune(origin_bucket, cutoff);
        if origin_bucket.len() >= PER_ORIGIN_LIMIT || inner.global.len() >= GLOBAL_LIMIT {
            return Ok(PairAttemptDecision::RateLimited);
        }
        Ok(PairAttemptDecision::Allowed)
    }

    async fn record_failure(&self, key: RateLimitKey) -> Result<(), StoreError> {
        let now = self.clock.now_millis().max(0);
        let cutoff = now.saturating_sub(WINDOW_MS);
        let mut inner = self.inner.lock().expect("pair limiter mutex poisoned");
        prune(&mut inner.global, cutoff);
        inner.global.push(now);
        let origin = origin_key(&key);
        let origin_bucket = inner.by_origin.entry(origin).or_default();
        prune(origin_bucket, cutoff);
        origin_bucket.push(now);
        Ok(())
    }
}

fn origin_key(key: &RateLimitKey) -> String {
    format!("{}:{}", key.route, key.origin)
}

fn prune(bucket: &mut Vec<i64>, cutoff: i64) {
    bucket.retain(|timestamp| *timestamp > cutoff);
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicI64, Ordering};

    struct FakeClock(AtomicI64);

    impl FakeClock {
        fn new(now: i64) -> Self {
            Self(AtomicI64::new(now))
        }

        fn set(&self, now: i64) {
            self.0.store(now, Ordering::SeqCst);
        }
    }

    impl Clock for FakeClock {
        fn now_millis(&self) -> i64 {
            self.0.load(Ordering::SeqCst)
        }
    }

    fn key(origin: &str) -> RateLimitKey {
        RateLimitKey {
            origin: origin.to_owned(),
            route: "/api/pair".to_owned(),
        }
    }

    #[tokio::test]
    async fn limits_five_failures_per_origin_per_minute_with_injected_clock() {
        let clock = Arc::new(FakeClock::new(1_000));
        let limiter = InMemoryPairAttemptLimiter::new(Arc::clone(&clock) as Arc<dyn Clock>);

        for _ in 0..PER_ORIGIN_LIMIT {
            assert_eq!(
                limiter.check(key("1.2.3.4")).await.unwrap(),
                PairAttemptDecision::Allowed
            );
            limiter.record_failure(key("1.2.3.4")).await.unwrap();
        }

        assert_eq!(
            limiter.check(key("1.2.3.4")).await.unwrap(),
            PairAttemptDecision::RateLimited
        );
        assert_eq!(
            limiter.check(key("1.2.3.5")).await.unwrap(),
            PairAttemptDecision::Allowed
        );

        clock.set(61_001);
        assert_eq!(
            limiter.check(key("1.2.3.4")).await.unwrap(),
            PairAttemptDecision::Allowed
        );
    }

    #[tokio::test]
    async fn limits_thirty_failures_globally_per_minute() {
        let clock = Arc::new(FakeClock::new(1_000));
        let limiter = InMemoryPairAttemptLimiter::new(clock as Arc<dyn Clock>);

        for n in 0..GLOBAL_LIMIT {
            let key = key(&format!("10.0.0.{n}"));
            assert_eq!(
                limiter.check(key.clone()).await.unwrap(),
                PairAttemptDecision::Allowed
            );
            limiter.record_failure(key).await.unwrap();
        }

        assert_eq!(
            limiter.check(key("10.0.0.99")).await.unwrap(),
            PairAttemptDecision::RateLimited
        );
    }
}