TheLanc3/rustdesk
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk.git synced 2026-04-05 23:01:29 +03:00
Code Issues Packages Projects Releases Wiki Activity

udp punch and ipv6 punch

Browse Source
This commit is contained in:
rustdesk
2025-06-12 21:32:28 +08:00
parent 05a812247a
commit 7792ac1481
12 changed files with 1151 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files

325
src/client.rs
View File

@@ -30,7 +30,9 @@ use uuid::Uuid;
use crate::{
check_port,
common::input::{MOUSE_BUTTON_LEFT, MOUSE_BUTTON_RIGHT, MOUSE_TYPE_DOWN, MOUSE_TYPE_UP},
create_symmetric_key_msg, decode_id_pk, get_rs_pk, is_keyboard_mode_supported, secure_tcp,
create_symmetric_key_msg, decode_id_pk, get_rs_pk, is_keyboard_mode_supported,
kcp_stream::KcpStream,
secure_tcp,
ui_interface::{get_builtin_option, use_texture_render},
ui_session_interface::{InvokeUiSession, Session},
};
@@ -40,7 +42,6 @@ pub use file_trait::FileManager;
#[cfg(not(feature = "flutter"))]
#[cfg(not(any(target_os = "android", target_os = "ios")))]
use hbb_common::tokio::sync::mpsc::UnboundedSender;
use hbb_common::tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver};
use hbb_common::{
allow_err,
anyhow::{anyhow, Context},
@@ -50,17 +51,23 @@ use hbb_common::{
READ_TIMEOUT, RELAY_PORT, RENDEZVOUS_PORT, RENDEZVOUS_SERVERS,
},
fs::JobType,
futures::future::{select_ok, FutureExt},
get_version_number, log,
message_proto::{option_message::BoolOption, *},
protobuf::{Message as _, MessageField},
rand,
rendezvous_proto::*,
sha2::{Digest, Sha256},
socket_client::{connect_tcp, connect_tcp_local, ipv4_to_ipv6},
socket_client::{connect_tcp, connect_tcp_local, ipv4_to_ipv6, new_direct_udp_for},
sodiumoxide::{base64, crypto::sign},
timeout,
tokio::{
self,
net::UdpSocket,
sync::{
mpsc::{unbounded_channel, UnboundedReceiver},
oneshot,
},
time::{interval, Duration, Instant},
},
AddrMangle, ResultType, Stream,
@@ -184,7 +191,10 @@ impl Client {
token: &str,
conn_type: ConnType,
interface: impl Interface,
) -> ResultType<((Stream, bool, Option<Vec<u8>>), (i32, String))> {
) -> ResultType<(
(Stream, bool, Option<Vec<u8>>, Option<KcpStream>),
(i32, String),
)> {
debug_assert!(peer == interface.get_id());
interface.update_direct(None);
interface.update_received(false);
@@ -208,7 +218,10 @@ impl Client {
token: &str,
conn_type: ConnType,
interface: impl Interface,
) -> ResultType<((Stream, bool, Option<Vec<u8>>), (i32, String))> {
) -> ResultType<(
(Stream, bool, Option<Vec<u8>>, Option<KcpStream>),
(i32, String),
)> {
if config::is_incoming_only() {
bail!("Incoming only mode");
}
@@ -220,6 +233,7 @@ impl Client {
.await?,
true,
None,
None,
),
(0, "".to_owned()),
));
@@ -231,6 +245,7 @@ impl Client {
connect_tcp_local(peer, None, CONNECT_TIMEOUT).await?,
true,
None,
None,
),
(0, "".to_owned()),
));
@@ -259,8 +274,29 @@ impl Client {
}
};
crate::test_ipv6().await;
let (stop_udp_tx, stop_udp_rx) = oneshot::channel::<()>();
let mut udp =
// no need to care about multiple rendezvous servers case, since it is acutally not used any more.
// Shared state for UDP NAT test result
if let Ok((socket, addr)) = new_direct_udp_for(&rendezvous_server).await {
let udp_port = Arc::new(Mutex::new(0));
let up_cloned = udp_port.clone();
let socket_cloned = socket.clone();
let func = async move {
allow_err!(test_udp_uat(socket_cloned, addr, up_cloned, stop_udp_rx).await);
};
tokio::spawn(func);
(Some(socket), Some(udp_port))
} else {
(None, None)
};
let mut start = Instant::now();
let mut socket = connect_tcp(&*rendezvous_server, CONNECT_TIMEOUT).await;
debug_assert!(!servers.contains(&rendezvous_server));
let rtt = start.elapsed();
log::debug!("TCP connection establishment time used: {:?}", rtt);
if socket.is_err() && !servers.is_empty() {
log::info!("try the other servers: {:?}", servers);
for server in servers {
@@ -280,40 +316,65 @@ impl Client {
let my_addr = socket.local_addr();
let mut signed_id_pk = Vec::new();
let mut relay_server = "".to_owned();
if !key.is_empty() && !token.is_empty() {
// mainly for the security of token
secure_tcp(&mut socket, key).await.map_err(|e| anyhow!("Failed to secure tcp: {}", e))?;
}
let start = std::time::Instant::now();
let mut peer_addr = Config::get_any_listen_addr(true);
let mut peer_nat_type = NatType::UNKNOWN_NAT;
let my_nat_type = crate::get_nat_type(100).await;
let mut is_local = false;
let mut feedback = 0;
let force_relay = interface.is_force_relay() || use_ws() || Config::is_proxy();
use hbb_common::protobuf::Enum;
let nat_type = if force_relay {
NatType::SYMMETRIC
} else {
NatType::from_i32(my_nat_type).unwrap_or(NatType::UNKNOWN_NAT)
};
if !key.is_empty() && !token.is_empty() {
// mainly for the security of token
secure_tcp(&mut socket, key)
.await
.map_err(|e| anyhow!("Failed to secure tcp: {}", e))?;
} else if let Some(udp) = udp.1.as_ref() {
let tm = Instant::now();
loop {
let port = *udp.lock().unwrap();
if port > 0 {
break;
}
// await for 0.5 RTT
if tm.elapsed() > rtt / 2 {
break;
}
hbb_common::sleep(0.001).await;
}
}
// Stop UDP NAT test task if still running
let _ = stop_udp_tx.send(());
let mut msg_out = RendezvousMessage::new();
let mut ipv6 = if let Some((socket, addr)) = crate::get_ipv6_socket().await {
(Some(socket), Some(addr))
} else {
(None, None)
};
let udp_nat_port = udp.1.map(|x| *x.lock().unwrap()).unwrap_or(0);
msg_out.set_punch_hole_request(PunchHoleRequest {
id: peer.to_owned(),
token: token.to_owned(),
nat_type: nat_type.into(),
licence_key: key.to_owned(),
conn_type: conn_type.into(),
version: crate::VERSION.to_owned(),
udp_port: udp_nat_port as _,
force_relay,
socket_addr_v6: ipv6.1.unwrap_or_default(),
..Default::default()
});
for i in 1..=3 {
log::info!("#{} punch attempt with {}, id: {}", i, my_addr, peer);
let mut msg_out = RendezvousMessage::new();
use hbb_common::protobuf::Enum;
let nat_type = if interface.is_force_relay() || use_ws() || Config::is_proxy() {
NatType::SYMMETRIC
} else {
NatType::from_i32(my_nat_type).unwrap_or(NatType::UNKNOWN_NAT)
};
msg_out.set_punch_hole_request(PunchHoleRequest {
id: peer.to_owned(),
token: token.to_owned(),
nat_type: nat_type.into(),
licence_key: key.to_owned(),
conn_type: conn_type.into(),
version: crate::VERSION.to_owned(),
..Default::default()
});
socket.send(&msg_out).await?;
// below timeout should not bigger than hbbs's connection timeout.
if let Some(msg_in) =
crate::get_next_nonkeyexchange_msg(&mut socket, Some(i * 6000)).await
crate::get_next_nonkeyexchange_msg(&mut socket, Some(i * 3000)).await
{
match msg_in.union {
Some(rendezvous_message::Union::PunchHoleResponse(ph)) => {
@@ -343,6 +404,23 @@ impl Client {
relay_server = ph.relay_server;
peer_addr = AddrMangle::decode(&ph.socket_addr);
feedback = ph.feedback;
let s = udp.0.take();
if ph.is_udp && s.is_some() {
if let Some(s) = s {
allow_err!(s.connect(peer_addr).await);
udp.0 = Some(s);
}
}
let s = ipv6.0.take();
if !ph.socket_addr_v6.is_empty() && s.is_some() {
let addr = AddrMangle::decode(&ph.socket_addr_v6);
if addr.port() > 0 {
if let Some(s) = s {
allow_err!(s.connect(addr).await);
ipv6.0 = Some(s);
}
}
}
log::info!("Hole Punched {} = {}", peer, peer_addr);
break;
}
@@ -353,20 +431,44 @@ impl Client {
start.elapsed(),
rr.relay_server
);
start = Instant::now();
let mut connect_futures = Vec::new();
if let Some(s) = ipv6.0 {
let addr = AddrMangle::decode(&rr.socket_addr_v6);
if addr.port() > 0 {
if s.connect(addr).await.is_ok() {
connect_futures.push(udp_nat_connect(s, "IPv6").boxed());
}
}
}
signed_id_pk = rr.pk().into();
let mut conn = Self::create_relay(
let fut = Self::create_relay(
peer,
rr.uuid,
rr.relay_server,
key,
conn_type,
my_addr.is_ipv4(),
)
.await?;
);
connect_futures.push(
async move {
let conn = fut.await?;
Ok((conn, None, "Relay"))
}
.boxed(),
);
// Run all connection attempts concurrently, return the first successful one
let (conn, kcp, typ) = match select_ok(connect_futures).await {
Ok(conn) => (Ok(conn.0 .0), conn.0 .1, conn.0 .2),
Err(e) => (Err(e), None, ""),
};
let mut conn = conn?;
feedback = rr.feedback;
log::info!("{:?} used to establish {typ} connection", start.elapsed());
let pk =
Self::secure_connection(peer, signed_id_pk, key, &mut conn).await?;
return Ok(((conn, false, pk), (feedback, rendezvous_server)));
return Ok(((conn, false, pk, kcp), (feedback, rendezvous_server)));
}
_ => {
log::error!("Unexpected protobuf msg received: {:?}", msg_in);
@@ -405,6 +507,8 @@ impl Client {
token,
conn_type,
interface,
udp.0,
ipv6.0,
)
.await?,
(feedback, rendezvous_server),
@@ -427,7 +531,9 @@ impl Client {
token: &str,
conn_type: ConnType,
interface: impl Interface,
) -> ResultType<(Stream, bool, Option<Vec<u8>>)> {
udp_socket_nat: Option<Arc<UdpSocket>>,
udp_socket_v6: Option<Arc<UdpSocket>>,
) -> ResultType<(Stream, bool, Option<Vec<u8>>, Option<KcpStream>)> {
let direct_failures = interface.get_lch().read().unwrap().direct_failures;
let mut connect_timeout = 0;
const MIN: u64 = 1000;
@@ -462,8 +568,28 @@ impl Client {
}
log::info!("peer address: {}, timeout: {}", peer, connect_timeout);
let start = std::time::Instant::now();
// NOTICE: Socks5 is be used event in intranet. Which may be not a good way.
let mut conn = connect_tcp_local(peer, Some(local_addr), connect_timeout).await;
let mut connect_futures = Vec::new();
let fut = connect_tcp_local(peer, Some(local_addr), connect_timeout);
connect_futures.push(
async move {
let conn = fut.await?;
Ok((conn, None, "TCP"))
}
.boxed(),
);
if let Some(udp_socket_nat) = udp_socket_nat {
connect_futures.push(udp_nat_connect(udp_socket_nat, "UDP").boxed());
}
if let Some(udp_socket_v6) = udp_socket_v6 {
connect_futures.push(udp_nat_connect(udp_socket_v6, "IPv6").boxed());
}
// Run all connection attempts concurrently, return the first successful one
let (mut conn, kcp, mut typ) = match select_ok(connect_futures).await {
Ok(conn) => (Ok(conn.0 .0), conn.0 .1, conn.0 .2),
Err(e) => (Err(e), None, ""),
};
let mut direct = !conn.is_err();
interface.update_direct(Some(direct));
if interface.is_force_relay() || conn.is_err() {
@@ -482,6 +608,7 @@ impl Client {
if let Err(e) = conn {
bail!("Failed to connect via relay server: {}", e);
}
typ = "Relay";
direct = false;
} else {
bail!("Failed to make direct connection to remote desktop");
@@ -493,9 +620,9 @@ impl Client {
interface.get_lch().write().unwrap().set_direct_failure(n);
}
let mut conn = conn?;
log::info!("{:?} used to establish connection", start.elapsed());
log::info!("{:?} used to establish {typ} connection", start.elapsed());
let pk = Self::secure_connection(peer_id, signed_id_pk, key, &mut conn).await?;
Ok((conn, direct, pk))
Ok((conn, direct, pk, kcp))
}
/// Establish secure connection with the server.
@@ -3707,3 +3834,125 @@ pub mod peer_online {
}
}
}
async fn test_udp_uat(
udp_socket: Arc<UdpSocket>,
server_addr: SocketAddr,
udp_port: Arc<Mutex<u16>>,
mut stop_udp_rx: oneshot::Receiver<()>,
) -> ResultType<()> {
let (tx, mut rx) = oneshot::channel::<_>();
tokio::spawn(async {
if let Ok(v) = crate::test_nat_ipv4().await {
tx.send(v).ok();
}
});
let start = Instant::now();
let mut msg_out = RendezvousMessage::new();
msg_out.set_test_nat_request(TestNatRequest {
..Default::default()
});
// Adaptive retry strategy that works within TCP RTT constraints
// Start with aggressive sending, then back off
let mut retry_interval = Duration::from_millis(20); // Start fast
const MAX_INTERVAL: Duration = Duration::from_millis(200);
let mut packets_sent = 0;
// Send initial burst to improve reliability
let data = msg_out.write_to_bytes()?;
for _ in 0..2 {
if let Err(e) = udp_socket.send_to(&data, server_addr).await {
log::warn!("Failed to send initial UDP NAT test packet: {}", e);
} else {
packets_sent += 1;
}
}
let mut last_send_time = Instant::now();
let mut buf = [0u8; 1024];
loop {
tokio::select! {
Ok((addr, server)) = &mut rx => {
*udp_port.lock().unwrap() = addr.port();
log::debug!("UDP NAT test received response from {}: {}", addr, server);
break;
}
_ = &mut stop_udp_rx => {
log::debug!("UDP NAT test received stop signal after {} packets", packets_sent);
break;
}
_ = hbb_common::sleep(retry_interval.as_secs_f32()) => {
// Adaptive retry: send fewer packets as time goes on
let elapsed = last_send_time.elapsed();
if elapsed >= retry_interval {
// Send single packet (not double) to reduce network load
if let Err(e) = udp_socket.send_to(&data, server_addr).await {
log::warn!("Failed to send UDP NAT test retry packet: {}", e);
} else {
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
);
last_send_time = Instant::now();
}
}
res = udp_socket.recv(&mut buf[..]) => {
match res {
Ok(n) => {
match RendezvousMessage::parse_from_bytes(&buf[0..n]) {
Ok(msg_in) => {
if let Some(rendezvous_message::Union::TestNatResponse(response)) = msg_in.union {
*udp_port.lock().unwrap() = response.port as u16;
break;
}
}
Err(e) => {
log::warn!("Failed to parse UDP NAT test response: {}", e);
}
}
}
Err(e) => {
log::warn!("UDP NAT test socket error: {}", e);
}
}
}
}
}
let final_port = *udp_port.lock().unwrap();
log::debug!(
"UDP NAT test to {:?} finished: time={:?}, port={}, packets_sent={}, success={}",
server_addr,
start.elapsed(),
final_port,
packets_sent,
final_port > 0
);
Ok(())
}
#[inline]
async fn udp_nat_connect(
socket: Arc<UdpSocket>,
typ: &'static str,
) -> ResultType<(Stream, Option<KcpStream>, &'static str)> {
crate::punch_udp(socket.clone(), false)
.await
.map_err(|err| {
log::debug!("{err}");
anyhow!(err)
})?;
let res = KcpStream::connect(socket, Duration::from_secs(CONNECT_TIMEOUT as _))
.await
.map_err(|err| {
log::debug!("Failed to connect KCP stream: {}", err);
anyhow!(err)
})?;
Ok((res.1, Some(res.0), typ))
}