fix(proxy): stop SOCKS5 optimistic sending deadlock

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,14 +164,19 @@ 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;
 
     // A SOCKS5 ProxyRes message. Successful, bound to 127.0.0.1:8080.
     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<NegotiationRes, SocksError<()>> = read_message(&mut conn, &mut buf).await;
+        assert_eq!(m.unwrap(), NegotiationRes(AuthMethod::UserPass));
+
+        let m: Result<AuthenticationRes, SocksError<()>> = 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<ProxyRes, SocksError<()>> = 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<NegotiationRes, SocksError<()>> = read_message(&mut conn, &mut buf).await;
+        assert_eq!(m.unwrap(), NegotiationRes(AuthMethod::UserPass));
+
+        let m: Result<AuthenticationRes, SocksError<()>> = read_message(&mut conn, &mut buf).await;
+        assert_eq!(m.unwrap(), AuthenticationRes(true));
+
+        let m: Result<ProxyRes, SocksError<()>> = 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);
+    }
 }

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,17 +528,42 @@ 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())
             .await
             .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))]