diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs index eaf7190..50b55b5 100644 --- a/src/rendezvous_server.rs +++ b/src/rendezvous_server.rs @@ -33,7 +33,7 @@ use hbb_common::{ use ipnetwork::Ipv4Network; use sodiumoxide::crypto::sign; use std::{ - collections::HashMap, + collections::{HashMap, VecDeque}, net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}, sync::atomic::{AtomicBool, AtomicUsize, Ordering}, sync::Arc, @@ -66,8 +66,24 @@ use once_cell::sync::Lazy; use tokio::sync::Mutex as TokioMutex; // differentiate if needed #[derive(Clone)] struct PunchReqEntry { tm: Instant, from_ip: String, to_ip: String, to_id: String } -static PUNCH_REQS: Lazy>> = Lazy::new(|| TokioMutex::new(Vec::new())); +static PUNCH_REQS: Lazy>> = Lazy::new(|| TokioMutex::new(VecDeque::new())); const PUNCH_REQ_DEDUPE_SEC: u64 = 60; +const PUNCH_REQ_MAX: usize = 10_000; + +async fn record_punch_req(from_ip: String, to_ip: String, to_id: String) { + let mut lock = PUNCH_REQS.lock().await; + let mut dup = false; + for e in lock.iter().rev().take(30) { // only check recent tail subset for speed + if e.from_ip == from_ip && e.to_id == to_id { + if e.tm.elapsed().as_secs() < PUNCH_REQ_DEDUPE_SEC { dup = true; } + break; + } + } + if !dup { + if lock.len() >= PUNCH_REQ_MAX { lock.pop_front(); } + lock.push_back(PunchReqEntry { tm: Instant::now(), from_ip, to_ip, to_id }); + } +} #[derive(Clone)] struct Inner { @@ -738,20 +754,12 @@ impl RendezvousServer { } // record punch hole request (from addr -> peer id/peer_addr) - { - let from_ip = try_into_v4(addr).ip().to_string(); - let to_ip = try_into_v4(peer_addr).ip().to_string(); - let to_id_clone = id.clone(); - let mut lock = PUNCH_REQS.lock().await; - let mut dup = false; - for e in lock.iter().rev().take(30) { // only check recent tail subset for speed - if e.from_ip == from_ip && e.to_id == to_id_clone { - if e.tm.elapsed().as_secs() < PUNCH_REQ_DEDUPE_SEC { dup = true; } - break; - } - } - if !dup { lock.push(PunchReqEntry { tm: Instant::now(), from_ip, to_ip, to_id: to_id_clone }); } - } + record_punch_req( + try_into_v4(addr).ip().to_string(), + try_into_v4(peer_addr).ip().to_string(), + id.clone(), + ) + .await; let mut msg_out = RendezvousMessage::new(); let peer_is_lan = self.is_lan(peer_addr); @@ -1418,4 +1426,18 @@ mod tests { let socket = create_udp_listener(Some(bind_addr), 0, 0).await.unwrap(); assert_eq!(socket.local_addr().unwrap().ip(), bind_addr); } + + #[hbb_common::tokio::test] + async fn punch_reqs_bounded() { + // 31 ids keep every request just outside the 30-entry dedupe window + for i in 0..PUNCH_REQ_MAX * 2 { + record_punch_req( + "1.2.3.4".to_owned(), + "5.6.7.8".to_owned(), + format!("peer{}", i % 31), + ) + .await; + } + assert_eq!(PUNCH_REQS.lock().await.len(), PUNCH_REQ_MAX); + } }