From 2efd4cf7c5f0969112a28bca34b15e1b5a37b3bc Mon Sep 17 00:00:00 2001 From: "glm-5.1" Date: Tue, 2 Jun 2026 10:49:07 +0000 Subject: [PATCH] Implement SOCKS5 server: local proxy forwarding through SSH channels Convert socks5.rs to directory module with protocol parsing and server implementation. Socks5Server binds to configurable address (default 127.0.0.1:1080), handles SOCKS5 handshake (no-auth), parses IPv4/IPv6/domain addresses, and proxies bidirectionally via SSH direct_tcpip channels. Domain names sent unresolved (SOCKS5h) to prevent DNS leaks (ADR-006). No logging of request targets per privacy requirements. --- crates/wraith-core/src/socks5.rs | 0 crates/wraith-core/src/socks5/mod.rs | 490 ++++++++++++++++++++++ crates/wraith-core/src/socks5/protocol.rs | 304 ++++++++++++++ 3 files changed, 794 insertions(+) delete mode 100644 crates/wraith-core/src/socks5.rs create mode 100644 crates/wraith-core/src/socks5/mod.rs create mode 100644 crates/wraith-core/src/socks5/protocol.rs diff --git a/crates/wraith-core/src/socks5.rs b/crates/wraith-core/src/socks5.rs deleted file mode 100644 index e69de29..0000000 diff --git a/crates/wraith-core/src/socks5/mod.rs b/crates/wraith-core/src/socks5/mod.rs new file mode 100644 index 0000000..12db855 --- /dev/null +++ b/crates/wraith-core/src/socks5/mod.rs @@ -0,0 +1,490 @@ +mod protocol; + +use std::net::SocketAddr; +use std::sync::Arc; + +use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; +use tokio::net::TcpListener; +use tokio::sync::Mutex; +use tracing::debug; + +use protocol::{Socks5Reply, Socks5Request, Socks5VersionMethod}; + +pub use protocol::Socks5Address; + +const DEFAULT_SOCKS5_ADDR: &str = "127.0.0.1:1080"; + +pub trait ChannelOpener: Send + Sync + 'static { + type Stream: AsyncRead + AsyncWrite + Unpin + Send + 'static; + + fn open_channel( + &self, + host: String, + port: u16, + ) -> impl std::future::Future> + Send; +} + +#[derive(Debug, thiserror::Error)] +pub enum ChannelOpenError { + #[error("session closed")] + SessionClosed, + #[error("channel open failed")] + ChannelOpenFailed, + #[error("connection refused")] + ConnectionRefused, +} + +pub struct Socks5Server { + listen_addr: SocketAddr, + channel_opener: Arc, +} + +impl Socks5Server { + pub fn new(channel_opener: C) -> Self { + Self::with_addr(channel_opener, DEFAULT_SOCKS5_ADDR) + } + + pub fn with_addr(channel_opener: C, addr: &str) -> Self { + let listen_addr: SocketAddr = addr + .parse() + .expect("invalid SOCKS5 listen address"); + Self { + listen_addr, + channel_opener: Arc::new(channel_opener), + } + } + + pub fn listen_addr(&self) -> SocketAddr { + self.listen_addr + } + + pub async fn run(self) -> Result<(), std::io::Error> { + let listener = TcpListener::bind(self.listen_addr).await?; + debug!("socks5 server listening on {}", self.listen_addr); + loop { + let (socket, _peer) = listener.accept().await?; + let opener = Arc::clone(&self.channel_opener); + tokio::spawn(async move { + if let Err(e) = handle_socks5_connection(socket, opener).await { + debug!("socks5 connection error: {e}"); + } + }); + } + } +} + +async fn handle_socks5_connection( + mut socket: S, + opener: Arc, +) -> Result<(), Socks5Error> +where + S: AsyncRead + AsyncWrite + Unpin, + C: ChannelOpener, +{ + let vm = Socks5VersionMethod::read_from(&mut socket).await?; + if vm.version != 0x05 { + return Err(Socks5Error::InvalidVersion(vm.version)); + } + if !vm.methods.contains(&0x00) { + let reply = [0x05, 0xFF]; + socket.write_all(&reply).await?; + socket.shutdown().await?; + return Err(Socks5Error::NoAcceptableAuth); + } + let reply = [0x05, 0x00]; + socket.write_all(&reply).await?; + + let request = Socks5Request::read_from(&mut socket).await?; + if request.version != 0x05 { + return Err(Socks5Error::InvalidVersion(request.version)); + } + if request.command != 0x01 { + send_error_reply(&mut socket, Socks5Reply::command_not_supported()).await?; + return Err(Socks5Error::UnsupportedCommand(request.command)); + } + + let (host, port) = match &request.address { + Socks5Address::Ipv4(addr) => (addr.to_string(), request.port), + Socks5Address::Ipv6(addr) => (addr.to_string(), request.port), + Socks5Address::Domain(name) => (name.clone(), request.port), + }; + + match opener.open_channel(host, port).await { + Ok(mut ssh_stream) => { + let bind_addr = Socks5Address::Ipv4(std::net::Ipv4Addr::UNSPECIFIED); + let reply = Socks5Reply::success(bind_addr, 0); + reply.write_to(&mut socket).await?; + tokio::io::copy_bidirectional(&mut socket, &mut ssh_stream).await?; + Ok(()) + } + Err(_) => { + send_error_reply(&mut socket, Socks5Reply::connection_refused()).await?; + Err(Socks5Error::ChannelOpenFailed) + } + } +} + +async fn send_error_reply( + socket: &mut S, + reply: Socks5Reply, +) -> Result<(), Socks5Error> { + reply.write_to(socket).await?; + let _ = socket.shutdown().await; + Ok(()) +} + +#[derive(Debug, thiserror::Error)] +pub enum Socks5Error { + #[error("invalid SOCKS version: {0}")] + InvalidVersion(u8), + #[error("no acceptable auth method")] + NoAcceptableAuth, + #[error("unsupported command: {0}")] + UnsupportedCommand(u8), + #[error("channel open failed")] + ChannelOpenFailed, + #[error("io error")] + Io(#[from] std::io::Error), +} + +pub struct HandleChannelOpener { + handle: Arc>>, +} + +impl HandleChannelOpener { + pub fn new(handle: russh::client::Handle) -> Self { + Self { + handle: Arc::new(Mutex::new(handle)), + } + } + + pub fn from_arc(handle: Arc>>) -> Self { + Self { handle } + } +} + +impl ChannelOpener for HandleChannelOpener { + type Stream = russh::ChannelStream; + + async fn open_channel(&self, host: String, port: u16) -> Result { + let handle = self.handle.lock().await; + if handle.is_closed() { + return Err(ChannelOpenError::SessionClosed); + } + let channel = handle + .channel_open_direct_tcpip(host, port as u32, "127.0.0.1", 0) + .await + .map_err(|_| ChannelOpenError::ChannelOpenFailed)?; + Ok(channel.into_stream()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::{duplex, AsyncReadExt, AsyncWriteExt, DuplexStream}; + + struct MockChannelOpener { + fail: bool, + } + + impl ChannelOpener for MockChannelOpener { + type Stream = DuplexStream; + + async fn open_channel( + &self, + _host: String, + _port: u16, + ) -> Result { + if self.fail { + Err(ChannelOpenError::ChannelOpenFailed) + } else { + let (client, _server) = duplex(4096); + Ok(client) + } + } + } + + fn build_socks5_greeting(methods: &[u8]) -> Vec { + let mut buf = vec![0x05, methods.len() as u8]; + buf.extend_from_slice(methods); + buf + } + + fn build_socks5_connect_ipv4(addr: [u8; 4], port: u16) -> Vec { + let mut buf = vec![0x05, 0x01, 0x00, 0x01]; + buf.extend_from_slice(&addr); + buf.extend_from_slice(&port.to_be_bytes()); + buf + } + + fn build_socks5_connect_domain(domain: &str, port: u16) -> Vec { + let mut buf = vec![0x05, 0x01, 0x00, 0x03]; + buf.push(domain.len() as u8); + buf.extend_from_slice(domain.as_bytes()); + buf.extend_from_slice(&port.to_be_bytes()); + buf + } + + fn build_socks5_connect_ipv6(addr: [u8; 16], port: u16) -> Vec { + let mut buf = vec![0x05, 0x01, 0x00, 0x04]; + buf.extend_from_slice(&addr); + buf.extend_from_slice(&port.to_be_bytes()); + buf + } + + async fn do_handshake(client: &mut DuplexStream) -> [u8; 2] { + client.write_all(&build_socks5_greeting(&[0x00])).await.unwrap(); + client.flush().await.unwrap(); + let mut resp = [0u8; 2]; + client.read_exact(&mut resp).await.unwrap(); + resp + } + + async fn do_connect_ipv4(client: &mut DuplexStream, addr: [u8; 4], port: u16) -> Vec { + client + .write_all(&build_socks5_connect_ipv4(addr, port)) + .await + .unwrap(); + client.flush().await.unwrap(); + let mut reply_buf = [0u8; 10]; + client.read_exact(&mut reply_buf).await.unwrap(); + reply_buf.to_vec() + } + + #[tokio::test] + async fn handshake_no_auth_method() { + let (mut client, server) = duplex(4096); + let opener = MockChannelOpener { fail: false }; + + let server_handle = tokio::spawn(async move { + handle_socks5_connection(server, Arc::new(opener)).await + }); + + let resp = do_handshake(&mut client).await; + assert_eq!(resp, [0x05, 0x00]); + + let reply_buf = do_connect_ipv4(&mut client, [127, 0, 0, 1], 80).await; + assert_eq!(reply_buf[0], 0x05); + assert_eq!(reply_buf[1], 0x00); + + drop(client); + let _ = server_handle.await; + } + + #[tokio::test] + async fn handshake_rejects_no_acceptable_method() { + let (mut client, server) = duplex(4096); + let opener = MockChannelOpener { fail: false }; + + let server_handle = tokio::spawn(async move { + handle_socks5_connection(server, Arc::new(opener)).await + }); + + client + .write_all(&build_socks5_greeting(&[0x02])) + .await + .unwrap(); + client.flush().await.unwrap(); + + let mut resp = [0u8; 2]; + client.read_exact(&mut resp).await.unwrap(); + assert_eq!(resp, [0x05, 0xFF]); + + drop(client); + let result = server_handle.await.unwrap(); + assert!(result.is_err()); + assert!(matches!( + result.unwrap_err(), + Socks5Error::NoAcceptableAuth + )); + } + + #[tokio::test] + async fn address_type_ipv4() { + let (mut client, server) = duplex(4096); + let opener = MockChannelOpener { fail: false }; + + let server_handle = tokio::spawn(async move { + handle_socks5_connection(server, Arc::new(opener)).await + }); + + do_handshake(&mut client).await; + let reply_buf = do_connect_ipv4(&mut client, [10, 0, 0, 1], 443).await; + assert_eq!(reply_buf[1], 0x00); + + drop(client); + let _ = server_handle.await; + } + + #[tokio::test] + async fn address_type_domain() { + let (mut client, server) = duplex(4096); + let opener = MockChannelOpener { fail: false }; + + let server_handle = tokio::spawn(async move { + handle_socks5_connection(server, Arc::new(opener)).await + }); + + do_handshake(&mut client).await; + + client + .write_all(&build_socks5_connect_domain("example.com", 443)) + .await + .unwrap(); + client.flush().await.unwrap(); + + let mut reply_buf = [0u8; 10]; + client.read_exact(&mut reply_buf).await.unwrap(); + assert_eq!(reply_buf[1], 0x00); + + drop(client); + let _ = server_handle.await; + } + + #[tokio::test] + async fn address_type_ipv6() { + let (mut client, server) = duplex(4096); + let opener = MockChannelOpener { fail: false }; + + let server_handle = tokio::spawn(async move { + handle_socks5_connection(server, Arc::new(opener)).await + }); + + do_handshake(&mut client).await; + + let ipv6_addr = [0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]; + client + .write_all(&build_socks5_connect_ipv6(ipv6_addr, 443)) + .await + .unwrap(); + client.flush().await.unwrap(); + + let mut reply_buf = [0u8; 10]; + client.read_exact(&mut reply_buf).await.unwrap(); + assert_eq!(reply_buf[0], 0x05); + assert_eq!(reply_buf[1], 0x00); + + drop(client); + let _ = server_handle.await; + } + + #[tokio::test] + async fn channel_open_failure_returns_socks5_error() { + let (mut client, server) = duplex(4096); + let opener = MockChannelOpener { fail: true }; + + let server_handle = tokio::spawn(async move { + handle_socks5_connection(server, Arc::new(opener)).await + }); + + do_handshake(&mut client).await; + let reply_buf = do_connect_ipv4(&mut client, [10, 0, 0, 1], 80).await; + assert_eq!(reply_buf[0], 0x05); + assert_eq!(reply_buf[1], 0x05); + + drop(client); + let _ = server_handle.await; + } + + #[tokio::test] + async fn unsupported_command_returns_error() { + let (mut client, server) = duplex(4096); + let opener = MockChannelOpener { fail: false }; + + let server_handle = tokio::spawn(async move { + handle_socks5_connection(server, Arc::new(opener)).await + }); + + do_handshake(&mut client).await; + + let mut bind_req = vec![0x05, 0x02, 0x00, 0x01]; + bind_req.extend_from_slice(&[127, 0, 0, 1]); + bind_req.extend_from_slice(&80u16.to_be_bytes()); + client.write_all(&bind_req).await.unwrap(); + client.flush().await.unwrap(); + + let mut reply_buf = [0u8; 10]; + client.read_exact(&mut reply_buf).await.unwrap(); + assert_eq!(reply_buf[1], 0x07); + + drop(client); + let _ = server_handle.await; + } + + #[tokio::test] + async fn bidirectional_proxy_flow() { + let (mut client_sock, server_sock) = duplex(4096); + let (ssh_client, mut ssh_server) = duplex(4096); + + let ssh_stream = Arc::new(Mutex::new(Some(ssh_client))); + + struct ProxyOpener { + stream: Arc>>, + } + + impl ChannelOpener for ProxyOpener { + type Stream = DuplexStream; + + async fn open_channel( + &self, + _host: String, + _port: u16, + ) -> Result { + self.stream + .lock() + .await + .take() + .ok_or(ChannelOpenError::ChannelOpenFailed) + } + } + + let opener = ProxyOpener { + stream: Arc::clone(&ssh_stream), + }; + + let server_handle = tokio::spawn(async move { + handle_socks5_connection(server_sock, Arc::new(opener)).await + }); + + do_handshake(&mut client_sock).await; + let reply_buf = do_connect_ipv4(&mut client_sock, [127, 0, 0, 1], 80).await; + assert_eq!(reply_buf[1], 0x00); + + let test_data = b"hello through tunnel"; + client_sock.write_all(test_data).await.unwrap(); + client_sock.flush().await.unwrap(); + + let mut received = vec![0u8; test_data.len()]; + AsyncReadExt::read_exact(&mut ssh_server, &mut received) + .await + .unwrap(); + assert_eq!(&received, test_data); + + let echo_data = b"response from tunnel"; + ssh_server.write_all(echo_data).await.unwrap(); + ssh_server.flush().await.unwrap(); + + let mut received_back = vec![0u8; echo_data.len()]; + client_sock.read_exact(&mut received_back).await.unwrap(); + assert_eq!(&received_back, echo_data); + + drop(client_sock); + drop(ssh_server); + let _ = server_handle.await; + } + + #[tokio::test] + async fn default_listen_address() { + let opener = MockChannelOpener { fail: false }; + let server = Socks5Server::new(opener); + assert_eq!(server.listen_addr(), "127.0.0.1:1080".parse().unwrap()); + } + + #[tokio::test] + async fn custom_listen_address() { + let opener = MockChannelOpener { fail: false }; + let server = Socks5Server::with_addr(opener, "127.0.0.1:9050"); + assert_eq!(server.listen_addr(), "127.0.0.1:9050".parse().unwrap()); + } +} \ No newline at end of file diff --git a/crates/wraith-core/src/socks5/protocol.rs b/crates/wraith-core/src/socks5/protocol.rs new file mode 100644 index 0000000..272c187 --- /dev/null +++ b/crates/wraith-core/src/socks5/protocol.rs @@ -0,0 +1,304 @@ +use std::net::{Ipv4Addr, Ipv6Addr}; + +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; + +#[derive(Debug, Clone, PartialEq)] +pub enum Socks5Address { + Ipv4(Ipv4Addr), + Ipv6(Ipv6Addr), + Domain(String), +} + +#[derive(Debug)] +pub struct Socks5VersionMethod { + pub version: u8, + pub methods: Vec, +} + +impl Socks5VersionMethod { + pub async fn read_from(reader: &mut R) -> std::io::Result { + let version = reader.read_u8().await?; + let nmethods = reader.read_u8().await?; + let mut methods = vec![0u8; nmethods as usize]; + reader.read_exact(&mut methods).await?; + Ok(Self { version, methods }) + } +} + +#[derive(Debug)] +pub struct Socks5Request { + pub version: u8, + pub command: u8, + pub address: Socks5Address, + pub port: u16, +} + +impl Socks5Request { + pub async fn read_from(reader: &mut R) -> std::io::Result { + let version = reader.read_u8().await?; + let command = reader.read_u8().await?; + let _rsv = reader.read_u8().await?; + let atyp = reader.read_u8().await?; + + let address = match atyp { + 0x01 => { + let mut octets = [0u8; 4]; + reader.read_exact(&mut octets).await?; + Socks5Address::Ipv4(Ipv4Addr::from(octets)) + } + 0x04 => { + let mut octets = [0u8; 16]; + reader.read_exact(&mut octets).await?; + Socks5Address::Ipv6(Ipv6Addr::from(octets)) + } + 0x03 => { + let len = reader.read_u8().await?; + let mut domain = vec![0u8; len as usize]; + reader.read_exact(&mut domain).await?; + Socks5Address::Domain(String::from_utf8_lossy(&domain).into_owned()) + } + _ => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("unsupported address type: {atyp}"), + )) + } + }; + + let port = reader.read_u16().await?; + + Ok(Self { + version, + command, + address, + port, + }) + } +} + +#[derive(Debug)] +pub struct Socks5Reply { + pub version: u8, + pub reply: u8, + pub address: Socks5Address, + pub port: u16, +} + +impl Socks5Reply { + pub fn success(address: Socks5Address, port: u16) -> Self { + Self { + version: 0x05, + reply: 0x00, + address, + port, + } + } + + pub fn connection_refused() -> Self { + Self { + version: 0x05, + reply: 0x05, + address: Socks5Address::Ipv4(Ipv4Addr::UNSPECIFIED), + port: 0, + } + } + + pub fn command_not_supported() -> Self { + Self { + version: 0x05, + reply: 0x07, + address: Socks5Address::Ipv4(Ipv4Addr::UNSPECIFIED), + port: 0, + } + } + + pub async fn write_to(&self, writer: &mut W) -> std::io::Result<()> { + writer.write_u8(self.version).await?; + writer.write_u8(self.reply).await?; + writer.write_u8(0x00).await?; + match &self.address { + Socks5Address::Ipv4(addr) => { + writer.write_u8(0x01).await?; + writer.write_all(&addr.octets()).await?; + } + Socks5Address::Ipv6(addr) => { + writer.write_u8(0x04).await?; + writer.write_all(&addr.octets()).await?; + } + Socks5Address::Domain(name) => { + writer.write_u8(0x03).await?; + writer.write_u8(name.len() as u8).await?; + writer.write_all(name.as_bytes()).await?; + } + } + writer.write_u16(self.port).await?; + writer.flush().await?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Cursor; + + #[tokio::test] + async fn parse_version_method_no_auth() { + let data = [0x05, 0x01, 0x00]; + let mut cursor = Cursor::new(&data[..]); + let vm = Socks5VersionMethod::read_from(&mut cursor).await.unwrap(); + assert_eq!(vm.version, 0x05); + assert_eq!(vm.methods, vec![0x00]); + } + + #[tokio::test] + async fn parse_version_method_multiple() { + let data = [0x05, 0x02, 0x00, 0x02]; + let mut cursor = Cursor::new(&data[..]); + let vm = Socks5VersionMethod::read_from(&mut cursor).await.unwrap(); + assert_eq!(vm.version, 0x05); + assert_eq!(vm.methods, vec![0x00, 0x02]); + } + + #[tokio::test] + async fn parse_request_ipv4() { + let mut data = vec![0x05, 0x01, 0x00, 0x01]; + data.extend_from_slice(&[10, 0, 0, 1]); + data.extend_from_slice(&443u16.to_be_bytes()); + let mut cursor = Cursor::new(&data[..]); + let req = Socks5Request::read_from(&mut cursor).await.unwrap(); + assert_eq!(req.version, 0x05); + assert_eq!(req.command, 0x01); + assert_eq!( + req.address, + Socks5Address::Ipv4(Ipv4Addr::new(10, 0, 0, 1)) + ); + assert_eq!(req.port, 443); + } + + #[tokio::test] + async fn parse_request_ipv6() { + let mut data = vec![0x05, 0x01, 0x00, 0x04]; + let octets: [u8; 16] = [0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]; + data.extend_from_slice(&octets); + data.extend_from_slice(&443u16.to_be_bytes()); + let mut cursor = Cursor::new(&data[..]); + let req = Socks5Request::read_from(&mut cursor).await.unwrap(); + assert_eq!(req.version, 0x05); + assert_eq!(req.command, 0x01); + assert!(matches!(req.address, Socks5Address::Ipv6(_))); + assert_eq!(req.port, 443); + } + + #[tokio::test] + async fn parse_request_domain() { + let domain = "example.com"; + let mut data = vec![0x05, 0x01, 0x00, 0x03]; + data.push(domain.len() as u8); + data.extend_from_slice(domain.as_bytes()); + data.extend_from_slice(&443u16.to_be_bytes()); + let mut cursor = Cursor::new(&data[..]); + let req = Socks5Request::read_from(&mut cursor).await.unwrap(); + assert_eq!(req.version, 0x05); + assert_eq!(req.command, 0x01); + assert_eq!(req.address, Socks5Address::Domain("example.com".to_string())); + assert_eq!(req.port, 443); + } + + #[tokio::test] + async fn parse_request_unsupported_address_type() { + let data = [0x05, 0x01, 0x00, 0x05]; + let mut cursor = Cursor::new(&data[..]); + let result = Socks5Request::read_from(&mut cursor).await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn reply_success_ipv4() { + let reply = Socks5Reply::success(Socks5Address::Ipv4(Ipv4Addr::UNSPECIFIED), 0); + let mut buf = Vec::new(); + reply.write_to(&mut buf).await.unwrap(); + assert_eq!(buf[0], 0x05); + assert_eq!(buf[1], 0x00); + assert_eq!(buf[2], 0x00); + assert_eq!(buf[3], 0x01); + } + + #[tokio::test] + async fn reply_connection_refused() { + let reply = Socks5Reply::connection_refused(); + let mut buf = Vec::new(); + reply.write_to(&mut buf).await.unwrap(); + assert_eq!(buf[0], 0x05); + assert_eq!(buf[1], 0x05); + } + + #[tokio::test] + async fn reply_command_not_supported() { + let reply = Socks5Reply::command_not_supported(); + let mut buf = Vec::new(); + reply.write_to(&mut buf).await.unwrap(); + assert_eq!(buf[0], 0x05); + assert_eq!(buf[1], 0x07); + } + + #[tokio::test] + async fn roundtrip_ipv4_reply() { + let reply = Socks5Reply::success(Socks5Address::Ipv4(Ipv4Addr::new(127, 0, 0, 1)), 1080); + let mut buf = Vec::new(); + reply.write_to(&mut buf).await.unwrap(); + + let mut cursor = Cursor::new(&buf[..]); + let version = cursor.read_u8().await.unwrap(); + let _reply_code = cursor.read_u8().await.unwrap(); + let _rsv = cursor.read_u8().await.unwrap(); + let atyp = cursor.read_u8().await.unwrap(); + assert_eq!(version, 0x05); + assert_eq!(atyp, 0x01); + let mut octets = [0u8; 4]; + cursor.read_exact(&mut octets).await.unwrap(); + assert_eq!(Ipv4Addr::from(octets), Ipv4Addr::new(127, 0, 0, 1)); + let port = cursor.read_u16().await.unwrap(); + assert_eq!(port, 1080); + } + + #[tokio::test] + async fn roundtrip_ipv6_reply() { + let addr = Ipv6Addr::new(0x2001, 0x0db8, 0, 0, 0, 0, 0, 1); + let reply = Socks5Reply::success(Socks5Address::Ipv6(addr), 443); + let mut buf = Vec::new(); + reply.write_to(&mut buf).await.unwrap(); + + let mut cursor = Cursor::new(&buf[..]); + let _version = cursor.read_u8().await.unwrap(); + let _reply_code = cursor.read_u8().await.unwrap(); + let _rsv = cursor.read_u8().await.unwrap(); + let atyp = cursor.read_u8().await.unwrap(); + assert_eq!(atyp, 0x04); + let mut octets = [0u8; 16]; + cursor.read_exact(&mut octets).await.unwrap(); + assert_eq!(Ipv6Addr::from(octets), addr); + let port = cursor.read_u16().await.unwrap(); + assert_eq!(port, 443); + } + + #[tokio::test] + async fn roundtrip_domain_reply() { + let reply = Socks5Reply::success(Socks5Address::Domain("example.com".to_string()), 8080); + let mut buf = Vec::new(); + reply.write_to(&mut buf).await.unwrap(); + + let mut cursor = Cursor::new(&buf[..]); + let _version = cursor.read_u8().await.unwrap(); + let _reply_code = cursor.read_u8().await.unwrap(); + let _rsv = cursor.read_u8().await.unwrap(); + let atyp = cursor.read_u8().await.unwrap(); + assert_eq!(atyp, 0x03); + let len = cursor.read_u8().await.unwrap(); + let mut domain = vec![0u8; len as usize]; + cursor.read_exact(&mut domain).await.unwrap(); + assert_eq!(String::from_utf8(domain).unwrap(), "example.com"); + let port = cursor.read_u16().await.unwrap(); + assert_eq!(port, 8080); + } +} \ No newline at end of file