diff --git a/Cargo.toml b/Cargo.toml index 98c3cf1b..e26719b6 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/client/legacy/connect/proxy/socks/mod.rs b/src/client/legacy/connect/proxy/socks/mod.rs index a6e43f49..9bc89e9d 100644 --- a/src/client/legacy/connect/proxy/socks/mod.rs +++ b/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,7 +44,7 @@ pub enum SerializeError { async fn read_message(mut conn: &mut T, buf: &mut BytesMut) -> Result> 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]; @@ -52,7 +52,8 @@ where 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> = 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> = read_message(&mut conn, &mut buf).await; + assert!(m.is_ok()); + assert_eq!(m.unwrap(), ProxyRes(Status::Success)); + + _writer.await.unwrap(); + } +} diff --git a/src/client/legacy/connect/proxy/socks/v4/messages.rs b/src/client/legacy/connect/proxy/socks/v4/messages.rs index bec8d081..c6421547 100644 --- a/src/client/legacy/connect/proxy/socks/v4/messages.rs +++ b/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 { + fn try_from(buf: &mut &[u8]) -> Result { 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)); + } +} diff --git a/src/client/legacy/connect/proxy/socks/v5/messages.rs b/src/client/legacy/connect/proxy/socks/v5/messages.rs index bdad5405..2c5fecc8 100644 --- a/src/client/legacy/connect/proxy/socks/v5/messages.rs +++ b/src/client/legacy/connect/proxy/socks/v5/messages.rs @@ -8,7 +8,7 @@ use std::net::SocketAddr; /// +----+----------+----------+ /// | 1 | 1 | 1 to 255 | /// +----+----------+----------+ -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub struct NegotiationReq<'a>(pub &'a AuthMethod); /// +----+--------+ @@ -16,7 +16,7 @@ pub struct NegotiationReq<'a>(pub &'a AuthMethod); /// +----+--------+ /// | 1 | 1 | /// +----+--------+ -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub struct NegotiationRes(pub AuthMethod); /// +----+------+----------+------+----------+ @@ -24,7 +24,7 @@ pub struct NegotiationRes(pub AuthMethod); /// +----+------+----------+------+----------+ /// | 1 | 1 | 1 to 255 | 1 | 1 to 255 | /// +----+------+----------+------+----------+ -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub struct AuthenticationReq<'a>(pub &'a str, pub &'a str); /// +----+--------+ @@ -32,7 +32,7 @@ pub struct AuthenticationReq<'a>(pub &'a str, pub &'a str); /// +----+--------+ /// | 1 | 1 | /// +----+--------+ -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub struct AuthenticationRes(pub bool); /// +----+-----+-------+------+----------+----------+ @@ -40,7 +40,7 @@ pub struct AuthenticationRes(pub bool); /// +----+-----+-------+------+----------+----------+ /// | 1 | 1 | X'00' | 1 | Variable | 2 | /// +----+-----+-------+------+----------+----------+ -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub struct ProxyReq<'a>(pub &'a Address); /// +----+-----+-------+------+----------+----------+ @@ -48,7 +48,7 @@ pub struct ProxyReq<'a>(pub &'a Address); /// +----+-----+-------+------+----------+----------+ /// | 1 | 1 | X'00' | 1 | Variable | 2 | /// +----+-----+-------+------+----------+----------+ -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub struct ProxyRes(pub Status); #[repr(u8)] @@ -59,7 +59,7 @@ pub enum AuthMethod { NoneAcceptable = 0xFF, } -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub enum Address { Socket(SocketAddr), Domain(String, u16), @@ -92,10 +92,10 @@ impl NegotiationReq<'_> { } } -impl TryFrom<&mut BytesMut> for NegotiationRes { +impl TryFrom<&mut &[u8]> for NegotiationRes { type Error = ParsingError; - fn try_from(buf: &mut BytesMut) -> Result { + fn try_from(buf: &mut &[u8]) -> Result { if buf.remaining() < 2 { return Err(ParsingError::Incomplete); } @@ -117,20 +117,20 @@ impl AuthenticationReq<'_> { buf.put_u8(0x01); // Version - buf.put_u8(self.0.len() as u8); // Username length (guarenteed to be 255 or less) + buf.put_u8(self.0.len() as u8); // Username length (guaranteed to be 255 or less) buf.put_slice(self.0.as_bytes()); // Username - buf.put_u8(self.1.len() as u8); // Password length (guarenteed to be 255 or less) + buf.put_u8(self.1.len() as u8); // Password length (guaranteed to be 255 or less) buf.put_slice(self.1.as_bytes()); // Password Ok(3 + self.0.len() + self.1.len()) } } -impl TryFrom<&mut BytesMut> for AuthenticationRes { +impl TryFrom<&mut &[u8]> for AuthenticationRes { type Error = ParsingError; - fn try_from(buf: &mut BytesMut) -> Result { + fn try_from(buf: &mut &[u8]) -> Result { if buf.remaining() < 2 { return Err(ParsingError::Incomplete); } @@ -168,10 +168,10 @@ impl ProxyReq<'_> { } } -impl TryFrom<&mut BytesMut> for ProxyRes { +impl TryFrom<&mut &[u8]> for ProxyRes { type Error = ParsingError; - fn try_from(buf: &mut BytesMut) -> Result { + fn try_from(buf: &mut &[u8]) -> Result { if buf.remaining() < 3 { return Err(ParsingError::Incomplete); } @@ -239,10 +239,10 @@ impl Address { } } -impl TryFrom<&mut BytesMut> for Address { +impl TryFrom<&mut &[u8]> for Address { type Error = ParsingError; - fn try_from(buf: &mut BytesMut) -> Result { + fn try_from(buf: &mut &[u8]) -> Result { if buf.remaining() < 2 { return Err(ParsingError::Incomplete); } @@ -263,17 +263,18 @@ impl TryFrom<&mut BytesMut> for Address { } // Domain 0x03 => { - let len = buf.get_u8(); + let len = buf.get_u8() as usize; if len == 0 { return Err(ParsingError::Other); - } else if buf.remaining() < (len as usize) + 2 { + } else if buf.remaining() < len + 2 { return Err(ParsingError::Incomplete); } - let domain = std::str::from_utf8(&buf[..len as usize]) + let domain = std::str::from_utf8(&buf.chunk()[..len]) .map_err(|_| ParsingError::Other)? .to_string(); + buf.advance(len); let port = buf.get_u16(); @@ -346,3 +347,156 @@ impl std::fmt::Display for Status { }) } } + +#[cfg(test)] +mod test { + use super::*; + use std::net::{Ipv4Addr, Ipv6Addr}; + + #[test] + fn negotiation_req_serialization() { + let expected = [ + 0x05, // protocol version + 0x01, // number of authentication methods: 1 + 0x00, // method: no authentication + ]; + + let mut buf = BytesMut::with_capacity(expected.len()); + let n = NegotiationReq(&AuthMethod::NoAuth) + .write_to_buf(&mut buf) + .unwrap(); + assert_eq!(n, buf.len()); + assert_eq!(&buf[..], &expected[..]); + } + + #[test] + fn negotiation_res_deserialization() { + let raw = [ + 0x05, // protocol version + 0x02, // selected method: username/password + ]; + + let mut view = &raw[..]; + + let res = NegotiationRes::try_from(&mut view).unwrap(); + assert_eq!(res, NegotiationRes(AuthMethod::UserPass)); + assert!(view.is_empty()); + } + + #[test] + fn authentication_req_serialization() { + let expected = [ + 0x01, // authentication version + 0x04, // username length: 4 + b'u', b's', b'e', b'r', // username + 0x04, // password length: 4 + b'p', b'a', b's', b's', // password + ]; + + let mut buf = BytesMut::with_capacity(expected.len()); + let n = AuthenticationReq("user", "pass") + .write_to_buf(&mut buf) + .unwrap(); + assert_eq!(n, buf.len()); + assert_eq!(&buf[..], &expected[..]); + } + + #[test] + fn authentication_res_deserialization() { + let raw = [ + 0x01, // authentication version + 0x00, // status: success + ]; + assert_eq!( + AuthenticationRes::try_from(&mut &raw[..]).unwrap(), + AuthenticationRes(true) + ); + + let raw = [ + 0x01, // authentication version + 0x01, // status: failure + ]; + assert_eq!( + AuthenticationRes::try_from(&mut &raw[..]).unwrap(), + AuthenticationRes(false) + ); + } + + #[test] + fn proxy_req_serialization() { + let expected = [ + 0x05, // protocol version + 0x01, // command: connect + 0x00, // reserved + 0x01, // address type: IPv4 + 127, 0, 0, 1, // destination address: 127.0.0.1 + 0x1F, 0x90, // destination port: 8080 + ]; + + let addr = Address::Socket(SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 8080)); + + let mut buf = BytesMut::with_capacity(expected.len()); + let n = ProxyReq(&addr).write_to_buf(&mut buf).unwrap(); + assert_eq!(n, buf.len()); + assert_eq!(&buf[..], &expected[..]); + } + + #[test] + fn proxy_res_deserialization() { + let raw = [ + 0x05, // protocol version + 0x00, // reply: success + 0x00, // reserved + 0x01, // address type: IPv4 + 127, 0, 0, 1, // bound address: 127.0.0.1 + 0x1F, 0x90, // bound port: 8080 + ]; + let mut view = &raw[..]; + + let res = ProxyRes::try_from(&mut view).unwrap(); + assert_eq!(res, ProxyRes(Status::Success)); + assert!(view.is_empty()); + } + + #[test] + fn serialization_would_overflow() { + let mut buf = BytesMut::with_capacity(2); + + let err = NegotiationReq(&AuthMethod::NoAuth) + .write_to_buf(&mut buf) + .unwrap_err(); + assert!(matches!(err, SerializeError::WouldOverflow)); + } + + fn assert_address_roundtrips(addr: Address) { + let mut buf = BytesMut::with_capacity(64); + + let n = addr.write_to_buf(&mut buf).unwrap(); + assert_eq!(n, buf.len()); + + let mut view = &buf[..]; + assert_eq!(Address::try_from(&mut view).unwrap(), addr); + assert!(view.is_empty(), "address bytes should be fully consumed"); + } + + #[test] + fn address_roundtrip_ipv4() { + assert_address_roundtrips(Address::Socket(SocketAddr::new( + Ipv4Addr::LOCALHOST.into(), + 8080, + ))); + } + + #[test] + fn address_roundtrip_ipv6() { + assert_address_roundtrips(Address::Socket(SocketAddr::new( + Ipv6Addr::LOCALHOST.into(), + 8080, + ))); + } + + #[test] + fn address_roundtrip_domain() { + assert_address_roundtrips(Address::Domain("example.com".into(), 8080)); + } +} diff --git a/src/client/legacy/connect/proxy/socks/v5/mod.rs b/src/client/legacy/connect/proxy/socks/v5/mod.rs index d1c6f229..71d4a255 100644 --- a/src/client/legacy/connect/proxy/socks/v5/mod.rs +++ b/src/client/legacy/connect/proxy/socks/v5/mod.rs @@ -1,7 +1,7 @@ mod errors; pub use errors::*; -mod messages; +pub(super) mod messages; use messages::*; use std::net::{IpAddr, SocketAddr, ToSocketAddrs}; @@ -65,7 +65,7 @@ impl SocksV5 { /// /// Username and Password must be maximum of 255 characters each. /// 0 length strings are allowed despite RFC prohibiting it. This is done for - /// compatablity with server implementations that use empty credentials + /// compatibility with server implementations that use empty credentials /// to allow returning error codes during IP authentication. pub fn with_auth(mut self, user: String, pass: String) -> Self { self.config.proxy_auth = Some((user, pass));