diff --git a/src/client/legacy/connect/proxy/socks/mod.rs b/src/client/legacy/connect/proxy/socks/mod.rs index 9bc89e9d..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) => { @@ -163,7 +164,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 +172,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 +209,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 9ff48d4e..8aa8070c 100644 --- a/tests/proxy.rs +++ b/tests/proxy.rs @@ -473,7 +473,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()) @@ -485,7 +486,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 @@ -514,6 +528,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()) @@ -521,10 +537,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))]