diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 488921f..28fb251 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -48,6 +48,7 @@ in the inherited process environment. | `ALWAYS_USE_RELAY` 🅴 | *(none)* | `N` | `Y` forces every session through a relay (disables direct/hole‑punched connections). At runtime, send `always-use-relay Y` or `always-use-relay N` to the `hbbs` [loopback console](#runtime-console). | | `DB_URL` 🅴 | *(none)* | `./db_v2.sqlite3` | Path/URL of the SQLite database file. See [Database](#database). | | `MAX_DATABASE_CONNECTIONS` 🅴 | *(none)* | `1` | Size of the SQLite connection pool. | +| `MAX_PUNCH_REQUESTS` 🅴 | *(none)* | `100000` | Number of punch-request records kept in memory for the `punch-requests` [console](#runtime-console) command; the oldest is dropped when full. `0` disables recording. | 🅴 = set through the inherited process environment. diff --git a/src/peer.rs b/src/peer.rs index 0f218df..9e4cbec 100644 --- a/src/peer.rs +++ b/src/peer.rs @@ -22,6 +22,9 @@ pub const IP_CHANGE_DUR: u64 = 180; pub const IP_CHANGE_DUR_X2: u64 = IP_CHANGE_DUR * 2; pub const DAY_SECONDS: u64 = 3600 * 24; pub const IP_BLOCK_DUR: u64 = 60; +// Same bound as the Pro server: a longer id is not a normal RustDesk id or hostname, +// and every peer entry and punch-request record keeps its own copy. +pub const MAX_PEER_ID_LEN: usize = 253; #[derive(Debug, Default, Serialize, Deserialize, Clone)] pub(crate) struct PeerInfo { diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs index eaf7190..daed76d 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,27 @@ 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; +static PUNCH_REQ_MAX: Lazy = + Lazy::new(|| get_arg("MAX_PUNCH_REQUESTS").parse().unwrap_or(100_000)); + +async fn record_punch_req(from_ip: String, to_ip: String, to_id: String) { + let max = *PUNCH_REQ_MAX; + if max == 0 { return; } + 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() >= max { lock.pop_front(); } + lock.push_back(PunchReqEntry { tm: Instant::now(), from_ip, to_ip, to_id }); + } +} #[derive(Clone)] struct Inner { @@ -354,7 +373,7 @@ impl RendezvousServer { match msg_in.union { Some(rendezvous_message::Union::RegisterPeer(rp)) => { // B registered - if !rp.id.is_empty() { + if !rp.id.is_empty() && rp.id.len() <= MAX_PEER_ID_LEN { log::trace!("New peer registered: {:?} {:?}", &rp.id, &addr); self.update_addr(rp.id, addr, socket).await?; if self.inner.serial > rp.serial { @@ -374,7 +393,7 @@ impl RendezvousServer { } let id = rk.id; let ip = addr.ip().to_string(); - if id.len() < 6 { + if id.len() < 6 || id.len() > MAX_PEER_ID_LEN { return send_rk_res(socket, addr, UUID_MISMATCH).await; } else if !self.check_ip_blocker(&ip, &id).await { return send_rk_res(socket, addr, TOO_FREQUENT).await; @@ -738,20 +757,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 +1429,26 @@ 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() { + let max = *PUNCH_REQ_MAX; + // 31 ids keep every request just outside the 30-entry dedupe window; + // one request at least, so MAX_PUNCH_REQUESTS=0 is exercised too + for i in 0..(max * 2).max(1) { + record_punch_req( + "1.2.3.4".to_owned(), + "5.6.7.8".to_owned(), + format!("peer{}", i % 31), + ) + .await; + } + let lock = PUNCH_REQS.lock().await; + assert_eq!(lock.len(), max); + if max > 0 { + // the oldest half is evicted and the newest kept + assert_eq!(lock.front().unwrap().to_id, format!("peer{}", max % 31)); + assert_eq!(lock.back().unwrap().to_id, format!("peer{}", (max * 2 - 1) % 31)); + } + } }