GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2026-09-12 18:44:25 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor: 删除部分多余的代码和注释

Browse Source
This commit is contained in:
mofeng-git
2026-05-01 17:31:04 +08:00
parent 74035f8e12
commit d8e7de74a6
165 changed files with 2960 additions and 9917 deletions
Download Patch File Download Diff File Expand all files Collapse all files

174
src/rustdesk/mod.rs
View File

@@ -1,17 +1,4 @@
//! RustDesk Protocol Integration Module
//!
//! This module implements the RustDesk client protocol, enabling One-KVM devices
//! to be accessed via standard RustDesk clients through existing hbbs/hbbr servers.
//!
//! ## Architecture
//!
//! - `config`: Configuration types for RustDesk settings
//! - `protocol`: Protobuf message wrappers and serialization
//! - `crypto`: NaCl cryptography (key generation, encryption, signatures)
//! - `rendezvous`: Communication with hbbs rendezvous server
//! - `connection`: Client session handling
//! - `frame_adapters`: Video/audio frame conversion to RustDesk format
//! - `hid_adapter`: RustDesk HID events to One-KVM conversion
//! RustDesk peer protocol (hbbs / hbbr).
pub mod bytes_codec;
pub mod config;
@@ -44,19 +31,13 @@ use self::connection::ConnectionManager;
use self::protocol::{make_local_addr, make_relay_response, make_request_relay};
use self::rendezvous::{AddrMangle, RendezvousMediator, RendezvousStatus};
/// Relay connection timeout
const RELAY_CONNECT_TIMEOUT_MS: u64 = 10_000;
/// RustDesk service status
#[derive(Debug, Clone, PartialEq)]
pub enum ServiceStatus {
/// Service is stopped
Stopped,
/// Service is starting
Starting,
/// Service is running and registered with rendezvous server
Running,
/// Service encountered an error
Error(String),
}
@@ -71,15 +52,8 @@ impl std::fmt::Display for ServiceStatus {
}
}
/// Default port for direct TCP connections (same as RustDesk)
const DIRECT_LISTEN_PORT: u16 = 21118;
/// RustDesk Service
///
/// Manages the RustDesk protocol integration, including:
/// - Registration with hbbs rendezvous server
/// - Accepting connections from RustDesk clients
/// - Streaming video/audio and receiving HID input
pub struct RustDeskService {
config: Arc<RwLock<RustDeskConfig>>,
status: Arc<RwLock<ServiceStatus>>,
@@ -95,7 +69,6 @@ pub struct RustDeskService {
}
impl RustDeskService {
/// Create a new RustDesk service instance
pub fn new(
config: RustDeskConfig,
video_manager: Arc<VideoStreamManager>,
@@ -120,42 +93,34 @@ impl RustDeskService {
}
}
/// Get the port for direct TCP connections
pub fn listen_port(&self) -> u16 {
*self.listen_port.read()
}
/// Get current service status
pub fn status(&self) -> ServiceStatus {
self.status.read().clone()
}
/// Get current configuration
pub fn config(&self) -> RustDeskConfig {
self.config.read().clone()
}
/// Update configuration
pub fn update_config(&self, config: RustDeskConfig) {
*self.config.write() = config;
}
/// Get rendezvous status
pub fn rendezvous_status(&self) -> Option<RendezvousStatus> {
self.rendezvous.read().as_ref().map(|r| r.status())
}
/// Get device ID
pub fn device_id(&self) -> String {
self.config.read().device_id.clone()
}
/// Get connection count
pub fn connection_count(&self) -> usize {
self.connection_manager.connection_count()
}
/// Start the RustDesk service
pub async fn start(&self) -> anyhow::Result<()> {
let config = self.config.read().clone();
@@ -181,74 +146,44 @@ impl RustDeskService {
config.rendezvous_addr()
);
// Initialize crypto
if let Err(e) = crypto::init() {
error!("Failed to initialize crypto: {}", e);
*self.status.write() = ServiceStatus::Error(e.to_string());
return Err(e.into());
}
// Create and start rendezvous mediator with relay callback
let mediator = Arc::new(RendezvousMediator::new(config.clone()));
// Set the keypair on connection manager (Curve25519 for encryption)
let keypair = mediator.ensure_keypair();
self.connection_manager.set_keypair(keypair);
// Set the signing keypair on connection manager (Ed25519 for SignedId)
let signing_keypair = mediator.ensure_signing_keypair();
self.connection_manager.set_signing_keypair(signing_keypair);
// Set the HID controller on connection manager
self.connection_manager.set_hid(self.hid.clone());
// Set the audio controller on connection manager for audio streaming
self.connection_manager.set_audio(self.audio.clone());
// Set the video manager on connection manager for video streaming
self.connection_manager
.set_video_manager(self.video_manager.clone());
*self.rendezvous.write() = Some(mediator.clone());
// Start TCP listener BEFORE the rendezvous mediator to ensure port is set correctly
// This prevents race condition where mediator starts registration with wrong port
let (tcp_handles, listen_port) = self.start_tcp_listener_with_port().await?;
*self.tcp_listener_handle.write() = Some(tcp_handles);
// Set the listen port on mediator before starting the registration loop
mediator.set_listen_port(listen_port);
// Create relay request handler
let connection_manager = self.connection_manager.clone();
let video_manager = self.video_manager.clone();
let hid = self.hid.clone();
let audio = self.audio.clone();
let service_config = self.config.clone();
// Set the punch callback on the mediator (tries P2P first, then relay)
let connection_manager_punch = self.connection_manager.clone();
let video_manager_punch = self.video_manager.clone();
let hid_punch = self.hid.clone();
let audio_punch = self.audio.clone();
let service_config_punch = self.config.clone();
mediator.set_punch_callback(Arc::new(
mediator.set_punch_callback(Arc::new({
let connection_manager = connection_manager.clone();
let service_config = service_config.clone();
move |peer_addr, rendezvous_addr, relay_server, uuid, socket_addr, device_id| {
let conn_mgr = connection_manager_punch.clone();
let video = video_manager_punch.clone();
let hid = hid_punch.clone();
let audio = audio_punch.clone();
let config = service_config_punch.clone();
let conn_mgr = connection_manager.clone();
let config = service_config.clone();
tokio::spawn(async move {
// Get relay_key from config (no public server fallback)
let relay_key = {
let cfg = config.read();
cfg.relay_key.clone().unwrap_or_default()
};
// Try P2P direct connection first
if let Some(addr) = peer_addr {
info!("Attempting P2P direct connection to {}", addr);
match punch::try_direct_connection(addr).await {
@@ -265,7 +200,7 @@ impl RustDeskService {
}
}
// Fall back to relay
let relay_key = rustdesk_relay_key(&config);
if let Err(e) = handle_relay_request(
&rendezvous_addr,
&relay_server,
@@ -274,34 +209,23 @@ impl RustDeskService {
&device_id,
&relay_key,
conn_mgr,
video,
hid,
audio,
)
.await
{
error!("Failed to handle relay request: {}", e);
}
});
},
));
}
}));
// Set the relay callback on the mediator
mediator.set_relay_callback(Arc::new(
mediator.set_relay_callback(Arc::new({
let connection_manager = connection_manager.clone();
let service_config = service_config.clone();
move |rendezvous_addr, relay_server, uuid, socket_addr, device_id| {
let conn_mgr = connection_manager.clone();
let video = video_manager.clone();
let hid = hid.clone();
let audio = audio.clone();
let config = service_config.clone();
tokio::spawn(async move {
// Get relay_key from config (no public server fallback)
let relay_key = {
let cfg = config.read();
cfg.relay_key.clone().unwrap_or_default()
};
let relay_key = rustdesk_relay_key(&config);
if let Err(e) = handle_relay_request(
&rendezvous_addr,
&relay_server,
@@ -310,19 +234,15 @@ impl RustDeskService {
&device_id,
&relay_key,
conn_mgr,
video,
hid,
audio,
)
.await
{
error!("Failed to handle relay request: {}", e);
}
});
},
));
}
}));
// Set the intranet callback on the mediator for same-LAN connections
let connection_manager2 = self.connection_manager.clone();
mediator.set_intranet_callback(Arc::new(
move |rendezvous_addr, peer_socket_addr, local_addr, relay_server, device_id| {
@@ -345,7 +265,6 @@ impl RustDeskService {
},
));
// Spawn rendezvous task
let status = self.status.clone();
let handle = tokio::spawn(async move {
loop {
@@ -357,7 +276,6 @@ impl RustDeskService {
Err(e) => {
error!("Rendezvous mediator error: {}", e);
*status.write() = ServiceStatus::Error(e.to_string());
// Wait before retry
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
*status.write() = ServiceStatus::Starting;
}
@@ -372,10 +290,7 @@ impl RustDeskService {
Ok(())
}
/// Start TCP listener for direct peer connections
/// Returns the join handle and the port that was bound
async fn start_tcp_listener_with_port(&self) -> anyhow::Result<(Vec<JoinHandle<()>>, u16)> {
// Try to bind to the default port, or find an available port
let (listeners, listen_port) = match self.bind_direct_listeners(DIRECT_LISTEN_PORT) {
Ok(result) => result,
Err(err) => {
@@ -453,7 +368,6 @@ impl RustDeskService {
Ok((listeners, listen_port))
}
/// Stop the RustDesk service
pub async fn stop(&self) -> anyhow::Result<()> {
if self.status() == ServiceStatus::Stopped {
return Ok(());
@@ -461,23 +375,18 @@ impl RustDeskService {
info!("Stopping RustDesk service");
// Send shutdown signal (this will stop the TCP listener)
let _ = self.shutdown_tx.send(());
// Close all connections
self.connection_manager.close_all();
// Stop rendezvous mediator
if let Some(mediator) = self.rendezvous.read().as_ref() {
mediator.stop();
}
// Wait for rendezvous task to finish
if let Some(handle) = self.rendezvous_handle.write().take() {
handle.abort();
}
// Wait for TCP listener task to finish
if let Some(handles) = self.tcp_listener_handle.write().take() {
for handle in handles {
handle.abort();
@@ -490,15 +399,12 @@ impl RustDeskService {
Ok(())
}
/// Restart the service with new configuration
pub async fn restart(&self, config: RustDeskConfig) -> anyhow::Result<()> {
self.stop().await?;
self.update_config(config);
self.start().await
}
/// Save keypair and UUID to config
/// Returns the updated config if changes were made
pub fn save_credentials(&self) -> Option<RustDeskConfig> {
if let Some(mediator) = self.rendezvous.read().as_ref() {
let kp = mediator.ensure_keypair();
@@ -506,7 +412,6 @@ impl RustDeskService {
let mut config = self.config.write();
let mut changed = false;
// Save encryption keypair (Curve25519)
let pk = kp.public_key_base64();
let sk = kp.secret_key_base64();
if config.public_key.as_ref() != Some(&pk) || config.private_key.as_ref() != Some(&sk) {
@@ -515,7 +420,6 @@ impl RustDeskService {
changed = true;
}
// Save signing keypair (Ed25519)
let signing_pk = skp.public_key_base64();
let signing_sk = skp.secret_key_base64();
if config.signing_public_key.as_ref() != Some(&signing_pk)
@@ -526,7 +430,6 @@ impl RustDeskService {
changed = true;
}
// Save UUID if it was newly generated
if mediator.uuid_needs_save() {
let mediator_config = mediator.config();
if let Some(uuid) = mediator_config.uuid {
@@ -545,21 +448,16 @@ impl RustDeskService {
None
}
/// Save keypair to config (deprecated, use save_credentials instead)
#[deprecated(note = "Use save_credentials instead")]
pub fn save_keypair(&self) {
let _ = self.save_credentials();
}
}
/// Handle relay request from rendezvous server
///
/// Correct flow based on RustDesk protocol:
/// 1. Connect to RENDEZVOUS server (not relay!)
/// 2. Send RelayResponse with client's socket_addr
/// 3. Connect to RELAY server
/// 4. Accept connection without waiting for response
#[allow(clippy::too_many_arguments)]
fn rustdesk_relay_key(config: &Arc<RwLock<RustDeskConfig>>) -> String {
config.read().relay_key.clone().unwrap_or_default()
}
async fn handle_relay_request(
rendezvous_addr: &str,
relay_server: &str,
@@ -568,16 +466,12 @@ async fn handle_relay_request(
device_id: &str,
relay_key: &str,
connection_manager: Arc<ConnectionManager>,
_video_manager: Arc<VideoStreamManager>,
_hid: Arc<HidController>,
_audio: Arc<AudioController>,
) -> anyhow::Result<()> {
info!(
"Handling relay request: rendezvous={}, relay={}, uuid={}",
rendezvous_addr, relay_server, uuid
);
// Step 1: Connect to RENDEZVOUS server and send RelayResponse
let rendezvous_socket_addr: SocketAddr = tokio::net::lookup_host(rendezvous_addr)
.await?
.next()
@@ -597,8 +491,7 @@ async fn handle_relay_request(
rendezvous_socket_addr
);
// Send RelayResponse to rendezvous server with client's socket_addr
// IMPORTANT: Include our device ID so rendezvous server can look up and sign our public key
// Rendezvous looks up our pk by device id (must set `id`, not raw pk on wire).
let relay_response = make_relay_response(uuid, socket_addr, relay_server, device_id);
let bytes = relay_response
.write_to_bytes()
@@ -606,10 +499,8 @@ async fn handle_relay_request(
bytes_codec::write_frame(&mut rendezvous_stream, &bytes).await?;
debug!("Sent RelayResponse to rendezvous server for uuid={}", uuid);
// Close rendezvous connection - we don't need to wait for response
drop(rendezvous_stream);
// Step 2: Connect to RELAY server and send RequestRelay to identify ourselves
let relay_addr: SocketAddr = tokio::net::lookup_host(relay_server)
.await?
.next()
@@ -624,9 +515,7 @@ async fn handle_relay_request(
info!("Connected to relay server at {}", relay_addr);
// Send RequestRelay to relay server with our uuid, licence_key, and peer's socket_addr
// The licence_key is required if the relay server is configured with -k option
// The socket_addr is CRITICAL - the relay server uses it to match us with the peer
// Relay pairs peers by uuid + mangled peer socket_addr (required when hbbr uses -k).
let request_relay = make_request_relay(uuid, relay_key, socket_addr);
let bytes = request_relay
.write_to_bytes()
@@ -634,10 +523,8 @@ async fn handle_relay_request(
bytes_codec::write_frame(&mut stream, &bytes).await?;
debug!("Sent RequestRelay to relay server for uuid={}", uuid);
// Decode peer address for logging
let peer_addr = rendezvous::AddrMangle::decode(socket_addr).unwrap_or(relay_addr);
// Step 3: Accept connection - relay server bridges the connection
connection_manager
.accept_connection(stream, peer_addr)
.await?;
@@ -649,14 +536,6 @@ async fn handle_relay_request(
Ok(())
}
/// Handle intranet/same-LAN connection request
///
/// When the server determines that the client and peer are on the same intranet
/// (same public IP or both on LAN), it sends FetchLocalAddr to the peer.
/// The peer must:
/// 1. Open a TCP connection to the rendezvous server
/// 2. Send LocalAddr with our local address
/// 3. Accept the peer connection over that same TCP stream
async fn handle_intranet_request(
rendezvous_addr: &str,
peer_socket_addr: &[u8],
@@ -670,11 +549,9 @@ async fn handle_intranet_request(
rendezvous_addr, local_addr, device_id
);
// Decode peer address for logging
let peer_addr = AddrMangle::decode(peer_socket_addr);
debug!("Peer address from FetchLocalAddr: {:?}", peer_addr);
// Connect to rendezvous server via TCP with timeout
let mut stream =
tokio::time::timeout(Duration::from_secs(5), TcpStream::connect(rendezvous_addr))
.await
@@ -685,7 +562,6 @@ async fn handle_intranet_request(
rendezvous_addr
);
// Build LocalAddr message with our local address (mangled)
let local_addr_bytes = AddrMangle::encode(local_addr);
let msg = make_local_addr(
peer_socket_addr,
@@ -698,24 +574,16 @@ async fn handle_intranet_request(
.write_to_bytes()
.map_err(|e| anyhow::anyhow!("Failed to encode: {}", e))?;
// Send LocalAddr using RustDesk's variable-length framing
bytes_codec::write_frame(&mut stream, &bytes).await?;
info!("Sent LocalAddr to rendezvous server, waiting for peer connection");
// Now the rendezvous server will forward this to the client,
// and the client will connect to us through this same TCP stream.
// The server proxies the connection between client and peer.
// Get peer address for logging/connection tracking
let effective_peer_addr = peer_addr.unwrap_or_else(|| {
// If we can't decode the peer address, use the rendezvous server address
rendezvous_addr
.parse()
.unwrap_or_else(|_| "0.0.0.0:0".parse().unwrap())
});
// Accept the connection - the stream is now a proxied connection to the client
connection_manager
.accept_connection(stream, effective_peer_addr)
.await?;