use std::net::SocketAddr;
use std::time::Duration;
use bytes::Bytes;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, UdpSocket};
use tokio::time::timeout;
use super::*;
use crate::rfc1928::consts;
const BOUND: Duration = Duration::from_secs(10);
async fn decode_inner(inner: &[u8]) -> (TargetAddr, Bytes) {
let atyp = inner[0];
let (mut client, mut server) = tokio::io::duplex(inner.len() + 16);
client
.write_all(&inner[1..])
.await
.expect("feed the inner bytes");
drop(client);
let addr = crate::rfc1928::util::target_addr::read_address(&mut server, atyp)
.await
.expect("the address decodes");
let addr_len = addr.to_be_bytes().expect("re-encode").len();
(addr, Bytes::copy_from_slice(&inner[addr_len..]))
}
async fn echo_tcp_server() -> SocketAddr {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let addr = listener.local_addr().expect("local_addr");
tokio::spawn(async move {
loop {
let Ok((mut stream, _)) = listener.accept().await else {
return;
};
tokio::spawn(async move {
let mut buf = [0u8; 4096];
loop {
match stream.read(&mut buf).await {
Ok(0) | Err(_) => return,
Ok(n) => {
if stream.write_all(&buf[..n]).await.is_err() {
return;
}
}
}
}
});
}
});
addr
}
async fn echo_udp_server() -> SocketAddr {
let sock = UdpSocket::bind("127.0.0.1:0").await.expect("bind");
let addr = sock.local_addr().expect("local_addr");
tokio::spawn(async move {
let mut buf = [0u8; 65535];
loop {
let Ok((n, from)) = sock.recv_from(&mut buf).await else {
return;
};
if sock.send_to(&buf[..n], from).await.is_err() {
return;
}
}
});
addr
}
#[tokio::test]
async fn tcp_dial_round_trips_through_the_boxed_halves() {
let server = echo_tcp_server().await;
let handle = tcp_dial()(TargetAddr::Ip(server))
.await
.expect("the dial connects");
let mut write = handle.write;
write.write_all(b"ping").await.expect("write");
write.flush().await.expect("flush");
let mut read = handle.read;
let mut buf = [0u8; 4];
timeout(BOUND, read.read_exact(&mut buf))
.await
.expect("echo within bound")
.expect("read");
assert_eq!(&buf, b"ping");
}
#[tokio::test]
async fn tcp_dial_resolves_domain_targets_here() {
let server = echo_tcp_server().await;
let handle = tcp_dial()(TargetAddr::Domain("localhost".to_owned(), server.port()))
.await
.expect("the domain target resolves and connects");
let mut write = handle.write;
write.write_all(b"pong").await.expect("write");
write.flush().await.expect("flush");
let mut read = handle.read;
let mut buf = [0u8; 4];
timeout(BOUND, read.read_exact(&mut buf))
.await
.expect("echo within bound")
.expect("read");
assert_eq!(&buf, b"pong");
}
#[tokio::test]
async fn tcp_dial_connect_refused_maps_faithfully() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let refused_port = listener.local_addr().expect("local_addr").port();
drop(listener);
let err = tcp_dial()(TargetAddr::Ip(SocketAddr::from((
[127, 0, 0, 1],
refused_port,
))))
.await
.expect_err("nothing listens there");
assert!(
matches!(err, SocksDialError::Refused(_)),
"connect-refused → Refused: {err}"
);
assert_eq!(
err.to_reply_error().as_u8(),
consts::SOCKS5_REPLY_CONNECTION_REFUSED
);
}
#[tokio::test]
async fn tcp_dial_unresolvable_domain_maps_to_host_unreachable() {
let err = tcp_dial()(TargetAddr::Domain("nonexistent.invalid".to_owned(), 443))
.await
.expect_err("the domain does not resolve");
assert!(
matches!(err, SocksDialError::HostUnreachable(_)),
"unresolvable → HostUnreachable: {err}"
);
assert_eq!(
err.to_reply_error().as_u8(),
consts::SOCKS5_REPLY_HOST_UNREACHABLE
);
}
#[tokio::test]
async fn udp_egress_round_trips_inner_bytes_with_sender_address() {
let echo = echo_udp_server().await;
let mut egress = udp_egress()().await.expect("the egress binds");
let dest = TargetAddr::Ip(echo);
egress
.send(dest.clone(), Bytes::from_static(b"query"))
.await
.expect("send");
let frame = timeout(BOUND, egress.next_reply())
.await
.expect("reply within bound")
.expect("recv")
.expect("a reply");
let mut expect = dest.to_be_bytes().expect("encode").to_vec();
expect.extend_from_slice(b"query");
assert_eq!(frame, Bytes::from(expect), "inner bytes, framed-once");
let (src, payload) = decode_inner(&frame).await;
assert_eq!(src, dest);
assert_eq!(&payload[..], b"query");
}
#[tokio::test]
async fn udp_egress_next_reply_parks_until_a_datagram_arrives() {
let echo = echo_udp_server().await;
let mut egress = udp_egress()().await.expect("the egress binds");
egress
.send(TargetAddr::Ip(echo), Bytes::from_static(b"park"))
.await
.expect("send");
let frame = timeout(BOUND, egress.next_reply())
.await
.expect("the reply arrives (the wait parked, the datagram resolved it)")
.expect("recv")
.expect("the echo arrives");
let (_, payload) = decode_inner(&frame).await;
assert_eq!(&payload[..], b"park");
egress
.send(TargetAddr::Ip(echo), Bytes::from_static(b"again"))
.await
.expect("send");
let mut frame = None;
let deadline = std::time::Instant::now() + BOUND;
while std::time::Instant::now() < deadline {
match timeout(Duration::from_millis(500), egress.next_reply()).await {
Ok(Ok(Some(f))) => {
frame = Some(f);
break;
}
Ok(Ok(None)) => continue,
Ok(Err(e)) => panic!("the second reply errored: {e}"),
Err(_) => continue,
}
}
let frame = frame.expect("the second echo arrives (the wait parked, the datagram resolved it)");
let (_, payload) = decode_inner(&frame).await;
assert_eq!(&payload[..], b"again");
}
#[tokio::test]
async fn udp_egress_round_trips_multiple_datagrams_in_order() {
let echo = echo_udp_server().await;
let mut egress = udp_egress()().await.expect("the egress binds");
for i in 0..5u8 {
egress
.send(TargetAddr::Ip(echo), Bytes::from(vec![i; 8]))
.await
.expect("send");
}
for i in 0..5u8 {
let frame = timeout(BOUND, egress.next_reply())
.await
.expect("reply within bound")
.expect("recv")
.expect("a reply");
let (src, payload) = decode_inner(&frame).await;
assert_eq!(src, TargetAddr::Ip(echo));
assert_eq!(&payload[..], &[i; 8][..], "in-order echo");
}
}
#[tokio::test]
async fn udp_egress_zero_length_payload_round_trips() {
let echo = echo_udp_server().await;
let mut egress = udp_egress()().await.expect("the egress binds");
egress
.send(TargetAddr::Ip(echo), Bytes::new())
.await
.expect("send the empty datagram");
let frame = timeout(BOUND, egress.next_reply())
.await
.expect("reply within bound")
.expect("recv")
.expect("the empty echo");
let (_, payload) = decode_inner(&frame).await;
assert!(payload.is_empty(), "the empty datagram is empty");
}
#[tokio::test]
async fn udp_egress_unresolvable_destination_maps_to_unreachable() {
let mut egress = udp_egress()().await.expect("the egress binds");
let err = egress
.send(
TargetAddr::Domain("nonexistent.invalid".to_owned(), 53),
Bytes::from_static(b"q"),
)
.await
.expect_err("the destination does not resolve");
assert!(matches!(err, EgressError::Unreachable(_)), "{err}");
assert_eq!(
err.to_reply_error().as_u8(),
consts::SOCKS5_REPLY_NETWORK_UNREACHABLE
);
}
#[tokio::test]
async fn bind_accept_reply1_real_addr_reply2_real_peer_and_round_trip() {
let accept = bind_accept([127, 0, 0, 1])(TargetAddr::Domain("ftp.test".to_owned(), 0))
.await
.expect("the factory binds an ephemeral port");
let bind_addr = accept.bind_addr();
assert!(
bind_addr.ip().is_loopback() && bind_addr.port() != 0,
"reply#1 carries the real bound address: {bind_addr}"
);
let client = timeout(BOUND, TcpStream::connect(bind_addr))
.await
.expect("dial within bound")
.expect("connect to the advertised address");
let mut handle = timeout(BOUND, accept.accept())
.await
.expect("accept within bound")
.expect("the accept resolves");
let peer = match handle.peer {
TargetAddr::Ip(sa) => sa,
other => panic!("an IP peer, got {other:?}"),
};
assert_eq!(
peer,
client.local_addr().expect("client local_addr"),
"reply#2's peer is the client's real address (byte-exact)"
);
let mut client_to_echo = client;
tokio::spawn(async move {
let mut buf = [0u8; 4096];
loop {
match client_to_echo.read(&mut buf).await {
Ok(0) | Err(_) => return,
Ok(n) => {
if client_to_echo.write_all(&buf[..n]).await.is_err() {
return;
}
}
}
}
});
let mut write = handle.write;
write.write_all(b"active").await.expect("write");
write.flush().await.expect("flush");
let mut buf = [0u8; 6];
timeout(BOUND, handle.read.read_exact(&mut buf))
.await
.expect("echo within bound")
.expect("read");
assert_eq!(&buf, b"active");
write.write_all(b"pass").await.expect("write");
write.flush().await.expect("flush");
let mut buf = [0u8; 4];
timeout(BOUND, handle.read.read_exact(&mut buf))
.await
.expect("round-trip within bound")
.expect("read");
assert_eq!(&buf, b"pass");
drop(write);
}
#[tokio::test]
async fn bind_accept_honors_a_requested_port_and_refuses_a_held_one() {
let holder = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let held = holder.local_addr().expect("local_addr");
let ok = bind_accept([127, 0, 0, 1])(TargetAddr::Ip(held))
.await
.expect_err("the port is held");
assert!(
matches!(
ok,
BindAcceptFactoryError::Shortage(_) | BindAcceptFactoryError::Refused(_)
),
"bind refused/held → Shortage/Refused: {ok}"
);
assert_eq!(
ok.to_reply_error().as_u8(),
consts::SOCKS5_REPLY_GENERAL_FAILURE,
"Shortage/Refused map to the RFC's catch-all"
);
let free = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let want = free.local_addr().expect("local_addr");
drop(free);
let accept = bind_accept([127, 0, 0, 1])(TargetAddr::Ip(want))
.await
.expect("the free port binds");
assert_eq!(accept.bind_addr().port(), want.port(), "the requested port");
assert_eq!(accept.bind_addr().ip().to_string(), "127.0.0.1");
}
#[tokio::test]
async fn bind_accept_drop_of_the_accept_future_closes_the_listener() {
let accept = bind_accept([127, 0, 0, 1])(TargetAddr::Domain("gone.test".to_owned(), 0))
.await
.expect("the factory binds");
let bind_addr = accept.bind_addr();
drop(accept);
let rebinder = TcpListener::bind(bind_addr).await;
match rebinder {
Ok(rebound) => {
assert_eq!(rebound.local_addr().expect("local_addr"), bind_addr);
}
Err(e) => panic!("the listener leaked (the port is still held): {e}"),
}
}
type DialCase = (ErrorKind, &'static dyn Fn(&SocksDialError) -> bool);
type BindCase = (ErrorKind, &'static dyn Fn(&BindAcceptFactoryError) -> bool);
#[test]
fn dial_error_of_maps_the_vocabulary_table() {
let what = "connect probe";
let cases: &[DialCase] = &[
(ErrorKind::ConnectionRefused, &|e| {
matches!(e, SocksDialError::Refused(_))
}),
(ErrorKind::NetworkUnreachable, &|e| {
matches!(e, SocksDialError::NetworkUnreachable(_))
}),
(ErrorKind::HostUnreachable, &|e| {
matches!(e, SocksDialError::HostUnreachable(_))
}),
(
ErrorKind::TimedOut,
&|e| matches!(e, SocksDialError::NetworkUnreachable(_)),
),
(
ErrorKind::PermissionDenied,
&|e| matches!(e, SocksDialError::General(_)),
),
(
ErrorKind::Interrupted,
&|e| matches!(e, SocksDialError::General(_)),
),
];
for (kind, check) in cases {
let err = std::io::Error::new(*kind, "table probe");
let mapped = super::dial_error_of(what, &err);
assert!(
check(&mapped),
"{kind:?} must map to the table's arm, got {mapped:?}"
);
let what_needle = format!("{what}: ");
assert!(
mapped.to_string().contains(&what_needle),
"the context rides the message: {mapped}"
);
}
}
#[test]
fn dial_vocabulary_reply_codes_table() {
use crate::egress::SocksDialError;
let cases: &[(&SocksDialError, u8)] = &[
(
&SocksDialError::Refused("x".into()),
consts::SOCKS5_REPLY_CONNECTION_REFUSED,
),
(
&SocksDialError::NetworkUnreachable("x".into()),
consts::SOCKS5_REPLY_NETWORK_UNREACHABLE,
),
(
&SocksDialError::HostUnreachable("x".into()),
consts::SOCKS5_REPLY_HOST_UNREACHABLE,
),
(
&SocksDialError::General("x".into()),
consts::SOCKS5_REPLY_GENERAL_FAILURE,
),
];
for (err, code) in cases {
assert_eq!(err.to_reply_error().as_u8(), *code, "{err:?}");
}
}
#[test]
fn bind_map_splits_the_io_kinds_table() {
let map = |kind: ErrorKind| {
let err = std::io::Error::new(kind, "bind table probe");
match err.kind() {
ErrorKind::PermissionDenied | ErrorKind::AddrInUse => {
BindAcceptFactoryError::Shortage(format!("tcp bind: {err}"))
}
ErrorKind::AddrNotAvailable => {
BindAcceptFactoryError::Refused(format!("tcp bind: {err}"))
}
_ => BindAcceptFactoryError::General(format!("tcp bind: {err}")),
}
};
let cases: &[BindCase] = &[
(ErrorKind::AddrInUse, &|e| {
matches!(e, BindAcceptFactoryError::Shortage(_))
}),
(ErrorKind::PermissionDenied, &|e| {
matches!(e, BindAcceptFactoryError::Shortage(_))
}),
(ErrorKind::AddrNotAvailable, &|e| {
matches!(e, BindAcceptFactoryError::Refused(_))
}),
(ErrorKind::Interrupted, &|e| {
matches!(e, BindAcceptFactoryError::General(_))
}),
];
for (kind, check) in cases {
let mapped = map(*kind);
assert!(
check(&mapped),
"{kind:?} must map to the table's arm, got {mapped:?}"
);
}
}
#[test]
fn bind_map_reply_codes_table() {
let cases: &[(&BindAcceptFactoryError, u8)] = &[
(
&BindAcceptFactoryError::Refused("x".into()),
consts::SOCKS5_REPLY_CONNECTION_REFUSED,
),
(
&BindAcceptFactoryError::Shortage("x".into()),
consts::SOCKS5_REPLY_GENERAL_FAILURE,
),
(
&BindAcceptFactoryError::General("x".into()),
consts::SOCKS5_REPLY_GENERAL_FAILURE,
),
];
for (err, code) in cases {
assert_eq!(err.to_reply_error().as_u8(), *code, "{err:?}");
}
}