Merge pull request #152 from 21pages/websocket

websocket
This commit is contained in:
RustDesk 2025-05-08 19:27:46 +08:00 committed by GitHub
commit 7839dcf4e4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 292 additions and 38 deletions

View File

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

View File

@ -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,

View File

@ -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,13 +147,6 @@ pub async fn connect_tcp_local<
local: Option<SocketAddr>,
ms_timeout: u64,
) -> 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?,
@ -141,7 +167,6 @@ pub async fn connect_tcp_local<
Ok(Stream::Tcp(
FramedStream::new(target, local, ms_timeout).await?,
))
}
}
#[inline]

View File

@ -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");
}
}