Merge pull request #137 from YinMo19/main

[enhance] Add websocket support.

Author: rustdesk

Cargo.toml

@@ -46,25 +46,36 @@ httparse = "1.10"
 base64 = "0.22"
 url = "2.5"
 sha2 = "0.10"
-tokio-tungstenite = "0.26"
+tokio-tungstenite = { version = "0.26.2" }
+tungstenite = { version = "0.26.2" }
+
 
 [target.'cfg(not(any(target_os = "android", target_os = "ios")))'.dependencies]
 mac_address = "1.1"
 default_net = { git = "https://github.com/rustdesk-org/default_net" }
 machine-uid = { git = "https://github.com/rustdesk-org/machine-uid" }
 [target.'cfg(not(any(target_os = "macos", target_os = "windows")))'.dependencies]
-tokio-rustls = { version = "0.26", features = ["logging", "tls12", "ring"], default-features = false }
+tokio-rustls = { version = "0.26", features = [
+    "logging",
+    "tls12",
+    "ring",
+], default-features = false }
 rustls-platform-verifier = "0.5"
 rustls-pki-types = "1.11"
 [target.'cfg(any(target_os = "macos", target_os = "windows"))'.dependencies]
-tokio-native-tls ="0.3"
+tokio-native-tls = "0.3"
 
 [build-dependencies]
 protobuf-codegen = { version = "3.7" }
 
 [target.'cfg(target_os = "windows")'.dependencies]
-winapi = { version = "0.3", features = ["winuser", "synchapi", "pdh", "memoryapi", "sysinfoapi"] }
+winapi = { version = "0.3", features = [
+    "winuser",
+    "synchapi",
+    "pdh",
+    "memoryapi",
+    "sysinfoapi",
+] }
 
 [target.'cfg(target_os = "macos")'.dependencies]
 osascript = "0.3"
-

src/lib.rs

@@ -57,8 +57,10 @@ pub use toml;
 pub use uuid;
 pub mod fingerprint;
 pub use flexi_logger;
+pub mod websocket;
+pub mod stream;
+pub use stream::Stream;
 
-pub type Stream = tcp::FramedStream;
 pub type SessionID = uuid::Uuid;
 
 #[inline]

src/proxy.rs

@@ -56,7 +56,7 @@ const MAXIMUM_RESPONSE_HEADERS: usize = 16;
 const DEFINE_TIME_OUT: u64 = 600;
 
 pub trait IntoUrl {
-    
+
     // Besides parsing as a valid `Url`, the `Url` must be a valid
     // `http::Uri`, in that it makes sense to use in a network request.
     fn into_url(self) -> Result<Url, ProxyError>;

src/socket_client.rs

@@ -2,7 +2,7 @@ use crate::{
     config::{Config, NetworkType},
     tcp::FramedStream,
     udp::FramedSocket,
-    ResultType,
+    websocket, ResultType, Stream,
 };
 use anyhow::Context;
 use std::net::SocketAddr;
@@ -102,7 +102,7 @@ pub async fn connect_tcp<
 >(
     target: T,
     ms_timeout: u64,
-) -> ResultType<FramedStream> {
+) -> ResultType<crate::Stream> {
     connect_tcp_local(target, None, ms_timeout).await
 }
 
@@ -113,19 +113,35 @@ pub async fn connect_tcp_local<
     target: T,
     local: Option<SocketAddr>,
     ms_timeout: u64,
-) -> ResultType<FramedStream> {
-    if let Some(conf) = Config::get_socks() {
-        return FramedStream::connect(target, local, &conf, ms_timeout).await;
-    }
-    if let Some(target) = target.resolve() {
-        if let Some(local) = local {
-            if local.is_ipv6() && target.is_ipv4() {
-                let target = query_nip_io(target).await?;
-                return FramedStream::new(target, Some(local), ms_timeout).await;
+) -> ResultType<Stream> {
+    let target_str = target.to_string();
+
+    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?,
+                    ));
+                }
             }
         }
+
+        Ok(Stream::Tcp(
+            FramedStream::new(target, local, ms_timeout).await?,
+        ))
     }
-    FramedStream::new(target, local, ms_timeout).await
 }
 
 #[inline]

src/stream.rs

@@ -0,0 +1,110 @@
+use crate::{config, tcp, websocket, ResultType};
+use sodiumoxide::crypto::secretbox::Key;
+use std::net::SocketAddr;
+
+// support Websocket and tcp.
+pub enum Stream {
+    WebSocket(websocket::WsFramedStream),
+    Tcp(tcp::FramedStream),
+}
+
+impl Stream {
+    #[inline]
+    pub fn set_send_timeout(&mut self, ms: u64) {
+        match self {
+            Stream::WebSocket(s) => s.set_send_timeout(ms),
+            Stream::Tcp(s) => s.set_send_timeout(ms),
+        }
+    }
+
+    #[inline]
+    pub fn set_raw(&mut self) {
+        match self {
+            Stream::WebSocket(s) => s.set_raw(),
+            Stream::Tcp(s) => s.set_raw(),
+        }
+    }
+
+    #[inline]
+    pub async fn send_bytes(&mut self, bytes: bytes::Bytes) -> ResultType<()> {
+        match self {
+            Stream::WebSocket(s) => s.send_bytes(bytes).await,
+            Stream::Tcp(s) => s.send_bytes(bytes).await,
+        }
+    }
+
+    #[inline]
+    pub async fn send_raw(&mut self, bytes: Vec<u8>) -> ResultType<()> {
+        match self {
+            Stream::WebSocket(s) => s.send_raw(bytes).await,
+            Stream::Tcp(s) => s.send_raw(bytes).await,
+        }
+    }
+
+    #[inline]
+    pub fn set_key(&mut self, key: Key) {
+        match self {
+            Stream::WebSocket(s) => s.set_key(key),
+            Stream::Tcp(s) => s.set_key(key),
+        }
+    }
+
+    #[inline]
+    pub fn is_secured(&self) -> bool {
+        match self {
+            Stream::WebSocket(s) => s.is_secured(),
+            Stream::Tcp(s) => s.is_secured(),
+        }
+    }
+
+    #[inline]
+    pub async fn next_timeout(
+        &mut self,
+        timeout: u64,
+    ) -> Option<Result<bytes::BytesMut, std::io::Error>> {
+        match self {
+            Stream::WebSocket(s) => s.next_timeout(timeout).await,
+            Stream::Tcp(s) => s.next_timeout(timeout).await,
+        }
+    }
+
+    /// establish connect from websocket
+    #[inline]
+    pub async fn connect_websocket(
+        url: impl AsRef<str>,
+        local_addr: Option<SocketAddr>,
+        proxy_conf: Option<&config::Socks5Server>,
+        timeout_ms: u64,
+    ) -> ResultType<Self> {
+        let ws_stream =
+            websocket::WsFramedStream::new(url, local_addr, proxy_conf, timeout_ms).await?;
+        log::debug!("WebSocket connection established");
+        Ok(Self::WebSocket(ws_stream))
+    }
+
+    /// send message
+    #[inline]
+    pub async fn send(&mut self, msg: &impl protobuf::Message) -> ResultType<()> {
+        match self {
+            Self::WebSocket(ws) => ws.send(msg).await,
+            Self::Tcp(tcp) => tcp.send(msg).await,
+        }
+    }
+
+    /// receive message
+    #[inline]
+    pub async fn next(&mut self) -> Option<Result<bytes::BytesMut, std::io::Error>> {
+        match self {
+            Self::WebSocket(ws) => ws.next().await,
+            Self::Tcp(tcp) => tcp.next().await,
+        }
+    }
+
+    #[inline]
+    pub fn local_addr(&self) -> SocketAddr {
+        match self {
+            Self::WebSocket(ws) => ws.local_addr(),
+            Self::Tcp(tcp) => tcp.local_addr(),
+        }
+    }
+}

src/tcp.rs

@@ -25,7 +25,7 @@ pub trait TcpStreamTrait: AsyncRead + AsyncWrite + Unpin {}
 pub struct DynTcpStream(pub(crate) Box<dyn TcpStreamTrait + Send + Sync>);
 
 #[derive(Clone)]
-pub struct Encrypt(Key, u64, u64);
+pub struct Encrypt(pub Key, pub u64, pub u64);
 
 pub struct FramedStream(
     pub(crate) Framed<DynTcpStream, BytesCodec>,