mirror of
https://github.com/rustdesk/rustdesk-server.git
synced 2026-10-12 03:34:26 +00:00
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:
co-authored by
Claude Fable 5.1
parent
a7736be5e4
commit
1cc5895ac0
+38
-16
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user