dogfood
Jump to…
snowcorepublic2 bookmarks
snowinitial commitqoxwzsukwmkx1mo
mainroot / crates / df-ssh / src / limits.rsHide BlameRaw
Matt W1//! Per-source-IP connection limits (spec §9).
Matt W2//!
Matt W3//! > **SSH.** Key auth only. […] Rate-limit connections per source IP.
Matt W4//!
Matt W5//! SSH is the cheapest thing on this host to attack: a connection costs the
Matt W6//! client one TCP handshake and costs the server a key exchange, a database
Matt W7//! lookup, and an Argon2-free but still real authentication path. Without a
Matt W8//! limit, one source can occupy every connection slot.
Matt W9//!
Matt W10//! Two limits, for the two shapes that takes:
Matt W11//!
Matt W12//! * a **rate** limit, so a source cannot open connections in a tight loop, and
Matt W13//! * a **concurrency** limit, so it cannot hold many open at once.
Matt W14//!
Matt W15//! Enforced at `new_client`, before the key exchange — the earliest point at
Matt W16//! which the peer is known, and before any work has been done on its behalf.
Matt W17
Matt W18use std::collections::HashMap;
Matt W19use std::net::IpAddr;
Matt W20use std::sync::{Arc, Mutex};
Matt W21use std::time::{Duration, Instant};
Matt W22
Matt W23/// Sustained new connections per second, per source.
Matt W24///
Matt W25/// A `jj git push` opens one connection. Even a script fetching every
Matt W26/// repository in a loop stays well under this.
Matt W27const REFILL_PER_SEC: f64 = 0.5;
Matt W28
Matt W29/// Burst above the sustained rate, for a client that clones several
Matt W30/// repositories at once.
Matt W31const BURST: f64 = 10.0;
Matt W32
Matt W33/// Simultaneous connections from one source.
Matt W34const MAX_CONCURRENT: u32 = 8;
Matt W35
Matt W36/// Sources idle longer than this are forgotten.
Matt W37const IDLE_EVICT: Duration = Duration::from_secs(600);
Matt W38
Matt W39/// Distinct sources tracked before new ones are refused. Reached only under a
Matt W40/// distributed flood.
Matt W41const MAX_TRACKED: usize = 20_000;
Matt W42
Matt W43#[derive(Clone, Default)]
Matt W44pub struct ConnectionLimits(Arc<Mutex<State>>);
Matt W45
Matt W46#[derive(Default)]
Matt W47struct State {
Matt W48 sources: HashMap<IpAddr, Source>,
Matt W49 last_sweep: Option<Instant>,
Matt W50}
Matt W51
Matt W52struct Source {
Matt W53 tokens: f64,
Matt W54 open: u32,
Matt W55 last: Instant,
Matt W56}
Matt W57
Matt W58impl ConnectionLimits {
Matt W59 pub fn new() -> Self {
Matt W60 Self::default()
Matt W61 }
Matt W62
Matt W63 /// Admit a connection from `peer`, or refuse it.
Matt W64 ///
Matt W65 /// The returned guard releases the concurrency slot when the connection's
Matt W66 /// handler is dropped — which happens however the connection ends, so a
Matt W67 /// client that disappears mid-session does not leak a slot.
Matt W68 pub fn admit(&self, peer: Option<std::net::SocketAddr>) -> Option<ConnectionGuard> {
Matt W69 // A connection whose peer cannot be determined is not attributable to a
Matt W70 // source, so it is admitted without a slot rather than refused: in
Matt W71 // practice this does not happen, and failing closed here would break
Matt W72 // every connection if it ever did.
Matt W73 let ip = peer?.ip();
Matt W74 self.admit_at(ip, Instant::now())
Matt W75 }
Matt W76
Matt W77 fn admit_at(&self, ip: IpAddr, now: Instant) -> Option<ConnectionGuard> {
Matt W78 let mut state = self.0.lock().ok()?;
Matt W79 state.sweep(now);
Matt W80
Matt W81 if state.sources.len() >= MAX_TRACKED && !state.sources.contains_key(&ip) {
Matt W82 return None;
Matt W83 }
Matt W84
Matt W85 let source = state.sources.entry(ip).or_insert(Source {
Matt W86 tokens: BURST,
Matt W87 open: 0,
Matt W88 last: now,
Matt W89 });
Matt W90
Matt W91 let elapsed = now.saturating_duration_since(source.last).as_secs_f64();
Matt W92 source.tokens = (source.tokens + elapsed * REFILL_PER_SEC).min(BURST);
Matt W93 source.last = now;
Matt W94
Matt W95 if source.open >= MAX_CONCURRENT || source.tokens < 1.0 {
Matt W96 return None;
Matt W97 }
Matt W98
Matt W99 source.tokens -= 1.0;
Matt W100 source.open += 1;
Matt W101
Matt W102 Some(ConnectionGuard { limits: self.clone(), ip })
Matt W103 }
Matt W104
Matt W105 fn release(&self, ip: IpAddr) {
Matt W106 if let Ok(mut state) = self.0.lock() {
Matt W107 if let Some(s) = state.sources.get_mut(&ip) {
Matt W108 s.open = s.open.saturating_sub(1);
Matt W109 }
Matt W110 }
Matt W111 }
Matt W112
Matt W113 #[cfg(test)]
Matt W114 fn tracked(&self) -> usize {
Matt W115 self.0.lock().unwrap().sources.len()
Matt W116 }
Matt W117}
Matt W118
Matt W119impl State {
Matt W120 fn sweep(&mut self, now: Instant) {
Matt W121 let due = self
Matt W122 .last_sweep
Matt W123 .is_none_or(|t| now.saturating_duration_since(t) >= Duration::from_secs(60));
Matt W124 if !due {
Matt W125 return;
Matt W126 }
Matt W127 self.last_sweep = Some(now);
Matt W128 // Never evict a source with a connection open; that would lose the count
Matt W129 // and let the concurrency cap be bypassed by long-lived sessions.
Matt W130 self.sources
Matt W131 .retain(|_, s| s.open > 0 || now.saturating_duration_since(s.last) < IDLE_EVICT);
Matt W132 }
Matt W133}
Matt W134
Matt W135/// Holds one connection slot for a source.
Matt W136pub struct ConnectionGuard {
Matt W137 limits: ConnectionLimits,
Matt W138 ip: IpAddr,
Matt W139}
Matt W140
Matt W141impl Drop for ConnectionGuard {
Matt W142 fn drop(&mut self) {
Matt W143 self.limits.release(self.ip);
Matt W144 }
Matt W145}
Matt W146
Matt W147#[cfg(test)]
Matt W148mod tests {
Matt W149 use super::*;
Matt W150
Matt W151 fn ip(s: &str) -> IpAddr {
Matt W152 s.parse().unwrap()
Matt W153 }
Matt W154
Matt W155 #[test]
Matt W156 fn concurrent_connections_are_capped() {
Matt W157 let l = ConnectionLimits::new();
Matt W158 let now = Instant::now();
Matt W159
Matt W160 let mut held = Vec::new();
Matt W161 for i in 0..MAX_CONCURRENT {
Matt W162 held.push(l.admit_at(ip("10.0.0.1"), now).expect_and_count(i));
Matt W163 }
Matt W164 assert!(
Matt W165 l.admit_at(ip("10.0.0.1"), now).is_none(),
Matt W166 "the concurrency cap must refuse the next connection"
Matt W167 );
Matt W168
Matt W169 drop(held.pop());
Matt W170 assert!(
Matt W171 l.admit_at(ip("10.0.0.1"), now).is_some(),
Matt W172 "closing a connection frees its slot"
Matt W173 );
Matt W174 }
Matt W175
Matt W176 #[test]
Matt W177 fn the_burst_is_spent_and_then_refills() {
Matt W178 let l = ConnectionLimits::new();
Matt W179 let start = Instant::now();
Matt W180
Matt W181 // Guards dropped immediately, so only the token bucket applies.
Matt W182 for _ in 0..BURST as usize {
Matt W183 assert!(l.admit_at(ip("10.0.0.2"), start).is_some());
Matt W184 }
Matt W185 assert!(l.admit_at(ip("10.0.0.2"), start).is_none(), "burst spent");
Matt W186
Matt W187 let later = start + Duration::from_secs(4);
Matt W188 assert!(
Matt W189 l.admit_at(ip("10.0.0.2"), later).is_some(),
Matt W190 "{REFILL_PER_SEC}/s should have refilled by then"
Matt W191 );
Matt W192 }
Matt W193
Matt W194 #[test]
Matt W195 fn sources_do_not_share_a_budget() {
Matt W196 let l = ConnectionLimits::new();
Matt W197 let now = Instant::now();
Matt W198 for _ in 0..BURST as usize {
Matt W199 let _ = l.admit_at(ip("10.0.0.3"), now);
Matt W200 }
Matt W201 assert!(l.admit_at(ip("10.0.0.3"), now).is_none());
Matt W202 assert!(
Matt W203 l.admit_at(ip("10.0.0.4"), now).is_some(),
Matt W204 "one source must not exhaust another's budget"
Matt W205 );
Matt W206 }
Matt W207
Matt W208 #[test]
Matt W209 fn a_peerless_connection_is_admitted_without_a_slot() {
Matt W210 let l = ConnectionLimits::new();
Matt W211 assert!(
Matt W212 l.admit(None).is_none(),
Matt W213 "no peer means no guard — the connection proceeds unmetered"
Matt W214 );
Matt W215 }
Matt W216
Matt W217 #[test]
Matt W218 fn idle_sources_are_forgotten() {
Matt W219 let l = ConnectionLimits::new();
Matt W220 let start = Instant::now();
Matt W221 drop(l.admit_at(ip("10.0.0.5"), start));
Matt W222 assert_eq!(l.tracked(), 1);
Matt W223
Matt W224 let later = start + IDLE_EVICT + Duration::from_secs(1);
Matt W225 let _g = l.admit_at(ip("10.0.0.6"), later);
Matt W226 assert_eq!(l.tracked(), 1, "the idle source should have been swept");
Matt W227 }
Matt W228
Matt W229 #[test]
Matt W230 fn a_backwards_clock_does_not_mint_tokens() {
Matt W231 let l = ConnectionLimits::new();
Matt W232 let now = Instant::now();
Matt W233 for _ in 0..BURST as usize {
Matt W234 let _ = l.admit_at(ip("10.0.0.7"), now);
Matt W235 }
Matt W236 let earlier = now.checked_sub(Duration::from_secs(600)).unwrap_or(now);
Matt W237 assert!(l.admit_at(ip("10.0.0.7"), earlier).is_none());
Matt W238 }
Matt W239
Matt W240 /// Small helper so the concurrency test reads clearly when it fails.
Matt W241 trait ExpectAndCount {
Matt W242 fn expect_and_count(self, i: u32) -> ConnectionGuard;
Matt W243 }
Matt W244 impl ExpectAndCount for Option<ConnectionGuard> {
Matt W245 fn expect_and_count(self, i: u32) -> ConnectionGuard {
Matt W246 self.unwrap_or_else(|| panic!("connection {i} should have been admitted"))
Matt W247 }
Matt W248 }
Matt W249}

249 lines · Rust