+1

Cargo.lock

View file Reviewed

··· 4749 4749 "rustix", 4750 4750 "tempfile", 4751 4751 "tokio", 4752 4752 + "tracing", 4752 4753 ] 4753 4754 4754 4755 [[package]]

+1

knot2/crates/knot-resource/Cargo.toml

View file Reviewed

··· 9 9 knot-types = { workspace = true } 10 10 rustix = { workspace = true, features = ["fs"] } 11 11 tokio = { workspace = true } 12 12 + tracing = { workspace = true } 12 13 13 14 [dev-dependencies] 14 15 tokio = { workspace = true, features = ["rt", "macros"] }

+129 -41

knot2/crates/knot-resource/src/admission.rs

View file Reviewed

··· 119 119 } 120 120 } 121 121 122 122 + const CONCENTRATED_REFUSALS: u32 = 1_024; 123 123 + 124 124 + #[derive(Default)] 125 125 + struct RefusalMajority { 126 126 + peer: Option<IpAddr>, 127 127 + votes: u32, 128 128 + reported: bool, 129 129 + } 130 130 + 131 131 + impl RefusalMajority { 132 132 + fn observe(&mut self, peer: IpAddr) -> Option<IpAddr> { 133 133 + match (self.peer == Some(peer), self.votes) { 134 134 + (true, _) => self.votes = self.votes.saturating_add(1), 135 135 + (false, 0) => { 136 136 + self.peer = Some(peer); 137 137 + self.votes = 1; 138 138 + } 139 139 + (false, _) => self.votes -= 1, 140 140 + } 141 141 + let crossed = self.votes >= CONCENTRATED_REFUSALS && !self.reported; 142 142 + self.reported |= crossed; 143 143 + crossed.then_some(peer) 144 144 + } 145 145 + } 146 146 + 122 147 struct Inner { 123 148 peers: HashMap<Option<IpAddr>, PeerState>, 124 149 global_inflight: usize, 125 150 last_sweep: UnixMicros, 151 151 + refusal_majority: RefusalMajority, 152 152 + } 153 153 + 154 154 + impl Inner { 155 155 + fn has_room_for( 156 156 + &mut self, 157 157 + peer: Option<IpAddr>, 158 158 + rate: Option<RateLimit>, 159 159 + now: UnixMicros, 160 160 + ) -> bool { 161 161 + if rate.is_none() || self.peers.contains_key(&peer) || self.peers.len() < MAX_TRACKED_PEERS 162 162 + { 163 163 + return true; 164 164 + } 165 165 + if now.get().saturating_sub(self.last_sweep.get()) >= SWEEP_INTERVAL_MICROS { 166 166 + self.last_sweep = now; 167 167 + self.peers 168 168 + .retain(|_, state| state.worth_tracking(rate, now)); 169 169 + } 170 170 + self.peers.len() < MAX_TRACKED_PEERS 171 171 + } 126 172 } 127 173 128 174 pub struct PreAuthLimiter { ··· 150 196 peers: HashMap::new(), 151 197 global_inflight: 0, 152 198 last_sweep: UnixMicros::new(0), 199 199 + refusal_majority: RefusalMajority::default(), 153 200 }), 154 201 } 155 202 } ··· 172 219 .global_inflight 173 220 .is_some_and(|limit| inner.global_inflight >= limit.get()); 174 221 175 175 - // Only a rate budget outlives the work it admitted, so only a rate 176 176 - // budget can pile up entries for peers that have gone away. Without one 177 177 - // the last operation removes the entry and the table is already bounded 178 178 - // by live concurrency, so capping it would shed callers over a table 179 179 - // that can't grow. 180 180 - if rate.is_some() 181 181 - && !inner.peers.contains_key(&peer) 182 182 - && inner.peers.len() >= MAX_TRACKED_PEERS 183 183 - { 184 184 - let sweep_due = 185 185 - now.get().saturating_sub(inner.last_sweep.get()) >= SWEEP_INTERVAL_MICROS; 186 186 - if sweep_due { 187 187 - inner.last_sweep = now; 188 188 - inner 189 189 - .peers 190 190 - .retain(|_, state| state.worth_tracking(rate, now)); 191 191 - } 192 192 - if inner.peers.len() >= MAX_TRACKED_PEERS { 193 193 - return Err(Refusal::Saturated); 194 194 - } 195 195 - } 196 196 - 197 222 let per_peer = self.config.per_peer_inflight; 198 198 - let decision = { 199 199 - let state = inner.peers.entry(peer).or_insert_with(|| PeerState { 200 200 - bucket: rate.map(|rate| Bucket::new(rate, now)), 201 201 - inflight: 0, 202 202 - }); 203 203 - let ready = match (rate, state.bucket.as_mut()) { 204 204 - (Some(rate), Some(bucket)) => bucket.replenish(rate, now), 205 205 - _ => true, 206 206 - }; 207 207 - let over_peer = per_peer.is_some_and(|limit| state.inflight >= limit.get()); 208 208 - match (ready, over_peer || over_global) { 209 209 - (false, _) => Err(Refusal::RateLimited), 210 210 - (_, true) => Err(Refusal::Saturated), 211 211 - (true, false) => { 212 212 - if let Some(bucket) = state.bucket.as_mut() { 213 213 - bucket.tokens -= 1; 223 223 + let decision = match inner.has_room_for(peer, rate, now) { 224 224 + false => Err(Refusal::Saturated), 225 225 + true => { 226 226 + let state = inner.peers.entry(peer).or_insert_with(|| PeerState { 227 227 + bucket: rate.map(|rate| Bucket::new(rate, now)), 228 228 + inflight: 0, 229 229 + }); 230 230 + let ready = match (rate, state.bucket.as_mut()) { 231 231 + (Some(rate), Some(bucket)) => bucket.replenish(rate, now), 232 232 + _ => true, 233 233 + }; 234 234 + let over_peer = per_peer.is_some_and(|limit| state.inflight >= limit.get()); 235 235 + match (ready, over_peer || over_global) { 236 236 + (false, _) => Err(Refusal::RateLimited), 237 237 + (_, true) => Err(Refusal::Saturated), 238 238 + (true, false) => { 239 239 + if let Some(bucket) = state.bucket.as_mut() { 240 240 + bucket.tokens -= 1; 241 241 + } 242 242 + state.inflight += 1; 243 243 + Ok(()) 214 244 } 215 215 - state.inflight += 1; 216 216 - Ok(()) 217 245 } 218 246 } 219 247 }; ··· 221 249 Err(refusal) => { 222 250 if inner.peers.get(&peer).is_some_and(PeerState::forgettable) { 223 251 inner.peers.remove(&peer); 252 252 + } 253 253 + let concentrated = peer.and_then(|peer| inner.refusal_majority.observe(peer)); 254 254 + drop(inner); 255 255 + if let Some(peer) = concentrated { 256 256 + tracing::warn!( 257 257 + %peer, 258 258 + "one address has taken {CONCENTRATED_REFUSALS} more of the pre-authentication limiter's refusals than every other address combined. If this address is a proxy, set xrpc.trusted_proxy_header to the header it forwards the client address in and add the address to xrpc.trusted_proxies, since every client behind a proxy will share its one rate-limit bucket. This warning reports the first such address only." 259 259 + ); 224 260 } 225 261 Err(refusal) 226 262 } ··· 295 331 } 296 332 297 333 fn peer(last: u8) -> Option<IpAddr> { 298 298 - Some(IpAddr::V4(Ipv4Addr::new(127, 0, 0, last))) 334 334 + Some(ip(last)) 299 335 } 300 336 301 337 fn rotating(index: u64) -> Option<IpAddr> { ··· 311 347 312 348 fn tracked(limiter: &Arc<PreAuthLimiter>) -> usize { 313 349 limiter.lock().peers.len() 350 350 + } 351 351 + 352 352 + fn ip(last: u8) -> IpAddr { 353 353 + IpAddr::V4(Ipv4Addr::new(127, 0, 0, last)) 354 354 + } 355 355 + 356 356 + #[test] 357 357 + fn only_an_address_taking_most_of_the_refusals_is_reported_and_only_once() { 358 358 + let reported = |rounds, peer: fn(u32) -> u8| { 359 359 + let mut majority = RefusalMajority::default(); 360 360 + (0..rounds) 361 361 + .filter_map(|round| majority.observe(ip(peer(round)))) 362 362 + .collect::<Vec<IpAddr>>() 363 363 + }; 364 364 + [ 365 365 + (CONCENTRATED_REFUSALS * 3, (|round| (round % 4 == 0) as u8) as fn(u32) -> u8, vec![ip(0)], 366 366 + "the warning reports that address once, since three refusals in every four come from it"), 367 367 + (CONCENTRATED_REFUSALS * 8, |round| (round % 251) as u8, vec![], 368 368 + "a knot under scattered load doesn't have a proxy to point the operator at"), 369 369 + (CONCENTRATED_REFUSALS - 1, |_| 1, vec![], 370 370 + "a burst under the threshold is ordinary rate limiting and doesn't point at a proxy"), 371 371 + (CONCENTRATED_REFUSALS, |_| 1, vec![ip(1)], 372 372 + "the threshold itself is where an address earns the warning"), 373 373 + ] 374 374 + .into_iter() 375 375 + .for_each(|(rounds, peer, expected, why)| { 376 376 + assert_eq!(reported(rounds, peer), expected, "{why}"); 377 377 + }); 314 378 } 315 379 316 380 #[test] ··· 492 556 assert!( 493 557 tracked <= MAX_TRACKED_PEERS, 494 558 "a flood of distinct source addresses mustn't grow the peer map past its limit, saw {tracked}" 559 559 + ); 560 560 + } 561 561 + 562 562 + #[test] 563 563 + fn the_majority_counter_sees_a_refusal_from_a_full_peer_table() { 564 564 + let limiter = limiter(LimitConfig { 565 565 + rate: rate(1, 1_000), 566 566 + per_peer_inflight: Some(PerPeerInflight::new(8)), 567 567 + global_inflight: Some(GlobalInflight::new(64)), 568 568 + }); 569 569 + (0..MAX_TRACKED_PEERS as u64).for_each(|index| { 570 570 + let _ = limiter.admit(rotating(index), at(0)); 571 571 + }); 572 572 + assert_eq!( 573 573 + limiter 574 574 + .admit(peer(201), at(SWEEP_INTERVAL_MICROS - 1)) 575 575 + .err(), 576 576 + Some(Refusal::Saturated) 577 577 + ); 578 578 + let inner = limiter.lock(); 579 579 + assert_eq!( 580 580 + (inner.refusal_majority.peer, inner.refusal_majority.votes), 581 581 + (peer(201), 1), 582 582 + "that refusal has to reach the counter like every other refusal, since a peer shed by a full table is what the warning most needs to report" 495 583 ); 496 584 } 497 585 }