From ef4a3c6f7cb594bc951423b1d9528a358ec65dd5 Mon Sep 17 00:00:00 2001 From: Joe Dye Date: Fri, 31 Jul 2026 17:33:13 +0100 Subject: [PATCH 1/2] tests for optimistic sending deadlock --- src/client/legacy/connect/proxy/socks/mod.rs | 68 +++++++++++++++++++- tests/proxy.rs | 43 ++++++++++++- 2 files changed, 108 insertions(+), 3 deletions(-) diff --git a/src/client/legacy/connect/proxy/socks/mod.rs b/src/client/legacy/connect/proxy/socks/mod.rs index 9bc89e9d..d0b298c6 100644 --- a/src/client/legacy/connect/proxy/socks/mod.rs +++ b/src/client/legacy/connect/proxy/socks/mod.rs @@ -163,7 +163,7 @@ mod test { use bytes::BytesMut; use tokio::io::AsyncWriteExt; - use super::v5::messages::{ProxyRes, Status}; + use super::v5::messages::{AuthMethod, AuthenticationRes, NegotiationRes, ProxyRes, Status}; use super::{SocksError, read_message}; use crate::rt::TokioIo; @@ -171,6 +171,11 @@ mod test { const SEG1: [u8; 4] = [0x05, 0x00, 0x00, 0x01]; const SEG2: [u8; 6] = [0x7F, 0x00, 0x00, 0x01, 0x1F, 0x90]; + // A SOCKS5 NegotiationRes message: username/password method selected. + const NEG_RES: [u8; 2] = [0x05, 0x02]; + // A SOCKS5 AuthenticationRes message: success. + const AUTH_RES: [u8; 2] = [0x01, 0x00]; + #[tokio::test] async fn it_works_in_one_read() { let (client, mut server) = tokio::io::duplex(SEG1.len() + SEG2.len()); @@ -203,4 +208,65 @@ mod test { _writer.await.unwrap(); } + + #[tokio::test(start_paused = true)] + async fn optimistic_sending_works_in_single_read() { + // Messages will arrive in a single read + let message = [&NEG_RES[..], &AUTH_RES[..], &SEG1[..], &SEG2[..]].concat(); + let (client, mut server) = tokio::io::duplex(message.len()); + server.write_all(&message).await.unwrap(); + + let mut conn = TokioIo::new(client); + let mut buf = BytesMut::new(); + + let m: Result> = read_message(&mut conn, &mut buf).await; + assert_eq!(m.unwrap(), NegotiationRes(AuthMethod::UserPass)); + + let m: Result> = tokio::time::timeout( + std::time::Duration::from_secs(1), + read_message(&mut conn, &mut buf), + ) + .await + .expect("second message should be parsed from the buffer, not read from the socket"); + assert_eq!(m.unwrap(), AuthenticationRes(true)); + + let m: Result> = tokio::time::timeout( + std::time::Duration::from_secs(1), + read_message(&mut conn, &mut buf), + ) + .await + .expect("third message should be parsed from the buffer, not read from the socket"); + assert_eq!(m.unwrap(), ProxyRes(Status::Success)); + + assert!(buf.is_empty(), "all handshake bytes should be consumed"); + drop(server); + } + + #[tokio::test] + async fn optimistic_sending_works_in_multiple_reads() { + // Bounded stream ensures message arrive in multiple reads + let message = [&NEG_RES[..], &AUTH_RES[..], &SEG1[..], &SEG2[..]].concat(); + let (client, mut server) = tokio::io::duplex(message.len() / 4); + let _writer = tokio::spawn(async move { + server.write_all(&message).await.unwrap(); + server + }); + + let mut conn = TokioIo::new(client); + let mut buf = BytesMut::new(); + + let m: Result> = read_message(&mut conn, &mut buf).await; + assert_eq!(m.unwrap(), NegotiationRes(AuthMethod::UserPass)); + + let m: Result> = read_message(&mut conn, &mut buf).await; + assert_eq!(m.unwrap(), AuthenticationRes(true)); + + let m: Result> = read_message(&mut conn, &mut buf).await; + assert_eq!(m.unwrap(), ProxyRes(Status::Success)); + + assert!(buf.is_empty(), "all handshake bytes should be consumed"); + + let server = _writer.await.unwrap(); + drop(server); + } } diff --git a/tests/proxy.rs b/tests/proxy.rs index 65e7e819..5899a923 100644 --- a/tests/proxy.rs +++ b/tests/proxy.rs @@ -425,7 +425,8 @@ async fn test_socks_v5_optimistic_works() { let proxy_addr = proxy_tcp.local_addr().expect("local_addr"); let proxy_dst = format!("http://{proxy_addr}").parse().expect("uri"); - let target_addr = std::net::SocketAddr::new([127, 0, 0, 1].into(), 1234); + let target_tcp = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let target_addr = target_tcp.local_addr().expect("local_addr"); let target_dst = format!("http://{target_addr}").parse().expect("uri"); let mut connector = SocksV5::new(proxy_dst, HttpConnector::new()) @@ -437,7 +438,20 @@ async fn test_socks_v5_optimistic_works() { // Will use `SocksV5` to establish proxy tunnel. // Will send "Hello World!" to the target and receive "Goodbye!" back. let t1 = tokio::spawn(async move { - let _ = connector.call(target_dst).await.expect("tunnel"); + let conn = tokio::time::timeout( + std::time::Duration::from_secs(5), + connector.call(target_dst), + ) + .await + .expect("handshake timed out") + .expect("tunnel"); + + let mut tcp = conn.into_inner(); + tcp.write_all(b"Hello World!").await.expect("write 1"); + + let mut buf = [0u8; 64]; + let n = tcp.read(&mut buf).await.expect("read 1"); + assert_eq!(&buf[..n], b"Goodbye!"); }); // Proxy @@ -466,6 +480,8 @@ async fn test_socks_v5_optimistic_works() { to_client.read_exact(&mut buf).await.expect("read"); assert_eq!(request.as_slice(), buf); + let mut to_target = TcpStream::connect(target_addr).await.expect("connect"); + // Send all handshake messages back to_client .write_all(response.as_slice()) @@ -473,10 +489,33 @@ async fn test_socks_v5_optimistic_works() { .expect("write"); to_client.flush().await.expect("flush"); + + let (from_client, from_target) = + tokio::io::copy_bidirectional(&mut to_client, &mut to_target) + .await + .expect("proxy"); + + assert_eq!(from_client, 12); + assert_eq!(from_target, 8) + }); + + // Target server + // + // Will accept connection from proxy server + // Will receive "Hello World!" from the client and return "Goodbye!" + let t3 = tokio::spawn(async move { + let (mut io, _) = target_tcp.accept().await.expect("accept"); + let mut buf = [0u8; 64]; + + let n = io.read(&mut buf).await.expect("read 1"); + assert_eq!(&buf[..n], b"Hello World!"); + + io.write_all(b"Goodbye!").await.expect("write 1"); }); t1.await.expect("task - client"); t2.await.expect("task - proxy"); + t3.await.expect("task - target"); } #[cfg(not(miri))] From c8b7aff6b57cd725e0db7af8578c08da57f92fe8 Mon Sep 17 00:00:00 2001 From: Joe Dye Date: Fri, 31 Jul 2026 17:06:44 +0100 Subject: [PATCH 2/2] parse existing bytes before reading from network, to prevent deadlock during optimisitic sending --- src/client/legacy/connect/proxy/socks/mod.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/client/legacy/connect/proxy/socks/mod.rs b/src/client/legacy/connect/proxy/socks/mod.rs index d0b298c6..caaff5ef 100644 --- a/src/client/legacy/connect/proxy/socks/mod.rs +++ b/src/client/legacy/connect/proxy/socks/mod.rs @@ -49,12 +49,11 @@ where let mut tmp = [0; 513]; loop { - let n = crate::rt::read(&mut conn, &mut tmp).await?; - buf.extend_from_slice(&tmp[..n]); - let mut view = &buf[..]; match M::try_from(&mut view) { Err(ParsingError::Incomplete) => { + let n = crate::rt::read(&mut conn, &mut tmp).await?; + if n == 0 { if buf.spare_capacity_mut().is_empty() { return Err(SocksError::Parsing(ParsingError::WouldOverflow)); @@ -66,6 +65,8 @@ where .into()); } } + + buf.extend_from_slice(&tmp[..n]); } Err(err) => return Err(err.into()), Ok(res) => {