alksocks/tests/local-providers.rs

//! The local providers' integration tests (feature `local`): the
//! three socket-backed seams driven through the producer's real
//! command-dispatch path (the harness's consumer-side topology) with
//! real targets: the CONNECT echo round-trip, the UDP datagram
//! round-trip through the producer's datagram stage (the POC #2
//! shape), the BIND accept end-to-end (real addresses in reply#1/
//! reply#2), the readiness contract re-pinned with the real-socket
//! egress (park-then-datagram, no spin), and the refusal mapping
//! (connect-refused → 0x05 through the producer's reply).
#![cfg(feature = "local")]
#![allow(dead_code)]

mod harness;

use std::net::SocketAddr;
use std::time::Duration;

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::time::timeout;

use alksocks::consumer::{AssociateSession, BindSession, Socks5ChannelSession};
use alksocks::error::SocksSessionError;
use alksocks::rfc1928::consts;
use alksocks::rfc1928::util::target_addr::{TargetAddr, ToTargetAddr};
use alksocks::SocksError as SocksErrorKind;

use harness::{both_scopes_identity, wire_consumer_side_with};

const BOUND: Duration = Duration::from_secs(10);

/// A session error carrying an RFC reply code (the refusal probe's
/// expectation shape: the mapped code rides the session error).
fn reply_code_of(err: &SocksSessionError) -> Option<u8> {
    match err {
        SocksSessionError::Rfc(SocksErrorKind::ReplyError(code)) => Some(code.as_u8()),
        _ => None,
    }
}

/// Both managers hold only channel 0 (the no-phantom-channel probe;
/// the connect suite's shape, re-pinned).
async fn no_leaked_channels(topo: &harness::ConsumerSideTopology) {
    let deadline = Duration::from_secs(5);
    let consumer_ids = topo.consumer.manager().channel_ids();
    assert!(
        consumer_ids.iter().all(|&id| id == 0),
        "consumer leaked channel entries: {consumer_ids:?}"
    );
    timeout(deadline, async {
        loop {
            let producer_ids = topo.producer_manager.channel_ids();
            if producer_ids.iter().all(|&id| id == 0) {
                return;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await
    .expect("the producer's handler-exit teardown reaps the entry");
}

/// Spawn an echo TCP server on an ephemeral loopback port.
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 (the datagram stage's far side).
async fn echo_udp_server() -> SocketAddr {
    let sock = tokio::net::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
}

/// The local providers wired as the producer's seams (the assembly
/// layer's shape; producer.md §Wiring's constructor names).
fn local_providers() -> (
    alksocks::egress::DialFn,
    alksocks::egress::EgressFactory,
    alksocks::egress::BindAcceptFactory,
) {
    (
        alksocks::tcp_dial(),
        alksocks::udp_egress(),
        alksocks::bind_accept([127, 0, 0, 1]),
    )
}

/// CONNECT through the local dial: the consumer session's bytes
/// round-trip through the real echo server; the domain target
/// (`localhost`) resolves producer-side (the ssh `-D` semantic,
/// ADR-012); the dial witness saw it unresolved.
#[tokio::test]
async fn connect_round_trips_through_a_real_target() {
    let server = echo_tcp_server().await;
    let (dial, egress, accept) = local_providers();
    let topo = wire_consumer_side_with(
        dial,
        Some(egress),
        Some(accept),
        Some(both_scopes_identity()),
    )
    .await;

    let mut session = Socks5ChannelSession::connect(
        topo.consumer.as_ref(),
        ("localhost", server.port()).to_target_addr().unwrap(),
    )
    .await
    .expect("the session connects (the domain resolved producer-side)");

    let (mut read, mut write) = session.split().expect("post-reply pipe");
    write.write_all(b"ping").await.expect("write");
    write.flush().await.expect("flush");
    let mut buf = [0u8; 4];
    timeout(BOUND, read.read_exact(&mut buf))
        .await
        .expect("echo within bound")
        .expect("read");
    assert_eq!(&buf, b"ping");

    session.close().await;
    no_leaked_channels(&topo).await;
}

/// A connect-refused target (a bound-then-dropped port): the dial
/// seam refuses with `Refused` → the producer's reply carries 0x05
/// → the session error surfaces the mapped code (the refusal rides
/// in-stream (no phantom channel).
#[tokio::test]
async fn connect_refused_surfaces_reply_code_0x05_through_the_producer() {
    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 (dial, egress, accept) = local_providers();
    let topo = wire_consumer_side_with(dial, Some(egress), Some(accept), None).await;

    let err = Socks5ChannelSession::connect(
        topo.consumer.as_ref(),
        TargetAddr::Ip(SocketAddr::from(([127, 0, 0, 1], refused_port))),
    )
    .await
    .expect_err("nothing listens there");

    assert_eq!(
        reply_code_of(&err),
        Some(consts::SOCKS5_REPLY_CONNECTION_REFUSED),
        "connect-refused rides the producer's reply as 0x05: {err:?}"
    );
    no_leaked_channels(&topo).await;
}

/// ASSOCIATE through the real-socket egress: the datagram stage's
/// framed round-trip reaches a real UDP echo (the POC #2 shape
/// through the producer's stage: framing once, inner bytes at the
/// egress, the sender's address reported back).
#[tokio::test]
async fn associate_round_trips_through_the_real_egress() {
    let echo = echo_udp_server().await;
    let (dial, egress, accept) = local_providers();
    let topo = wire_consumer_side_with(dial, Some(egress), Some(accept), None).await;

    let mut assoc = AssociateSession::associate(topo.consumer.as_ref())
        .await
        .expect("the association associates");

    let target = TargetAddr::Ip(echo);
    for i in 0..3u8 {
        assoc
            .send_datagram(&target, &[i; 12])
            .await
            .expect("send datagram");
    }
    for i in 0..3u8 {
        let (src, payload) = timeout(BOUND, assoc.recv_datagram())
            .await
            .expect("recv within bound")
            .expect("a datagram (not eof)")
            .expect("a datagram");
        assert_eq!(src, target, "the echo's address is the sender");
        assert_eq!(&payload[..], &[i; 12][..], "in-order echo");
    }

    assoc.shutdown().await.expect("half-close");
    no_leaked_channels(&topo).await;
}

/// The readiness contract re-pinned with the real-socket egress,
/// through the producer's datagram stage: the association stays
/// parked (the egress's `next_reply` parks on `readable()`) while no
/// datagram is in flight; the consumer's `send` after a quiet
/// window still round-trips (no spin starved the channel branch; a
/// hot `Ok(None)` egress would have wedged the stage's biased
/// select into livelock and the second datagram would never
/// complete).
#[tokio::test]
async fn the_readiness_contract_holds_through_the_producer_stage() {
    let echo = echo_udp_server().await;
    let (dial, egress, accept) = local_providers();
    let topo = wire_consumer_side_with(dial, Some(egress), Some(accept), None).await;

    let mut assoc = AssociateSession::associate(topo.consumer.as_ref())
        .await
        .expect("the association associates");
    let target = TargetAddr::Ip(echo);

    // A datagram now, then a quiet window (the stage parked on the
    // real socket's readiness), then a datagram after; the second
    // round-trip proves the stage never wedged.
    assoc
        .send_datagram(&target, b"first")
        .await
        .expect("send the first");
    let (_, payload) = timeout(BOUND, assoc.recv_datagram())
        .await
        .expect("recv within bound")
        .expect("a datagram")
        .expect("a datagram");
    assert_eq!(&payload[..], b"first");

    tokio::time::sleep(Duration::from_millis(150)).await;

    assoc
        .send_datagram(&target, b"second")
        .await
        .expect("send the second (the stage was parked, not spinning)");
    let (_, payload) = timeout(BOUND, assoc.recv_datagram())
        .await
        .expect("recv within bound")
        .expect("a datagram")
        .expect("a datagram");
    assert_eq!(&payload[..], b"second");

    assoc.shutdown().await.expect("half-close");
    no_leaked_channels(&topo).await;
}

/// BIND end-to-end with the local accept: reply#1 carries the real
/// listener address; the app server (the in-test client) dials it;
/// reply#2 carries the real peer (byte-exact); the data plane
/// round-trips through the producer's `pump_bidi`.
#[tokio::test]
async fn bind_end_to_end_with_real_addresses_and_round_trip() {
    let (dial, egress, accept) = local_providers();
    let topo = wire_consumer_side_with(
        dial,
        Some(egress),
        Some(accept),
        Some(both_scopes_identity()),
    )
    .await;

    let mut session = BindSession::bind(
        topo.consumer.as_ref(),
        ("bind.test", 0u16).to_target_addr().unwrap(),
    )
    .await
    .expect("the session binds");

    // Reply#1: the real bound address (the listener's local_addr).
    let bind_addr = session.reply1().expect("reply#1's handle").bind_addr;
    let TargetAddr::Ip(sa) = bind_addr else {
        panic!("an IP BND.ADDR, got {bind_addr:?}")
    };
    assert!(
        sa.ip().is_loopback() && sa.port() != 0,
        "reply#1 carries the real listener address: {sa}"
    );

    // The app server dials the advertised address; a byte arrives
    // through the data plane (the quiet wait's accept resolved).
    let mut client = timeout(BOUND, tokio::net::TcpStream::connect(sa))
        .await
        .expect("dial within bound")
        .expect("connect to the advertised address");

    let mut plane = timeout(BOUND, session.await_inbound())
        .await
        .expect("reply#2 within bound")
        .expect("the plane exists")
        .expect("the accept resolved");
    let peer = match plane.peer {
        TargetAddr::Ip(peer_sa) => peer_sa,
        other => panic!("an IP peer, got {other:?}"),
    };
    assert_eq!(peer.ip(), sa.ip(), "the peer rides the same listener");
    assert_ne!(
        peer.port(),
        sa.port(),
        "the peer is the client, not the listener"
    );

    // The data plane: bytes through the channel reach the dialing
    // client and back (the producer's `pump_bidi` against the real
    // accepted halves).
    plane.write_all(b"inbound").await.expect("write");
    plane.flush().await.expect("flush");
    let mut buf = [0u8; 7];
    timeout(BOUND, client.read_exact(&mut buf))
        .await
        .expect("read within bound")
        .expect("read");
    assert_eq!(&buf, b"inbound");
    client.write_all(b"outbound").await.expect("echo");
    client.flush().await.expect("flush");
    let mut buf = [0u8; 8];
    timeout(BOUND, plane.read_exact(&mut buf))
        .await
        .expect("read within bound")
        .expect("read");
    assert_eq!(&buf, b"outbound");

    // Both ends close (the producer's pump exits when the accepted
    // halves' peer is gone; a real TCP stream stays open until its
    // owner closes it, unlike the harness duplexes).
    drop(plane);
    drop(client);
    assert!(session.close().await, "the adopted channel entry existed");
    no_leaked_channels(&topo).await;
}

/// The accept-future-owns-its-listener contract, through the
/// producer's quiet wait: a client-gone BIND session (the consumer
/// vanishes mid-quiet-wait) tears the listener down with the session;
/// the advertised address stops accepting (the drop-witness probe:
/// a fresh re-bind of the same port succeeds, the listener was
/// dropped, not leaked).
#[tokio::test]
async fn client_gone_mid_quiet_wait_closes_the_listener() {
    let (dial, egress, accept) = local_providers();
    let topo = wire_consumer_side_with(
        dial,
        Some(egress),
        Some(accept),
        Some(both_scopes_identity()),
    )
    .await;

    let mut session = BindSession::bind(
        topo.consumer.as_ref(),
        ("gone.test", 0u16).to_target_addr().unwrap(),
    )
    .await
    .expect("the session binds");
    let TargetAddr::Ip(sa) = session.reply1().expect("reply#1's handle").bind_addr else {
        panic!("an IP BND.ADDR")
    };

    // The client-gone teardown: drop the session mid-quiet-wait (the
    // channel's write half EOFs; the producer's quiet wait exits on
    // the EOF sentinel and the accept future drops, taking the
    // listener with it).
    drop(session);

    // The producer's teardown reaps the channel entry (poll past the
    // reap race), then the advertised address must be re-bindable;
    // the listener closed with the quiet wait's drop.
    timeout(Duration::from_secs(5), async {
        loop {
            if topo
                .producer_manager
                .channel_ids()
                .iter()
                .all(|&id| id == 0)
            {
                return;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await
    .expect("the producer's entry reaped");

    // A held port would refuse the re-bind; the dropped listener
    // frees it. (The OS may hand the port to the re-bind; the probe
    // is that it is free either way.)
    let mut rebound = None;
    for _ in 0..50 {
        match tokio::net::TcpListener::bind(sa).await {
            Ok(listener) => {
                rebound = Some(listener);
                break;
            }
            Err(_) => tokio::time::sleep(Duration::from_millis(20)).await,
        }
    }
    assert!(
        rebound.is_some(),
        "the listener closed with the client-gone teardown (the port freed): {sa}"
    );
    no_leaked_channels(&topo).await;
}