diff --git a/crates/lanspread-peer/src/services/discovery.rs b/crates/lanspread-peer/src/services/discovery.rs index ea31518..e0ef382 100644 --- a/crates/lanspread-peer/src/services/discovery.rs +++ b/crates/lanspread-peer/src/services/discovery.rs @@ -28,6 +28,11 @@ use crate::{ }; const MAX_ACTIVE_DISCOVERY_CANDIDATES: usize = 64; +/// Upper bound on active-or-cooling candidates advertised from one source IP. +/// mDNS lets a single LAN host claim arbitrarily many peer IDs and ports, so +/// without this bound one host could occupy every discovery slot and starve +/// genuinely new peers. Legitimate hosts run one peer, occasionally a few. +const MAX_DISCOVERY_CANDIDATES_PER_SOURCE_IP: usize = 8; const MAX_PENDING_MDNS_SERVICES: usize = 64; const DISCOVERY_CANDIDATE_COOLDOWN: Duration = Duration::from_secs(5); @@ -60,6 +65,21 @@ impl RecentCandidates { } } + /// Counts unexpired recent candidates from `ip` that are not also in the + /// active set, so the caller can sum both without double counting. + fn count_from_ip_outside( + &mut self, + ip: std::net::IpAddr, + active: &HashSet, + now: tokio::time::Instant, + ) -> usize { + self.expire(now); + self.entries + .iter() + .filter(|(endpoint, _)| endpoint.addr.ip() == ip && !active.contains(endpoint)) + .count() + } + #[cfg(test)] fn len(&self) -> usize { self.entries.len() @@ -297,6 +317,23 @@ fn candidate_conflicts(active: &HashSet, candidate: PeerEndpoint) .any(|endpoint| endpoint.peer_id == candidate.peer_id || endpoint.addr == candidate.addr) } +/// Returns whether admitting `candidate` would exceed the per-source-IP +/// discovery budget across active negotiations and cooling-down attempts. +fn source_ip_is_saturated( + active: &HashSet, + recent: &mut RecentCandidates, + candidate: PeerEndpoint, + now: tokio::time::Instant, +) -> bool { + let ip = candidate.addr.ip(); + let active_from_ip = active + .iter() + .filter(|endpoint| endpoint.addr.ip() == ip) + .count(); + active_from_ip + recent.count_from_ip_outside(ip, active, now) + >= MAX_DISCOVERY_CANDIDATES_PER_SOURCE_IP +} + fn candidate_is_admissible( active: &HashSet, recent: &mut RecentCandidates, @@ -305,6 +342,7 @@ fn candidate_is_admissible( ) -> bool { active.len() < MAX_ACTIVE_DISCOVERY_CANDIDATES && !candidate_conflicts(active, candidate) + && !source_ip_is_saturated(active, recent, candidate, now) && recent.try_record(candidate, now) } @@ -475,6 +513,7 @@ mod tests { DISCOVERY_CANDIDATE_COOLDOWN, DiscoveryWorker, MAX_ACTIVE_DISCOVERY_CANDIDATES, + MAX_DISCOVERY_CANDIDATES_PER_SOURCE_IP, MdnsPeerInfo, RecentCandidates, candidate_conflicts, @@ -580,15 +619,72 @@ mod tests { assert_eq!(recent.len(), 1); } + fn endpoint_at(seed: u8, ip: [u8; 4], port: u16) -> PeerEndpoint { + PeerEndpoint::new(PeerId::from_bytes([seed; 32]), SocketAddr::from((ip, port))) + } + + #[test] + fn one_source_ip_cannot_occupy_more_than_its_candidate_budget() { + let now = tokio::time::Instant::now(); + let mut active = HashSet::new(); + let mut recent = RecentCandidates::default(); + let flood_ip = [192, 168, 1, 66]; + + // A single host rotating peer IDs and ports fills only its own budget. + for index in 0..MAX_DISCOVERY_CANDIDATES_PER_SOURCE_IP { + let seed = u8::try_from(index + 1).expect("test index fits u8"); + let candidate = endpoint_at(seed, flood_ip, 20000 + u16::from(seed)); + assert!(candidate_is_admissible( + &active, + &mut recent, + candidate, + now + )); + active.insert(candidate); + } + assert!(!candidate_is_admissible( + &active, + &mut recent, + endpoint_at(100, flood_ip, 30000), + now + )); + + // Other hosts are unaffected while the flood host is saturated. + assert!(candidate_is_admissible( + &active, + &mut recent, + endpoint_at(101, [192, 168, 1, 67], 42424), + now + )); + + // Completed negotiations keep counting while they cool down ... + let completed = endpoint_at(1, flood_ip, 20001); + active.remove(&completed); + assert!(!candidate_is_admissible( + &active, + &mut recent, + endpoint_at(102, flood_ip, 30001), + now + )); + // ... and free the budget once the cooldown expires. + let after_cooldown = now + DISCOVERY_CANDIDATE_COOLDOWN; + assert!(candidate_is_admissible( + &active, + &mut recent, + endpoint_at(102, flood_ip, 30001), + after_cooldown + )); + } + #[test] fn expiring_recent_entries_never_bypasses_the_independent_active_cap() { let now = tokio::time::Instant::now(); + // Spread the candidates over distinct source hosts so the per-IP + // budget does not interfere with the global active cap under test. let mut active = (0..MAX_ACTIVE_DISCOVERY_CANDIDATES) .map(|index| { - endpoint( - u8::try_from(index + 1).expect("test index fits u8"), - 15000 + u16::try_from(index).expect("test index fits u16"), - ) + let seed = u8::try_from(index + 1).expect("test index fits u8"); + endpoint_at(seed, [10, 0, seed, 1], 15000 + u16::from(seed)) }) .collect::>(); let mut recent = RecentCandidates::default(); @@ -597,7 +693,7 @@ mod tests { } let after_cooldown = now + DISCOVERY_CANDIDATE_COOLDOWN; - let next = endpoint(100, 16000); + let next = endpoint_at(100, [10, 0, 100, 1], 16000); assert!(!candidate_is_admissible( &active, &mut recent,