mirror of
https://github.com/rustdesk/rustdesk-server.git
synced 2026-10-22 16:27:22 +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:
1 parent
a7736be5e4
commit
1cc5895ac0
1 file changed
+38
-16
+38
-16
@@ -33,7 +33,7 @@ use hbb_common::{
|
|||||||
use ipnetwork::Ipv4Network;
|
use ipnetwork::Ipv4Network;
|
||||||
use sodiumoxide::crypto::sign;
|
use sodiumoxide::crypto::sign;
|
||||||
use std::{
|
use std::{
|
||||||
collections::HashMap,
|
collections::{HashMap, VecDeque},
|
||||||
net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
|
net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
|
||||||
sync::atomic::{AtomicBool, AtomicUsize, Ordering},
|
sync::atomic::{AtomicBool, AtomicUsize, Ordering},
|
||||||
sync::Arc,
|
sync::Arc,
|
||||||
@@ -66,8 +66,24 @@ use once_cell::sync::Lazy;
|
|||||||
use tokio::sync::Mutex as TokioMutex; // differentiate if needed
|
use tokio::sync::Mutex as TokioMutex; // differentiate if needed
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
struct PunchReqEntry { tm: Instant, from_ip: String, to_ip: String, to_id: String }
|
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_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)]
|
#[derive(Clone)]
|
||||||
struct Inner {
|
struct Inner {
|
||||||
@@ -738,20 +754,12 @@ impl RendezvousServer {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// record punch hole request (from addr -> peer id/peer_addr)
|
// record punch hole request (from addr -> peer id/peer_addr)
|
||||||
{
|
record_punch_req(
|
||||||
let from_ip = try_into_v4(addr).ip().to_string();
|
try_into_v4(addr).ip().to_string(),
|
||||||
let to_ip = try_into_v4(peer_addr).ip().to_string();
|
try_into_v4(peer_addr).ip().to_string(),
|
||||||
let to_id_clone = id.clone();
|
id.clone(),
|
||||||
let mut lock = PUNCH_REQS.lock().await;
|
)
|
||||||
let mut dup = false;
|
.await;
|
||||||
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 }); }
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut msg_out = RendezvousMessage::new();
|
let mut msg_out = RendezvousMessage::new();
|
||||||
let peer_is_lan = self.is_lan(peer_addr);
|
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();
|
let socket = create_udp_listener(Some(bind_addr), 0, 0).await.unwrap();
|
||||||
assert_eq!(socket.local_addr().unwrap().ip(), bind_addr);
|
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