dogfood
Jump to…
snowcorepublic2 bookmarks
snowinitial commitqoxwzsukwmkx1mo
mainroot / crates / df-ssh / src / limits.rsBlameRaw
1//! Per-source-IP connection limits (spec §9).
2//!
3//! > **SSH.** Key auth only. […] Rate-limit connections per source IP.
4//!
5//! SSH is the cheapest thing on this host to attack: a connection costs the
6//! client one TCP handshake and costs the server a key exchange, a database
7//! lookup, and an Argon2-free but still real authentication path. Without a
8//! limit, one source can occupy every connection slot.
9//!
10//! Two limits, for the two shapes that takes:
11//!
12//! * a **rate** limit, so a source cannot open connections in a tight loop, and
13//! * a **concurrency** limit, so it cannot hold many open at once.
14//!
15//! Enforced at `new_client`, before the key exchange — the earliest point at
16//! which the peer is known, and before any work has been done on its behalf.
17
18use std::collections::HashMap;
19use std::net::IpAddr;
20use std::sync::{Arc, Mutex};
21use std::time::{Duration, Instant};
22
23/// Sustained new connections per second, per source.
24///
25/// A `jj git push` opens one connection. Even a script fetching every
26/// repository in a loop stays well under this.
27const REFILL_PER_SEC: f64 = 0.5;
28
29/// Burst above the sustained rate, for a client that clones several
30/// repositories at once.
31const BURST: f64 = 10.0;
32
33/// Simultaneous connections from one source.
34const MAX_CONCURRENT: u32 = 8;
35
36/// Sources idle longer than this are forgotten.
37const IDLE_EVICT: Duration = Duration::from_secs(600);
38
39/// Distinct sources tracked before new ones are refused. Reached only under a
40/// distributed flood.
41const MAX_TRACKED: usize = 20_000;
42
43#[derive(Clone, Default)]
44pub struct ConnectionLimits(Arc<Mutex<State>>);
45
46#[derive(Default)]
47struct State {
48 sources: HashMap<IpAddr, Source>,
49 last_sweep: Option<Instant>,
50}
51
52struct Source {
53 tokens: f64,
54 open: u32,
55 last: Instant,
56}
57
58impl ConnectionLimits {
59 pub fn new() -> Self {
60 Self::default()
61 }
62
63 /// Admit a connection from `peer`, or refuse it.
64 ///
65 /// The returned guard releases the concurrency slot when the connection's
66 /// handler is dropped — which happens however the connection ends, so a
67 /// client that disappears mid-session does not leak a slot.
68 pub fn admit(&self, peer: Option<std::net::SocketAddr>) -> Option<ConnectionGuard> {
69 // A connection whose peer cannot be determined is not attributable to a
70 // source, so it is admitted without a slot rather than refused: in
71 // practice this does not happen, and failing closed here would break
72 // every connection if it ever did.
73 let ip = peer?.ip();
74 self.admit_at(ip, Instant::now())
75 }
76
77 fn admit_at(&self, ip: IpAddr, now: Instant) -> Option<ConnectionGuard> {
78 let mut state = self.0.lock().ok()?;
79 state.sweep(now);
80
81 if state.sources.len() >= MAX_TRACKED && !state.sources.contains_key(&ip) {
82 return None;
83 }
84
85 let source = state.sources.entry(ip).or_insert(Source {
86 tokens: BURST,
87 open: 0,
88 last: now,
89 });
90
91 let elapsed = now.saturating_duration_since(source.last).as_secs_f64();
92 source.tokens = (source.tokens + elapsed * REFILL_PER_SEC).min(BURST);
93 source.last = now;
94
95 if source.open >= MAX_CONCURRENT || source.tokens < 1.0 {
96 return None;
97 }
98
99 source.tokens -= 1.0;
100 source.open += 1;
101
102 Some(ConnectionGuard { limits: self.clone(), ip })
103 }
104
105 fn release(&self, ip: IpAddr) {
106 if let Ok(mut state) = self.0.lock() {
107 if let Some(s) = state.sources.get_mut(&ip) {
108 s.open = s.open.saturating_sub(1);
109 }
110 }
111 }
112
113 #[cfg(test)]
114 fn tracked(&self) -> usize {
115 self.0.lock().unwrap().sources.len()
116 }
117}
118
119impl State {
120 fn sweep(&mut self, now: Instant) {
121 let due = self
122 .last_sweep
123 .is_none_or(|t| now.saturating_duration_since(t) >= Duration::from_secs(60));
124 if !due {
125 return;
126 }
127 self.last_sweep = Some(now);
128 // Never evict a source with a connection open; that would lose the count
129 // and let the concurrency cap be bypassed by long-lived sessions.
130 self.sources
131 .retain(|_, s| s.open > 0 || now.saturating_duration_since(s.last) < IDLE_EVICT);
132 }
133}
134
135/// Holds one connection slot for a source.
136pub struct ConnectionGuard {
137 limits: ConnectionLimits,
138 ip: IpAddr,
139}
140
141impl Drop for ConnectionGuard {
142 fn drop(&mut self) {
143 self.limits.release(self.ip);
144 }
145}
146
147#[cfg(test)]
148mod tests {
149 use super::*;
150
151 fn ip(s: &str) -> IpAddr {
152 s.parse().unwrap()
153 }
154
155 #[test]
156 fn concurrent_connections_are_capped() {
157 let l = ConnectionLimits::new();
158 let now = Instant::now();
159
160 let mut held = Vec::new();
161 for i in 0..MAX_CONCURRENT {
162 held.push(l.admit_at(ip("10.0.0.1"), now).expect_and_count(i));
163 }
164 assert!(
165 l.admit_at(ip("10.0.0.1"), now).is_none(),
166 "the concurrency cap must refuse the next connection"
167 );
168
169 drop(held.pop());
170 assert!(
171 l.admit_at(ip("10.0.0.1"), now).is_some(),
172 "closing a connection frees its slot"
173 );
174 }
175
176 #[test]
177 fn the_burst_is_spent_and_then_refills() {
178 let l = ConnectionLimits::new();
179 let start = Instant::now();
180
181 // Guards dropped immediately, so only the token bucket applies.
182 for _ in 0..BURST as usize {
183 assert!(l.admit_at(ip("10.0.0.2"), start).is_some());
184 }
185 assert!(l.admit_at(ip("10.0.0.2"), start).is_none(), "burst spent");
186
187 let later = start + Duration::from_secs(4);
188 assert!(
189 l.admit_at(ip("10.0.0.2"), later).is_some(),
190 "{REFILL_PER_SEC}/s should have refilled by then"
191 );
192 }
193
194 #[test]
195 fn sources_do_not_share_a_budget() {
196 let l = ConnectionLimits::new();
197 let now = Instant::now();
198 for _ in 0..BURST as usize {
199 let _ = l.admit_at(ip("10.0.0.3"), now);
200 }
201 assert!(l.admit_at(ip("10.0.0.3"), now).is_none());
202 assert!(
203 l.admit_at(ip("10.0.0.4"), now).is_some(),
204 "one source must not exhaust another's budget"
205 );
206 }
207
208 #[test]
209 fn a_peerless_connection_is_admitted_without_a_slot() {
210 let l = ConnectionLimits::new();
211 assert!(
212 l.admit(None).is_none(),
213 "no peer means no guard — the connection proceeds unmetered"
214 );
215 }
216
217 #[test]
218 fn idle_sources_are_forgotten() {
219 let l = ConnectionLimits::new();
220 let start = Instant::now();
221 drop(l.admit_at(ip("10.0.0.5"), start));
222 assert_eq!(l.tracked(), 1);
223
224 let later = start + IDLE_EVICT + Duration::from_secs(1);
225 let _g = l.admit_at(ip("10.0.0.6"), later);
226 assert_eq!(l.tracked(), 1, "the idle source should have been swept");
227 }
228
229 #[test]
230 fn a_backwards_clock_does_not_mint_tokens() {
231 let l = ConnectionLimits::new();
232 let now = Instant::now();
233 for _ in 0..BURST as usize {
234 let _ = l.admit_at(ip("10.0.0.7"), now);
235 }
236 let earlier = now.checked_sub(Duration::from_secs(600)).unwrap_or(now);
237 assert!(l.admit_at(ip("10.0.0.7"), earlier).is_none());
238 }
239
240 /// Small helper so the concurrency test reads clearly when it fails.
241 trait ExpectAndCount {
242 fn expect_and_count(self, i: u32) -> ConnectionGuard;
243 }
244 impl ExpectAndCount for Option<ConnectionGuard> {
245 fn expect_and_count(self, i: u32) -> ConnectionGuard {
246 self.unwrap_or_else(|| panic!("connection {i} should have been admitted"))
247 }
248 }
249}

249 lines · Rust