Skip to main content

pwr_server/
auth.rs

1//! Authentication rate limiting for pwr-server.
2//!
3//! Tracks failed authentication attempts per IP address and rejects
4//! connections that exceed the configured threshold within a time window.
5//! Uses a simple in-memory map with periodic cleanup.
6
7use std::collections::HashMap;
8use std::net::IpAddr;
9use std::time::Instant;
10
11/// Maximum failed authentication attempts per IP before temporary ban.
12const MAX_ATTEMPTS: usize = 5;
13
14/// Time window for counting failed attempts (seconds).
15const WINDOW_SECS: u64 = 60;
16
17/// Ban duration after exceeding the attempt limit (seconds).
18const BAN_DURATION_SECS: u64 = 300;
19
20/// Entry in the rate limiting table.
21#[derive(Debug, Clone)]
22struct RateLimitEntry {
23    /// Number of failed attempts in the current window.
24    attempts: usize,
25    /// When the first attempt in this window occurred.
26    window_start: Instant,
27    /// If banned, the time when the ban expires.
28    banned_until: Option<Instant>,
29}
30
31/// Tracks authentication attempts per source IP address.
32#[derive(Debug, Default)]
33pub struct RateLimiter {
34    entries: HashMap<IpAddr, RateLimitEntry>,
35}
36
37impl RateLimiter {
38    /// Create a new rate limiter.
39    pub fn new() -> Self {
40        Self {
41            entries: HashMap::new(),
42        }
43    }
44
45    /// Check whether an authentication attempt from this IP should be allowed.
46    ///
47    /// Returns `true` if the attempt is permitted, `false` if the
48    /// IP is currently banned or has exceeded the attempt limit.
49    pub fn check_attempt(&mut self, ip: IpAddr) -> bool {
50        let now = Instant::now();
51        self.cleanup_expired(now);
52
53        let entry = self.entries.entry(ip).or_insert(RateLimitEntry {
54            attempts: 0,
55            window_start: now,
56            banned_until: None,
57        });
58
59        // Check if banned
60        if let Some(banned_until) = entry.banned_until {
61            if now < banned_until {
62                log::warn!("Rejected auth attempt from banned IP: {}", ip);
63                return false;
64            }
65            // Ban expired, reset
66            entry.banned_until = None;
67            entry.attempts = 0;
68            entry.window_start = now;
69        }
70
71        // Reset window if expired
72        if now.duration_since(entry.window_start).as_secs() > WINDOW_SECS {
73            entry.attempts = 0;
74            entry.window_start = now;
75        }
76
77        entry.attempts += 1;
78
79        // Check if this attempt exceeds the limit
80        if entry.attempts > MAX_ATTEMPTS {
81            entry.banned_until = Some(now + std::time::Duration::from_secs(BAN_DURATION_SECS));
82            log::warn!(
83                "IP {} banned for {} seconds after {} failed attempts",
84                ip,
85                BAN_DURATION_SECS,
86                entry.attempts
87            );
88            return false;
89        }
90
91        // Reset attempts on successful auth (called externally)
92        true
93    }
94
95    /// Reset the counter for an IP after a successful authentication.
96    pub fn record_success(&mut self, ip: IpAddr) {
97        if let Some(entry) = self.entries.get_mut(&ip) {
98            entry.attempts = 0;
99            entry.banned_until = None;
100        }
101    }
102
103    /// Remove expired entries to prevent unbounded memory growth.
104    fn cleanup_expired(&mut self, now: Instant) {
105        self.entries.retain(|_ip, entry| {
106            let expired = now.duration_since(entry.window_start).as_secs() > WINDOW_SECS + BAN_DURATION_SECS;
107            if let Some(banned_until) = entry.banned_until {
108                !expired || now < banned_until
109            } else {
110                !expired
111            }
112        });
113    }
114
115    /// Return the number of tracked IPs (for monitoring).
116    pub fn tracked_count(&self) -> usize {
117        self.entries.len()
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    fn test_ip() -> IpAddr {
126        "192.168.1.100".parse().unwrap()
127    }
128
129    #[test]
130    fn test_allows_first_five_attempts() {
131        let mut limiter = RateLimiter::new();
132        let ip = test_ip();
133
134        for _ in 0..5 {
135            assert!(limiter.check_attempt(ip));
136        }
137    }
138
139    #[test]
140    fn test_blocks_sixth_attempt() {
141        let mut limiter = RateLimiter::new();
142        let ip = test_ip();
143
144        for _ in 0..5 {
145            assert!(limiter.check_attempt(ip));
146        }
147        // Sixth attempt should be blocked
148        assert!(!limiter.check_attempt(ip));
149    }
150
151    #[test]
152    fn test_success_resets_counter() {
153        let mut limiter = RateLimiter::new();
154        let ip = test_ip();
155
156        // Three failed
157        for _ in 0..3 {
158            limiter.check_attempt(ip);
159        }
160        // Success resets
161        limiter.record_success(ip);
162
163        // Should have a fresh window of 5
164        for _ in 0..5 {
165            assert!(limiter.check_attempt(ip));
166        }
167    }
168
169    #[test]
170    fn test_different_ips_independent() {
171        let mut limiter = RateLimiter::new();
172        let ip1: IpAddr = "10.0.0.1".parse().unwrap();
173        let ip2: IpAddr = "10.0.0.2".parse().unwrap();
174
175        // Exhaust ip1
176        for _ in 0..5 {
177            limiter.check_attempt(ip1);
178        }
179        assert!(!limiter.check_attempt(ip1));
180
181        // ip2 should still be fine
182        assert!(limiter.check_attempt(ip2));
183    }
184}