fix(hbbs): bound the punch-request record store

PUNCH_REQS grew without limit; its dedupe scan only covers the newest 30
entries, so cycling 31 online peer ids from one source appended a record
per request for the lifetime of the process. Keep at most 10,000 entries,
evicting the oldest.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PkpVy2ufxgEjKd3SjXPnJ8
This commit is contained in:
rustdesk
2026-10-06 17:29:01 +08:00
co-authored by Claude Fable 5.1
parent a7736be5e4
commit 1cc5895ac0
+38 -16
View File
@@ -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<TokioMutex<Vec<PunchReqEntry>>> = Lazy::new(|| TokioMutex::new(Vec::new()));
static PUNCH_REQS: Lazy<TokioMutex<VecDeque<PunchReqEntry>>> = 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);
}
}