fix(proxy): SOCKS split message parsing fix

Cargo.toml

@@ -40,7 +40,7 @@ hyper = { version = "1.4.0", features = ["full"] }
 bytes = "1"
 futures-util = { version = "0.3.16", default-features = false, features = ["alloc"] }
 http-body-util = "0.1.0"
-tokio = { version = "1", features = ["macros", "test-util", "signal", "net"] }
+tokio = { version = "1", features = ["macros", "test-util", "signal", "net", "io-util"] }
 tokio-test = "0.4"
 tower-test = "0.4"
 pretty_env_logger = "0.5"

src/client/legacy/connect/proxy/socks/mod.rs

@@ -9,7 +9,7 @@ use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
-use bytes::BytesMut;
+use bytes::{Buf, BytesMut};
 
 use hyper::rt::Read;
 
@@ -44,15 +44,16 @@ pub enum SerializeError {
 async fn read_message<T, M, C>(mut conn: &mut T, buf: &mut BytesMut) -> Result<M, SocksError<C>>
 where
     T: Read + Unpin,
-    M: for<'a> TryFrom<&'a mut BytesMut, Error = ParsingError>,
+    M: for<'a, 'b> TryFrom<&'a mut &'b [u8], Error = ParsingError>,
 {
     let mut tmp = [0; 513];
 
     loop {
         let n = crate::rt::read(&mut conn, &mut tmp).await?;
         buf.extend_from_slice(&tmp[..n]);
 
-        match M::try_from(buf) {
+        let mut view = &buf[..];
+        match M::try_from(&mut view) {
             Err(ParsingError::Incomplete) => {
                 if n == 0 {
                     if buf.spare_capacity_mut().is_empty() {
@@ -67,7 +68,11 @@ where
                 }
             }
             Err(err) => return Err(err.into()),
-            Ok(res) => return Ok(res),
+            Ok(res) => {
+                let consumed = buf.len() - view.len();
+                buf.advance(consumed);
+                return Ok(res);
+            }
         }
     }
 }
@@ -152,3 +157,50 @@ where
         self.project().fut.poll(cx)
     }
 }
+
+#[cfg(all(test, feature = "tokio"))]
+mod test {
+    use bytes::BytesMut;
+    use tokio::io::AsyncWriteExt;
+
+    use super::v5::messages::{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];
+
+    #[tokio::test]
+    async fn it_works_in_one_read() {
+        let (client, mut server) = tokio::io::duplex(SEG1.len() + SEG2.len());
+        server.write_all(&SEG1).await.unwrap();
+        server.write_all(&SEG2).await.unwrap();
+
+        let mut conn = TokioIo::new(client);
+        let mut buf = BytesMut::new();
+
+        let m: Result<ProxyRes, SocksError<()>> = read_message(&mut conn, &mut buf).await;
+        assert!(m.is_ok());
+        assert_eq!(m.unwrap(), ProxyRes(Status::Success));
+    }
+
+    #[tokio::test]
+    async fn it_works_in_multiple_reads() {
+        // Bounded stream ensures message arrives in two reads
+        let (client, mut server) = tokio::io::duplex(SEG1.len());
+        let _writer = tokio::spawn(async move {
+            server.write_all(&SEG1).await.unwrap();
+            server.write_all(&SEG2).await.unwrap();
+        });
+
+        let mut conn = TokioIo::new(client);
+        let mut buf = BytesMut::new();
+
+        let m: Result<ProxyRes, SocksError<()>> = read_message(&mut conn, &mut buf).await;
+        assert!(m.is_ok());
+        assert_eq!(m.unwrap(), ProxyRes(Status::Success));
+
+        _writer.await.unwrap();
+    }
+}

src/client/legacy/connect/proxy/socks/v4/messages.rs

@@ -1,6 +1,6 @@
 use super::super::{ParsingError, SerializeError};
 
-use bytes::{Buf, BufMut, BytesMut};
+use bytes::{Buf, BufMut};
 use std::net::SocketAddrV4;
 
 /// +-----+-----+----+----+----+----+----+----+-------------+------+------------+------+
@@ -9,8 +9,8 @@ use std::net::SocketAddrV4;
 /// |  1  |  1  |    2    |         4         |   Variable  |  1   |  Variable  |   1  |
 /// +-----+-----+----+----+----+----+----+----+-------------+------+------------+------+
 ///                                                                ^^^^^^^^^^^^^^^^^^^^^
-///                                                      optional: only do IP is 0.0.0.X
-#[derive(Debug)]
+///                                                   optional: only do if IP is 0.0.0.X
+#[derive(Debug, PartialEq)]
 pub struct Request<'a>(pub &'a Address);
 
 /// +-----+-----+----+----+----+----+----+----+
@@ -20,10 +20,10 @@ pub struct Request<'a>(pub &'a Address);
 /// +-----+-----+----+----+----+----+----+----+
 ///             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ///              ignore: only for SOCKSv4 BIND
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub struct Response(pub Status);
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub enum Address {
     Socket(SocketAddrV4),
     Domain(String, u16),
@@ -80,10 +80,10 @@ impl Request<'_> {
     }
 }
 
-impl TryFrom<&mut BytesMut> for Response {
+impl TryFrom<&mut &[u8]> for Response {
     type Error = ParsingError;
 
-    fn try_from(buf: &mut BytesMut) -> Result<Self, Self::Error> {
+    fn try_from(buf: &mut &[u8]) -> Result<Self, Self::Error> {
         if buf.remaining() < 8 {
             return Err(ParsingError::Incomplete);
         }
@@ -129,3 +129,83 @@ impl std::fmt::Display for Status {
         })
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use bytes::BytesMut;
+    use std::net::Ipv4Addr;
+
+    #[test]
+    fn request_serialization_with_socket() {
+        let expected = [
+            0x04, // protocol version
+            0x01, // command: connect
+            0x1F, 0x90, // destination port: 8080
+            127, 0, 0, 1,    // destination address: 127.0.0.1
+            0x00, // userid (empty)
+            0x00, // null terminator
+        ];
+
+        let addr = Address::Socket(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080));
+        let mut buf = BytesMut::with_capacity(expected.len());
+        let n = Request(&addr).write_to_buf(&mut buf).unwrap();
+        assert_eq!(n, buf.len());
+        assert_eq!(&buf[..], &expected[..]);
+    }
+
+    #[test]
+    fn request_serialization_with_domain() {
+        let expected = [
+            0x04, // protocol version
+            0x01, // command: connect
+            0x1F, 0x90, // destination port: 8080
+            0x00, 0x00, 0x00, 0xFF, // invalid IP: signals that a domain follows (SOCKS4a)
+            0x00, // userid (empty)
+            0x00, // null terminator
+            b'e', b'x', b'a', b'm', b'p', b'l', b'e', b'.', b'c', b'o', b'm', // domain
+            0x00, // null terminator
+        ];
+
+        let addr = Address::Domain("example.com".into(), 8080);
+        let mut buf = BytesMut::with_capacity(expected.len());
+        let n = Request(&addr).write_to_buf(&mut buf).unwrap();
+        assert_eq!(n, buf.len());
+        assert_eq!(&buf[..], &expected[..]);
+    }
+
+    #[test]
+    fn response_deserialization() {
+        let raw = [
+            0x00, // reply version
+            90,   // status: request granted
+            0x1F, 0x90, // port: 8080 (ignored, only used for BIND)
+            127, 0, 0, 1, // address: 127.0.0.1 (ignored, only used for BIND)
+        ];
+        let mut view = &raw[..];
+
+        let res = Response::try_from(&mut view).unwrap();
+        assert_eq!(res, Response(Status::Success));
+        assert!(view.is_empty());
+    }
+
+    #[test]
+    fn response_incomplete() {
+        let raw = [
+            0x00, // reply version
+            90,   // status: request granted
+            0x00, // truncated mid-port
+        ];
+
+        let err = Response::try_from(&mut &raw[..]).unwrap_err();
+        assert!(matches!(err, ParsingError::Incomplete));
+    }
+
+    #[test]
+    fn request_serialization_would_overflow() {
+        let addr = Address::Socket(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080));
+        let mut short = [0u8; 5];
+        let err = Request(&addr).write_to_buf(&mut short[..]).unwrap_err();
+        assert!(matches!(err, SerializeError::WouldOverflow));
+    }
+}