alksocks/src/local/tests.rs

//! The local providers' unit tests: real sockets, including the echo
//! TCP server for [`super::tcp_dial`] round-trips, a real UDP echo for
//! the egress ([`super::LocalEgress`], the POC #2 shape, including
//! the readiness contract's park-then-datagram shape), a real BIND
//! accept ([`super::bind_accept`], reply#1/reply#2 carrying the real
//! addresses), and the refusal mappings table-tested against the
//! landed error types' `to_reply_error()`.

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);

/// Decode the inner bytes' `(addr, payload)` shape: the forked
/// `read_address` codec applied to the egress's reply (an honest
/// decode through the codec, not a hand-parse: the inner bytes are
/// `ATYP + addr + port + data`; the codec reads exactly the address
/// field, the payload tail is un-read surplus).
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..]))
}

/// Spawn an echo TCP server on an ephemeral loopback port; returns
/// its address (the dial target).
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
}

/// A UDP echo server: datagrams echo back from their sender address
/// (the egress test's far side; the POC #2 shape).
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
}

// --- tcp_dial -------------------------------------------------------------

/// The happy path: dial the echo server, round-trip bytes through the
/// boxed halves (ownership rides the handle; the halves are
/// Send+Sync+Unpin; the boxing compiles, pinned here).
#[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");
}

/// The `to_target_addr`-shaped target (the `impl ToTargetAddr`
/// consumer path): a domain that resolves here (`localhost`; the
/// `-D` semantic: resolution is the provider's job) reaches the same
/// echo server; the seam received it unresolved and dialed the
/// resolved address.
#[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");
}

/// An unroutable loopback port refuses: `ConnectionRefused` (0x05)
/// per the io error kind; the refusal rides the seam's error
/// vocabulary (`to_reply_error()`), table-pinned through the mapped
/// code.
#[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
    );
}

/// A domain that does not resolve → `HostUnreachable` (0x04): the
/// resolution failure is the dial seam's error, faithfully mapped.
#[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
    );
}

// --- udp_egress ------------------------------------------------------------

/// The datagram round-trip through a real UDP echo (the POC #2
/// shape): `send` (payload-only) → the echo → `next_reply`'s inner
/// bytes carry the sender's address + the data (no length prefix, as
/// framing happens once; the wire layer adds it downstream).
#[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");

    // The inner bytes: `to_be_bytes() + data`, no length prefix.
    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");

    // The sender's address decodes as the echo's (the RFC's
    // per-datagram source semantics; the address rides the frame).
    let (src, payload) = decode_inner(&frame).await;
    assert_eq!(src, dest);
    assert_eq!(&payload[..], b"query");
}

/// The readiness contract with the real socket: `next_reply` parks
/// (no reply is ready → `Pending`); the stage's select! branch would
/// stay parked; a datagram arriving resolves it. Pinned by racing the
/// wait against a `sleep`: no reply → not resolved by the bound; the
/// datagram → resolved. Never a hot `Ok(None)` (the send-then-drain
/// ordering above already proves the resolution; this proves the
/// park).
#[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");

    // A `Pending`-honoring wait stays pending across a no-datagram
    // window: if `next_reply` resolved hot with `Ok(None)`, this
    // would return within the bound with a spurious shape instead of
    // timing out.
    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");

    // The socket drained: a second wait parks again (no reply is
    // ready), resolved only by a fresh datagram.
    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");
}

/// Multiple datagrams in order (no flow table; per-datagram
/// addressing does the demultiplexing; the socket is the relay
/// handle).
#[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");
    }
}

/// A zero-length payload is a legal empty datagram (the F-2
/// invariant at payload level) and round-trips through the real
/// socket.
#[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");
}

/// An undecodable destination (an unresolvable domain) fails the
/// `send` with `Unreachable` → `NetworkUnreachable` (0x03); the
/// resolve failure is the egress's error vocabulary.
#[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
    );
}

// --- bind_accept ------------------------------------------------------------

/// The faithful RFC shape (ADR-007 variant 1): reply#1's `bind_addr`
/// is the listener's real `local_addr`; a client dialing it is
/// accepted; the accepted handle's `peer` is the client's real
/// address (byte-exact); the data plane round-trips.
#[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}"
    );

    // The app server dials the advertised address (the RFC's role
    // for BND.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)"
    );

    // The client echoes what it receives (the round-trip's far side:
    // the app server the RFC models, i.e. active-mode FTP's data
    // listener echo).
    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;
                    }
                }
            }
        }
    });

    // The data plane: bytes written to the handle reach the client;
    // the client echoes back; the handle's read half receives them.
    // The client echoes what it receives (the round-trip's far
    // side), so both directions ride the accepted halves.
    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);
}

/// The requested port is honored when non-zero (FTP active mode):
/// the listener binds exactly the requested port, and a held port
/// refuses the factory (`Shortage`/`Refused` per the io error; the
/// bind failure is reply#1's vocabulary).
#[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"
    );

    // A free port is honored exactly (the listener is the factory's
    // resource, dropped with the accept future).
    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");
}

/// The accept future owns its listener end-to-end: dropping the
/// `BindAccept` (the quiet wait's drop, i.e. the client-gone teardown)
/// closes the listener; the advertised address stops accepting. A
/// fresh factory on the same (now free) port re-binds.
#[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);

    // The listener closed with the future: a fresh bind of the same
    // ephemeral port succeeds (the OS may hand the port to either
    // the reuse or a fresh allocation; loopback ephemeral binds
    // succeed either way; the probe is that the re-bind is possible
    // at all, i.e. the listener was dropped, not leaked).
    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}"),
    }
}

// --- the io-kind → vocabulary mapping tables (review 003 D-3) --------------

type DialCase = (ErrorKind, &'static dyn Fn(&SocksDialError) -> bool);
type BindCase = (ErrorKind, &'static dyn Fn(&BindAcceptFactoryError) -> bool);

/// The dial map is table-tested (`dial_error_of` is a pure function
/// of `io::ErrorKind`; review 003 D-3): a refactor that swaps an arm
/// into the `_` catch-all fails here. The two loopback-unproducible
/// kinds (`NetworkUnreachable`, `HostUnreachable`, `TimedOut`,
/// `PermissionDenied`) ride synthetic errors of the right kind.
#[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(_))
        }),
        (
            // TimedOut shares NetworkUnreachable's arm (an io timeout
            // is not an RFC-codeable shape of its own).
            ErrorKind::TimedOut,
            &|e| matches!(e, SocksDialError::NetworkUnreachable(_)),
        ),
        (
            // PermissionDenied is a General (the catch-all's
            // ruleset-adjacent shape; NotAllowed is never synthesized
            // by this provider).
            ErrorKind::PermissionDenied,
            &|e| matches!(e, SocksDialError::General(_)),
        ),
        (
            // The `_` catch-all: any unmapped kind is a 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}"
        );
    }
}

/// The dial vocabulary's reply codes stay RFC-faithful (the mapped
/// surface the producer replies with; that is the split's whole purpose).
#[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:?}");
    }
}

/// The bind map's closure is table-tested (review 003 D-3: the tests
/// drove `AddrInUse` only through real sockets): `PermissionDenied`
/// → `Shortage`, `AddrNotAvailable` → `Refused`, the `_` catch-all →
/// `General`, and `AddrInUse` → `Shortage` (the previously-pinned
/// arm, in-table for the full split).
#[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:?}"
        );
    }
}

/// The bind map's reply codes stay faithful: `Refused` → 0x05, the
/// shortage/catch-all shapes → 0x01 (the POC #5 mapping, table).
#[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:?}");
    }
}