lee/rustdesk
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk.git synced 2026-08-27 08:36:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

udp: make the punch prove itself, and keep the listener answering

Browse Source 
punch_udp sent a zero-length datagram and called the hole open on whatever
arrived next. The rendezvous NAT test's own leftover replies satisfy that
immediately - connect() does not flush the receive queue - so the retry loop
never ran and success meant nothing. The dead socket then cost KCP its full
timeout to rediscover, which is how a failed punch came to take 18 seconds.

Probes now carry a magic and a 64-bit transaction id, and both ends answer
each other's probes, so returning is a fact: a reply echoing our own id is the
one thing that proves the pair carries traffic both ways. With failure now
distinguishable from 'not yet', the window drops from 20s to 3s.

Two asymmetries fall out of that:

Only the connector stops on its own acknowledgement, because only it has
something to send next. An acknowledgement proves our probe came back, not
that the peer's probe was answered - and after punch_udp returns nothing
answers probes any more, since KCP's io loop drops anything shorter than its
header. A listener that stopped there would go mute while a peer whose own
probe or answer was lost - the normal state of a hole still opening - kept
probing an endpoint that works, until it timed out.

So the listener stops on the peer's first real packet instead, and hands that
packet to KcpStream::accept as its init_packet: its arrival proves the pair as
well as an acknowledgement would, and KCP never retransmits its SYN.
This commit is contained in:
rustdesk
2026-08-27 14:49:34 +08:00
parent e50d620889
commit 81703e2a82
2 changed files with 80 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common.rs
+78 -29
View File
@@ -2552,50 +2552,87 @@ pub async fn test_ipv6() -> Option<tokio::task::JoinHandle<()>> {
}))
}
// A punch packet carries a magic and a transaction id so a reply can be *proven* to answer this
// probe. The punch it replaces sent a zero-length datagram and called the hole open on whatever
// arrived next - which the rendezvous NAT test's own leftover replies satisfied instantly, so the
// retry loop below never actually ran and its success meant nothing.
const PUNCH_PROBE: [u8; 4] = *b"RDP?";
const PUNCH_ACK: [u8; 4] = *b"RDP!";
const PUNCH_PACKET_LEN: usize = 12;
fn punch_packet(tag: &[u8; 4], tid: u64) -> [u8; PUNCH_PACKET_LEN] {
let mut packet = [0u8; PUNCH_PACKET_LEN];
packet[..4].copy_from_slice(tag);
packet[4..].copy_from_slice(&tid.to_le_bytes());
packet
}
fn punch_tid(packet: &[u8], tag: &[u8; 4]) -> Option<u64> {
if packet.len() != PUNCH_PACKET_LEN || packet[..4] != tag[..] {
return None;
}
packet[4..].try_into().ok().map(u64::from_le_bytes)
}
/// Punch until one of our own probes is acknowledged. Both ends run this identically - each
/// probes, each answers the other's probes - and each returns only once a reply carrying its own
/// transaction id comes back, the one thing that proves the pair carries traffic both ways.
///
/// Returning is therefore a fact rather than a guess, which is what lets the caller stop instead
/// of handing a dead socket to a transport whose only way to discover the truth is to time out.
///
/// A datagram that is neither probe nor acknowledgement is returned rather than dropped: it means
/// the peer finished first and is already speaking KCP, whose SYN is never retransmitted.
///
/// Only the connector stops on its own acknowledgement, because only it has something to send
/// next. An acknowledgement proves our probe came back, not that the peer's probe was answered -
/// and after this returns nothing answers probes any more, since KCP's io loop drops anything
/// shorter than its header. A listener that stopped here would go mute while a peer whose own
/// probe or answer was lost - the normal state of a hole that is still opening - kept probing an
/// endpoint that works, until it timed out. So the listener stops on the peer's first real packet.
pub async fn punch_udp(
socket: Arc<UdpSocket>,
listen: bool,
) -> ResultType<Option<bytes::BytesMut>> {
let tid = ((hbb_common::time_based_rand() as u64) << 32) | hbb_common::time_based_rand() as u64;
let probe = punch_packet(&PUNCH_PROBE, tid);
let mut data = [0u8; 1500];
// `connect` does not flush the receive queue, so the NAT test's extra replies are still in it.
while socket.try_recv(&mut data).is_ok() {}
let mut retry_interval = Duration::from_millis(20);
const MAX_INTERVAL: Duration = Duration::from_millis(200);
const MAX_TIME: Duration = Duration::from_secs(20);
let mut packets_sent = 0;
// Both ends start within one rendezvous round trip of each other and the acknowledgement is
// one peer round trip, so a pair that has not answered in this long is not going to. The old
// 20s came from having no way to tell "not yet" from "never".
const MAX_TIME: Duration = Duration::from_secs(3);
let mut probes_sent = 0u32;
let mut probes_seen = 0u32;
let mut acked = false;
let mut recv_errors = 0u32;
socket.send(&[]).await.ok();
packets_sent += 1;
socket.send(&probe).await.ok();
probes_sent += 1;
let mut last_send_time = Instant::now();
let tm = Instant::now();
let mut data = [0u8; 1500];
loop {
tokio::select! {
_ = hbb_common::sleep(retry_interval.as_secs_f32()) => {
if tm.elapsed() > MAX_TIME {
bail!("UDP punch is timed out, stop sending packets after {:?} packets, {} recv errors absorbed", packets_sent, recv_errors);
bail!("UDP punch is timed out, {probes_sent} probes sent, {probes_seen} probes received, acked: {acked}, {recv_errors} recv errors absorbed");
}
let elapsed = last_send_time.elapsed();
if elapsed >= retry_interval {
socket.send(&[]).await.ok();
packets_sent += 1;
// Exponentially increase interval to reduce network pressure
retry_interval = std::cmp::min(
Duration::from_millis((retry_interval.as_millis() as f64 * 1.5) as u64),
MAX_INTERVAL
);
if last_send_time.elapsed() >= retry_interval {
socket.send(&probe).await.ok();
probes_sent += 1;
retry_interval = std::cmp::min(retry_interval.mul_f64(1.5), MAX_INTERVAL);
last_send_time = Instant::now();
}
}
res = socket.recv(&mut data) => match res {
Err(e) => {
// While the hole is still forming, ICMP unreachable from the peer's NAT
// is expected and surfaces as ConnectionReset/Refused on a connected
// socket (notably 10054 on Windows). Treat it as loss and keep punching;
// MAX_TIME above still bounds the whole attempt.
// Log only the first: this retries every 10ms for up to MAX_TIME, so one
// line per occurrence would write thousands into the log file per punch.
// The count is reported once at the end.
// ICMP unreachable from the peer's NAT is expected while the hole forms and
// surfaces here as ConnectionReset/Refused; treat it as loss, MAX_TIME bounds
// the attempt. Log only the first - this retries every 10ms.
recv_errors += 1;
if recv_errors == 1 {
log::debug!("UDP punch recv error (treated as loss): {e}");
@@ -2603,14 +2640,26 @@ pub async fn punch_udp(
hbb_common::sleep(0.01).await;
}
Ok(n) => {
// log::debug!("UDP punch succeeded after sending {} packets after {:?}", packets_sent, tm.elapsed());
if listen {
if n == 0 {
continue;
let ack = punch_tid(&data[..n], &PUNCH_ACK);
if ack == Some(tid) {
if !listen {
log::debug!(
"UDP punch confirmed in {:?}, {probes_sent} probes sent, {probes_seen} received",
tm.elapsed()
);
return Ok(None);
}
acked = true;
} else if let Some(peer_tid) = punch_tid(&data[..n], &PUNCH_PROBE) {
probes_seen += 1;
socket.send(&punch_packet(&PUNCH_ACK, peer_tid)).await.ok();
} else if ack.is_none() && n > 0 {
log::debug!(
"UDP punch confirmed by {n} bytes of peer data in {:?}, {probes_sent} probes sent",
tm.elapsed()
);
return Ok(Some(bytes::BytesMut::from(&data[..n])));
}
return Ok(None);
}
}
}
src/rendezvous_mediator.rs
+2 -2
View File
@@ -1192,11 +1192,11 @@ async fn udp_nat_listen(
let socket_cloned = socket.clone();
let func = async {
socket.connect(peer_addr).await?;
let res = crate::punch_udp(socket.clone(), true).await?;
let init_packet = crate::punch_udp(socket.clone(), true).await?;
let stream = crate::kcp_stream::KcpStream::accept(
socket,
Duration::from_millis(CONNECT_TIMEOUT as _),
res,
init_packet,
)
.await?;
crate::server::create_tcp_connection(server, stream.1, peer_addr_v4, true, meta).await?;