diff --git a/Cargo.toml b/Cargo.toml index 00e30f2..0689f62 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -46,8 +46,6 @@ httparse = "1.10" base64 = "0.22" url = "2.5" sha2 = "0.10" -tokio-tungstenite = { version = "0.26" } -tungstenite = { version = "0.26" } whoami = "1.5" [target.'cfg(not(any(target_os = "android", target_os = "ios")))'.dependencies] @@ -63,8 +61,13 @@ tokio-rustls = { version = "0.26", features = [ ], default-features = false } rustls-platform-verifier = "0.5" rustls-pki-types = "1.11" +tokio-tungstenite = { version = "0.26", features = ["rustls-tls-native-roots", "rustls-tls-webpki-roots"] } +tungstenite = { version = "0.26", features = ["rustls-tls-native-roots", "rustls-tls-webpki-roots"] } + [target.'cfg(any(target_os = "macos", target_os = "windows"))'.dependencies] tokio-native-tls = "0.3" +tokio-tungstenite = { version = "0.26", features = ["native-tls"] } +tungstenite = { version = "0.26", features = ["native-tls"] } [build-dependencies] protobuf-codegen = { version = "3.7" } diff --git a/src/config.rs b/src/config.rs index 06790aa..38facc1 100644 --- a/src/config.rs +++ b/src/config.rs @@ -103,6 +103,8 @@ pub const RS_PUB_KEY: &str = "OeVuKk5nlHiXp+APNn0Y3pC1Iwpwn44JGqrQCsWqmBw="; pub const RENDEZVOUS_PORT: i32 = 21116; pub const RELAY_PORT: i32 = 21117; +pub const WS_RENDEZVOUS_PORT: i32 = 21118; +pub const WS_RELAY_PORT: i32 = 21119; macro_rules! serde_field_string { ($default_func:ident, $de_func:ident, $default_expr:expr) => { @@ -2293,6 +2295,11 @@ pub fn option2bool(option: &str, value: &str) -> bool { } } +pub fn use_ws() -> bool { + let option = keys::OPTION_ALLOW_WEBSOCKET; + option2bool(option, &Config::get_option(option)) +} + pub mod keys { pub const OPTION_VIEW_ONLY: &str = "view_only"; pub const OPTION_SHOW_MONITORS_TOOLBAR: &str = "show_monitors_toolbar"; @@ -2368,6 +2375,7 @@ pub mod keys { pub const OPTION_CUSTOM_RENDEZVOUS_SERVER: &str = "custom-rendezvous-server"; pub const OPTION_API_SERVER: &str = "api-server"; pub const OPTION_KEY: &str = "key"; + pub const OPTION_ALLOW_WEBSOCKET: &str = "allow-websocket"; pub const OPTION_PRESET_ADDRESS_BOOK_NAME: &str = "preset-address-book-name"; pub const OPTION_PRESET_ADDRESS_BOOK_TAG: &str = "preset-address-book-tag"; pub const OPTION_ENABLE_DIRECTX_CAPTURE: &str = "enable-directx-capture"; @@ -2388,6 +2396,7 @@ pub mod keys { pub const OPTION_HIDE_SERVER_SETTINGS: &str = "hide-server-settings"; pub const OPTION_HIDE_PROXY_SETTINGS: &str = "hide-proxy-settings"; pub const OPTION_HIDE_REMOTE_PRINTER_SETTINGS: &str = "hide-remote-printer-settings"; + pub const OPTION_HIDE_WEBSOCKET_SETTINGS: &str = "hide-websocket-settings"; pub const OPTION_HIDE_USERNAME_ON_CARD: &str = "hide-username-on-card"; pub const OPTION_HIDE_HELP_CARDS: &str = "hide-help-cards"; pub const OPTION_DEFAULT_CONNECT_PASSWORD: &str = "default-connect-password"; @@ -2529,6 +2538,7 @@ pub mod keys { OPTION_CUSTOM_RENDEZVOUS_SERVER, OPTION_API_SERVER, OPTION_KEY, + OPTION_ALLOW_WEBSOCKET, OPTION_PRESET_ADDRESS_BOOK_NAME, OPTION_PRESET_ADDRESS_BOOK_TAG, OPTION_ENABLE_DIRECTX_CAPTURE, @@ -2549,6 +2559,7 @@ pub mod keys { OPTION_HIDE_SERVER_SETTINGS, OPTION_HIDE_PROXY_SETTINGS, OPTION_HIDE_REMOTE_PRINTER_SETTINGS, + OPTION_HIDE_WEBSOCKET_SETTINGS, OPTION_HIDE_USERNAME_ON_CARD, OPTION_HIDE_HELP_CARDS, OPTION_DEFAULT_CONNECT_PASSWORD, diff --git a/src/socket_client.rs b/src/socket_client.rs index 60695af..ee37eb0 100644 --- a/src/socket_client.rs +++ b/src/socket_client.rs @@ -2,7 +2,8 @@ use crate::{ config::{Config, NetworkType}, tcp::FramedStream, udp::FramedSocket, - websocket, ResultType, Stream, + websocket::{self, check_ws, is_ws_endpoint}, + ResultType, Stream, }; use anyhow::Context; use std::net::SocketAddr; @@ -49,6 +50,30 @@ pub fn increase_port(host: T, offset: i32) -> String { host } +pub fn split_host_port(host: T) -> Option<(String, i32)> { + let host = host.to_string(); + if crate::is_ipv6_str(&host) { + if host.starts_with('[') { + let tmp: Vec<&str> = host.split("]:").collect(); + if tmp.len() == 2 { + let port: i32 = tmp[1].parse().unwrap_or(0); + if port > 0 { + return Some((format!("{}]", tmp[0]), port)); + } + } + } + } else if host.contains(':') { + let tmp: Vec<&str> = host.split(':').collect(); + if tmp.len() == 2 { + let port: i32 = tmp[1].parse().unwrap_or(0); + if port > 0 { + return Some((tmp[0].to_string(), port)); + } + } + } + None +} + pub fn test_if_valid_server(host: &str, test_with_proxy: bool) -> String { let host = check_port(host, 0); use std::net::ToSocketAddrs; @@ -95,6 +120,7 @@ impl IsResolvedSocketAddr for &str { } } +// This function checks if the target is a websocket endpoint and connects accordingly. #[inline] pub async fn connect_tcp< 't, @@ -103,9 +129,16 @@ pub async fn connect_tcp< target: T, ms_timeout: u64, ) -> ResultType { + let target_str = check_ws(&target.to_string()); + if is_ws_endpoint(&target_str) { + return Ok(Stream::WebSocket( + websocket::WsFramedStream::new(target_str, None, None, ms_timeout).await?, + )); + } connect_tcp_local(target, None, ms_timeout).await } +// This function connects directly to the target without checking for websocket endpoints. pub async fn connect_tcp_local< 't, T: IntoTargetAddr<'t> + ToSocketAddrs + IsResolvedSocketAddr + std::fmt::Display, @@ -114,34 +147,26 @@ pub async fn connect_tcp_local< local: Option, ms_timeout: u64, ) -> ResultType { - let target_str = target.to_string(); + if let Some(conf) = Config::get_socks() { + return Ok(Stream::Tcp( + FramedStream::connect(target, local, &conf, ms_timeout).await?, + )); + } - if target_str.starts_with("ws://") || target_str.starts_with("wss://") { - Ok(Stream::WebSocket( - websocket::WsFramedStream::new(target_str, local, None, ms_timeout).await?, - )) - } else { - if let Some(conf) = Config::get_socks() { - return Ok(Stream::Tcp( - FramedStream::connect(target, local, &conf, ms_timeout).await?, - )); - } - - if let Some(target_addr) = target.resolve() { - if let Some(local_addr) = local { - if local_addr.is_ipv6() && target_addr.is_ipv4() { - let resolved_target = query_nip_io(target_addr).await?; - return Ok(Stream::Tcp( - FramedStream::new(resolved_target, Some(local_addr), ms_timeout).await?, - )); - } + if let Some(target_addr) = target.resolve() { + if let Some(local_addr) = local { + if local_addr.is_ipv6() && target_addr.is_ipv4() { + let resolved_target = query_nip_io(target_addr).await?; + return Ok(Stream::Tcp( + FramedStream::new(resolved_target, Some(local_addr), ms_timeout).await?, + )); } } - - Ok(Stream::Tcp( - FramedStream::new(target, local, ms_timeout).await?, - )) } + + Ok(Stream::Tcp( + FramedStream::new(target, local, ms_timeout).await?, + )) } #[inline] diff --git a/src/websocket.rs b/src/websocket.rs index 25f2f29..4e9e318 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -1,5 +1,9 @@ use crate::{ - config::Socks5Server, protobuf::Message, sodiumoxide::crypto::secretbox::Key, tcp::Encrypt, + config::{use_ws, Config, Socks5Server, RELAY_PORT, RENDEZVOUS_PORT}, + protobuf::Message, + socket_client::split_host_port, + sodiumoxide::crypto::secretbox::Key, + tcp::Encrypt, ResultType, }; use bytes::{Bytes, BytesMut}; @@ -33,7 +37,6 @@ impl WsFramedStream { let url_str = url.as_ref(); // to-do: websocket proxy. - log::info!("{:?}", url_str); let request = url_str .into_client_request() @@ -44,6 +47,10 @@ impl WsFramedStream { let addr = match stream.get_ref() { MaybeTlsStream::Plain(tcp) => tcp.peer_addr()?, + #[cfg(any(target_os = "macos", target_os = "windows"))] + MaybeTlsStream::NativeTls(tls) => tls.get_ref().get_ref().get_ref().peer_addr()?, + #[cfg(not(any(target_os = "macos", target_os = "windows")))] + MaybeTlsStream::Rustls(tls) => tls.get_ref().0.peer_addr()?, _ => return Err(Error::new(ErrorKind::Other, "Unsupported stream type").into()), }; @@ -126,14 +133,11 @@ impl WsFramedStream { #[inline] pub async fn next(&mut self) -> Option> { - log::debug!("Waiting for next message"); - while let Some(msg) = self.stream.next().await { - log::debug!("receive msg: {:?}", msg); let msg = match msg { Ok(msg) => msg, Err(e) => { - log::debug!("{}", e); + log::error!("{}", e); return Some(Err(Error::new( ErrorKind::Other, format!("WebSocket protocol error: {}", e), @@ -141,10 +145,8 @@ impl WsFramedStream { } }; - log::debug!("Received message type: {}", msg.to_string()); match msg { WsMessage::Binary(data) => { - log::info!("Received binary data ({} bytes)", data.len()); let mut bytes = BytesMut::from(&data[..]); if let Some(key) = self.encrypt.as_mut() { if let Err(err) = key.dec(&mut bytes) { @@ -154,16 +156,13 @@ impl WsFramedStream { return Some(Ok(bytes)); } WsMessage::Text(text) => { - log::debug!("Received text message, converting to binary"); let bytes = BytesMut::from(text.as_bytes()); return Some(Ok(bytes)); } WsMessage::Close(_) => { - log::info!("Received close frame"); return None; } _ => { - log::debug!("Unhandled message type: {}", msg.to_string()); continue; } } @@ -180,3 +179,219 @@ impl WsFramedStream { } } } + +pub fn is_ws_endpoint(endpoint: &str) -> bool { + endpoint.starts_with("ws://") || endpoint.starts_with("wss://") +} + +/** + * Core function to convert an endpoint to WebSocket format + * + * Converts between different address formats: + * 1. IPv4 address with/without port -> ws://ipv4:port + * 2. IPv6 address with/without port -> ws://[ipv6]:port + * 3. Domain with/without port -> ws(s)://domain/ws/path + * + * @param endpoint The endpoint to convert + * @return The converted WebSocket endpoint + */ +pub fn check_ws(endpoint: &str) -> String { + if !use_ws() { + return endpoint.to_string(); + } + + if endpoint.is_empty() { + return endpoint.to_string(); + } + + if is_ws_endpoint(endpoint) { + return endpoint.to_string(); + } + + let Some((endpoint_host, endpoint_port)) = split_host_port(endpoint) else { + debug_assert!(false, "endpoint doesn't have port"); + return endpoint.to_string(); + }; + + let custom_rendezvous_server = Config::get_rendezvous_server(); + let relay_server = Config::get_option("relay-server"); + let rendezvous_port = split_host_port(&custom_rendezvous_server) + .map(|(_, p)| p) + .unwrap_or(RENDEZVOUS_PORT); + let relay_port = split_host_port(&relay_server) + .map(|(_, p)| p) + .unwrap_or(RELAY_PORT); + + let (relay, dst_port) = if endpoint_port == rendezvous_port { + // rendezvous + (false, endpoint_port + 2) + } else if endpoint_port == rendezvous_port - 1 { + // online + (false, endpoint_port + 3) + } else if endpoint_port == relay_port || endpoint_port == rendezvous_port + 1 { + // relay + // https://github.com/rustdesk/rustdesk/blob/6ffbcd1375771f2482ec4810680623a269be70f1/src/rendezvous_mediator.rs#L615 + // https://github.com/rustdesk/rustdesk-server/blob/235a3c326ceb665e941edb50ab79faa1208f7507/src/relay_server.rs#L83, based on relay port. + (true, endpoint_port + 2) + } else { + // fallback relay + // for controlling side, relay server is passed from the controlled side, not related to local config. + (true, endpoint_port + 2) + }; + + let (address, is_domain) = if crate::is_ip_str(endpoint) { + (format!("{}:{}", endpoint_host, dst_port), false) + } else { + let domain_path = if relay { "/ws/relay" } else { "/ws/id" }; + (format!("{}{}", endpoint_host, domain_path), true) + }; + let protocol = if is_domain { + let api_server = Config::get_option("api-server"); + if api_server.starts_with("https") { + "wss" + } else { + "ws" + } + } else { + "ws" + }; + format!("{}://{}", protocol, address) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{keys, Config}; + + #[test] + fn test_check_ws() { + // enable websocket + Config::set_option(keys::OPTION_ALLOW_WEBSOCKET.to_string(), "Y".to_string()); + + // not set custom-rendezvous-server + Config::set_option("custom-rendezvous-server".to_string(), "".to_string()); + Config::set_option("relay-server".to_string(), "".to_string()); + Config::set_option("api-server".to_string(), "".to_string()); + assert_eq!(check_ws("127.0.0.1:21115"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21116"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21117"), "ws://127.0.0.1:21119"); + assert_eq!(check_ws("rustdesk.com:21115"), "ws://rustdesk.com/ws/id"); + assert_eq!(check_ws("rustdesk.com:21116"), "ws://rustdesk.com/ws/id"); + assert_eq!(check_ws("rustdesk.com:21117"), "ws://rustdesk.com/ws/relay"); + // set relay-server without port + Config::set_option("relay-server".to_string(), "127.0.0.1".to_string()); + Config::set_option( + "api-server".to_string(), + "https://api.rustdesk.com".to_string(), + ); + assert_eq!( + check_ws("[0:0:0:0:0:0:0:1]:21115"), + "ws://[0:0:0:0:0:0:0:1]:21118" + ); + assert_eq!( + check_ws("[0:0:0:0:0:0:0:1]:21116"), + "ws://[0:0:0:0:0:0:0:1]:21118" + ); + assert_eq!( + check_ws("[0:0:0:0:0:0:0:1]:21117"), + "ws://[0:0:0:0:0:0:0:1]:21119" + ); + assert_eq!(check_ws("rustdesk.com:21115"), "wss://rustdesk.com/ws/id"); + assert_eq!(check_ws("rustdesk.com:21116"), "wss://rustdesk.com/ws/id"); + assert_eq!( + check_ws("rustdesk.com:21117"), + "wss://rustdesk.com/ws/relay" + ); + // set relay-server with default port + Config::set_option("relay-server".to_string(), "127.0.0.1:21117".to_string()); + assert_eq!(check_ws("127.0.0.1:21115"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21116"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21117"), "ws://127.0.0.1:21119"); + // set relay-server with custom port + Config::set_option("relay-server".to_string(), "127.0.0.1:34567".to_string()); + assert_eq!(check_ws("rustdesk.com:21115"), "wss://rustdesk.com/ws/id"); + assert_eq!(check_ws("rustdesk.com:21116"), "wss://rustdesk.com/ws/id"); + assert_eq!( + check_ws("rustdesk.com:34567"), + "wss://rustdesk.com/ws/relay" + ); + + // set custom-rendezvous-server without port + Config::set_option( + "custom-rendezvous-server".to_string(), + "127.0.0.1".to_string(), + ); + Config::set_option("relay-server".to_string(), "".to_string()); + Config::set_option("api-server".to_string(), "".to_string()); + assert_eq!(check_ws("127.0.0.1:21115"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21116"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21117"), "ws://127.0.0.1:21119"); + // set relay-server without port + Config::set_option("relay-server".to_string(), "127.0.0.1".to_string()); + assert_eq!(check_ws("127.0.0.1:21115"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21116"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21117"), "ws://127.0.0.1:21119"); + // set relay-server with default port + Config::set_option("relay-server".to_string(), "127.0.0.1:21117".to_string()); + assert_eq!(check_ws("127.0.0.1:21115"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21116"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21117"), "ws://127.0.0.1:21119"); + // set relay-server with custom port + Config::set_option("relay-server".to_string(), "127.0.0.1:34567".to_string()); + assert_eq!(check_ws("127.0.0.1:21115"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21116"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:34567"), "ws://127.0.0.1:34569"); + + // set custom-rendezvous-server without default port + Config::set_option( + "custom-rendezvous-server".to_string(), + "127.0.0.1".to_string(), + ); + Config::set_option("relay-server".to_string(), "".to_string()); + Config::set_option("api-server".to_string(), "".to_string()); + assert_eq!(check_ws("127.0.0.1:21115"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21116"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21117"), "ws://127.0.0.1:21119"); + // set relay-server without port + Config::set_option("relay-server".to_string(), "127.0.0.1".to_string()); + assert_eq!(check_ws("127.0.0.1:21115"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21116"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21117"), "ws://127.0.0.1:21119"); + // set relay-server with default port + Config::set_option("relay-server".to_string(), "127.0.0.1:21117".to_string()); + assert_eq!(check_ws("127.0.0.1:21115"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21116"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21117"), "ws://127.0.0.1:21119"); + // set relay-server with custom port + Config::set_option("relay-server".to_string(), "127.0.0.1:34567".to_string()); + assert_eq!(check_ws("127.0.0.1:21115"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:21116"), "ws://127.0.0.1:21118"); + assert_eq!(check_ws("127.0.0.1:34567"), "ws://127.0.0.1:34569"); + + // set custom-rendezvous-server with custom port + Config::set_option( + "custom-rendezvous-server".to_string(), + "127.0.0.1:23456".to_string(), + ); + Config::set_option("relay-server".to_string(), "".to_string()); + Config::set_option("api-server".to_string(), "".to_string()); + assert_eq!(check_ws("127.0.0.1:23455"), "ws://127.0.0.1:23458"); + assert_eq!(check_ws("127.0.0.1:23456"), "ws://127.0.0.1:23458"); + assert_eq!(check_ws("127.0.0.1:23457"), "ws://127.0.0.1:23459"); + // set relay-server without port + Config::set_option("relay-server".to_string(), "127.0.0.1".to_string()); + assert_eq!(check_ws("127.0.0.1:23455"), "ws://127.0.0.1:23458"); + assert_eq!(check_ws("127.0.0.1:23456"), "ws://127.0.0.1:23458"); + assert_eq!(check_ws("127.0.0.1:21117"), "ws://127.0.0.1:21119"); + // set relay-server with default port + Config::set_option("relay-server".to_string(), "127.0.0.1:21117".to_string()); + assert_eq!(check_ws("127.0.0.1:23455"), "ws://127.0.0.1:23458"); + assert_eq!(check_ws("127.0.0.1:23456"), "ws://127.0.0.1:23458"); + assert_eq!(check_ws("127.0.0.1:21117"), "ws://127.0.0.1:21119"); + // set relay-server with custom port + Config::set_option("relay-server".to_string(), "127.0.0.1:34567".to_string()); + assert_eq!(check_ws("127.0.0.1:23455"), "ws://127.0.0.1:23458"); + assert_eq!(check_ws("127.0.0.1:23456"), "ws://127.0.0.1:23458"); + assert_eq!(check_ws("127.0.0.1:34567"), "ws://127.0.0.1:34569"); + } +}