diff --git a/wstunnel/src/config.rs b/wstunnel/src/config.rs index 161f99c7..f453b5d0 100644 --- a/wstunnel/src/config.rs +++ b/wstunnel/src/config.rs @@ -72,6 +72,20 @@ pub struct Client { ))] pub connection_retry_max_backoff: Duration, + /// Timeout used when establishing connections for tunnel traffic. + /// This includes connecting to the wstunnel server and reverse tunnel endpoint connects. + /// For reverse SOCKS5, client/server values can differ; the shorter side effectively limits wait time. + #[cfg_attr(feature = "clap", arg( + long, + value_name = "DURATION(s|m|h)", + default_value = "3s", + value_parser = parsers::parse_duration_sec, + alias = "timeout-connect-sec", + env = "WSTUNNEL_TIMEOUT_CONNECT", + verbatim_doc_comment + ))] + pub timeout_connect: Duration, + /// When using reverse tunnel, the client will try to always keep a connection to the server to await for new tunnels /// This delay is the maximum of time the client will wait before trying to reconnect to the server in case of failure. /// The client follows an exponential backoff strategy until it reaches this maximum delay @@ -301,6 +315,19 @@ pub struct Server { #[cfg_attr(feature = "clap", arg(long, default_value = "false", verbatim_doc_comment))] pub websocket_mask_frame: bool, + /// Timeout used when establishing upstream connections and waiting for reverse SOCKS5 connect handshakes. + /// For reverse SOCKS5, client/server values can differ; the shorter side effectively limits wait time. + #[cfg_attr(feature = "clap", arg( + long, + value_name = "DURATION(s|m|h)", + default_value = "3s", + value_parser = parsers::parse_duration_sec, + alias = "timeout-connect-sec", + env = "WSTUNNEL_TIMEOUT_CONNECT", + verbatim_doc_comment + ))] + pub timeout_connect: Duration, + /// Dns resolver to use to lookup ips of domain name /// This option is not going to work if you use transparent proxy /// Can be specified multiple time diff --git a/wstunnel/src/lib.rs b/wstunnel/src/lib.rs index abb54dd4..c290cfe9 100644 --- a/wstunnel/src/lib.rs +++ b/wstunnel/src/lib.rs @@ -201,7 +201,7 @@ pub async fn create_client( http_headers: args.http_headers.into_iter().filter(|(k, _)| k != HOST).collect(), http_headers_file: args.http_headers_file, http_header_host: host_header, - timeout_connect: Duration::from_secs(10), + timeout_connect: args.timeout_connect, websocket_ping_frequency: args .websocket_ping_frequency .or(Some(Duration::from_secs(30))) @@ -555,7 +555,7 @@ async fn run_server_impl(args: Server, executor: impl TokioExecutorRef) -> anyho .websocket_ping_frequency .or(Some(Duration::from_secs(30))) .filter(|d| d.as_secs() > 0), - timeout_connect: Duration::from_secs(10), + timeout_connect: args.timeout_connect, websocket_mask_frame: args.websocket_mask_frame, tls: tls_config, dns_resolver: DnsResolver::new_from_urls( diff --git a/wstunnel/src/protocols/socks5/tcp_server.rs b/wstunnel/src/protocols/socks5/tcp_server.rs index c994922c..5739cd79 100644 --- a/wstunnel/src/protocols/socks5/tcp_server.rs +++ b/wstunnel/src/protocols/socks5/tcp_server.rs @@ -1,16 +1,18 @@ use super::udp_server::{Socks5UdpStream, Socks5UdpStreamWriter}; use crate::tunnel::LocalProtocol; use anyhow::Context; -use fast_socks5::Socks5Command; use fast_socks5::server::Socks5ServerProtocol; use fast_socks5::util::target_addr::TargetAddr; +use fast_socks5::{ReplyError, Socks5Command, consts}; use futures_util::{Stream, StreamExt, stream}; use std::io::{Error, IoSlice}; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::task::Poll; use std::time::Duration; -use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, ReadBuf}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf}; use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf}; use tokio::net::{TcpListener, TcpStream}; use tokio::select; @@ -23,6 +25,119 @@ use url::Host; /// hold it: a real client sends its handshake immediately. const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10); +/// Withholds the SOCKS5 reply produced by [`Socks5ServerProtocol`]. +/// +/// A CONNECT reply must not reach the client before the upstream connection has +/// actually been established: a client that is told `Succeeded` cannot tell a +/// working tunnel from one that fails immediately afterwards, and never sees the +/// real status code. The protocol type owns the client socket and only returns +/// it from `reply_success`, so the reply it writes is produced and discarded +/// here. The real reply is sent once the connect result is known, by +/// [`Socks5WriteHalf::send_reply_if_needed`]. +/// +/// Writes pass straight through until [`WithheldReplyHandle::arm`] is called, so +/// the greeting and authentication exchange behave exactly as before. +struct WithheldReply { + inner: T, + armed: Arc, +} + +#[derive(Clone)] +struct WithheldReplyHandle(Arc); + +impl WithheldReplyHandle { + /// Discard everything written from here on. + fn arm(&self) { + self.0.store(true, Ordering::Release); + } +} + +impl WithheldReply { + fn new(inner: T) -> (Self, WithheldReplyHandle) { + let armed = Arc::new(AtomicBool::new(false)); + ( + Self { + inner, + armed: Arc::clone(&armed), + }, + WithheldReplyHandle(armed), + ) + } + + fn into_inner(self) -> T { + self.inner + } + + fn is_armed(&self) -> bool { + self.armed.load(Ordering::Acquire) + } +} + +impl AsyncRead for WithheldReply { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for WithheldReply { + fn poll_write(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>, buf: &[u8]) -> Poll> { + if self.is_armed() { + return Poll::Ready(Ok(buf.len())); + } + Pin::new(&mut self.get_mut().inner).poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll> { + if self.is_armed() { + return Poll::Ready(Ok(())); + } + Pin::new(&mut self.get_mut().inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll> { + if self.is_armed() { + return Poll::Ready(Ok(())); + } + Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) + } +} + +/// Build a raw SOCKS5 reply packet. +/// +/// Upstream dropped this helper when the CONNECT path moved to +/// `Socks5ServerProtocol`, whose reply methods consume the protocol object. The +/// deferred reply is sent after the socket has already been split, so the bytes +/// have to be assembled directly. +fn new_reply(error: &ReplyError, sock_addr: SocketAddr) -> Vec { + let (addr_type, mut ip_oct, mut port) = match sock_addr { + SocketAddr::V4(sock) => ( + consts::SOCKS5_ADDR_TYPE_IPV4, + sock.ip().octets().to_vec(), + sock.port().to_be_bytes().to_vec(), + ), + SocketAddr::V6(sock) => ( + consts::SOCKS5_ADDR_TYPE_IPV6, + sock.ip().octets().to_vec(), + sock.port().to_be_bytes().to_vec(), + ), + }; + + let mut reply = vec![ + consts::SOCKS5_VERSION, + error.as_u8(), // transform the error into byte code + 0x00, // reserved + addr_type, // address type (ipv4, v6, domain) + ]; + reply.append(&mut ip_oct); + reply.append(&mut port); + + reply +} + #[allow(clippy::type_complexity)] pub struct Socks5Listener { socks_server: Pin> + Send>>, @@ -34,19 +149,22 @@ pub enum Socks5ReadHalf { } pub enum Socks5WriteHalf { - Tcp(OwnedWriteHalf), + Tcp { + writer: OwnedWriteHalf, + pending_reply: bool, + }, Udp(Socks5UdpStreamWriter), } pub enum Socks5Stream { - Tcp(TcpStream), + Tcp { stream: TcpStream, pending_reply: bool }, Udp((Socks5UdpStream, Socks5UdpStreamWriter)), } impl Socks5Stream { pub fn local_protocol(&self) -> LocalProtocol { match self { - Self::Tcp(_) => LocalProtocol::Tcp { proxy_protocol: false }, // TODO: Implement proxy protocol + Self::Tcp { .. } => LocalProtocol::Tcp { proxy_protocol: false }, // TODO: Implement proxy protocol Self::Udp(s) => LocalProtocol::Udp { timeout: s.0.watchdog_deadline.as_ref().map(|x| x.period()), }, @@ -55,9 +173,15 @@ impl Socks5Stream { pub fn into_split(self) -> (Socks5ReadHalf, Socks5WriteHalf) { match self { - Self::Tcp(s) => { - let (r, w) = s.into_split(); - (Socks5ReadHalf::Tcp(r), Socks5WriteHalf::Tcp(w)) + Self::Tcp { stream, pending_reply } => { + let (r, w) = stream.into_split(); + ( + Socks5ReadHalf::Tcp(r), + Socks5WriteHalf::Tcp { + writer: w, + pending_reply, + }, + ) } Self::Udp((r, w)) => (Socks5ReadHalf::Udp(r), Socks5WriteHalf::Udp(w)), } @@ -122,6 +246,9 @@ pub async fn run_server( // Authenticate the connection, bounding the handshake read so a // silent client cannot hold the accept loop (see HANDSHAKE_TIMEOUT). + // See WithheldReply: the CONNECT reply must not go out until the + // upstream connect result is known. + let (socket, withhold) = WithheldReply::new(socket); let proto = if let Some((ref username, ref password)) = credentials { let username = username.clone(); let password = password.clone(); @@ -179,13 +306,14 @@ pub async fn run_server( // Special case for UDP Associate where we return the bind addr of the udp server if matches!(cmd, Socks5Command::UDPAssociate) { - let mut cnx = match proto.reply_success(bind).await { + let cnx = match proto.reply_success(bind).await { Ok(cnx) => cnx, Err(err) => { warn!("Cannot reply to socks5 udp client: {}", err); continue; } }; + let mut cnx = cnx.into_inner(); tasks.spawn(async move { let mut buf = [0u8; 8]; loop { @@ -199,6 +327,11 @@ pub async fn run_server( continue; }; + // Withhold the CONNECT reply. The protocol type owns the socket and + // only hands it back from reply_success, so let it produce the reply + // and drop it: the real one is sent once the tunnel connect result + // is known, carrying the true status code. + withhold.arm(); let cnx = match proto .reply_success(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 0)) .await @@ -211,7 +344,13 @@ pub async fn run_server( }; return Some(( - Ok((Socks5Stream::Tcp(cnx), (host, port))), + Ok(( + Socks5Stream::Tcp { + stream: cnx.into_inner(), + pending_reply: true, + }, + (host, port), + )), (listener, udp_server, tasks, credentials), )); } @@ -226,6 +365,32 @@ pub async fn run_server( } impl Unpin for Socks5Stream {} + +impl Socks5WriteHalf { + pub(crate) async fn send_reply_if_needed(&mut self, error: ReplyError) -> anyhow::Result<()> { + let should_reply = match self { + Self::Tcp { pending_reply, .. } => { + if *pending_reply { + *pending_reply = false; + true + } else { + false + } + } + Self::Udp(_) => false, + }; + + if should_reply { + let bind_addr = match error { + ReplyError::Succeeded => SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0), + _ => SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0), + }; + self.write_all(&new_reply(&error, bind_addr)).await?; + } + + Ok(()) + } +} impl AsyncRead for Socks5ReadHalf { fn poll_read( self: Pin<&mut Self>, @@ -242,21 +407,21 @@ impl AsyncRead for Socks5ReadHalf { impl AsyncWrite for Socks5WriteHalf { fn poll_write(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>, buf: &[u8]) -> Poll> { match self.get_mut() { - Self::Tcp(s) => Pin::new(s).poll_write(cx, buf), + Self::Tcp { writer, .. } => Pin::new(writer).poll_write(cx, buf), Self::Udp(s) => Pin::new(s).poll_write(cx, buf), } } fn poll_flush(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll> { match self.get_mut() { - Self::Tcp(s) => Pin::new(s).poll_flush(cx), + Self::Tcp { writer, .. } => Pin::new(writer).poll_flush(cx), Self::Udp(s) => Pin::new(s).poll_flush(cx), } } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll> { match self.get_mut() { - Self::Tcp(s) => Pin::new(s).poll_shutdown(cx), + Self::Tcp { writer, .. } => Pin::new(writer).poll_shutdown(cx), Self::Udp(s) => Pin::new(s).poll_shutdown(cx), } } @@ -267,15 +432,70 @@ impl AsyncWrite for Socks5WriteHalf { bufs: &[IoSlice<'_>], ) -> Poll> { match self.get_mut() { - Self::Tcp(s) => Pin::new(s).poll_write_vectored(cx, bufs), + Self::Tcp { writer, .. } => Pin::new(writer).poll_write_vectored(cx, bufs), Self::Udp(s) => Pin::new(s).poll_write_vectored(cx, bufs), } } fn is_write_vectored(&self) -> bool { match self { - Self::Tcp(s) => s.is_write_vectored(), + Self::Tcp { writer, .. } => writer.is_write_vectored(), Self::Udp(s) => s.is_write_vectored(), } } } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::AsyncReadExt; + use tokio::net::TcpListener; + use tokio::time::{Duration, timeout}; + + async fn tcp_writer_with_client(pending_reply: bool) -> (tokio::net::TcpStream, Socks5WriteHalf) { + let listener = TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)) + .await + .unwrap(); + let client = tokio::net::TcpStream::connect(listener.local_addr().unwrap()) + .await + .unwrap(); + let (server, _) = listener.accept().await.unwrap(); + let (_, writer) = server.into_split(); + + (client, Socks5WriteHalf::Tcp { writer, pending_reply }) + } + + #[tokio::test] + async fn deferred_socks5_reply_is_sent_only_when_requested() { + let (mut client, mut writer) = tcp_writer_with_client(true).await; + + assert!(timeout(Duration::from_millis(30), client.read_u8()).await.is_err()); + + writer.send_reply_if_needed(ReplyError::Succeeded).await.unwrap(); + let mut reply = [0u8; 10]; + client.read_exact(&mut reply).await.unwrap(); + + assert_eq!(reply[0], consts::SOCKS5_VERSION); + assert_eq!(reply[1], ReplyError::Succeeded.as_u8()); + } + + #[tokio::test] + async fn deferred_socks5_reply_is_one_shot() { + let (mut client, mut writer) = tcp_writer_with_client(true).await; + + writer.send_reply_if_needed(ReplyError::GeneralFailure).await.unwrap(); + let mut reply = [0u8; 10]; + client.read_exact(&mut reply).await.unwrap(); + assert_eq!(reply[1], ReplyError::GeneralFailure.as_u8()); + + writer.send_reply_if_needed(ReplyError::Succeeded).await.unwrap(); + assert!(timeout(Duration::from_millis(30), client.read_u8()).await.is_err()); + } + + #[tokio::test] + async fn socks5_reply_is_skipped_when_not_pending() { + let (mut client, mut writer) = tcp_writer_with_client(false).await; + writer.send_reply_if_needed(ReplyError::Succeeded).await.unwrap(); + assert!(timeout(Duration::from_millis(30), client.read_u8()).await.is_err()); + } +} diff --git a/wstunnel/src/test_integrations.rs b/wstunnel/src/test_integrations.rs index 50f6d71a..943fa30c 100644 --- a/wstunnel/src/test_integrations.rs +++ b/wstunnel/src/test_integrations.rs @@ -6,7 +6,7 @@ use crate::restrictions::types; use crate::restrictions::types::{AllowConfig, MatchConfig, RestrictionConfig, RestrictionsRules}; use crate::somark::SoMark; use crate::tunnel::client::{TlsClientConfig, WsClient, WsClientConfig}; -use crate::tunnel::listeners::{TcpTunnelListener, UdpTunnelListener}; +use crate::tunnel::listeners::{Socks5TunnelListener, TcpTunnelListener, UdpTunnelListener}; use crate::tunnel::server::{TlsServerConfig, WsServer, WsServerConfig}; use crate::tunnel::transport::{TransportAddr, TransportScheme}; use bytes::BytesMut; @@ -18,6 +18,7 @@ use rstest::{fixture, rstest}; use scopeguard::defer; use serial_test::serial; use std::collections::{BTreeSet, HashMap}; +use std::io::{Error, ErrorKind}; use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -216,6 +217,46 @@ fn no_restrictions() -> RestrictionsRules { } } +async fn socks5_connect( + proxy_addr: SocketAddr, + target_addr: SocketAddr, +) -> std::io::Result<(tokio::net::TcpStream, u8)> { + let mut stream = tokio::net::TcpStream::connect(proxy_addr).await?; + stream.write_all(&[0x05, 0x01, 0x00]).await?; + + let mut auth_reply = [0u8; 2]; + stream.read_exact(&mut auth_reply).await?; + if auth_reply != [0x05, 0x00] { + return Err(Error::new( + ErrorKind::InvalidData, + format!("unexpected socks auth reply: {auth_reply:?}"), + )); + } + + let SocketAddr::V4(target_addr) = target_addr else { + return Err(Error::new(ErrorKind::InvalidInput, "test helper only supports IPv4 targets")); + }; + let mut connect_req = [0u8; 10]; + connect_req[0] = 0x05; + connect_req[1] = 0x01; + connect_req[2] = 0x00; + connect_req[3] = 0x01; + connect_req[4..8].copy_from_slice(&target_addr.ip().octets()); + connect_req[8..10].copy_from_slice(&target_addr.port().to_be_bytes()); + + stream.write_all(&connect_req).await?; + let mut reply = [0u8; 10]; + stream.read_exact(&mut reply).await?; + if reply[0] != 0x05 { + return Err(Error::new( + ErrorKind::InvalidData, + format!("unexpected socks connect reply: {reply:?}"), + )); + } + + Ok((stream, reply[1])) +} + #[rstest] #[timeout(Duration::from_secs(10))] #[tokio::test] @@ -356,6 +397,103 @@ async fn test_tcp_tunnel_webtransport( assert_eq!(&buf[..6], b"world!"); } +#[rstest] +#[timeout(Duration::from_secs(15))] +#[tokio::test] +#[serial] +async fn test_socks5_tunnel_connect_success_webtransport( + server_webtransport: WsServer, + no_restrictions: RestrictionsRules, + dns_resolver: DnsResolver, +) { + let tunnel_listen = free_addr().0; + let endpoint_listen = free_addr().0; + + let server_port = server_webtransport.config.bind.port(); + let server_h = tokio::spawn(server_webtransport.serve(no_restrictions)); + defer! { server_h.abort(); }; + + let client = client_webtransport(server_port, dns_resolver).await; + let socks_listener = Socks5TunnelListener::new(tunnel_listen, None, None).await.unwrap(); + tokio::spawn(async move { client.run_tunnel(socks_listener).await.unwrap() }); + + let mut endpoint_listener = protocols::tcp::run_server(endpoint_listen, false).await.unwrap(); + let (mut socks_client, connect_status) = socks5_connect(tunnel_listen, endpoint_listen).await.unwrap(); + assert_eq!(connect_status, 0x00, "SOCKS CONNECT should succeed"); + + socks_client.write_all(b"Hello").await.unwrap(); + let mut endpoint = endpoint_listener.next().await.unwrap().unwrap(); + let mut buf = BytesMut::new(); + endpoint.read_buf(&mut buf).await.unwrap(); + assert_eq!(&buf[..5], b"Hello"); + buf.clear(); + + endpoint.write_all(b"world!").await.unwrap(); + socks_client.read_buf(&mut buf).await.unwrap(); + assert_eq!(&buf[..6], b"world!"); +} + +#[rstest] +#[timeout(Duration::from_secs(15))] +#[tokio::test] +#[serial] +async fn test_socks5_tunnel_connect_failure_webtransport( + server_webtransport: WsServer, + no_restrictions: RestrictionsRules, + dns_resolver: DnsResolver, +) { + let tunnel_listen = free_addr().0; + // Reserved but never bound, so the connect is guaranteed to be refused. + let unreachable_target = free_addr().0; + + let server_port = server_webtransport.config.bind.port(); + let server_h = tokio::spawn(server_webtransport.serve(no_restrictions)); + defer! { server_h.abort(); }; + + let client = client_webtransport(server_port, dns_resolver).await; + let socks_listener = Socks5TunnelListener::new(tunnel_listen, None, None).await.unwrap(); + tokio::spawn(async move { client.run_tunnel(socks_listener).await.unwrap() }); + + let (_client, connect_status) = socks5_connect(tunnel_listen, unreachable_target).await.unwrap(); + assert_eq!(connect_status, 0x01, "SOCKS CONNECT should report failure"); +} + +#[rstest] +#[timeout(Duration::from_secs(10))] +#[tokio::test] +#[serial] +async fn test_socks5_tunnel_connect_success( + server_no_tls: WsServer, + no_restrictions: RestrictionsRules, + dns_resolver: DnsResolver, +) { + let tunnel_listen = free_addr().0; + let endpoint_listen = free_addr().0; + + let server_port = server_no_tls.config.bind.port(); + let server_h = tokio::spawn(server_no_tls.serve(no_restrictions)); + defer! { server_h.abort(); }; + + let client_ws = client_ws(server_port, dns_resolver).await; + let socks_listener = Socks5TunnelListener::new(tunnel_listen, None, None).await.unwrap(); + tokio::spawn(async move { client_ws.run_tunnel(socks_listener).await.unwrap() }); + + let mut endpoint_listener = protocols::tcp::run_server(endpoint_listen, false).await.unwrap(); + let (mut socks_client, connect_status) = socks5_connect(tunnel_listen, endpoint_listen).await.unwrap(); + assert_eq!(connect_status, 0x00, "SOCKS CONNECT should succeed"); + + socks_client.write_all(b"Hello").await.unwrap(); + let mut endpoint = endpoint_listener.next().await.unwrap().unwrap(); + let mut buf = BytesMut::new(); + endpoint.read_buf(&mut buf).await.unwrap(); + assert_eq!(&buf[..5], b"Hello"); + buf.clear(); + + endpoint.write_all(b"world!").await.unwrap(); + socks_client.read_buf(&mut buf).await.unwrap(); + assert_eq!(&buf[..6], b"world!"); +} + #[rstest] #[timeout(Duration::from_secs(15))] #[tokio::test] @@ -441,3 +579,28 @@ async fn test_udp_tunnel_webtransport( // client.read_buf(&mut buf).await.unwrap(); // assert_eq!(&buf[..6], b"world!"); //} + +#[rstest] +#[timeout(Duration::from_secs(10))] +#[tokio::test] +#[serial] +async fn test_socks5_tunnel_connect_failure( + server_no_tls: WsServer, + no_restrictions: RestrictionsRules, + dns_resolver: DnsResolver, +) { + let tunnel_listen = free_addr().0; + // Reserved but never bound, so the connect is guaranteed to be refused. + let unreachable_target = free_addr().0; + + let server_port = server_no_tls.config.bind.port(); + let server_h = tokio::spawn(server_no_tls.serve(no_restrictions)); + defer! { server_h.abort(); }; + + let client_ws = client_ws(server_port, dns_resolver).await; + let socks_listener = Socks5TunnelListener::new(tunnel_listen, None, None).await.unwrap(); + tokio::spawn(async move { client_ws.run_tunnel(socks_listener).await.unwrap() }); + + let (_client, connect_status) = socks5_connect(tunnel_listen, unreachable_target).await.unwrap(); + assert_eq!(connect_status, 0x01, "SOCKS CONNECT should report failure"); +} diff --git a/wstunnel/src/tunnel/client/client.rs b/wstunnel/src/tunnel/client/client.rs index 608e873d..f6b544ba 100644 --- a/wstunnel/src/tunnel/client/client.rs +++ b/wstunnel/src/tunnel/client/client.rs @@ -1,4 +1,5 @@ use crate::executor::{DefaultTokioExecutor, TokioExecutorRef}; +use crate::protocols::socks5::Socks5WriteHalf; use crate::protocols::tls; use crate::tunnel; use crate::tunnel::RemoteAddr; @@ -6,14 +7,18 @@ use crate::tunnel::client::WsClientConfig; use crate::tunnel::client::cnx_pool::WsConnection; use crate::tunnel::connectors::TunnelConnector; use crate::tunnel::listeners::TunnelListener; +use crate::tunnel::reverse_socks5::{HANDSHAKE_CONNECT_FAIL, HANDSHAKE_CONNECT_OK}; use crate::tunnel::tls_reloader::TlsReloader; -use crate::tunnel::transport::io::{TunnelReader, TunnelWriter}; +use crate::tunnel::transport::io::{TunnelReader, TunnelWrite, TunnelWriter}; use crate::tunnel::transport::webtransport::WebTransportEndpoint; use crate::tunnel::transport::{TransportScheme, jwt_token_to_tunnel}; use anyhow::{Context, anyhow}; +use bytes::BufMut; +use fast_socks5::ReplyError; use futures_util::pin_mut; use hyper::header::COOKIE; use log::debug; +use std::any::Any; use std::cmp::min; use std::sync::Arc; use std::time::Duration; @@ -36,6 +41,28 @@ pub struct WsClient { } impl WsClient { + async fn send_socks5_reply_if_needed( + writer: &mut W, + error: ReplyError, + ) -> anyhow::Result<()> { + let Some(socks5_writer) = (writer as &mut dyn Any).downcast_mut::() else { + return Ok(()); + }; + + socks5_writer.send_reply_if_needed(error).await + } + + async fn send_reverse_socks5_handshake(writer: &mut TunnelWriter, success: bool) -> anyhow::Result<()> { + let byte = if success { + HANDSHAKE_CONNECT_OK + } else { + HANDSHAKE_CONNECT_FAIL + }; + writer.buf_mut().put_u8(byte); + writer.write().await?; + Ok(()) + } + pub async fn new( config: WsClientConfig, connection_min_idle: u32, @@ -108,33 +135,49 @@ impl WsClient { ) -> anyhow::Result<()> where R: AsyncRead + Send + 'static, - W: AsyncWrite + Send + 'static, + W: AsyncWrite + Send + 'static + Any, { + let (local_rx, mut local_tx) = duplex_stream; + // Connect to server with the correct protocol let (ws_rx, ws_tx, response) = match self.config.remote_addr.scheme() { TransportScheme::Ws | TransportScheme::Wss => { - tunnel::transport::websocket::connect(request_id, self, remote_cfg) - .await - .map(|(r, w, response)| (TunnelReader::Websocket(r), TunnelWriter::Websocket(w), response))? + match tunnel::transport::websocket::connect(request_id, self, remote_cfg).await { + Ok((r, w, response)) => (TunnelReader::Websocket(r), TunnelWriter::Websocket(w), response), + Err(err) => { + let _ = Self::send_socks5_reply_if_needed(&mut local_tx, ReplyError::GeneralFailure).await; + return Err(err); + } + } } TransportScheme::Http | TransportScheme::Https => { - tunnel::transport::http2::connect(request_id, self, remote_cfg) - .await - .map(|(r, w, response)| (TunnelReader::Http2(r), TunnelWriter::Http2(w), response))? + match tunnel::transport::http2::connect(request_id, self, remote_cfg).await { + Ok((r, w, response)) => (TunnelReader::Http2(r), TunnelWriter::Http2(w), response), + Err(err) => { + let _ = Self::send_socks5_reply_if_needed(&mut local_tx, ReplyError::GeneralFailure).await; + return Err(err); + } + } } - TransportScheme::Wts => tunnel::transport::webtransport::connect(request_id, self, remote_cfg) - .await - .map(|(r, w, response)| { - ( + TransportScheme::Wts => { + match tunnel::transport::webtransport::connect(request_id, self, remote_cfg).await { + Ok((r, w, response)) => ( TunnelReader::WebTransport(Box::new(r)), TunnelWriter::WebTransport(Box::new(w)), response, - ) - })?, + ), + Err(err) => { + let _ = Self::send_socks5_reply_if_needed(&mut local_tx, ReplyError::GeneralFailure).await; + return Err(err); + } + } + } }; debug!("Server response: {response:?}"); - let (local_rx, local_tx) = duplex_stream; + + Self::send_socks5_reply_if_needed(&mut local_tx, ReplyError::Succeeded).await?; + let (close_tx, close_rx) = oneshot::channel::<()>(); // Forward local tx to websocket tx @@ -200,6 +243,8 @@ impl WsClient { } let mut reconnect_delay = new_reconnect_delay(self.reverse_tunnel_connection_retry_max_backoff); + let is_reverse_socks5 = matches!(remote_addr.protocol, tunnel::LocalProtocol::ReverseSocks5 { .. }); + loop { let client = self.clone(); let request_id = Uuid::now_v7(); @@ -210,7 +255,7 @@ impl WsClient { remote = format!("{}:{}", remote_addr.host, remote_addr.port) ); // Correctly configure tunnel cfg - let (ws_rx, ws_tx, response) = match client.config.remote_addr.scheme() { + let (ws_rx, mut ws_tx, response) = match client.config.remote_addr.scheme() { TransportScheme::Ws | TransportScheme::Wss => { match tunnel::transport::websocket::connect(request_id, &client, &remote_addr) .instrument(span.clone()) @@ -273,14 +318,29 @@ impl WsClient { port: jwt.claims.rp, }); + let need_reverse_socks5_handshake = is_reverse_socks5 + && matches!(remote.as_ref().map(|r| &r.protocol), Some(tunnel::LocalProtocol::Tcp { .. })); + let (local_rx, local_tx) = match connector.connect(&remote).instrument(span.clone()).await { Ok(s) => s, Err(err) => { event!(parent: &span, Level::ERROR, "Cannot connect to {remote:?}: {err:?}"); + if need_reverse_socks5_handshake { + let _ = Self::send_reverse_socks5_handshake(&mut ws_tx, false).await; + let _ = ws_tx.close().await; + } continue; } }; + if need_reverse_socks5_handshake + && let Err(err) = Self::send_reverse_socks5_handshake(&mut ws_tx, true).await + { + event!(parent: &span, Level::ERROR, "Cannot send reverse socks5 handshake: {err:?}"); + let _ = ws_tx.close().await; + continue; + } + let (close_tx, close_rx) = oneshot::channel::<()>(); self.executor.spawn({ let ping_frequency = client.config.websocket_ping_frequency; diff --git a/wstunnel/src/tunnel/mod.rs b/wstunnel/src/tunnel/mod.rs index 45fb302d..d0f3d1af 100644 --- a/wstunnel/src/tunnel/mod.rs +++ b/wstunnel/src/tunnel/mod.rs @@ -2,6 +2,7 @@ pub mod ca_reloader; pub mod client; pub mod connectors; pub mod listeners; +pub(crate) mod reverse_socks5; pub mod server; mod tls_reloader; pub mod transport; diff --git a/wstunnel/src/tunnel/reverse_socks5.rs b/wstunnel/src/tunnel/reverse_socks5.rs new file mode 100644 index 00000000..a6d0fbd3 --- /dev/null +++ b/wstunnel/src/tunnel/reverse_socks5.rs @@ -0,0 +1,72 @@ +use std::future::Future; +use std::io; +use std::io::ErrorKind; +use std::time::Duration; + +pub(crate) const HANDSHAKE_CONNECT_OK: u8 = 0; +pub(crate) const HANDSHAKE_CONNECT_FAIL: u8 = 1; + +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub(crate) enum ReverseSocks5ConnectResult { + Connected, + Failed, +} + +pub(crate) async fn read_reverse_socks5_connect_result( + read_handshake_byte: impl Future>, + timeout: Duration, +) -> io::Result { + match tokio::time::timeout(timeout, read_handshake_byte).await { + Ok(Ok(HANDSHAKE_CONNECT_OK)) => Ok(ReverseSocks5ConnectResult::Connected), + Ok(Ok(_)) => Ok(ReverseSocks5ConnectResult::Failed), + Ok(Err(err)) => Err(err), + Err(_) => Err(io::Error::new(ErrorKind::TimedOut, "reverse socks5 handshake timeout")), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::time::sleep; + + #[tokio::test] + async fn reverse_socks5_handshake_reports_connected() { + let ret = read_reverse_socks5_connect_result(async { Ok(HANDSHAKE_CONNECT_OK) }, Duration::from_millis(50)) + .await + .unwrap(); + assert_eq!(ret, ReverseSocks5ConnectResult::Connected); + } + + #[tokio::test] + async fn reverse_socks5_handshake_reports_failed() { + let ret = read_reverse_socks5_connect_result(async { Ok(HANDSHAKE_CONNECT_FAIL) }, Duration::from_millis(50)) + .await + .unwrap(); + assert_eq!(ret, ReverseSocks5ConnectResult::Failed); + } + + #[tokio::test] + async fn reverse_socks5_handshake_times_out() { + let err = read_reverse_socks5_connect_result( + async { + sleep(Duration::from_millis(25)).await; + Ok(HANDSHAKE_CONNECT_OK) + }, + Duration::from_millis(5), + ) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::TimedOut); + } + + #[tokio::test] + async fn reverse_socks5_handshake_propagates_io_error() { + let err = read_reverse_socks5_connect_result( + async { Err(io::Error::new(ErrorKind::ConnectionAborted, "oops")) }, + Duration::from_millis(50), + ) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::ConnectionAborted); + } +} diff --git a/wstunnel/src/tunnel/server/handler_http2.rs b/wstunnel/src/tunnel/server/handler_http2.rs index 65de37d9..5c33f1a9 100644 --- a/wstunnel/src/tunnel/server/handler_http2.rs +++ b/wstunnel/src/tunnel/server/handler_http2.rs @@ -1,10 +1,14 @@ use crate::executor::TokioExecutorRef; use crate::restrictions::types::RestrictionsRules; +use crate::tunnel::LocalProtocol; +use crate::tunnel::reverse_socks5::{ReverseSocks5ConnectResult, read_reverse_socks5_connect_result}; use crate::tunnel::server::WsServer; +use crate::tunnel::server::send_socks5_reply_if_needed; use crate::tunnel::server::utils::{HttpResponse, bad_request, inject_cookie}; use crate::tunnel::transport; use crate::tunnel::transport::http2::{Http2TunnelRead, Http2TunnelWrite}; use bytes::Bytes; +use fast_socks5::ReplyError; use futures_util::StreamExt; use http_body_util::combinators::BoxBody; use http_body_util::{BodyStream, Either, StreamBody}; @@ -24,7 +28,7 @@ pub(super) async fn http_server_upgrade( client_addr: SocketAddr, mut req: Request, ) -> HttpResponse { - let (remote_addr, local_rx, local_tx, need_cookie) = match server + let (remote_addr, local_rx, mut local_tx, need_cookie, reverse_socks5) = match server .handle_tunnel_request(restrictions, restrict_path_prefix, client_addr, &req) .await { @@ -44,15 +48,41 @@ pub(super) async fn http_server_upgrade( .body(Either::Right(body)) .expect("bug: failed to build response"); - let (close_tx, close_rx) = oneshot::channel::<()>(); + let reverse_socks5_handshake_timeout = server.config.timeout_connect; + let reverse_socks5_tcp = reverse_socks5 && matches!(&remote_addr.protocol, LocalProtocol::Tcp { .. }); + let executor = server.executor.clone(); server.executor.spawn( - transport::io::propagate_remote_to_local(local_tx, Http2TunnelRead::new(ws_rx, None), close_rx) - .instrument(Span::current()), - ); + async move { + let mut ws_rx = Http2TunnelRead::new(ws_rx, None); - server.executor.spawn( - transport::io::propagate_local_to_remote(local_rx, Http2TunnelWrite::new(ws_tx), close_tx, None) - .instrument(Span::current()), + if reverse_socks5_tcp { + match read_reverse_socks5_connect_result(ws_rx.read_handshake_byte(), reverse_socks5_handshake_timeout) + .await + { + Ok(ReverseSocks5ConnectResult::Connected) => { + if send_socks5_reply_if_needed(&mut local_tx, ReplyError::Succeeded) + .await + .is_err() + { + return; + } + } + Ok(ReverseSocks5ConnectResult::Failed) | Err(_) => { + let _ = send_socks5_reply_if_needed(&mut local_tx, ReplyError::GeneralFailure).await; + return; + } + } + } + + let (close_tx, close_rx) = oneshot::channel::<()>(); + executor + .spawn(transport::io::propagate_remote_to_local(local_tx, ws_rx, close_rx).instrument(Span::current())); + + let _ = transport::io::propagate_local_to_remote(local_rx, Http2TunnelWrite::new(ws_tx), close_tx, None) + .instrument(Span::current()) + .await; + } + .instrument(Span::current()), ); if need_cookie && inject_cookie(&mut response, &remote_addr).is_err() { diff --git a/wstunnel/src/tunnel/server/handler_websocket.rs b/wstunnel/src/tunnel/server/handler_websocket.rs index adb0495a..47812c79 100644 --- a/wstunnel/src/tunnel/server/handler_websocket.rs +++ b/wstunnel/src/tunnel/server/handler_websocket.rs @@ -1,9 +1,14 @@ use crate::executor::TokioExecutorRef; use crate::restrictions::types::RestrictionsRules; +use crate::tunnel::LocalProtocol; +use crate::tunnel::reverse_socks5::{ReverseSocks5ConnectResult, read_reverse_socks5_connect_result}; use crate::tunnel::server::WsServer; +use crate::tunnel::server::send_socks5_reply_if_needed; use crate::tunnel::server::utils::{HttpResponse, bad_request, inject_cookie}; use crate::tunnel::transport; +use crate::tunnel::transport::io::TunnelWrite; use crate::tunnel::transport::websocket::mk_websocket_tunnel; +use fast_socks5::ReplyError; use fastwebsockets::Role; use http_body_util::Either; use http_body_util::combinators::BoxBody; @@ -28,7 +33,7 @@ pub(super) async fn ws_server_upgrade( } let mask_frame = server.config.websocket_mask_frame; - let (remote_addr, local_rx, local_tx, need_cookie) = match server + let (remote_addr, local_rx, mut local_tx, need_cookie, reverse_socks5) = match server .handle_tunnel_request(restrictions, restrict_path_prefix, client_addr, &req) .await { @@ -44,10 +49,12 @@ pub(super) async fn ws_server_upgrade( } }; + let reverse_socks5_handshake_timeout = server.config.timeout_connect; + let reverse_socks5_tcp = reverse_socks5 && matches!(&remote_addr.protocol, LocalProtocol::Tcp { .. }); let executor = server.executor.clone(); server.executor.spawn( async move { - let (ws_rx, ws_tx) = match fut.await { + let (mut ws_rx, mut ws_tx) = match fut.await { Ok(ws) => match mk_websocket_tunnel(ws, Role::Server, mask_frame) { Ok(ws) => ws, Err(err) => { @@ -60,6 +67,31 @@ pub(super) async fn ws_server_upgrade( return Err(anyhow::Error::from(err)); } }; + if reverse_socks5_tcp { + match read_reverse_socks5_connect_result(ws_rx.read_handshake_byte(), reverse_socks5_handshake_timeout) + .await + { + Ok(ReverseSocks5ConnectResult::Connected) => { + if let Err(err) = send_socks5_reply_if_needed(&mut local_tx, ReplyError::Succeeded).await { + error!("Cannot reply to socks5 client: {err:?}"); + let _ = ws_tx.close().await; + return Err(err); + } + } + Ok(ReverseSocks5ConnectResult::Failed) => { + let _ = send_socks5_reply_if_needed(&mut local_tx, ReplyError::GeneralFailure).await; + let _ = ws_tx.close().await; + return Ok(()); + } + Err(err) => { + warn!("Reverse socks5 handshake failed: {err}"); + let _ = send_socks5_reply_if_needed(&mut local_tx, ReplyError::GeneralFailure).await; + let _ = ws_tx.close().await; + return Err(err.into()); + } + } + } + let (close_tx, close_rx) = oneshot::channel::<()>(); executor diff --git a/wstunnel/src/tunnel/server/handler_webtransport.rs b/wstunnel/src/tunnel/server/handler_webtransport.rs index cc4d0d07..20fd53c2 100644 --- a/wstunnel/src/tunnel/server/handler_webtransport.rs +++ b/wstunnel/src/tunnel/server/handler_webtransport.rs @@ -1,7 +1,10 @@ use crate::executor::TokioExecutorRef; use crate::protocols::tls; use crate::restrictions::types::RestrictionsRules; +use crate::tunnel::LocalProtocol; +use crate::tunnel::reverse_socks5::{ReverseSocks5ConnectResult, read_reverse_socks5_connect_result}; use crate::tunnel::server::WsServer; +use crate::tunnel::server::send_socks5_reply_if_needed; use crate::tunnel::tls_reloader::TlsReloader; use crate::tunnel::transport; use crate::tunnel::transport::tunnel_to_jwt_token; @@ -10,6 +13,7 @@ use crate::tunnel::transport::webtransport::{ }; use anyhow::{Context, anyhow}; use arc_swap::ArcSwap; +use fast_socks5::ReplyError; use hyper::{Request, StatusCode}; use std::any::Any; use std::sync::Arc; @@ -136,7 +140,9 @@ async fn handle_session( .handle_tunnel_request(restrictions, restrict_path, client_addr, &http_request) .await; - let (remote_addr, local_rx, local_tx, need_preamble) = match tunnel { + // `reverse_socks5` gates the deferred SOCKS5 reply on the reverse connect handshake, the same + // way the websocket and http2 handlers do. + let (remote_addr, local_rx, mut local_tx, need_preamble, reverse_socks5) = match tunnel { Ok(tunnel) => tunnel, Err(response) => { let status = response.status(); @@ -179,15 +185,38 @@ async fn handle_session( } } + let mut tunnel_rx = WebTransportTunnelRead::new(recv, session.clone()); + + // Wait for the client to report the outcome of its local connect before answering the SOCKS + // client, so a refused target is reported as a failure instead of a premature success. The + // preamble above travels the other way, on `send`, so the two never interleave. + if reverse_socks5 && matches!(&remote_addr.protocol, LocalProtocol::Tcp { .. }) { + match read_reverse_socks5_connect_result(tunnel_rx.read_handshake_byte(), server.config.timeout_connect).await { + Ok(ReverseSocks5ConnectResult::Connected) => { + if let Err(err) = send_socks5_reply_if_needed(&mut local_tx, ReplyError::Succeeded).await { + error!("Cannot reply to socks5 client: {err:?}"); + session.close(StatusCode::INTERNAL_SERVER_ERROR.as_u16() as u32, b"socks5 reply failed"); + return Err(err); + } + } + Ok(ReverseSocks5ConnectResult::Failed) => { + let _ = send_socks5_reply_if_needed(&mut local_tx, ReplyError::GeneralFailure).await; + session.close(StatusCode::BAD_GATEWAY.as_u16() as u32, b"reverse socks5 connect failed"); + return Ok(()); + } + Err(err) => { + warn!("Reverse socks5 handshake failed: {err}"); + let _ = send_socks5_reply_if_needed(&mut local_tx, ReplyError::GeneralFailure).await; + session.close(StatusCode::GATEWAY_TIMEOUT.as_u16() as u32, b"reverse socks5 handshake failed"); + return Err(err.into()); + } + } + } + let (close_tx, close_rx) = oneshot::channel::<()>(); - server.executor.spawn( - transport::io::propagate_remote_to_local( - local_tx, - WebTransportTunnelRead::new(recv, session.clone()), - close_rx, - ) - .instrument(Span::current()), - ); + server + .executor + .spawn(transport::io::propagate_remote_to_local(local_tx, tunnel_rx, close_rx).instrument(Span::current())); server.executor.spawn( transport::io::propagate_local_to_remote(local_rx, WebTransportTunnelWrite::new(send, session), close_tx, None) diff --git a/wstunnel/src/tunnel/server/mod.rs b/wstunnel/src/tunnel/server/mod.rs index b7e564d0..89682f73 100644 --- a/wstunnel/src/tunnel/server/mod.rs +++ b/wstunnel/src/tunnel/server/mod.rs @@ -4,8 +4,11 @@ mod handler_websocket; mod handler_webtransport; mod reverse_tunnel; mod server; +mod socks5_reply; mod utils; pub use server::TlsServerConfig; pub use server::WsServer; pub use server::WsServerConfig; + +pub(crate) use socks5_reply::{AnyAsyncWrite, send_socks5_reply_if_needed}; diff --git a/wstunnel/src/tunnel/server/server.rs b/wstunnel/src/tunnel/server/server.rs index dd190fef..1c96f1be 100644 --- a/wstunnel/src/tunnel/server/server.rs +++ b/wstunnel/src/tunnel/server/server.rs @@ -7,6 +7,7 @@ use crate::restrictions::types::{RestrictionConfig, RestrictionsRules}; use crate::somark::SoMark; use crate::tunnel::connectors::{TcpTunnelConnector, TunnelConnector, UdpTunnelConnector}; use crate::tunnel::listeners::{HttpProxyTunnelListener, Socks5TunnelListener, TcpTunnelListener, UdpTunnelListener}; +use crate::tunnel::server::AnyAsyncWrite; use crate::tunnel::server::handler_http2::http_server_upgrade; use crate::tunnel::server::handler_websocket::ws_server_upgrade; use crate::tunnel::server::handler_webtransport; @@ -37,7 +38,7 @@ use std::path::PathBuf; use std::pin::Pin; use std::sync::{Arc, LazyLock}; use std::time::Duration; -use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; +use tokio::io::{AsyncRead, AsyncWriteExt}; use tokio::net::TcpListener; use tokio_rustls::TlsAcceptor; use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer}; @@ -99,7 +100,8 @@ impl WsServer { ( RemoteAddr, Pin>, - Pin>, + Pin>, + bool, bool, ), HttpResponse, @@ -144,6 +146,7 @@ impl WsServer { info!("Tunnel accepted due to matched restriction: {}", restriction.name); let req_protocol = remote.protocol.clone(); + let reverse_socks5 = matches!(req_protocol, LocalProtocol::ReverseSocks5 { .. }); let inject_cookie = req_protocol.is_dynamic_reverse_tunnel(); let tunnel = self .exec_tunnel(restriction, remote, client_addr) @@ -155,7 +158,7 @@ impl WsServer { let (remote_addr, local_rx, local_tx) = tunnel; info!("connected to {:?} {}:{}", req_protocol, remote_addr.host, remote_addr.port); - Ok((remote_addr, local_rx, local_tx, inject_cookie)) + Ok((remote_addr, local_rx, local_tx, inject_cookie, reverse_socks5)) } async fn exec_tunnel( @@ -163,14 +166,14 @@ impl WsServer { restriction: &RestrictionConfig, remote: RemoteAddr, client_address: SocketAddr, - ) -> anyhow::Result<(RemoteAddr, Pin>, Pin>)> { + ) -> anyhow::Result<(RemoteAddr, Pin>, Pin>)> { match remote.protocol { LocalProtocol::Udp { timeout, .. } => { let connector = UdpTunnelConnector::new( &remote.host, remote.port, self.config.socket_so_mark, - timeout.unwrap_or(Duration::from_secs(10)), + timeout.unwrap_or(self.config.timeout_connect), &self.config.dns_resolver, ); let (rx, tx) = match &self.config.http_proxy { @@ -185,7 +188,7 @@ impl WsServer { &remote.host, remote.port, self.config.socket_so_mark, - Duration::from_secs(10), + self.config.timeout_connect, &self.config.dns_resolver, ); let (rx, mut tx) = match &self.config.http_proxy { diff --git a/wstunnel/src/tunnel/server/socks5_reply.rs b/wstunnel/src/tunnel/server/socks5_reply.rs new file mode 100644 index 00000000..cf1de05a --- /dev/null +++ b/wstunnel/src/tunnel/server/socks5_reply.rs @@ -0,0 +1,29 @@ +use crate::protocols::socks5::Socks5WriteHalf; +use fast_socks5::ReplyError; +use std::any::Any; +use std::pin::Pin; +use tokio::io::AsyncWrite; + +pub(crate) trait AnyAsyncWrite: AsyncWrite + Send + Unpin { + fn as_any_mut(&mut self) -> &mut dyn Any; +} + +impl AnyAsyncWrite for T +where + T: AsyncWrite + Send + Unpin + Any + 'static, +{ + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +pub(crate) async fn send_socks5_reply_if_needed( + writer: &mut Pin>, + error: ReplyError, +) -> anyhow::Result<()> { + let Some(socks5_writer) = writer.as_mut().get_mut().as_any_mut().downcast_mut::() else { + return Ok(()); + }; + + socks5_writer.send_reply_if_needed(error).await +} diff --git a/wstunnel/src/tunnel/transport/http2.rs b/wstunnel/src/tunnel/transport/http2.rs index a7f6acf1..b4d7748e 100644 --- a/wstunnel/src/tunnel/transport/http2.rs +++ b/wstunnel/src/tunnel/transport/http2.rs @@ -28,11 +28,16 @@ use uuid::Uuid; pub struct Http2TunnelRead { inner: BodyStream, cnx_poller: Option, + prefetched: BytesMut, } impl Http2TunnelRead { - pub const fn new(inner: BodyStream, cnx_poller: Option) -> Self { - Self { inner, cnx_poller } + pub fn new(inner: BodyStream, cnx_poller: Option) -> Self { + Self { + inner, + cnx_poller, + prefetched: BytesMut::new(), + } } } @@ -46,6 +51,14 @@ impl Drop for Http2TunnelRead { impl TunnelRead for Http2TunnelRead { async fn copy(&mut self, mut writer: impl AsyncWrite + Unpin + Send) -> Result<(), io::Error> { + if !self.prefetched.is_empty() { + let data = self.prefetched.split().freeze(); + return match writer.write_all(data.as_ref()).await { + Ok(_) => Ok(()), + Err(err) => Err(io::Error::new(ErrorKind::ConnectionAborted, err)), + }; + } + loop { match self.inner.next().await { Some(Ok(frame)) => match frame.into_data() { @@ -69,6 +82,40 @@ impl TunnelRead for Http2TunnelRead { } } +impl Http2TunnelRead { + pub async fn read_handshake_byte(&mut self) -> Result { + if !self.prefetched.is_empty() { + let data = self.prefetched.split_to(1); + return Ok(data[0]); + } + + loop { + match self.inner.next().await { + Some(Ok(frame)) => match frame.into_data() { + Ok(data) => { + if data.is_empty() { + continue; + } + let first = data[0]; + if data.len() > 1 { + self.prefetched.extend_from_slice(&data[1..]); + } + return Ok(first); + } + Err(err) => { + warn!("{err:?}"); + continue; + } + }, + Some(Err(err)) => { + return Err(io::Error::new(ErrorKind::ConnectionAborted, err)); + } + None => return Err(io::Error::new(ErrorKind::BrokenPipe, "closed")), + } + } + } +} + pub struct Http2TunnelWrite { inner: mpsc::Sender, buf: BytesMut, diff --git a/wstunnel/src/tunnel/transport/websocket.rs b/wstunnel/src/tunnel/transport/websocket.rs index b10c68ef..98d54a31 100644 --- a/wstunnel/src/tunnel/transport/websocket.rs +++ b/wstunnel/src/tunnel/transport/websocket.rs @@ -161,6 +161,7 @@ pub struct WebsocketTunnelRead { inner: WebSocketRead, pending_operations: Sender>, notify_pending_ops: Arc, + prefetched: BytesMut, } impl WebsocketTunnelRead { @@ -172,6 +173,7 @@ impl WebsocketTunnelRead { inner: ws, pending_operations: tx, notify_pending_ops: notify.clone(), + prefetched: BytesMut::new(), }, (rx, notify), ) @@ -186,6 +188,14 @@ fn frame_reader(_: Frame<'_>) -> futures_util::future::Ready> impl TunnelRead for WebsocketTunnelRead { async fn copy(&mut self, mut writer: impl AsyncWrite + Unpin + Send) -> Result<(), io::Error> { loop { + if !self.prefetched.is_empty() { + let data = self.prefetched.split().freeze(); + return match writer.write_all(data.as_ref()).await { + Ok(_) => Ok(()), + Err(err) => Err(io::Error::new(ErrorKind::ConnectionAborted, err)), + }; + } + let msg = match self.inner.read_frame(&mut frame_reader).await { Ok(msg) => msg, Err(err) => return Err(io::Error::new(ErrorKind::ConnectionAborted, err)), @@ -234,6 +244,66 @@ impl TunnelRead for WebsocketTunnelRead { } } +impl WebsocketTunnelRead { + pub async fn read_handshake_byte(&mut self) -> Result { + if !self.prefetched.is_empty() { + let data = self.prefetched.split_to(1); + return Ok(data[0]); + } + + loop { + let msg = match self.inner.read_frame(&mut frame_reader).await { + Ok(msg) => msg, + Err(err) => return Err(io::Error::new(ErrorKind::ConnectionAborted, err)), + }; + + match msg.opcode { + OpCode::Continuation | OpCode::Text | OpCode::Binary => { + let payload = msg.payload.as_ref(); + if payload.is_empty() { + continue; + } + let first = payload[0]; + if payload.len() > 1 { + self.prefetched.extend_from_slice(&payload[1..]); + } + return Ok(first); + } + OpCode::Close => { + let _ = self + .pending_operations + .send(Frame::close(CloseCode::Normal.into(), &[])) + .await; + self.notify_pending_ops.notify_waiters(); + return Err(io::Error::new(ErrorKind::NotConnected, "websocket close")); + } + OpCode::Ping => { + if self + .pending_operations + .send(Frame::new(true, msg.opcode, None, Payload::Owned(msg.payload.to_owned()))) + .await + .is_err() + { + return Err(io::Error::new(ErrorKind::ConnectionAborted, "cannot send ping")); + } + self.notify_pending_ops.notify_waiters(); + } + OpCode::Pong => { + if self + .pending_operations + .send(Frame::pong(Payload::Borrowed(&[]))) + .await + .is_err() + { + return Err(io::Error::new(ErrorKind::ConnectionAborted, "cannot send pong")); + } + self.notify_pending_ops.notify_waiters(); + } + } + } + } +} + pub async fn connect( request_id: Uuid, client: &WsClient, diff --git a/wstunnel/src/tunnel/transport/webtransport.rs b/wstunnel/src/tunnel/transport/webtransport.rs index 589bdf3c..2b7975ac 100644 --- a/wstunnel/src/tunnel/transport/webtransport.rs +++ b/wstunnel/src/tunnel/transport/webtransport.rs @@ -153,6 +153,19 @@ impl WebTransportTunnelRead { _session: session, } } + + /// Read the single reverse SOCKS5 connect-result byte. + /// + /// Unlike the websocket and http2 tunnels, which read a whole frame at a time and have to + /// stash its remainder, a QUIC stream is a byte stream: taking one byte leaves the rest of + /// the payload queued, so no prefetch buffer is needed here. + pub async fn read_handshake_byte(&mut self) -> Result { + let mut byte = [0u8; 1]; + match self.inner.read_exact(&mut byte).await { + Ok(()) => Ok(byte[0]), + Err(err) => Err(io::Error::new(ErrorKind::ConnectionAborted, err)), + } + } } impl TunnelRead for WebTransportTunnelRead {