diff --git a/src/runtime/remote_access.rs b/src/runtime/remote_access.rs index 6b468faa..91a598c8 100644 --- a/src/runtime/remote_access.rs +++ b/src/runtime/remote_access.rs @@ -22,6 +22,9 @@ use super::ConfigApplyOptions; pub struct RustDeskRuntimeStatus { pub service_status: String, pub rendezvous_status: Option, + pub connection_count: usize, + pub listening: bool, + pub listen_port: Option, } pub struct RemoteAccessCoordinator { @@ -145,10 +148,16 @@ impl RemoteAccessCoordinator { Some(service) => RustDeskRuntimeStatus { service_status: service.status().to_string(), rendezvous_status: service.rendezvous_status().map(|status| status.to_string()), + connection_count: service.connection_count(), + listening: service.is_listening(), + listen_port: service.is_listening().then(|| service.listen_port()), }, None => RustDeskRuntimeStatus { service_status: "not_initialized".to_string(), rendezvous_status: None, + connection_count: 0, + listening: false, + listen_port: None, }, } } diff --git a/src/rustdesk/bytes_codec.rs b/src/rustdesk/bytes_codec.rs index d00896ab..a45dc2bd 100644 --- a/src/rustdesk/bytes_codec.rs +++ b/src/rustdesk/bytes_codec.rs @@ -1,7 +1,7 @@ //! Variable-length TCP framing (RustDesk wire format). use bytes::{Buf, BufMut, Bytes, BytesMut}; -use std::io; +use std::io::{self, IoSlice}; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; const MAX_PACKET_LENGTH: usize = 0x3FFFFFFF; @@ -53,6 +53,18 @@ fn decode_header(first_byte: u8, header_bytes: &[u8]) -> (usize, usize) { } pub async fn read_frame(reader: &mut R) -> io::Result { + read_frame_with_limit(reader, MAX_PACKET_LENGTH).await +} + +/// Read one framed message while enforcing a caller-selected allocation limit. +/// +/// Network-facing protocol stages should use a substantially smaller limit than +/// the wire format's theoretical maximum so an untrusted peer cannot force a +/// huge allocation by sending only a length header. +pub async fn read_frame_with_limit( + reader: &mut R, + max_packet_length: usize, +) -> io::Result { let mut first_byte = [0u8; 1]; reader.read_exact(&mut first_byte).await?; @@ -65,10 +77,10 @@ pub async fn read_frame(reader: &mut R) -> io::Result MAX_PACKET_LENGTH { + if msg_len > max_packet_length { return Err(io::Error::new( io::ErrorKind::InvalidData, - "Message too large", + format!("Message too large: {msg_len} bytes exceeds {max_packet_length}-byte limit"), )); } @@ -86,6 +98,53 @@ pub async fn write_frame(writer: &mut W, data: &[u8]) -> Ok(()) } +/// Write a frame without copying its payload into a second contiguous buffer. +/// TCP writers normally send the small header and payload in one vectored write. +pub async fn write_frame_vectored( + writer: &mut W, + data: &[u8], +) -> io::Result<()> { + let len = data.len(); + let mut header = [0u8; 4]; + let header_len = if len <= 0x3F { + header[0] = (len << 2) as u8; + 1 + } else if len <= 0x3FFF { + header[..2].copy_from_slice(&(((len << 2) as u16) | 0x1).to_le_bytes()); + 2 + } else if len <= 0x3FFFFF { + let value = ((len << 2) as u32) | 0x2; + header[0] = (value & 0xFF) as u8; + header[1] = ((value >> 8) & 0xFF) as u8; + header[2] = ((value >> 16) & 0xFF) as u8; + 3 + } else if len <= MAX_PACKET_LENGTH { + header.copy_from_slice(&(((len << 2) as u32) | 0x3).to_le_bytes()); + 4 + } else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "Message too large", + )); + }; + + let slices = [IoSlice::new(&header[..header_len]), IoSlice::new(data)]; + let written = writer.write_vectored(&slices).await?; + if written == 0 { + return Err(io::Error::new( + io::ErrorKind::WriteZero, + "failed to write RustDesk frame", + )); + } + if written < header_len { + writer.write_all(&header[written..header_len]).await?; + writer.write_all(data).await?; + } else { + writer.write_all(&data[written - header_len..]).await?; + } + Ok(()) +} + pub async fn write_frame_buffered( writer: &mut W, data: &[u8], @@ -281,4 +340,32 @@ mod tests { let decoded = codec.decode(&mut buf).unwrap().unwrap(); assert_eq!(decoded.len(), 100000); } + + #[tokio::test] + async fn read_limit_rejects_length_before_allocating_payload() { + let encoded = encode_frame(&vec![0u8; 1024]).unwrap(); + let (mut writer, mut reader) = tokio::io::duplex(encoded.len()); + tokio::spawn(async move { + writer.write_all(&encoded).await.unwrap(); + }); + + let error = read_frame_with_limit(&mut reader, 128).await.unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + } + + #[tokio::test] + async fn vectored_writer_round_trips() { + let payload = vec![0x5a; 100_000]; + let (mut writer, mut reader) = tokio::io::duplex(payload.len() + 4); + let expected = payload.clone(); + let send = tokio::spawn(async move { + write_frame_vectored(&mut writer, &payload).await.unwrap(); + }); + + let decoded = read_frame_with_limit(&mut reader, expected.len()) + .await + .unwrap(); + send.await.unwrap(); + assert_eq!(decoded.as_ref(), expected.as_slice()); + } } diff --git a/src/rustdesk/connection.rs b/src/rustdesk/connection.rs index ed51cd92..dafffe7f 100644 --- a/src/rustdesk/connection.rs +++ b/src/rustdesk/connection.rs @@ -1,16 +1,17 @@ //! Incoming RustDesk TCP sessions (handshake, AV, input). +use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; -use bytes::{Bytes, BytesMut}; +use bytes::Bytes; use parking_lot::RwLock; use protobuf::Message as ProtobufMessage; use sodiumoxide::crypto::box_; use tokio::net::tcp::OwnedWriteHalf; use tokio::net::TcpStream; -use tokio::sync::{broadcast, mpsc, Mutex}; +use tokio::sync::{mpsc, watch}; use tracing::{debug, error, info, warn}; use crate::audio::AudioController; @@ -21,7 +22,7 @@ use crate::video::codec::BitratePreset; use crate::video::codec_constraints::{encoder_codec_to_id, encoder_codec_to_video_codec}; use crate::video::stream_manager::VideoStreamManager; -use super::bytes_codec::{read_frame, write_frame, write_frame_buffered}; +use super::bytes_codec::{read_frame_with_limit, write_frame_vectored}; use super::config::RustDeskConfig; use super::crypto::{self, KeyPair, SigningKeyPair}; use super::frame_adapters::{AudioFrameAdapter, VideoCodec, VideoFrameAdapter}; @@ -41,6 +42,14 @@ const DEFAULT_SCREEN_HEIGHT: u32 = 1080; /// Default mouse event throttle interval (16ms ≈ 60Hz) const DEFAULT_MOUSE_THROTTLE_MS: u64 = 16; +/// Limit work retained for unauthenticated peers. +const MAX_CONNECTIONS: usize = 8; +const MAX_UNAUTHENTICATED_PACKET_LENGTH: usize = 256 * 1024; +const MAX_AUTHENTICATED_PACKET_LENGTH: usize = 8 * 1024 * 1024; +const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(15); +const MAX_PASSWORD_ATTEMPTS: u8 = 5; +const CONNECTION_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(3); + /// Advertised RustDesk version for client compatibility. const RUSTDESK_COMPAT_VERSION: &str = "1.4.5"; // Advertised platform for RustDesk clients. This affects which UI options are shown. @@ -140,10 +149,8 @@ pub struct Connection { /// Screen dimensions for mouse coordinate conversion screen_width: u32, screen_height: u32, - /// Message sender to connection handler - tx: mpsc::UnboundedSender, /// Shutdown signal - shutdown_tx: broadcast::Sender<()>, + shutdown_tx: watch::Sender, /// Video streaming task handle video_task: Option>, /// Audio streaming task handle @@ -152,14 +159,10 @@ pub struct Connection { session_key: Option, /// Encryption enabled flag encryption_enabled: bool, - /// Encryption sequence number (for nonce generation) - enc_seqnum: u64, /// Decryption sequence number (for nonce generation) dec_seqnum: u64, /// Negotiated video codec (after client capability exchange) negotiated_codec: Option, - /// Video frame sender for restarting video after codec switch - video_frame_tx: Option>, /// Input event throttler to prevent HID device EAGAIN errors input_throttler: InputThrottler, /// Last measured round-trip delay in milliseconds (for TestDelay responses) @@ -170,21 +173,18 @@ pub struct Connection { last_caps_lock: bool, /// Whether relative mouse mode is currently active for this connection relative_mouse_active: bool, + /// Latest throttled move; flushed on the next HID interval. + pending_mouse_move: Option, /// Server-configured RustDesk video codec. configured_codec: VideoEncoderType, + /// Failed authentication attempts on this TCP session. + password_attempts: u8, } -/// Messages sent to connection handler -#[derive(Debug)] -pub enum ConnectionMessage { - /// Send video frame - VideoFrame(Bytes), - /// Send audio frame - AudioFrame(Bytes), - /// Send cursor data - CursorData(Bytes), - /// Close connection - Close, +enum WriterCommand { + SetSessionKey(secretbox::Key), + Send { data: Bytes, encrypt: bool }, + Shutdown, } /// Messages received from client @@ -214,9 +214,8 @@ impl Connection { hid: Option>, audio: Option>, video_manager: Option>, - ) -> (Self, mpsc::UnboundedReceiver) { - let (tx, rx) = mpsc::unbounded_channel(); - let (shutdown_tx, _) = broadcast::channel(1); + ) -> Self { + let (shutdown_tx, _) = watch::channel(false); // Generate fresh Curve25519 keypair for this connection // This is used for encrypting the symmetric key exchange @@ -227,7 +226,7 @@ impl Connection { super::config::RustDeskCodec::H265 => VideoEncoderType::H265, }; - let conn = Self { + Self { id, device_id: config.device_id.clone(), peer_id: String::new(), @@ -244,25 +243,22 @@ impl Connection { video_manager, screen_width: DEFAULT_SCREEN_WIDTH, screen_height: DEFAULT_SCREEN_HEIGHT, - tx, shutdown_tx, video_task: None, audio_task: None, session_key: None, encryption_enabled: false, - enc_seqnum: 0, dec_seqnum: 0, negotiated_codec: None, - video_frame_tx: None, input_throttler: InputThrottler::new(), last_delay: 0, last_test_delay_sent: None, last_caps_lock: false, relative_mouse_active: false, + pending_mouse_move: None, configured_codec, - }; - - (conn, rx) + password_attempts: 0, + } } /// Get connection ID @@ -280,9 +276,8 @@ impl Connection { &self.peer_id } - /// Get message sender - pub fn sender(&self) -> mpsc::UnboundedSender { - self.tx.clone() + pub fn shutdown_sender(&self) -> watch::Sender { + self.shutdown_tx.clone() } /// Handle an incoming TCP connection @@ -292,11 +287,24 @@ impl Connection { peer_addr: SocketAddr, ) -> anyhow::Result<()> { info!("New connection from {}", peer_addr); + stream.set_nodelay(true)?; *self.state.write() = ConnectionState::Handshaking; let (mut reader, writer) = stream.into_split(); - let writer = Arc::new(Mutex::new(writer)); let mut shutdown_rx = self.shutdown_tx.subscribe(); + if *shutdown_rx.borrow() { + *self.state.write() = ConnectionState::Closed; + return Ok(()); + } + + // Keep socket writes out of the input loop. A congested video path must + // never prevent us from reading keyboard and mouse events. + let (control_tx, control_rx) = mpsc::channel::(32); + let (video_tx, video_rx) = mpsc::channel::(1); + let (audio_tx, audio_rx) = mpsc::channel::(8); + let mut writer_task = tokio::spawn(run_connection_writer( + writer, control_rx, video_rx, audio_rx, + )); match self.mode { ConnectionMode::Secure => { @@ -306,7 +314,7 @@ impl Connection { .write_to_bytes() .map_err(|e| anyhow::anyhow!("Failed to encode SignedId: {}", e))?; debug!("Sending SignedId with device_id={}", self.device_id); - self.send_framed_arc(&writer, &signed_id_bytes).await?; + self.send_framed(&control_tx, &signed_id_bytes).await?; } ConnectionMode::DirectIp => { // Standard RustDesk direct-IP clients do not perform the signed-ID handshake. @@ -316,34 +324,36 @@ impl Connection { .write_to_bytes() .map_err(|e| anyhow::anyhow!("Failed to encode Hash: {}", e))?; debug!("Sending password challenge for direct IP connection"); - self.send_framed_arc(&writer, &hash_bytes).await?; + self.send_framed(&control_tx, &hash_bytes).await?; } } - // Channel for receiving video frames to send (bounded to provide backpressure) - let (video_tx, mut video_rx) = mpsc::channel::(4); let mut video_streaming = false; - - // Channel for receiving audio frames to send (bounded to provide backpressure) - let (audio_tx, mut audio_rx) = mpsc::channel::(8); let mut audio_streaming = false; // Timer for sending TestDelay to measure round-trip latency // RustDesk clients display this delay information let mut test_delay_interval = tokio::time::interval(Duration::from_secs(1)); test_delay_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut mouse_flush_interval = + tokio::time::interval(Duration::from_millis(DEFAULT_MOUSE_THROTTLE_MS)); + mouse_flush_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - // Pre-allocated buffer for framing (reused across sends to reduce allocations) - // Typical H264 frame is 10-100KB, pre-allocate 128KB - let mut frame_buf = BytesMut::with_capacity(128 * 1024); + let handshake_deadline = tokio::time::sleep(HANDSHAKE_TIMEOUT); + tokio::pin!(handshake_deadline); loop { + let packet_limit = if self.state() == ConnectionState::Active { + MAX_AUTHENTICATED_PACKET_LENGTH + } else { + MAX_UNAUTHENTICATED_PACKET_LENGTH + }; tokio::select! { // Read framed message from client using RustDesk's variable-length encoding - result = read_frame(&mut reader) => { + result = read_frame_with_limit(&mut reader, packet_limit) => { match result { Ok(msg_buf) => { - if let Err(e) = self.handle_message_arc(&msg_buf, &writer, &video_tx, &mut video_streaming, &audio_tx, &mut audio_streaming).await { + if let Err(e) = self.handle_message_arc(&msg_buf, &control_tx, &video_tx, &mut video_streaming, &audio_tx, &mut audio_streaming).await { error!("Error handling message: {}", e); break; } @@ -363,61 +373,38 @@ impl Connection { } } - // Send video frames (encrypted if session key is set) - // Optimized path: inline encryption and use pre-allocated buffer - Some(frame_data) = video_rx.recv() => { - let send_result = if let Some(ref key) = self.session_key { - // Encrypt the frame - self.enc_seqnum += 1; - let nonce = Self::get_nonce(self.enc_seqnum); - let ciphertext = secretbox::seal(&frame_data, &nonce, key); - // Send using pre-allocated buffer - let mut w = writer.lock().await; - write_frame_buffered(&mut *w, &ciphertext, &mut frame_buf).await - } else { - // No encryption, send plain - let mut w = writer.lock().await; - write_frame_buffered(&mut *w, &frame_data, &mut frame_buf).await - }; - - if let Err(e) = send_result { - error!("Error sending video frame: {}", e); - break; - } - } - - // Send audio frames (encrypted if session key is set) - Some(frame_data) = audio_rx.recv() => { - let send_result = if let Some(ref key) = self.session_key { - // Encrypt the frame - self.enc_seqnum += 1; - let nonce = Self::get_nonce(self.enc_seqnum); - let ciphertext = secretbox::seal(&frame_data, &nonce, key); - let mut w = writer.lock().await; - write_frame_buffered(&mut *w, &ciphertext, &mut frame_buf).await - } else { - // No encryption, send plain - let mut w = writer.lock().await; - write_frame_buffered(&mut *w, &frame_data, &mut frame_buf).await - }; - - if let Err(e) = send_result { - error!("Error sending audio frame: {}", e); - break; - } - } - // Send TestDelay periodically to measure latency _ = test_delay_interval.tick() => { if self.state() == ConnectionState::Active && self.last_test_delay_sent.is_none() { - if let Err(e) = self.send_test_delay(&writer).await { + if let Err(e) = self.send_test_delay(&control_tx).await { warn!("Failed to send TestDelay: {}", e); } } } + _ = mouse_flush_interval.tick(), if self.pending_mouse_move.is_some() => { + if let Some(mouse_event) = self.pending_mouse_move.take() { + self.send_mouse_event_to_hid(&mouse_event).await; + self.input_throttler.mark_mouse_sent(); + } + } + + _ = &mut handshake_deadline, if self.state() != ConnectionState::Active => { + warn!("RustDesk handshake timed out for {}", peer_addr); + break; + } + + result = &mut writer_task => { + match result { + Ok(Ok(())) => debug!("RustDesk writer stopped for {}", peer_addr), + Ok(Err(error)) => warn!("RustDesk writer failed for {}: {}", peer_addr, error), + Err(error) => warn!("RustDesk writer task failed for {}: {}", peer_addr, error), + } + break; + } + // Shutdown signal - _ = shutdown_rx.recv() => { + _ = shutdown_rx.changed() => { info!("Connection shutdown requested"); break; } @@ -434,19 +421,33 @@ impl Connection { task.abort(); } + let _ = control_tx.try_send(WriterCommand::Shutdown); + if !writer_task.is_finished() { + match tokio::time::timeout(Duration::from_secs(1), &mut writer_task).await { + Ok(_) => {} + Err(_) => { + writer_task.abort(); + let _ = writer_task.await; + } + } + } + *self.state.write() = ConnectionState::Closed; Ok(()) } - /// Send framed message using Arc> with RustDesk's variable-length encoding - async fn send_framed_arc( + async fn send_framed( &self, - writer: &Arc>, + writer: &mpsc::Sender, data: &[u8], ) -> anyhow::Result<()> { - let mut w = writer.lock().await; - write_frame(&mut *w, data).await?; - Ok(()) + writer + .send(WriterCommand::Send { + data: Bytes::copy_from_slice(data), + encrypt: false, + }) + .await + .map_err(|_| anyhow::anyhow!("RustDesk writer is closed")) } /// Generate nonce from sequence number (RustDesk format) @@ -458,22 +459,18 @@ impl Connection { /// Send encrypted framed message if encryption is enabled /// RustDesk uses sequence-based nonce, NOT nonce prefix in message - async fn send_encrypted_arc( - &mut self, - writer: &Arc>, + async fn send_encrypted( + &self, + writer: &mpsc::Sender, data: &[u8], ) -> anyhow::Result<()> { - if let Some(ref key) = self.session_key { - // Increment encryption sequence number - self.enc_seqnum += 1; - let nonce = Self::get_nonce(self.enc_seqnum); - // Encrypt the message - RustDesk only sends ciphertext, no nonce prefix - let ciphertext = secretbox::seal(data, &nonce, key); - self.send_framed_arc(writer, &ciphertext).await - } else { - // No encryption, send plain - self.send_framed_arc(writer, data).await - } + writer + .send(WriterCommand::Send { + data: Bytes::copy_from_slice(data), + encrypt: self.session_key.is_some(), + }) + .await + .map_err(|_| anyhow::anyhow!("RustDesk writer is closed")) } /// Handle incoming message with Arc writer @@ -481,7 +478,7 @@ impl Connection { async fn handle_message_arc( &mut self, data: &[u8], - writer: &Arc>, + writer: &mpsc::Sender, video_tx: &mpsc::Sender, video_streaming: &mut bool, audio_tx: &mpsc::Sender, @@ -534,8 +531,6 @@ impl Connection { // Handle login and start video/audio streaming if successful if self.handle_login_request_arc(&lr, writer).await? { - // Store video_tx for potential codec switching - self.video_frame_tx = Some(video_tx.clone()); // Start video streaming if !*video_streaming { self.start_video_streaming(video_tx.clone()); @@ -605,7 +600,7 @@ impl Connection { async fn handle_login_request_arc( &mut self, lr: &LoginRequest, - writer: &Arc>, + writer: &mpsc::Sender, ) -> anyhow::Result { info!( "Login request from {} ({}), password_len={}", @@ -627,19 +622,23 @@ impl Connection { let response_bytes = error_response .write_to_bytes() .map_err(|e| anyhow::anyhow!("Failed to encode: {}", e))?; - self.send_encrypted_arc(writer, &response_bytes).await?; + self.send_encrypted(writer, &response_bytes).await?; // Don't close connection - wait for retry with password return Ok(false); } // Verify the password if !self.verify_password(&lr.password) { + self.password_attempts = self.password_attempts.saturating_add(1); warn!("Wrong password from {}", lr.my_id); let error_response = self.create_login_error_response("Wrong Password"); let response_bytes = error_response .write_to_bytes() .map_err(|e| anyhow::anyhow!("Failed to encode: {}", e))?; - self.send_encrypted_arc(writer, &response_bytes).await?; + self.send_encrypted(writer, &response_bytes).await?; + if self.password_attempts >= MAX_PASSWORD_ATTEMPTS { + anyhow::bail!("Too many failed RustDesk password attempts"); + } // Don't close connection - wait for retry with correct password return Ok(false); } @@ -647,6 +646,7 @@ impl Connection { // Password valid or no password required info!("Login successful for {}", lr.my_id); + self.password_attempts = 0; *self.state.write() = ConnectionState::Active; // Select the best available video codec @@ -659,7 +659,7 @@ impl Connection { let response_bytes = response .write_to_bytes() .map_err(|e| anyhow::anyhow!("Failed to encode: {}", e))?; - self.send_encrypted_arc(writer, &response_bytes).await?; + self.send_encrypted(writer, &response_bytes).await?; Ok(true) } @@ -702,7 +702,7 @@ impl Connection { async fn handle_misc_arc( &mut self, misc: &Misc, - _writer: &Arc>, + _writer: &mpsc::Sender, ) -> anyhow::Result<()> { match &misc.union { Some(misc::Union::SwitchDisplay(sd)) => { @@ -890,7 +890,7 @@ impl Connection { async fn handle_peer_public_key( &mut self, pk: &PublicKey, - writer: &Arc>, + writer: &mpsc::Sender, ) -> anyhow::Result<()> { // RustDesk's PublicKey message has two parts: // - asymmetric_value: The peer's temporary Curve25519 public key (32 bytes) @@ -913,6 +913,10 @@ impl Connection { ) { Ok(session_key) => { info!("Session key negotiated successfully"); + writer + .send(WriterCommand::SetSessionKey(session_key.clone())) + .await + .map_err(|_| anyhow::anyhow!("RustDesk writer is closed"))?; self.session_key = Some(session_key); self.encryption_enabled = true; } @@ -950,7 +954,7 @@ impl Connection { "Sending Hash message for password authentication (encrypted={})", self.encryption_enabled ); - self.send_encrypted_arc(writer, &hash_bytes).await?; + self.send_encrypted(writer, &hash_bytes).await?; Ok(()) } @@ -963,7 +967,7 @@ impl Connection { async fn handle_signed_id( &mut self, si: &SignedId, - writer: &Arc>, + writer: &mpsc::Sender, ) -> anyhow::Result<()> { // The SignedId contains a signed IdPk message // Try to parse the IdPk from the signed data @@ -1005,7 +1009,7 @@ impl Connection { let signed_id_bytes = signed_id_msg .write_to_bytes() .map_err(|e| anyhow::anyhow!("Failed to encode: {}", e))?; - self.send_framed_arc(writer, &signed_id_bytes).await?; + self.send_framed(writer, &signed_id_bytes).await?; Ok(()) } @@ -1033,7 +1037,9 @@ impl Connection { ); // Try comparison with double hash - if provided == expected_hash.as_slice() { + if provided.len() == expected_hash.len() + && sodiumoxide::utils::memcmp(provided, expected_hash.as_slice()) + { debug!("Password verified with double hash"); return true; } @@ -1041,7 +1047,10 @@ impl Connection { // Keep the legacy single-hash fallback only inside the encrypted ID-service path. // It has no per-connection challenge and must not be accepted on direct IP access. let expected_hash_single = crypto::hash_password(&self.password, &self.password_salt); - if self.mode == ConnectionMode::Secure && provided == expected_hash_single.as_slice() { + if self.mode == ConnectionMode::Secure + && provided.len() == expected_hash_single.len() + && sodiumoxide::utils::memcmp(provided, expected_hash_single.as_slice()) + { debug!("Password verified with single hash"); return true; } @@ -1058,7 +1067,7 @@ impl Connection { } /// Create login response with dynamically detected encoder capabilities - async fn create_login_response(&self, success: bool) -> HbbMessage { + async fn create_login_response(&mut self, success: bool) -> HbbMessage { if success { // Dynamically detect available encoders let registry = EncoderRegistry::global(); @@ -1089,6 +1098,10 @@ impl Connection { display_height = height; } } + // Use the same geometry for both the advertised display and HID + // absolute-coordinate conversion. + self.screen_width = display_width.max(1); + self.screen_height = display_height.max(1); let mut display_info = DisplayInfo::new(); display_info.x = 0; @@ -1171,7 +1184,7 @@ impl Connection { async fn handle_test_delay( &mut self, td: &TestDelay, - writer: &Arc>, + writer: &mpsc::Sender, ) -> anyhow::Result<()> { if td.from_client { // Client initiated the delay test, respond with the same time @@ -1187,7 +1200,7 @@ impl Connection { let data = response .write_to_bytes() .map_err(|e| anyhow::anyhow!("Failed to encode: {}", e))?; - self.send_encrypted_arc(writer, &data).await?; + self.send_encrypted(writer, &data).await?; debug!( "TestDelay response sent: time={}, last_delay={}ms", @@ -1214,7 +1227,10 @@ impl Connection { /// The client will echo this back, allowing us to calculate RTT. /// The measured delay is then included in future TestDelay messages /// for the client to display. - async fn send_test_delay(&mut self, writer: &Arc>) -> anyhow::Result<()> { + async fn send_test_delay( + &mut self, + writer: &mpsc::Sender, + ) -> anyhow::Result<()> { // Get current time in milliseconds since epoch let time_ms = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -1233,7 +1249,7 @@ impl Connection { let data = msg .write_to_bytes() .map_err(|e| anyhow::anyhow!("Failed to encode: {}", e))?; - self.send_encrypted_arc(writer, &data).await?; + self.send_encrypted(writer, &data).await?; // Record when we sent this, so we can calculate RTT when client echoes back self.last_test_delay_sent = Some(Instant::now()); @@ -1332,46 +1348,131 @@ impl Connection { // For pure move events, apply throttling if is_pure_move && !self.input_throttler.should_send_mouse_move() { - // Skip this move event to prevent HID EAGAIN + // Coalesce moves instead of losing the final pointer position. + self.pending_mouse_move = Some(me.clone()); return Ok(()); } + // Preserve input ordering when a button/scroll event follows a move + // that was waiting for the next HID poll interval. + if !is_pure_move { + if let Some(pending) = self.pending_mouse_move.take() { + self.send_mouse_event_to_hid(&pending).await; + } + } + debug!("Mouse event: x={}, y={}, mask={}", me.x, me.y, me.mask); - // Convert RustDesk mouse event to One-KVM mouse events - let mouse_events = convert_mouse_event(me, self.screen_width, self.screen_height); - - // Send to HID controller if available - if let Some(ref hid) = self.hid { - for event in mouse_events { - if let Err(e) = hid.send_mouse(event).await { - warn!("Failed to send mouse event: {}", e); - } - } + self.send_mouse_event_to_hid(me).await; + if self.hid.is_some() { // Mark that we sent a mouse event (for non-move events) if !is_pure_move { self.input_throttler.mark_mouse_sent(); } - } else { - debug!("HID controller not available, skipping mouse event"); } Ok(()) } + async fn send_mouse_event_to_hid(&self, event: &MouseEvent) { + let mouse_events = convert_mouse_event(event, self.screen_width, self.screen_height); + if let Some(ref hid) = self.hid { + for mouse_event in mouse_events { + if let Err(error) = hid.send_mouse(mouse_event).await { + warn!("Failed to send mouse event: {}", error); + } + } + } else { + debug!("HID controller not available, skipping mouse event"); + } + } + /// Close the connection pub fn close(&self) { - let _ = self.shutdown_tx.send(()); + self.shutdown_tx.send_replace(true); *self.state.write() = ConnectionState::Closed; } } +async fn run_connection_writer( + mut writer: OwnedWriteHalf, + mut control_rx: mpsc::Receiver, + mut video_rx: mpsc::Receiver, + mut audio_rx: mpsc::Receiver, +) -> anyhow::Result<()> { + let mut session_key: Option = None; + let mut sequence = 0u64; + + loop { + let (data, encrypt) = tokio::select! { + command = control_rx.recv() => { + match command { + Some(WriterCommand::SetSessionKey(key)) => { + session_key = Some(key); + continue; + } + Some(WriterCommand::Send { data, encrypt }) => (data, encrypt), + Some(WriterCommand::Shutdown) | None => break, + } + } + frame = audio_rx.recv() => { + match frame { + Some(data) => (data, session_key.is_some()), + None => continue, + } + } + frame = video_rx.recv() => { + match frame { + Some(data) => (data, session_key.is_some()), + None => continue, + } + } + }; + + let encrypted; + let payload = if encrypt { + if let Some(key) = session_key.as_ref() { + sequence = sequence.wrapping_add(1); + let nonce = Connection::get_nonce(sequence); + encrypted = secretbox::seal(&data, &nonce, key); + encrypted.as_slice() + } else { + data.as_ref() + } + } else { + data.as_ref() + }; + + tokio::time::timeout( + Duration::from_secs(10), + write_frame_vectored(&mut writer, payload), + ) + .await + .map_err(|_| anyhow::anyhow!("RustDesk socket write timed out"))??; + } + + Ok(()) +} + /// Lightweight connection info for tracking active connections pub struct ConnectionInfo { /// Connection ID pub id: u32, /// Connection state (shared with Connection) pub state: Arc>, + shutdown_tx: watch::Sender, + abort_handle: Option, +} + +struct ConnectionCleanup { + id: u32, + connections: Arc>>, +} + +impl Drop for ConnectionCleanup { + fn drop(&mut self) { + self.connections.write().remove(&self.id); + } } impl ConnectionInfo { @@ -1384,7 +1485,7 @@ impl ConnectionInfo { /// Connection manager pub struct ConnectionManager { /// Active connection info - connections: Arc>>>>, + connections: Arc>>, /// Next connection ID next_id: Arc>, /// Configuration @@ -1405,7 +1506,7 @@ impl ConnectionManager { /// Create a new connection manager pub fn new(config: RustDeskConfig) -> Self { Self { - connections: Arc::new(RwLock::new(Vec::new())), + connections: Arc::new(RwLock::new(HashMap::new())), next_id: Arc::new(RwLock::new(1)), config: Arc::new(RwLock::new(config)), keypair: Arc::new(RwLock::new(None)), @@ -1483,12 +1584,29 @@ impl ConnectionManager { .await } + pub async fn accept_listener_connection( + &self, + stream: TcpStream, + peer_addr: SocketAddr, + ) -> anyhow::Result { + let mode = match self.config.read().mode { + super::config::RustDeskMode::Id => ConnectionMode::Secure, + super::config::RustDeskMode::DirectIp => ConnectionMode::DirectIp, + }; + self.accept_connection_with_mode(stream, peer_addr, mode) + .await + } + async fn accept_connection_with_mode( &self, stream: TcpStream, peer_addr: SocketAddr, mode: ConnectionMode, ) -> anyhow::Result { + if self.connection_count() >= MAX_CONNECTIONS { + anyhow::bail!("RustDesk connection limit ({MAX_CONNECTIONS}) reached"); + } + let id = { let mut next = self.next_id.write(); let id = *next; @@ -1504,7 +1622,7 @@ impl ConnectionManager { let hid = self.hid.read().clone(); let audio = self.audio.read().clone(); let video_manager = self.video_manager.read().clone(); - let (mut conn, _rx) = Connection::new( + let mut conn = Connection::new( id, &config, mode, @@ -1516,16 +1634,34 @@ impl ConnectionManager { // Track connection state for external access let state = conn.state.clone(); - self.connections - .write() - .push(Arc::new(RwLock::new(ConnectionInfo { id, state }))); + let shutdown_tx = conn.shutdown_sender(); + { + let mut connections = self.connections.write(); + if connections.len() >= MAX_CONNECTIONS { + anyhow::bail!("RustDesk connection limit ({MAX_CONNECTIONS}) reached"); + } + connections.insert( + id, + ConnectionInfo { + id, + state, + shutdown_tx, + abort_handle: None, + }, + ); + } // Spawn connection handler - Connection is moved, not locked - tokio::spawn(async move { + let connections = self.connections.clone(); + let task = tokio::spawn(async move { + let _cleanup = ConnectionCleanup { id, connections }; if let Err(e) = conn.handle_tcp(stream, peer_addr).await { error!("Connection {} error: {}", id, e); } }); + if let Some(connection) = self.connections.write().get_mut(&id) { + connection.abort_handle = Some(task.abort_handle()); + } Ok(id) } @@ -1535,11 +1671,47 @@ impl ConnectionManager { self.connections.read().len() } - /// Mark all connections as closed (actual connection tasks will detect this) - pub fn close_all(&self) { - let connections = self.connections.read(); - for conn_info in connections.iter() { - *conn_info.read().state.write() = ConnectionState::Closed; + /// Cancel all sessions and wait briefly for their TCP tasks to exit. + pub async fn close_all(&self) { + let senders = self + .connections + .read() + .values() + .map(|connection| connection.shutdown_tx.clone()) + .collect::>(); + for sender in senders { + sender.send_replace(true); + } + + let wait_for_empty = async { + while !self.connections.read().is_empty() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }; + if tokio::time::timeout(CONNECTION_SHUTDOWN_TIMEOUT, wait_for_empty) + .await + .is_err() + { + warn!( + "Timed out waiting for {} RustDesk connection(s) to close", + self.connection_count() + ); + let abort_handles = self + .connections + .read() + .values() + .filter_map(|connection| connection.abort_handle.clone()) + .collect::>(); + for handle in abort_handles { + handle.abort(); + } + + let abort_wait = async { + while !self.connections.read().is_empty() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }; + let _ = tokio::time::timeout(Duration::from_secs(1), abort_wait).await; } } } @@ -1557,7 +1729,7 @@ async fn run_video_streaming( video_manager: Arc, video_tx: mpsc::Sender, state: Arc>, - shutdown_tx: broadcast::Sender<()>, + shutdown_tx: watch::Sender, negotiated_codec: VideoEncoderType, ) -> anyhow::Result<()> { use crate::video::codec::VideoCodecType; @@ -1590,6 +1762,9 @@ async fn run_video_streaming( let mut video_adapter = VideoFrameAdapter::new(codec); let mut shutdown_rx = shutdown_tx.subscribe(); + if *shutdown_rx.borrow() { + return Ok(()); + } let mut encoded_count: u64 = 0; let mut last_log_time = Instant::now(); let mut waiting_for_keyframe = true; @@ -1652,7 +1827,7 @@ async fn run_video_streaming( tokio::select! { biased; - _ = shutdown_rx.recv() => { + _ = shutdown_rx.changed() => { debug!("Shutdown signal received, stopping video for connection {}", conn_id); break 'subscribe_loop; } @@ -1699,10 +1874,26 @@ async fn run_video_streaming( frame.pts_ms as u64, ); - // Send to connection (backpressure instead of dropping) - if video_tx.send(msg_bytes).await.is_err() { - debug!("Video channel closed for connection {}", conn_id); - break 'subscribe_loop; + // Never queue a run of stale frames behind a slow socket. + // If the one-frame queue is full, wait for a fresh keyframe + // before resuming so the decoder cannot receive a broken GOP. + match video_tx.try_send(msg_bytes) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(_)) => { + waiting_for_keyframe = true; + let now = Instant::now(); + if now.duration_since(last_keyframe_request) >= Duration::from_millis(200) { + if let Err(error) = video_manager.request_keyframe().await { + debug!("Failed to request recovery keyframe for connection {}: {}", conn_id, error); + } + last_keyframe_request = now; + } + continue; + } + Err(mpsc::error::TrySendError::Closed(_)) => { + debug!("Video channel closed for connection {}", conn_id); + break 'subscribe_loop; + } } last_sequence = Some(frame.sequence); @@ -1738,12 +1929,15 @@ async fn run_audio_streaming( audio_controller: Arc, audio_tx: mpsc::Sender, state: Arc>, - shutdown_tx: broadcast::Sender<()>, + shutdown_tx: watch::Sender, ) -> anyhow::Result<()> { // Audio format: 48kHz stereo Opus let mut audio_adapter = AudioFrameAdapter::new(48000, 2); let mut shutdown_rx = shutdown_tx.subscribe(); + if *shutdown_rx.borrow() { + return Ok(()); + } let mut frame_count: u64 = 0; let mut last_log_time = Instant::now(); @@ -1798,7 +1992,7 @@ async fn run_audio_streaming( tokio::select! { biased; - _ = shutdown_rx.recv() => { + _ = shutdown_rx.changed() => { debug!("Shutdown signal received, stopping audio for connection {}", conn_id); break 'subscribe_loop; } @@ -1820,10 +2014,14 @@ async fn run_audio_streaming( // Convert OpusFrame to RustDesk AudioFrame message let msg_bytes = audio_adapter.encode_opus_bytes(&opus_frame.data); - // Send to connection (blocks if channel is full, providing backpressure) - if audio_tx.send(msg_bytes).await.is_err() { - debug!("Audio channel closed for connection {}", conn_id); - break 'subscribe_loop; + // Audio is real-time data; dropping an old packet is preferable + // to accumulating seconds of latency. + match audio_tx.try_send(msg_bytes) { + Ok(()) | Err(mpsc::error::TrySendError::Full(_)) => {} + Err(mpsc::error::TrySendError::Closed(_)) => { + debug!("Audio channel closed for connection {}", conn_id); + break 'subscribe_loop; + } } frame_count += 1; @@ -1860,7 +2058,7 @@ mod tests { device_password: "fixed-password".to_string(), ..Default::default() }; - let (connection, _rx) = Connection::new( + let connection = Connection::new( 1, &config, mode, @@ -1892,7 +2090,7 @@ mod tests { async move { server_connection.handle_tcp(server_stream, peer_addr).await }, ); - let bytes = read_frame(&mut client_stream) + let bytes = read_frame_with_limit(&mut client_stream, MAX_UNAUTHENTICATED_PACKET_LENGTH) .await .expect("client should receive the first frame"); drop(client_stream); @@ -1944,4 +2142,80 @@ mod tests { let message = first_server_message(ConnectionMode::Secure).await; assert!(matches!(message.union, Some(message::Union::SignedId(_)))); } + + #[tokio::test] + async fn connection_manager_close_all_cancels_and_removes_sessions() { + let config = RustDeskConfig { + enabled: true, + mode: super::super::config::RustDeskMode::DirectIp, + ..Default::default() + }; + let manager = ConnectionManager::new(config); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let client = tokio::spawn(async move { TcpStream::connect(address).await.unwrap() }); + let (server, peer) = listener.accept().await.unwrap(); + let _client = client.await.unwrap(); + + manager + .accept_direct_connection(server, peer) + .await + .unwrap(); + assert_eq!(manager.connection_count(), 1); + manager.close_all().await; + assert_eq!(manager.connection_count(), 0); + } + + #[tokio::test] + async fn connection_manager_enforces_session_limit() { + let config = RustDeskConfig { + enabled: true, + mode: super::super::config::RustDeskMode::DirectIp, + ..Default::default() + }; + let manager = ConnectionManager::new(config); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let mut clients = Vec::new(); + + for _ in 0..MAX_CONNECTIONS { + let client = tokio::spawn(async move { TcpStream::connect(address).await.unwrap() }); + let (server, peer) = listener.accept().await.unwrap(); + clients.push(client.await.unwrap()); + manager + .accept_direct_connection(server, peer) + .await + .unwrap(); + } + let extra_client = tokio::spawn(async move { TcpStream::connect(address).await.unwrap() }); + let (extra_server, extra_peer) = listener.accept().await.unwrap(); + clients.push(extra_client.await.unwrap()); + assert!(manager + .accept_direct_connection(extra_server, extra_peer) + .await + .is_err()); + + manager.close_all().await; + assert_eq!(manager.connection_count(), 0); + } + + #[tokio::test] + async fn password_attempts_are_bounded_per_session() { + let mut connection = connection(ConnectionMode::DirectIp); + let (writer, _reader) = mpsc::channel(32); + let mut login = LoginRequest::new(); + login.my_id = "attacker".to_string(); + login.password = vec![0u8; 32].into(); + + for _ in 1..MAX_PASSWORD_ATTEMPTS { + assert!(!connection + .handle_login_request_arc(&login, &writer) + .await + .unwrap()); + } + assert!(connection + .handle_login_request_arc(&login, &writer) + .await + .is_err()); + } } diff --git a/src/rustdesk/frame_adapters.rs b/src/rustdesk/frame_adapters.rs index 93dc6f67..de57bc4c 100644 --- a/src/rustdesk/frame_adapters.rs +++ b/src/rustdesk/frame_adapters.rs @@ -91,6 +91,12 @@ impl VideoFrameAdapter { return data; } + // Parameter sets are relevant only on random-access frames. Avoid a + // full Annex-B/AVCC scan on every delta frame in every client session. + if !is_keyframe { + return data; + } + let (sps, pps) = crate::video::codec::h264_bitstream::extract_sps_pps(&data); let mut has_sps = false; let mut has_pps = false; diff --git a/src/rustdesk/mod.rs b/src/rustdesk/mod.rs index ff70b042..3e443641 100644 --- a/src/rustdesk/mod.rs +++ b/src/rustdesk/mod.rs @@ -17,7 +17,7 @@ use std::time::Duration; use parking_lot::RwLock; use protobuf::Message; use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::broadcast; +use tokio::sync::{broadcast, Semaphore}; use tokio::task::JoinHandle; use tracing::{debug, error, info, warn}; @@ -33,6 +33,7 @@ use self::rendezvous::{AddrMangle, RendezvousMediator, RendezvousStatus}; const RELAY_CONNECT_TIMEOUT_MS: u64 = 10_000; const SERVICE_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(2); +const MAX_PENDING_CONNECTION_ATTEMPTS: usize = 8; #[derive(Debug, Clone, PartialEq)] pub enum ServiceStatus { @@ -59,6 +60,7 @@ pub struct RustDeskService { rendezvous: Arc>>>, rendezvous_handle: Arc>>>, tcp_listener_handle: Arc>>>>, + listener_start_lock: Arc>, listen_port: Arc>, connection_manager: Arc, video_manager: Arc, @@ -84,6 +86,7 @@ impl RustDeskService { rendezvous: Arc::new(RwLock::new(None)), rendezvous_handle: Arc::new(RwLock::new(None)), tcp_listener_handle: Arc::new(RwLock::new(None)), + listener_start_lock: Arc::new(tokio::sync::Mutex::new(())), listen_port: Arc::new(RwLock::new(direct_access_port)), connection_manager, video_manager, @@ -130,7 +133,7 @@ impl RustDeskService { self.status() == ServiceStatus::Running } - pub async fn start(&self) -> anyhow::Result<()> { + pub async fn start(self: &Arc) -> anyhow::Result<()> { let config = self.config.read().clone(); if !config.enabled { @@ -168,14 +171,16 @@ impl RustDeskService { .set_video_manager(self.video_manager.clone()); if config.mode == RustDeskMode::DirectIp { - let (tcp_handles, listen_port) = match self.start_tcp_listener_with_port().await { + let listen_port = match self + .ensure_tcp_listener(config.direct_access_port, false) + .await + { Ok(result) => result, Err(err) => { *self.status.write() = ServiceStatus::Error(err.to_string()); return Err(err); } }; - *self.tcp_listener_handle.write() = Some(tcp_handles); *self.listen_port.write() = listen_port; *self.status.write() = ServiceStatus::Running; return Ok(()); @@ -193,14 +198,23 @@ impl RustDeskService { let connection_manager = self.connection_manager.clone(); let service_config = self.config.clone(); + let connection_attempts = Arc::new(Semaphore::new(MAX_PENDING_CONNECTION_ATTEMPTS)); mediator.set_punch_callback(Arc::new({ let connection_manager = connection_manager.clone(); let service_config = service_config.clone(); + let connection_attempts = connection_attempts.clone(); move |peer_addr, rendezvous_addr, relay_server, uuid, socket_addr, device_id| { let conn_mgr = connection_manager.clone(); let config = service_config.clone(); + let attempts = connection_attempts.clone(); tokio::spawn(async move { + let Ok(_permit) = attempts.try_acquire_owned() else { + warn!( + "Dropping RustDesk punch request: too many pending connection attempts" + ); + return; + }; if let Some(addr) = peer_addr { info!("Attempting P2P direct connection to {}", addr); match punch::try_direct_connection(addr).await { @@ -238,10 +252,18 @@ impl RustDeskService { mediator.set_relay_callback(Arc::new({ let connection_manager = connection_manager.clone(); let service_config = service_config.clone(); + let connection_attempts = connection_attempts.clone(); move |rendezvous_addr, relay_server, uuid, socket_addr, device_id| { let conn_mgr = connection_manager.clone(); let config = service_config.clone(); + let attempts = connection_attempts.clone(); tokio::spawn(async move { + let Ok(_permit) = attempts.try_acquire_owned() else { + warn!( + "Dropping RustDesk relay request: too many pending connection attempts" + ); + return; + }; let relay_key = rustdesk_relay_key(&config); if let Err(e) = handle_relay_request( &rendezvous_addr, @@ -260,19 +282,39 @@ impl RustDeskService { } })); - let connection_manager2 = self.connection_manager.clone(); + let weak_service = Arc::downgrade(self); + let intranet_attempts = connection_attempts.clone(); mediator.set_intranet_callback(Arc::new( - move |rendezvous_addr, peer_socket_addr, local_addr, relay_server, device_id| { - let conn_mgr = connection_manager2.clone(); - + move |rendezvous_addr, peer_socket_addr, local_ip, relay_server, device_id| { + let weak_service = weak_service.clone(); + let attempts = intranet_attempts.clone(); tokio::spawn(async move { + let Ok(_permit) = attempts.try_acquire_owned() else { + warn!("Dropping RustDesk intranet request: too many pending connection attempts"); + return; + }; + let Some(service) = weak_service.upgrade() else { + return; + }; + let preferred_port = service.config.read().direct_access_port; + let listen_port = match service + .ensure_tcp_listener(preferred_port, true) + .await + { + Ok(port) => port, + Err(error) => { + error!("Failed to start on-demand RustDesk listener: {}", error); + return; + } + }; + let local_addr = SocketAddr::new(local_ip, listen_port); if let Err(e) = handle_intranet_request( &rendezvous_addr, &peer_socket_addr, local_addr, &relay_server, &device_id, - conn_mgr, + service.connection_manager.clone(), ) .await { @@ -304,9 +346,30 @@ impl RustDeskService { Ok(()) } - async fn start_tcp_listener_with_port(&self) -> anyhow::Result<(Vec>, u16)> { - let direct_access_port = self.config.read().direct_access_port; - let (listeners, listen_port) = self.bind_direct_listeners(direct_access_port)?; + async fn ensure_tcp_listener( + self: &Arc, + preferred_port: u16, + allow_ephemeral_fallback: bool, + ) -> anyhow::Result { + let _guard = self.listener_start_lock.lock().await; + if self.tcp_listener_handle.read().is_some() { + return Ok(*self.listen_port.read()); + } + if self.status() == ServiceStatus::Stopped { + anyhow::bail!("RustDesk service stopped before listener could start"); + } + + let (listeners, listen_port) = match self.bind_direct_listeners(preferred_port) { + Ok(result) => result, + Err(error) if allow_ephemeral_fallback => { + warn!( + "RustDesk port {} unavailable for on-demand listening: {}; using an ephemeral port", + preferred_port, error + ); + self.bind_direct_listeners(0)? + } + Err(error) => return Err(error), + }; *self.listen_port.write() = listen_port; @@ -326,15 +389,13 @@ impl RustDeskService { match result { Ok((stream, peer_addr)) => { info!("Accepted direct connection from {}", peer_addr); - let conn_mgr = conn_mgr.clone(); - tokio::spawn(async move { - if let Err(e) = conn_mgr.accept_direct_connection(stream, peer_addr).await { - error!("Failed to handle direct connection from {}: {}", peer_addr, e); - } - }); + if let Err(e) = conn_mgr.accept_listener_connection(stream, peer_addr).await { + warn!("Rejected direct connection from {}: {}", peer_addr, e); + } } Err(e) => { error!("TCP accept error: {}", e); + tokio::time::sleep(Duration::from_millis(100)).await; } } } @@ -348,7 +409,8 @@ impl RustDeskService { handles.push(handle); } - Ok((handles, listen_port)) + *self.tcp_listener_handle.write() = Some(handles); + Ok(listen_port) } fn bind_direct_listeners(&self, port: u16) -> anyhow::Result<(Vec, u16)> { @@ -382,8 +444,8 @@ impl RustDeskService { info!("Stopping RustDesk service"); let _ = self.shutdown_tx.send(()); - - self.connection_manager.close_all(); + let _listener_guard = self.listener_start_lock.lock().await; + *self.status.write() = ServiceStatus::Stopped; if let Some(mediator) = self.rendezvous.read().as_ref() { mediator.stop(); @@ -401,13 +463,15 @@ impl RustDeskService { } } + // No listener can admit a new session after this point. + self.connection_manager.close_all().await; + *self.rendezvous.write() = None; - *self.status.write() = ServiceStatus::Stopped; Ok(()) } - pub async fn restart(&self, config: RustDeskConfig) -> anyhow::Result<()> { + pub async fn restart(self: &Arc, config: RustDeskConfig) -> anyhow::Result<()> { self.stop().await?; self.update_config(config); self.start().await diff --git a/src/rustdesk/rendezvous.rs b/src/rustdesk/rendezvous.rs index 6147c535..ab5f919b 100644 --- a/src/rustdesk/rendezvous.rs +++ b/src/rustdesk/rendezvous.rs @@ -128,7 +128,7 @@ pub type RelayCallback = Arc, String) + S pub type PunchCallback = Arc, String, String, String, Vec, String) + Send + Sync>; -pub type IntranetCallback = Arc, SocketAddr, String, String) + Send + Sync>; +pub type IntranetCallback = Arc, IpAddr, String, String) + Send + Sync>; pub struct RendezvousMediator { config: Arc>, @@ -143,7 +143,6 @@ pub struct RendezvousMediator { relay_callback: Arc>>, punch_callback: Arc>>, intranet_callback: Arc>>, - listen_port: Arc>, shutdown_tx: broadcast::Sender<()>, } @@ -166,23 +165,10 @@ impl RendezvousMediator { relay_callback: Arc::new(RwLock::new(None)), punch_callback: Arc::new(RwLock::new(None)), intranet_callback: Arc::new(RwLock::new(None)), - listen_port: Arc::new(RwLock::new(21118)), shutdown_tx, } } - pub fn set_listen_port(&self, port: u16) { - let old_port = *self.listen_port.read(); - if old_port != port { - *self.listen_port.write() = port; - self.increment_serial(); - } - } - - pub fn listen_port(&self) -> u16 { - *self.listen_port.read() - } - pub fn increment_serial(&self) { let mut serial = self.serial.write(); *serial = serial.wrapping_add(1); @@ -430,7 +416,9 @@ impl RendezvousMediator { ) -> anyhow::Result<()> { let id = self.device_id(); - let local_addrs = get_local_addresses(); + let local_addrs = tokio::task::spawn_blocking(get_local_addresses) + .await + .map_err(|error| anyhow::anyhow!("Failed to inspect local addresses: {error}"))?; if local_addrs.is_empty() { debug!("No local addresses available for LocalAddr response"); return Ok(()); @@ -439,21 +427,18 @@ impl RendezvousMediator { let config = self.config.read().clone(); let rendezvous_addr = config.rendezvous_addr(); - let listen_port = self.listen_port(); - let local_ip = local_addrs[0]; - let local_sock_addr = SocketAddr::new(local_ip, listen_port); info!( - "FetchLocalAddr: calling intranet callback with local_addr={}, rendezvous={}", - local_sock_addr, rendezvous_addr + "FetchLocalAddr: requesting an on-demand listener for {}, rendezvous={}", + local_ip, rendezvous_addr ); if let Some(callback) = self.intranet_callback.read().as_ref() { callback( rendezvous_addr, peer_socket_addr.to_vec(), - local_sock_addr, + local_ip, relay_server.to_string(), id, ); diff --git a/src/web/handlers/config/rustdesk.rs b/src/web/handlers/config/rustdesk.rs index 2af7df96..6801a6c8 100644 --- a/src/web/handlers/config/rustdesk.rs +++ b/src/web/handlers/config/rustdesk.rs @@ -47,6 +47,9 @@ async fn current_status( config: RustDeskConfigResponse::from(&config), service_status: runtime.service_status, rendezvous_status: runtime.rendezvous_status, + connection_count: runtime.connection_count, + listening: runtime.listening, + listen_port: runtime.listen_port, } } @@ -86,6 +89,9 @@ pub struct RustDeskStatusResponse { pub config: RustDeskConfigResponse, pub service_status: String, pub rendezvous_status: Option, + pub connection_count: usize, + pub listening: bool, + pub listen_port: Option, } pub async fn get_rustdesk_config( diff --git a/web/src/api/config.ts b/web/src/api/config.ts index a6da99b4..f54b9628 100644 --- a/web/src/api/config.ts +++ b/web/src/api/config.ts @@ -236,6 +236,9 @@ export interface RustDeskStatusResponse { config: RustDeskConfigResponse service_status: string rendezvous_status: string | null + connection_count: number + listening: boolean + listen_port: number | null } export interface RustDeskConfigUpdate { diff --git a/web/src/i18n/en-US.ts b/web/src/i18n/en-US.ts index ce266394..1a8b5876 100644 --- a/web/src/i18n/en-US.ts +++ b/web/src/i18n/en-US.ts @@ -1079,8 +1079,6 @@ export default { modeDirectIpDesc: 'Listen on the device only, without connecting to an ID or relay server', directAccessPort: 'Direct Access Port', directAccessPortInvalid: 'The direct access port must be between 1 and 65535', - directAccessWarningTitle: 'Direct IP access is not end-to-end encrypted', - directAccessWarningDesc: 'Use it only on a trusted LAN or encrypted VPN. Do not expose this port directly to the internet.', rendezvousServer: 'ID Server', rendezvousServerPlaceholder: 'hbbs.example.com:21116', rendezvousServerRequired: 'Enter the RustDesk ID server', @@ -1088,7 +1086,6 @@ export default { relayServerPlaceholder: 'hbbr.example.com:21117', relayKey: 'Relay Key', codec: 'Codec', - deviceInfo: 'Device Info', deviceId: 'Device ID', devicePassword: 'Device Password', showPassword: 'Show Password', diff --git a/web/src/i18n/zh-CN.ts b/web/src/i18n/zh-CN.ts index a625f869..37b544dd 100644 --- a/web/src/i18n/zh-CN.ts +++ b/web/src/i18n/zh-CN.ts @@ -1078,8 +1078,6 @@ export default { modeDirectIpDesc: '仅监听设备端口,不连接 ID 或中继服务器', directAccessPort: '直连端口', directAccessPortInvalid: '直连端口必须在 1 到 65535 之间', - directAccessWarningTitle: 'IP 直连不提供端到端加密', - directAccessWarningDesc: '仅建议在可信局域网或加密 VPN 中使用,请勿将此端口直接暴露到公网。', rendezvousServer: 'ID 服务器', rendezvousServerPlaceholder: 'hbbs.example.com:21116', rendezvousServerRequired: '请填写 RustDesk ID 服务器', @@ -1087,7 +1085,6 @@ export default { relayServerPlaceholder: 'hbbr.example.com:21117', relayKey: '中继密钥', codec: '编码格式', - deviceInfo: '设备信息', deviceId: '设备 ID', devicePassword: '设备密码', showPassword: '显示密码', diff --git a/web/src/views/SettingsView.vue b/web/src/views/SettingsView.vue index 8f65c6f5..8562c038 100644 --- a/web/src/views/SettingsView.vue +++ b/web/src/views/SettingsView.vue @@ -5166,18 +5166,11 @@ watch(isWindows, () => {

{{ rustdeskValidationMessage }}

- - - {{ t('extensions.rustdesk.directAccessWarningTitle') }} - {{ t('extensions.rustdesk.directAccessWarningDesc') }} -
-

{{ t('extensions.rustdesk.deviceInfo') }}

-