feat: add opt-in WebSocket proxy support

Support WS and WSS over SOCKS5, HTTP, and HTTPS proxies, including custom
ports. Use the configured proxy only when allow-websocket-proxy is enabled;
explicit proxy configuration takes precedence. Propagate proxy failures
without retrying a direct connection.

Add regression coverage for proxy routing, authentication, TLS, custom
ports, direct defaults, and failure handling.

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

fix: flush WebSocket proxy tunnels before native TLS reads

Flush buffered HTTP and HTTPS tunnel writes before NativeTls reads so WSS handshakes can complete on macOS. Apply the adapter only to NativeTls WebSocket connections through HTTP or HTTPS proxies.

Extend the existing proxy round-trip test with HTTPS and SOCKS5 NativeTls cases. The websocket_proxy integration test passes in the rustdesk workspace.

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

refactor: simplify WebSocket proxy configuration

Select the configured proxy inside WsFramedStream::new based on the WebSocket proxy option and remove the explicit proxy argument. Update the socket client and drop the unused Stream::connect_websocket wrapper.

Adapt integration tests to configured proxies and cover custom target ports and WebSocket close handling.
This commit is contained in:
21pages committed 2026-10-10 11:50:09 +08:00
1 parent 94ec710894
commit 3aade011c0
7 files changed
+627 -70

No files matched your search

+6
View File
@@ -2842,6 +2842,11 @@ pub fn use_ws() -> bool {
option2bool(option, &Config::get_option(option)) option2bool(option, &Config::get_option(option))
} }
pub fn ws_use_proxy() -> bool {
let option = keys::OPTION_ALLOW_WEBSOCKET_PROXY;
option2bool(option, &Config::get_option(option))
}
pub fn allow_insecure_tls_fallback() -> bool { pub fn allow_insecure_tls_fallback() -> bool {
let option = keys::OPTION_ALLOW_INSECURE_TLS_FALLBACK; let option = keys::OPTION_ALLOW_INSECURE_TLS_FALLBACK;
option2bool(option, &Config::get_option(option)) option2bool(option, &Config::get_option(option))
@@ -2869,6 +2874,7 @@ pub mod keys {
pub const OPTION_ALLOW_NUMERNIC_ONE_TIME_PASSWORD: &str = "allow-numeric-one-time-password"; pub const OPTION_ALLOW_NUMERNIC_ONE_TIME_PASSWORD: &str = "allow-numeric-one-time-password";
pub const OPTION_DIRECT_SERVER: &str = "direct-server"; pub const OPTION_DIRECT_SERVER: &str = "direct-server";
pub const OPTION_ALLOW_WEBSOCKET: &str = "allow-websocket"; pub const OPTION_ALLOW_WEBSOCKET: &str = "allow-websocket";
pub const OPTION_ALLOW_WEBSOCKET_PROXY: &str = "allow-websocket-proxy";
pub const OPTION_TRACKPAD_SPEED: &str = "trackpad-speed"; pub const OPTION_TRACKPAD_SPEED: &str = "trackpad-speed";
pub const OPTION_REGISTER_DEVICE: &str = "register-device"; pub const OPTION_REGISTER_DEVICE: &str = "register-device";
pub const OPTION_RELAY_SERVER: &str = "relay-server"; pub const OPTION_RELAY_SERVER: &str = "relay-server";
+1 -1
View File
@@ -140,7 +140,7 @@ pub async fn connect_tcp<
let target_str = check_ws(&target.to_string()); let target_str = check_ws(&target.to_string());
if is_ws_endpoint(&target_str) { if is_ws_endpoint(&target_str) {
return Ok(Stream::WebSocket( return Ok(Stream::WebSocket(
websocket::WsFramedStream::new(target_str, None, None, ms_timeout).await?, websocket::WsFramedStream::new(target_str, None, ms_timeout).await?,
)); ));
} }
connect_tcp_local(target, None, ms_timeout).await connect_tcp_local(target, None, ms_timeout).await
+1 -15
View File
@@ -1,4 +1,4 @@
use crate::{bail, config, tcp, websocket, ResultType}; use crate::{bail, tcp, websocket, ResultType};
#[cfg(feature = "webrtc")] #[cfg(feature = "webrtc")]
use crate::webrtc; use crate::webrtc;
use sodiumoxide::crypto::secretbox::Key; use sodiumoxide::crypto::secretbox::Key;
@@ -245,20 +245,6 @@ impl Stream {
} }
} }
/// 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 /// send message
#[inline] #[inline]
pub async fn send(&mut self, msg: &impl protobuf::Message) -> ResultType<()> { pub async fn send(&mut self, msg: &impl protobuf::Message) -> ResultType<()> {
+183 -54
View File
@@ -1,40 +1,84 @@
use crate::{ use crate::{
config::{ config::{
keys::OPTION_RELAY_SERVER, use_ws, Config, Socks5Server, RELAY_PORT, RENDEZVOUS_PORT, keys::OPTION_RELAY_SERVER, use_ws, ws_use_proxy, Config, Socks5Server, RELAY_PORT,
RENDEZVOUS_PORT,
}, },
protobuf::Message, protobuf::Message,
socket_client::split_host_port, socket_client::split_host_port,
sodiumoxide::crypto::secretbox::Key, sodiumoxide::crypto::secretbox::Key,
tcp::{Encrypt, KxTranscript}, tcp::{DynTcpStream, Encrypt, KxTranscript},
tls::{get_cached_tls_accept_invalid_cert, get_cached_tls_type, upsert_tls_cache, TlsType}, tls::{get_cached_tls_accept_invalid_cert, get_cached_tls_type, upsert_tls_cache, TlsType},
ResultType, ResultType,
}; };
use anyhow::bail; use anyhow::bail;
use async_recursion::async_recursion; use async_recursion::async_recursion;
use bytes::{Bytes, BytesMut}; use bytes::{Bytes, BytesMut};
use futures::{SinkExt, StreamExt}; use futures::{future::Either, SinkExt, StreamExt};
use std::{ use std::{
io::{Error, ErrorKind}, io::{self, Error, ErrorKind},
net::SocketAddr, net::SocketAddr,
pin::Pin,
sync::Arc, sync::Arc,
task::{ready, Context, Poll},
time::Duration, time::Duration,
}; };
use tokio::{net::TcpStream, time::timeout}; use tokio::{
io::{AsyncRead, AsyncWrite, ReadBuf},
net::TcpStream,
time::timeout,
};
use tokio_native_tls::native_tls::TlsConnector; use tokio_native_tls::native_tls::TlsConnector;
use tokio_tungstenite::{ use tokio_tungstenite::{
connect_async_tls_with_config, tungstenite::protocol::Message as WsMessage, Connector, client_async_tls_with_config, connect_async_tls_with_config,
MaybeTlsStream, WebSocketStream, tungstenite::protocol::Message as WsMessage, Connector, MaybeTlsStream, WebSocketStream,
}; };
use tungstenite::client::IntoClientRequest; use tungstenite::client::IntoClientRequest;
use tungstenite::protocol::Role;
/// tungstenite's own defaults. `set_max_packet_length` clamps to these, so `usize::MAX` puts the /// tungstenite's own defaults. `set_max_packet_length` clamps to these, so `usize::MAX` puts the
/// transport back exactly where it started rather than above it. /// transport back exactly where it started rather than above it.
const DEFAULT_MAX_MESSAGE_SIZE: usize = 64 << 20; const DEFAULT_MAX_MESSAGE_SIZE: usize = 64 << 20;
const DEFAULT_MAX_FRAME_SIZE: usize = 16 << 20; const DEFAULT_MAX_FRAME_SIZE: usize = 16 << 20;
// macOS/iOS Native TLS may wait for a response without flushing the proxy tunnel's write buffer.
struct FlushBeforeRead(DynTcpStream);
impl AsyncRead for FlushBeforeRead {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
ready!(Pin::new(&mut self.0).poll_flush(cx))?;
Pin::new(&mut self.0).poll_read(cx, buf)
}
}
impl AsyncWrite for FlushBeforeRead {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.0).poll_write(cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.0).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.0).poll_shutdown(cx)
}
}
// Keep the original non-proxy connection path while supporting proxy streams.
type WsStream = Either<
WebSocketStream<MaybeTlsStream<TcpStream>>,
WebSocketStream<MaybeTlsStream<DynTcpStream>>,
>;
pub struct WsFramedStream { pub struct WsFramedStream {
stream: WebSocketStream<MaybeTlsStream<TcpStream>>, stream: WsStream,
addr: SocketAddr, addr: SocketAddr,
encrypt: Option<Encrypt>, encrypt: Option<Encrypt>,
send_timeout: u64, send_timeout: u64,
@@ -72,16 +116,18 @@ impl WsFramedStream {
async fn connect( async fn connect(
url: &str, url: &str,
local_addr: Option<SocketAddr>,
proxy_conf: Option<&Socks5Server>,
ms_timeout: u64, ms_timeout: u64,
) -> ResultType<WebSocketStream<MaybeTlsStream<TcpStream>>> { ) -> ResultType<(WsStream, SocketAddr)> {
// to-do: websocket proxy.
let tls_type = get_cached_tls_type(url); let tls_type = get_cached_tls_type(url);
let is_tls_type_cached = tls_type.is_some(); let is_tls_type_cached = tls_type.is_some();
let tls_type = tls_type.unwrap_or(TlsType::Rustls); let tls_type = tls_type.unwrap_or(TlsType::Rustls);
let danger_accept_invalid_cert = get_cached_tls_accept_invalid_cert(&url); let danger_accept_invalid_cert = get_cached_tls_accept_invalid_cert(&url);
Self::try_connect( Self::try_connect(
url, url,
local_addr,
proxy_conf,
ms_timeout, ms_timeout,
tls_type, tls_type,
is_tls_type_cached, is_tls_type_cached,
@@ -94,28 +140,83 @@ impl WsFramedStream {
#[async_recursion] #[async_recursion]
async fn try_connect( async fn try_connect(
url: &str, url: &str,
local_addr: Option<SocketAddr>,
proxy_conf: Option<&Socks5Server>,
ms_timeout: u64, ms_timeout: u64,
tls_type: TlsType, tls_type: TlsType,
is_tls_type_cached: bool, is_tls_type_cached: bool,
danger_accept_invalid_cert: Option<bool>, danger_accept_invalid_cert: Option<bool>,
original_danger_accept_invalid_certs: Option<bool>, original_danger_accept_invalid_certs: Option<bool>,
) -> ResultType<WebSocketStream<MaybeTlsStream<TcpStream>>> { ) -> ResultType<(WsStream, SocketAddr)> {
let ws_config = None; let ws_config = None;
let disable_nagle = false;
let request = url let request = url
.into_client_request() .into_client_request()
.map_err(|e| Error::new(ErrorKind::Other, e))?; .map_err(|e| Error::new(ErrorKind::Other, e))?;
let connector = let connector =
Self::get_connector(&tls_type, danger_accept_invalid_cert.unwrap_or(false))?; Self::get_connector(&tls_type, danger_accept_invalid_cert.unwrap_or(false))?;
match timeout( let result = if let Some(proxy_conf) = proxy_conf {
Duration::from_millis(ms_timeout), let host = request.uri().host().ok_or_else(|| {
connect_async_tls_with_config(request, ws_config, disable_nagle, connector), Error::new(ErrorKind::InvalidInput, "Invalid WebSocket URL: no host")
) })?;
.await? let port = request
{ .uri()
Ok((ws_stream, _)) => { .port_u16()
.or_else(|| match request.uri().scheme_str() {
Some("wss") => Some(443),
Some("ws") => Some(80),
_ => None,
})
.ok_or_else(|| {
Error::new(ErrorKind::InvalidInput, "Invalid WebSocket URL scheme")
})?;
let target = format!("{host}:{port}");
timeout(Duration::from_millis(ms_timeout), async {
let stream = crate::tcp::FramedStream::connect(
target.as_str(),
local_addr,
proxy_conf,
ms_timeout,
)
.await?;
let addr = stream.1;
let mut stream = stream.0.into_inner();
if cfg!(any(target_os = "macos", target_os = "ios"))
&& url.starts_with("wss://")
&& matches!(tls_type, TlsType::NativeTls)
&& url::Url::parse(&proxy_conf.proxy)
.is_ok_and(|url| matches!(url.scheme(), "http" | "https"))
{
stream = DynTcpStream(Box::new(FlushBeforeRead(stream)));
}
Ok::<_, anyhow::Error>(
client_async_tls_with_config(request, stream, ws_config, connector)
.await
.map(|(stream, _)| (stream, addr)),
)
})
.await??
.map(Either::Right)
} else {
let disable_nagle = false;
timeout(
Duration::from_millis(ms_timeout),
connect_async_tls_with_config(request, ws_config, disable_nagle, connector),
)
.await?
.map(|(stream, _)| Either::Left(stream))
};
match result {
Ok(stream) => {
upsert_tls_cache(url, tls_type, danger_accept_invalid_cert.unwrap_or(false)); upsert_tls_cache(url, tls_type, danger_accept_invalid_cert.unwrap_or(false));
Ok(ws_stream) match stream {
Either::Left(stream) => {
// MaybeTlsStream::get_ref() handles Plain, NativeTls and Rustls,
// returning the underlying TcpStream for peer_addr().
let addr = stream.get_ref().get_ref().peer_addr()?;
Ok((Either::Left(stream), addr))
}
Either::Right((stream, addr)) => Ok((Either::Right(stream), addr)),
}
} }
Err(e) => match (tls_type, is_tls_type_cached, danger_accept_invalid_cert) { Err(e) => match (tls_type, is_tls_type_cached, danger_accept_invalid_cert) {
(TlsType::Rustls, _, None) => { (TlsType::Rustls, _, None) => {
@@ -126,6 +227,8 @@ impl WsFramedStream {
); );
Self::try_connect( Self::try_connect(
url, url,
local_addr,
proxy_conf,
ms_timeout, ms_timeout,
tls_type, tls_type,
is_tls_type_cached, is_tls_type_cached,
@@ -142,6 +245,8 @@ impl WsFramedStream {
); );
Self::try_connect( Self::try_connect(
url, url,
local_addr,
proxy_conf,
ms_timeout, ms_timeout,
TlsType::NativeTls, TlsType::NativeTls,
is_tls_type_cached, is_tls_type_cached,
@@ -158,6 +263,8 @@ impl WsFramedStream {
); );
Self::try_connect( Self::try_connect(
url, url,
local_addr,
proxy_conf,
ms_timeout, ms_timeout,
tls_type, tls_type,
is_tls_type_cached, is_tls_type_cached,
@@ -189,17 +296,16 @@ impl WsFramedStream {
pub async fn new<T: AsRef<str>>( pub async fn new<T: AsRef<str>>(
url: T, url: T,
_local_addr: Option<SocketAddr>, local_addr: Option<SocketAddr>,
_proxy_conf: Option<&Socks5Server>,
ms_timeout: u64, ms_timeout: u64,
) -> ResultType<Self> { ) -> ResultType<Self> {
let stream = Self::connect(url.as_ref(), ms_timeout).await?; let proxy_conf = if ws_use_proxy() {
let addr = match stream.get_ref() { Config::get_socks()
MaybeTlsStream::Plain(tcp) => tcp.peer_addr()?, } else {
MaybeTlsStream::NativeTls(tls) => tls.get_ref().get_ref().get_ref().peer_addr()?, None
MaybeTlsStream::Rustls(tls) => tls.get_ref().0.peer_addr()?,
_ => return Err(Error::new(ErrorKind::Other, "Unsupported stream type").into()),
}; };
let (stream, addr) =
Self::connect(url.as_ref(), local_addr, proxy_conf.as_ref(), ms_timeout).await?;
let ws = Self { let ws = Self {
stream, stream,
@@ -222,24 +328,14 @@ impl WsFramedStream {
/// fork growing the read buffer a chunk at a time rather than to the length a frame declares. /// fork growing the read buffer a chunk at a time rather than to the length a frame declares.
#[inline] #[inline]
pub fn set_max_packet_length(&mut self, n: usize) { pub fn set_max_packet_length(&mut self, n: usize) {
self.stream.set_config(|c| { let set_config = |c: &mut tungstenite::protocol::WebSocketConfig| {
c.max_message_size = Some(n.min(DEFAULT_MAX_MESSAGE_SIZE)); c.max_message_size = Some(n.min(DEFAULT_MAX_MESSAGE_SIZE));
c.max_frame_size = Some(n.min(DEFAULT_MAX_FRAME_SIZE)); c.max_frame_size = Some(n.min(DEFAULT_MAX_FRAME_SIZE));
}); };
} match &mut self.stream {
Either::Left(stream) => stream.set_config(set_config),
#[inline] Either::Right(stream) => stream.set_config(set_config),
pub async fn from_tcp_stream(stream: TcpStream, addr: SocketAddr) -> ResultType<Self> { }
let ws_stream =
WebSocketStream::from_raw_socket(MaybeTlsStream::Plain(stream), Role::Client, None)
.await;
Ok(Self {
stream: ws_stream,
addr,
encrypt: None,
send_timeout: 0,
})
} }
#[inline] #[inline]
@@ -439,6 +535,20 @@ mod tests {
use super::*; use super::*;
use crate::config::{keys, Config}; use crate::config::{keys, Config};
use tokio::{io::AsyncWriteExt, net::TcpListener}; use tokio::{io::AsyncWriteExt, net::TcpListener};
use tungstenite::protocol::Role;
async fn ws_from_tcp(stream: TcpStream, addr: SocketAddr) -> WsFramedStream {
let ws_stream =
WebSocketStream::from_raw_socket(MaybeTlsStream::Plain(stream), Role::Client, None)
.await;
WsFramedStream {
stream: Either::Left(ws_stream),
addr,
encrypt: None,
send_timeout: 0,
}
}
#[tokio::test(flavor = "current_thread")] #[tokio::test(flavor = "current_thread")]
async fn test_secured_stream_rejects_plaintext_text() { async fn test_secured_stream_rejects_plaintext_text() {
@@ -455,7 +565,7 @@ mod tests {
}); });
let tcp = TcpStream::connect(addr).await.unwrap(); let tcp = TcpStream::connect(addr).await.unwrap();
let mut client = WsFramedStream::from_tcp_stream(tcp, addr).await.unwrap(); let mut client = ws_from_tcp(tcp, addr).await;
client.set_key(Key([0x42; sodiumoxide::crypto::secretbox::KEYBYTES])); client.set_key(Key([0x42; sodiumoxide::crypto::secretbox::KEYBYTES]));
let result = client.next().await; let result = client.next().await;
@@ -629,7 +739,7 @@ mod tests {
let addr = listener.local_addr().unwrap(); let addr = listener.local_addr().unwrap();
let client = TcpStream::connect(addr).await.unwrap(); let client = TcpStream::connect(addr).await.unwrap();
let (server, _) = listener.accept().await.unwrap(); let (server, _) = listener.accept().await.unwrap();
let ws = WsFramedStream::from_tcp_stream(client, addr).await.unwrap(); let ws = ws_from_tcp(client, addr).await;
(ws, server) (ws, server)
} }
@@ -640,7 +750,10 @@ mod tests {
ws.set_max_packet_length(CAP); ws.set_max_packet_length(CAP);
// One byte over, header only: refused before there is any payload to buffer. // One byte over, header only: refused before there is any payload to buffer.
server.write_all(&ws_binary_frame(CAP + 1, false)).await.unwrap(); server
.write_all(&ws_binary_frame(CAP + 1, false))
.await
.unwrap();
match timeout(Duration::from_secs(5), ws.next()).await { match timeout(Duration::from_secs(5), ws.next()).await {
Ok(Some(Err(e))) => assert!(e.to_string().contains("Message too long"), "{}", e), Ok(Some(Err(e))) => assert!(e.to_string().contains("Message too long"), "{}", e),
Ok(other) => panic!( Ok(other) => panic!(
@@ -657,18 +770,28 @@ mod tests {
let (mut ws, mut server) = ws_loopback().await; let (mut ws, mut server) = ws_loopback().await;
ws.set_max_packet_length(CAP); ws.set_max_packet_length(CAP);
server.write_all(&ws_binary_frame(8 * 1024, true)).await.unwrap(); server
.write_all(&ws_binary_frame(8 * 1024, true))
.await
.unwrap();
let got = ws.next().await.unwrap().unwrap(); let got = ws.next().await.unwrap().unwrap();
assert_eq!(got.len(), 8 * 1024, "a message under the cap still arrives"); assert_eq!(got.len(), 8 * 1024, "a message under the cap still arrives");
ws.set_max_packet_length(usize::MAX); ws.set_max_packet_length(usize::MAX);
server.write_all(&ws_binary_frame(200_000, true)).await.unwrap(); server
.write_all(&ws_binary_frame(200_000, true))
.await
.unwrap();
let got = timeout(Duration::from_secs(5), ws.next()) let got = timeout(Duration::from_secs(5), ws.next())
.await .await
.expect("next() hung after the cap was lifted") .expect("next() hung after the cap was lifted")
.unwrap() .unwrap()
.unwrap(); .unwrap();
assert_eq!(got.len(), 200_000, "lifting the cap lets a large message through again"); assert_eq!(
got.len(),
200_000,
"lifting the cap lets a large message through again"
);
} }
// Two frames each under the cap that reassemble to a message over it: the frame bound lets // Two frames each under the cap that reassemble to a message over it: the frame bound lets
@@ -680,8 +803,14 @@ mod tests {
ws.set_max_packet_length(CAP); ws.set_max_packet_length(CAP);
// Binary with FIN clear, then a continuation with FIN set: 12 KiB each, 24 KiB together. // Binary with FIN clear, then a continuation with FIN set: 12 KiB each, 24 KiB together.
server.write_all(&ws_frame(0x02, 12 * 1024, true)).await.unwrap(); server
server.write_all(&ws_frame(0x80, 12 * 1024, true)).await.unwrap(); .write_all(&ws_frame(0x02, 12 * 1024, true))
.await
.unwrap();
server
.write_all(&ws_frame(0x80, 12 * 1024, true))
.await
.unwrap();
match timeout(Duration::from_secs(5), ws.next()).await { match timeout(Duration::from_secs(5), ws.next()).await {
Ok(Some(Err(e))) => assert!(e.to_string().contains("Message too long"), "{}", e), Ok(Some(Err(e))) => assert!(e.to_string().contains("Message too long"), "{}", e),
Ok(other) => panic!( Ok(other) => panic!(
+12
View File
@@ -0,0 +1,12 @@
-----BEGIN CERTIFICATE-----
MIIB1TCCAXugAwIBAgIUU5BrjQyAn16xm5+tfdJlWneyQA8wCgYIKoZIzj0EAwIw
IzEhMB8GA1UEAwwYd2Vic29ja2V0LXRhcmdldC5pbnZhbGlkMCAXDTI2MDkwODEx
NTg0NloYDzIxMjYwODE1MTE1ODQ2WjAjMSEwHwYDVQQDDBh3ZWJzb2NrZXQtdGFy
Z2V0LmludmFsaWQwWTATBgcqhkjOPQIBBggqhkjOPQMBBwNCAATGz2rNvO/J5IZc
R2tAP18dE/cDqaqhSI/GrG9wsRKRDypx6C1AlhHp67+xsBxfxGX8qP+2IYqxPeEf
OaXTNmRuo4GKMIGHMB0GA1UdDgQWBBSQaINp0ts4Dxa2+7VJqavc7dykejAfBgNV
HSMEGDAWgBSQaINp0ts4Dxa2+7VJqavc7dykejAPBgNVHRMBAf8EBTADAQH/MDQG
A1UdEQQtMCuCGHdlYnNvY2tldC10YXJnZXQuaW52YWxpZIIJbG9jYWxob3N0hwR/
AAABMAoGCCqGSM49BAMCA0gAMEUCIQCfNhZYsn3MB0781bnGmxY5y+qvq1bl1xgx
6hB2jeCiwQIgeJjyi0sD3kg9IwExk1tvY1CqME8S3vX6B3A6OihqYUg=
-----END CERTIFICATE-----
+5
View File
@@ -0,0 +1,5 @@
-----BEGIN PRIVATE KEY-----
MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgfMmTjDYIQrnprR5R
3NkH2woEGhwRtkkOSoFqi9yAeN6hRANCAATGz2rNvO/J5IZcR2tAP18dE/cDqaqh
SI/GrG9wsRKRDypx6C1AlhHp67+xsBxfxGX8qP+2IYqxPeEfOaXTNmRu
-----END PRIVATE KEY-----
+419
View File
@@ -0,0 +1,419 @@
use futures::{SinkExt, StreamExt};
use hbb_common::{
config::{keys, Socks5Server, APP_NAME, OVERWRITE_SETTINGS},
socket_client,
tcp::DynTcpStream,
tls::{upsert_tls_cache, TlsType},
websocket::WsFramedStream,
Stream,
};
use rustls_pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer};
use std::{sync::Arc, time::Duration};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
time::timeout,
};
use tokio_rustls::{rustls, TlsAcceptor};
use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};
const TIMEOUT: u64 = 5_000;
const PAYLOAD: &[u8] = b"RustDesk websocket proxy test";
fn tls_acceptor() -> TlsAcceptor {
// This self-signed identity is only used by the local mock servers.
let cert = CertificateDer::from_pem_slice(include_bytes!("data/websocket-cert.pem")).unwrap();
let key = PrivateKeyDer::from_pem_slice(include_bytes!("data/websocket-key.pem")).unwrap();
let config = rustls::ServerConfig::builder_with_provider(Arc::new(
rustls::crypto::ring::default_provider(),
))
.with_safe_default_protocol_versions()
.unwrap()
.with_no_client_auth()
.with_single_cert(vec![cert], key)
.unwrap();
TlsAcceptor::from(Arc::new(config))
}
async fn read_headers(stream: &mut DynTcpStream) -> String {
let mut headers = Vec::new();
while !headers.ends_with(b"\r\n\r\n") {
assert!(headers.len() < 4096);
headers.push(stream.read_u8().await.unwrap());
}
String::from_utf8(headers).unwrap()
}
async fn echo(stream: DynTcpStream, path: &str, host: &str) {
let mut ws =
tokio_tungstenite::accept_hdr_async(stream, |request: &Request, response: Response| {
assert_eq!(request.uri().to_string(), path);
assert_eq!(request.headers()["host"], host);
assert!(!request.headers().contains_key("proxy-authorization"));
Ok(response)
})
.await
.unwrap();
let message = ws.next().await.unwrap().unwrap();
assert!(message.is_binary());
assert_eq!(message.into_data().as_ref(), PAYLOAD);
ws.send(PAYLOAD.to_vec().into()).await.unwrap();
ws.close(None).await.unwrap();
}
async fn exchange(mut stream: Stream) {
stream.send_bytes(PAYLOAD.to_vec().into()).await.unwrap();
let reply = stream.next_timeout(TIMEOUT).await.unwrap().unwrap();
assert_eq!(reply.as_ref(), PAYLOAD);
if matches!(&stream, Stream::WebSocket(_)) {
assert!(timeout(Duration::from_millis(TIMEOUT), stream.next())
.await
.unwrap()
.is_none());
}
}
fn option(key: &str, value: &str) {
OVERWRITE_SETTINGS
.write()
.unwrap()
.insert(key.into(), value.into());
}
async fn proxy_round_trip(
scheme: &str,
url: &str,
tls_type: TlsType,
auth: bool,
configured: bool,
) {
option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "Y");
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let proxy = Socks5Server {
proxy: format!("{scheme}://{}", listener.local_addr().unwrap()),
username: if auth { "user".into() } else { String::new() },
password: if auth { "pass".into() } else { String::new() },
};
upsert_tls_cache(&proxy.proxy, TlsType::Rustls, true);
upsert_tls_cache(url, tls_type, true);
let target = url::Url::parse(url).unwrap();
let host = target.host_str().unwrap().to_owned();
let port = target.port_or_known_default().unwrap();
let authority = match target.port() {
Some(port) => format!("{host}:{port}"),
None => host.clone(),
};
let path = format!(
"{}{}",
target.path(),
target.query().map(|q| format!("?{q}")).unwrap_or_default()
);
let wss = target.scheme() == "wss";
let scheme = scheme.to_owned();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut stream = DynTcpStream(Box::new(stream));
if scheme == "https" {
stream = DynTcpStream(Box::new(tls_acceptor().accept(stream).await.unwrap()));
}
if scheme == "socks5" {
let mut greeting = [0; 3];
stream.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [5, 1, 0]);
stream.write_all(&[5, 0]).await.unwrap();
let mut request = [0; 5];
stream.read_exact(&mut request).await.unwrap();
assert_eq!(&request[..4], &[5, 1, 0, 3]);
let mut name = vec![0; request[4] as usize];
stream.read_exact(&mut name).await.unwrap();
assert_eq!(name, host.as_bytes());
assert_eq!(stream.read_u16().await.unwrap(), port);
stream
.write_all(&[5, 0, 0, 1, 127, 0, 0, 1, 0, 0])
.await
.unwrap();
} else {
let headers = read_headers(&mut stream).await;
assert!(headers.starts_with(&format!("CONNECT {host}:{port} HTTP/1.1\r\n")));
assert!(headers.contains(&format!("Host: {host}:{port}\r\n")));
assert_eq!(
headers.contains("Proxy-Authorization: Basic dXNlcjpwYXNz\r\n"),
auth
);
stream
.write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
.await
.unwrap();
}
stream.flush().await.unwrap();
if wss {
let tls = tls_acceptor().accept(stream).await.unwrap();
assert_eq!(tls.get_ref().1.server_name(), Some(host.as_str()));
stream = DynTcpStream(Box::new(tls));
}
echo(stream, &path, &authority).await;
});
option(keys::OPTION_PROXY_URL, &proxy.proxy);
option(keys::OPTION_PROXY_USERNAME, &proxy.username);
option(keys::OPTION_PROXY_PASSWORD, &proxy.password);
let stream = if configured {
option(keys::OPTION_ALLOW_WEBSOCKET, "Y");
option("custom-rendezvous-server", target.host_str().unwrap());
option(
"api-server",
if wss {
"https://websocket-target.invalid"
} else {
"http://websocket-target.invalid"
},
);
socket_client::connect_tcp("websocket-target.invalid:21116", TIMEOUT)
.await
.unwrap()
} else {
WsFramedStream::new(url, None, TIMEOUT)
.await
.map(Stream::WebSocket)
.unwrap()
};
exchange(stream).await;
timeout(Duration::from_millis(TIMEOUT), server)
.await
.unwrap()
.unwrap();
}
async fn rejected_or_stalled_proxy(stalled: bool, configured: bool) {
let direct = TcpListener::bind("127.0.0.1:0").await.unwrap();
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let proxy = Socks5Server {
proxy: format!("http://{}", listener.local_addr().unwrap()),
..Default::default()
};
option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "Y");
option(keys::OPTION_PROXY_URL, &proxy.proxy);
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut stream = DynTcpStream(Box::new(stream));
read_headers(&mut stream).await;
if !stalled {
stream
.write_all(b"HTTP/1.1 407 Proxy Authentication Required\r\n\r\n")
.await
.unwrap();
}
let mut byte = [0];
assert_eq!(stream.read(&mut byte).await.unwrap(), 0);
});
let result = timeout(Duration::from_millis(TIMEOUT), async {
let url = format!("wss://{}/ws/id", direct.local_addr().unwrap());
if configured {
socket_client::connect_tcp(url, 200).await
} else {
WsFramedStream::new(url, None, 200)
.await
.map(Stream::WebSocket)
}
})
.await
.unwrap();
let error = result.err().expect("the proxy must fail");
if stalled {
assert!(error
.downcast_ref::<tokio::time::error::Elapsed>()
.is_some());
} else {
assert!(error.to_string().contains("407"), "{}", error);
}
assert!(timeout(Duration::from_millis(50), direct.accept())
.await
.is_err());
timeout(Duration::from_millis(TIMEOUT), server)
.await
.unwrap()
.unwrap();
}
async fn direct_round_trip(setting: &str, configured: bool, wss: bool) {
option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, setting);
let unused_proxy = TcpListener::bind("127.0.0.1:0").await.unwrap();
if configured {
option(
keys::OPTION_PROXY_URL,
&format!("http://{}", unused_proxy.local_addr().unwrap()),
);
} else {
OVERWRITE_SETTINGS
.write()
.unwrap()
.remove(keys::OPTION_PROXY_URL);
}
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let url = format!("{}://{addr}/ws/id", if wss { "wss" } else { "ws" });
upsert_tls_cache(&url, TlsType::Rustls, true);
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut stream = DynTcpStream(Box::new(stream));
if wss {
stream = DynTcpStream(Box::new(tls_acceptor().accept(stream).await.unwrap()));
}
echo(stream, "/ws/id", &addr.to_string()).await;
});
let stream = socket_client::connect_tcp(url, TIMEOUT).await.unwrap();
exchange(stream).await;
timeout(Duration::from_millis(TIMEOUT), server)
.await
.unwrap()
.unwrap();
assert!(timeout(Duration::from_millis(50), unused_proxy.accept())
.await
.is_err());
}
async fn tcp_proxy_round_trip() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let proxy = Socks5Server {
proxy: format!("http://{}", listener.local_addr().unwrap()),
..Default::default()
};
option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "N");
option(keys::OPTION_PROXY_URL, &proxy.proxy);
option(keys::OPTION_PROXY_USERNAME, "");
option(keys::OPTION_PROXY_PASSWORD, "");
let server = tokio::spawn(async move {
let (stream, addr) = listener.accept().await.unwrap();
let mut stream = DynTcpStream(Box::new(stream));
let headers = read_headers(&mut stream).await;
assert!(headers.starts_with("CONNECT tcp-target.invalid:21116 HTTP/1.1\r\n"));
stream
.write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
.await
.unwrap();
let mut stream = hbb_common::tcp::FramedStream::from(stream, addr);
let message = stream.next().await.unwrap().unwrap();
assert_eq!(message.as_ref(), PAYLOAD);
stream.send_bytes(message.freeze()).await.unwrap();
});
let stream = socket_client::connect_tcp_local("tcp-target.invalid:21116", None, TIMEOUT)
.await
.unwrap();
exchange(stream).await;
timeout(Duration::from_millis(TIMEOUT), server)
.await
.unwrap()
.unwrap();
}
#[tokio::test]
async fn websocket_proxy_connections() {
// Keep configuration changes inside this integration test process.
*APP_NAME.write().unwrap() = format!("RustDesk WebSocket Test {}", std::process::id());
option(keys::OPTION_ALLOW_INSECURE_TLS_FALLBACK, "Y");
for (scheme, url, tls_type, auth, configured) in [
(
"http",
"ws://websocket-target.invalid/ws/id",
TlsType::Plain,
false,
true,
),
(
"http",
"wss://websocket-target.invalid/ws/id",
TlsType::Rustls,
true,
true,
),
(
"http",
"wss://websocket-target.invalid:8443/ws/relay?test=1",
TlsType::NativeTls,
true,
false,
),
(
"https",
"ws://websocket-target.invalid/ws/id",
TlsType::Plain,
true,
false,
),
(
"https",
"wss://websocket-target.invalid/ws/relay",
TlsType::Rustls,
true,
false,
),
(
"https",
"wss://websocket-target.invalid/ws/relay",
TlsType::NativeTls,
true,
false,
),
(
"socks5",
"ws://websocket-target.invalid/ws/id",
TlsType::Plain,
false,
false,
),
(
"socks5",
"wss://websocket-target.invalid/ws/relay",
TlsType::Rustls,
false,
false,
),
(
"socks5",
"wss://websocket-target.invalid/ws/relay",
TlsType::NativeTls,
false,
false,
),
] {
proxy_round_trip(scheme, url, tls_type, auth, configured).await;
}
for scheme in ["http", "https", "socks5"] {
for (url, tls_type) in [
(
"ws://websocket-target.invalid:8080/ws/id?test=1",
TlsType::Plain,
),
(
"wss://websocket-target.invalid:8443/ws/relay?test=1",
TlsType::Rustls,
),
(
"wss://websocket-target.invalid:8443/ws/relay?test=1",
TlsType::NativeTls,
),
] {
proxy_round_trip(scheme, url, tls_type, scheme != "socks5", false).await;
}
}
for configured in [false, true] {
rejected_or_stalled_proxy(false, configured).await;
rejected_or_stalled_proxy(true, configured).await;
}
tcp_proxy_round_trip().await;
for wss in [false, true] {
direct_round_trip("", true, wss).await;
direct_round_trip("N", true, wss).await;
direct_round_trip("Y", false, wss).await;
}
option(keys::OPTION_ALLOW_WEBSOCKET_PROXY, "N");
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
drop(listener);
for scheme in ["wss", "ws"] {
let url = format!("{scheme}://{addr}/ws/id");
let error = WsFramedStream::new(url, None, TIMEOUT).await.err().unwrap();
assert!(error.to_string().starts_with("WebSocket error:"));
}
}