1
//! Per-IP and per-account throttles for authentication attempts.
2
//!
3
//! The daemon keeps all this in process memory — the plan explicitly
4
//! does not require multi-host consistency. Pure state + an injected
5
//! clock so tests can drive the time axis deterministically.
6

            
7
use std::collections::HashMap;
8
use std::net::IpAddr;
9
use std::time::{Duration, Instant};
10

            
11
/// Configuration knobs applied to every `RateLimiter`.
12
#[derive(Debug, Clone, Copy)]
13
pub struct Config {
14
    /// Maximum *password* authentication attempts an IP may make inside
15
    /// one window before further attempts are refused.
16
    pub max_attempts_per_ip: u32,
17
    /// Maximum *publickey* authentication attempts an IP may make
18
    /// inside one window. Looser than the password bucket because a
19
    /// single ssh client legitimately offers several identities per
20
    /// connection.
21
    pub max_pubkey_attempts_per_ip: u32,
22
    /// Width of the sliding window, measured from the first attempt.
23
    pub window: Duration,
24
    /// After this many consecutive *password* failures on one
25
    /// account, the account enters lockout.
26
    pub account_lockout_threshold: u32,
27
    /// How long an account stays locked out once the threshold is
28
    /// tripped.
29
    pub account_lockout_duration: Duration,
30
}
31

            
32
impl Default for Config {
33
    fn default() -> Self {
34
        Self {
35
            max_attempts_per_ip: 6,
36
            max_pubkey_attempts_per_ip: 30,
37
            window: Duration::from_mins(1),
38
            account_lockout_threshold: 3,
39
            account_lockout_duration: Duration::from_mins(15),
40
        }
41
    }
42
}
43

            
44
/// Abstraction over "what time is it now" so tests don't hang on
45
/// real clocks. The daemon uses `InstantClock`; tests use
46
/// `MockClock`.
47
pub trait Clock {
48
    fn now(&self) -> Instant;
49
}
50

            
51
/// Default clock that delegates to `std::time::Instant::now`.
52
#[derive(Debug, Default, Clone, Copy)]
53
pub struct InstantClock;
54

            
55
impl Clock for InstantClock {
56
2
    fn now(&self) -> Instant {
57
2
        Instant::now()
58
2
    }
59
}
60

            
61
/// Outcome of a rate-limit check, returned by [`RateLimiter::check_ip`]
62
/// and [`RateLimiter::check_account`].
63
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
64
pub enum Decision {
65
    /// Proceed with the authentication attempt.
66
    Allow,
67
    /// Refuse the attempt. The attacker should see a generic
68
    /// failure — not "you're rate limited", which leaks information.
69
    Deny,
70
}
71

            
72
#[derive(Debug, Default)]
73
struct IpState {
74
    /// Timestamps of attempts that still fall inside the active
75
    /// window. Prunded on every observation.
76
    attempts: Vec<Instant>,
77
}
78

            
79
#[derive(Debug, Default)]
80
struct AccountState {
81
    consecutive_failures: u32,
82
    locked_until: Option<Instant>,
83
}
84

            
85
/// In-memory rate-limit + lockout tracker. Not `Send`/`Sync` on its
86
/// own — the daemon wraps it in a `tokio::sync::Mutex`.
87
pub struct RateLimiter<C: Clock> {
88
    config: Config,
89
    clock: C,
90
    ip: HashMap<IpAddr, IpState>,
91
    pubkey_ip: HashMap<IpAddr, IpState>,
92
    account: HashMap<String, AccountState>,
93
}
94

            
95
impl<C: Clock> RateLimiter<C> {
96
    #[must_use]
97
10
    pub fn new(config: Config, clock: C) -> Self {
98
10
        Self {
99
10
            config,
100
10
            clock,
101
10
            ip: HashMap::new(),
102
10
            pubkey_ip: HashMap::new(),
103
10
            account: HashMap::new(),
104
10
        }
105
10
    }
106

            
107
    /// Consider a fresh authentication attempt from `ip`. Records
108
    /// the attempt regardless of the return value so the window
109
    /// covers bursts that never reach the server's auth callbacks.
110
19
    pub fn check_ip(&mut self, ip: IpAddr) -> Decision {
111
19
        Self::tick_bucket(
112
19
            &mut self.ip,
113
19
            ip,
114
19
            self.clock.now(),
115
19
            self.config.window,
116
19
            self.config.max_attempts_per_ip,
117
        )
118
19
    }
119

            
120
    /// Same as [`check_ip`] but for the publickey bucket — separate
121
    /// counter so a normal client offering multiple identities does
122
    /// not trip the password limiter.
123
8
    pub fn check_pubkey_ip(&mut self, ip: IpAddr) -> Decision {
124
8
        Self::tick_bucket(
125
8
            &mut self.pubkey_ip,
126
8
            ip,
127
8
            self.clock.now(),
128
8
            self.config.window,
129
8
            self.config.max_pubkey_attempts_per_ip,
130
        )
131
8
    }
132

            
133
27
    fn tick_bucket(
134
27
        bucket: &mut HashMap<IpAddr, IpState>,
135
27
        ip: IpAddr,
136
27
        now: Instant,
137
27
        window: Duration,
138
27
        limit: u32,
139
27
    ) -> Decision {
140
27
        let entry = bucket.entry(ip).or_default();
141
42
        entry.attempts.retain(|t| now.duration_since(*t) < window);
142
27
        let seen = u32::try_from(entry.attempts.len()).unwrap_or(u32::MAX);
143
27
        if seen >= limit {
144
7
            return Decision::Deny;
145
20
        }
146
20
        entry.attempts.push(now);
147
20
        Decision::Allow
148
27
    }
149

            
150
    /// Should a password attempt against `account` be allowed right
151
    /// now? Pair with [`record_password_failure`] / [`reset_account`]
152
    /// once the auth backend has adjudicated the attempt.
153
6
    pub fn check_account(&mut self, account: &str) -> Decision {
154
6
        let now = self.clock.now();
155
6
        let state = self.account.entry(account.to_string()).or_default();
156
6
        if let Some(until) = state.locked_until
157
3
            && now < until
158
        {
159
2
            return Decision::Deny;
160
4
        }
161
4
        if state.locked_until.is_some() {
162
1
            state.locked_until = None;
163
1
            state.consecutive_failures = 0;
164
3
        }
165
4
        Decision::Allow
166
6
    }
167

            
168
6
    pub fn record_password_failure(&mut self, account: &str) {
169
6
        let now = self.clock.now();
170
6
        let state = self.account.entry(account.to_string()).or_default();
171
6
        state.consecutive_failures += 1;
172
6
        if state.consecutive_failures >= self.config.account_lockout_threshold {
173
2
            state.locked_until = Some(now + self.config.account_lockout_duration);
174
4
        }
175
6
    }
176

            
177
    /// Reset the consecutive-failure counter for an account — call
178
    /// this on a successful authentication.
179
1
    pub fn reset_account(&mut self, account: &str) {
180
1
        if let Some(state) = self.account.get_mut(account) {
181
1
            state.consecutive_failures = 0;
182
1
            state.locked_until = None;
183
1
        }
184
1
    }
185
}
186

            
187
#[cfg(test)]
188
mod tests {
189
    use super::*;
190
    use std::cell::Cell;
191
    use std::net::Ipv4Addr;
192

            
193
    /// Clock whose `now()` value the test drives by hand.
194
    pub struct MockClock {
195
        now: Cell<Instant>,
196
    }
197

            
198
    impl MockClock {
199
8
        pub fn new() -> Self {
200
8
            Self {
201
8
                now: Cell::new(Instant::now()),
202
8
            }
203
8
        }
204
2
        pub fn advance(&self, by: Duration) {
205
2
            self.now.set(self.now.get() + by);
206
2
        }
207
    }
208

            
209
    impl Clock for &MockClock {
210
37
        fn now(&self) -> Instant {
211
37
            self.now.get()
212
37
        }
213
    }
214

            
215
8
    fn cfg() -> Config {
216
8
        Config {
217
8
            max_attempts_per_ip: 3,
218
8
            max_pubkey_attempts_per_ip: 5,
219
8
            window: Duration::from_secs(10),
220
8
            account_lockout_threshold: 2,
221
8
            account_lockout_duration: Duration::from_mins(1),
222
8
        }
223
8
    }
224

            
225
20
    fn loopback() -> IpAddr {
226
20
        IpAddr::V4(Ipv4Addr::LOCALHOST)
227
20
    }
228

            
229
    #[test]
230
1
    fn ip_allows_up_to_the_limit_then_denies() {
231
1
        let clock = MockClock::new();
232
1
        let mut rl = RateLimiter::new(cfg(), &clock);
233
1
        assert_eq!(rl.check_ip(loopback()), Decision::Allow);
234
1
        assert_eq!(rl.check_ip(loopback()), Decision::Allow);
235
1
        assert_eq!(rl.check_ip(loopback()), Decision::Allow);
236
1
        assert_eq!(
237
1
            rl.check_ip(loopback()),
238
            Decision::Deny,
239
            "fourth attempt in-window must be denied"
240
        );
241
1
    }
242

            
243
    #[test]
244
1
    fn ip_window_expires_and_resets_allowance() {
245
1
        let clock = MockClock::new();
246
1
        let mut rl = RateLimiter::new(cfg(), &clock);
247
3
        for _ in 0..3 {
248
3
            rl.check_ip(loopback());
249
3
        }
250
1
        assert_eq!(rl.check_ip(loopback()), Decision::Deny);
251
1
        clock.advance(Duration::from_secs(11));
252
1
        assert_eq!(
253
1
            rl.check_ip(loopback()),
254
            Decision::Allow,
255
            "after window expiry the limiter must forget"
256
        );
257
1
    }
258

            
259
    #[test]
260
1
    fn account_locks_after_threshold_failures() {
261
1
        let clock = MockClock::new();
262
1
        let mut rl = RateLimiter::new(cfg(), &clock);
263
1
        assert_eq!(rl.check_account("alice"), Decision::Allow);
264
1
        rl.record_password_failure("alice");
265
1
        assert_eq!(rl.check_account("alice"), Decision::Allow);
266
1
        rl.record_password_failure("alice");
267
1
        assert_eq!(
268
1
            rl.check_account("alice"),
269
            Decision::Deny,
270
            "second failure trips the lockout"
271
        );
272
1
    }
273

            
274
    #[test]
275
1
    fn account_unlocks_after_duration() {
276
1
        let clock = MockClock::new();
277
1
        let mut rl = RateLimiter::new(cfg(), &clock);
278
1
        rl.record_password_failure("bob");
279
1
        rl.record_password_failure("bob");
280
1
        assert_eq!(rl.check_account("bob"), Decision::Deny);
281
1
        clock.advance(Duration::from_secs(61));
282
1
        assert_eq!(rl.check_account("bob"), Decision::Allow);
283
1
    }
284

            
285
    #[test]
286
1
    fn reset_clears_accumulated_failures() {
287
1
        let clock = MockClock::new();
288
1
        let mut rl = RateLimiter::new(cfg(), &clock);
289
1
        rl.record_password_failure("carol");
290
1
        rl.reset_account("carol");
291
1
        rl.record_password_failure("carol");
292
1
        assert_eq!(
293
1
            rl.check_account("carol"),
294
            Decision::Allow,
295
            "counter must be cleared by reset"
296
        );
297
1
    }
298

            
299
    #[test]
300
1
    fn pubkey_bucket_is_independent_of_password_bucket() {
301
1
        let clock = MockClock::new();
302
1
        let mut rl = RateLimiter::new(cfg(), &clock);
303
3
        for _ in 0..3 {
304
3
            rl.check_ip(loopback());
305
3
        }
306
1
        assert_eq!(
307
1
            rl.check_ip(loopback()),
308
            Decision::Deny,
309
            "password bucket exhausted"
310
        );
311
1
        assert_eq!(
312
1
            rl.check_pubkey_ip(loopback()),
313
            Decision::Allow,
314
            "pubkey bucket must not share state with the password bucket"
315
        );
316
1
    }
317

            
318
    #[test]
319
1
    fn pubkey_bucket_enforces_its_own_limit() {
320
1
        let clock = MockClock::new();
321
1
        let mut rl = RateLimiter::new(cfg(), &clock);
322
1
        for _ in 0..5 {
323
5
            assert_eq!(rl.check_pubkey_ip(loopback()), Decision::Allow);
324
        }
325
1
        assert_eq!(
326
1
            rl.check_pubkey_ip(loopback()),
327
            Decision::Deny,
328
            "sixth pubkey attempt in-window must be denied"
329
        );
330
1
    }
331

            
332
    #[test]
333
1
    fn different_ips_are_tracked_independently() {
334
1
        let clock = MockClock::new();
335
1
        let mut rl = RateLimiter::new(cfg(), &clock);
336
1
        let a: IpAddr = "10.0.0.1".parse().unwrap();
337
1
        let b: IpAddr = "10.0.0.2".parse().unwrap();
338
3
        for _ in 0..3 {
339
3
            rl.check_ip(a);
340
3
        }
341
1
        assert_eq!(rl.check_ip(a), Decision::Deny);
342
1
        assert_eq!(
343
1
            rl.check_ip(b),
344
            Decision::Allow,
345
            "a different IP must not inherit another's denial"
346
        );
347
1
    }
348
}