websocket

Signed-off-by: 21pages <sunboeasy@gmail.com>

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" }

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,

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<T: std::string::ToString>(host: T, offset: i32) -> String {
     host
 }
 
+pub fn split_host_port<T: std::string::ToString>(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<crate::Stream> {
+    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<SocketAddr>,
     ms_timeout: u64,
 ) -> ResultType<Stream> {
-    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://") {
+    if let Some(target_addr) = target.resolve() {
-        Ok(Stream::WebSocket(
+        if let Some(local_addr) = local {
-            websocket::WsFramedStream::new(target_str, local, None, ms_timeout).await?,
+            if local_addr.is_ipv6() && target_addr.is_ipv4() {
-        ))
+                let resolved_target = query_nip_io(target_addr).await?;
-    } else {
+                return Ok(Stream::Tcp(
-        if let Some(conf) = Config::get_socks() {
+                    FramedStream::new(resolved_target, Some(local_addr), ms_timeout).await?,
-            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?,
-                    ));
-                }
             }
         }
-
-        Ok(Stream::Tcp(
-            FramedStream::new(target, local, ms_timeout).await?,
-        ))
     }
+
+    Ok(Stream::Tcp(
+        FramedStream::new(target, local, ms_timeout).await?,
+    ))
 }
 
 #[inline]

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<Result<BytesMut, Error>> {
-        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");
+    }
+}