alksocks/tests/consumer-associate.rs

//! The consumer UDP ASSOCIATE integration tests — `AssociateSession`
//! end to end through the harness's consumer-side topology (the
//! producer's real command-dispatch path over channels; the framed
//! datagram stage, ADR-003, driven by the wrapper-aware client, POC
//! #2's shape).
//!
//! Covered: associate → sentinel reply (received, never used for
//! addressing) → send/recv round-trips (multiple in-order datagrams,
//! one association); the split-datagram persistence test (reader
//! state persists across `recv_datagram` calls — the POC's
//! 2-byte-prefix chunk finding); empty-payload datagram round-trip
//! (F-2, never EOF); shutdown → the producer's stage ends and a
//! fresh associate succeeds immediately; Drop teardown (drop
//! mid-stream → the producer's stage sees EOF, the connection stays
//! healthy); the ASSOCIATE refusal surfaced as a typed error with no
//! phantom channel; CONNECT + ASSOCIATE on the same registration
//! (the regression guard extended).

mod harness;

use tokio::io::{AsyncReadExt, AsyncWriteExt};

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

use harness::{echo_dial, echo_egress, refusing_egress, wire_consumer_side_with};

/// The no-leaked-channels probe (the connect/bind suites' shape,
/// re-pinned for the ASSOCIATE session): both managers hold only
/// channel 0 — the producer's entry reaps by its open-op wrapper;
/// the probe polls past the reap race.
async fn no_leaked_channels(topo: &harness::ConsumerSideTopology) {
    let deadline = std::time::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:?}"
    );
    tokio::time::timeout(deadline, async {
        loop {
            let producer_ids = topo.producer_manager.channel_ids();
            if producer_ids.iter().all(|&id| id == 0) {
                return;
            }
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        }
    })
    .await
    .expect("the producer's handler-exit teardown reaps the entry");
}

/// The happy path: `associate` → the reply's sentinel `BND.ADDR` is
/// received (RFC fidelity) and never used for addressing → send/recv
/// round-trips through the echo egress — multiple datagrams, in
/// order, on one association; the establisher witnessed the opener
/// identity (CF-006).
#[tokio::test]
async fn associate_round_trips_multiple_datagrams_one_association() {
    let topo = wire_consumer_side_with(
        echo_dial(),
        Some(echo_egress()),
        None,
        Some(harness::consumer_identity()),
    )
    .await;

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

    // The reply's BND.ADDR is the producer's sentinel (127.0.0.1:0 —
    // ADR-003): received for RFC fidelity, never used for addressing.
    let relay = assoc
        .relay_addr()
        .expect("the reply carried BND.ADDR")
        .into_string_and_port();
    assert_eq!(relay.0, "127.0.0.1", "the in-band sentinel BND.ADDR");
    assert_eq!(relay.1, 0);

    let target = TargetAddr::Ip("10.0.0.1:53".parse().unwrap());
    for i in 0..5u8 {
        let payload = vec![i; 10];
        assoc
            .send_datagram(&target, &payload)
            .await
            .expect("send datagram");
    }
    for i in 0..5u8 {
        let (src, payload) =
            tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
                .await
                .expect("recv within timeout")
                .expect("a datagram (not eof)")
                .expect("a datagram");
        assert_eq!(src, target, "the echo reports the destination");
        assert_eq!(payload, bytes::Bytes::from(vec![i; 10]), "in-order echo");
    }

    // The association is alive: another datagram round-trips.
    assoc.send_datagram(&target, b"alive").await.expect("send");
    let (_, payload) =
        tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
            .await
            .expect("recv timed out")
            .expect("a datagram (not eof)")
            .expect("a datagram");
    assert_eq!(&payload[..], b"alive");

    assert!(
        topo.consumer.manager().has_channel(assoc.channel_id()),
        "the adopted channel entry exists while associated"
    );

    let reaped = assoc.shutdown().await.expect("half-close");
    assert!(reaped, "the adopted channel entry existed");
    no_leaked_channels(&topo).await;
}

/// The split-datagram persistence test (POC #2's finding — the
/// reader state must persist across `recv_datagram` calls): a large
/// datagram's frame (under the u16 joint bound but well over the 16
/// KB read buffer) spans multiple channel reads AND multiple
/// `recv_datagram` calls — the first call returns the earlier small
/// datagram while the large one is mid-flight, the next call
/// reassembles it across its internal reads — and a trailing small
/// datagram batched in the same read batch drains through the
/// pending queue.
#[tokio::test]
async fn recv_datagram_state_persists_across_calls() {
    let topo = wire_consumer_side_with(echo_dial(), Some(echo_egress()), None, None).await;

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

    // The order tells the story: "first" completes early; the 60 KiB
    // datagram (60007 inner bytes — legal, but over the 16 KB read
    // buffer) is mid-flight when the first call returns; "third"
    // batches after it.
    assoc
        .send_datagram(&target, b"first")
        .await
        .expect("send the first datagram");
    let big = vec![7u8; 60 * 1024];
    assoc
        .send_datagram(&target, &big)
        .await
        .expect("send the big datagram");
    assoc
        .send_datagram(&target, b"third")
        .await
        .expect("send the third datagram");

    let (src, first) =
        tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
            .await
            .expect("recv timed out")
            .expect("a datagram (not eof)")
            .expect("a datagram");
    assert_eq!(src, target);
    assert_eq!(&first[..], b"first", "the first datagram, in order");

    // The large datagram reassembles across the call boundary (its
    // mid-frame state persisted into this call) — the POC finding's
    // shape.
    let (src, second) =
        tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
            .await
            .expect("recv timed out")
            .expect("a datagram (not eof)")
            .expect("a datagram");
    assert_eq!(src, target);
    assert_eq!(&second[..], &big[..], "the split datagram reassembled");

    // The trailing datagram rode the same read batches: the pending
    // queue carries it to this call.
    let (_, third) =
        tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
            .await
            .expect("recv timed out")
            .expect("a datagram (not eof)")
            .expect("a datagram");
    assert_eq!(&third[..], b"third", "the batched datagram drained");

    let reaped = assoc.shutdown().await.expect("half-close");
    assert!(reaped);
    no_leaked_channels(&topo).await;
}

/// An empty-payload datagram is a legal datagram — never EOF (the
/// F-2 invariant, end to end through the echo egress), and the
/// association survives it (pre → empty → post, in order).
#[tokio::test]
async fn empty_payload_datagram_round_trips_and_survives() {
    let topo = wire_consumer_side_with(echo_dial(), Some(echo_egress()), None, None).await;

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

    assoc
        .send_datagram(&target, b"pre")
        .await
        .expect("send pre");
    assoc
        .send_datagram(&target, b"")
        .await
        .expect("an empty payload is a legal datagram");
    assoc
        .send_datagram(&target, b"post")
        .await
        .expect("send post");

    let (_, pre) = tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
        .await
        .expect("recv timed out")
        .expect("a datagram (not eof)")
        .expect("a datagram");
    assert_eq!(&pre[..], b"pre");

    let (src, empty) =
        tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
            .await
            .expect("recv timed out")
            .expect("a datagram (not eof)")
            .expect("a datagram");
    assert_eq!(src, target);
    assert!(
        empty.is_empty(),
        "the empty-payload datagram round-trips (never EOF)"
    );

    let (_, post) = tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
        .await
        .expect("recv timed out")
        .expect("a datagram (not eof)")
        .expect("a datagram");
    assert_eq!(&post[..], b"post");

    let reaped = assoc.shutdown().await.expect("half-close");
    assert!(reaped);
    no_leaked_channels(&topo).await;
}

/// Shutdown half-closes the stream: the mux's implicit-EOF sentinel
/// reaches the producer's stage as EOF and the association ends
/// (association lifetime = channel lifetime) — and a fresh associate
/// succeeds immediately after (the POC's test 4, from the consumer
/// side; the same registration is unpoisoned).
#[tokio::test]
async fn shutdown_ends_the_association_and_a_fresh_associate_succeeds() {
    let topo = wire_consumer_side_with(echo_dial(), Some(echo_egress()), None, None).await;

    let mut assoc = AssociateSession::associate(topo.consumer.as_ref())
        .await
        .expect("the association associates");
    let channel_id = assoc.channel_id();
    let target = TargetAddr::Ip("10.0.0.4:53".parse().unwrap());
    assoc
        .send_datagram(&target, b"before")
        .await
        .expect("send before shutdown");
    let (_, echoed) =
        tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
            .await
            .expect("recv timed out")
            .expect("a datagram (not eof)")
            .expect("a datagram");
    assert_eq!(&echoed[..], b"before");

    // Half-close: the producer's stage exits on the EOF sentinel.
    assoc.shutdown().await.expect("shutdown");
    assert!(
        !topo.consumer.manager().has_channel(channel_id),
        "the adopted channel entry was reaped"
    );

    // The connection is healthy: a fresh associate round-trips.
    let mut fresh = AssociateSession::associate(topo.consumer.as_ref())
        .await
        .expect("fresh associate after the shutdown");
    fresh
        .send_datagram(&target, b"fresh")
        .await
        .expect("send after fresh associate");
    let (_, echoed) =
        tokio::time::timeout(std::time::Duration::from_secs(10), fresh.recv_datagram())
            .await
            .expect("recv timed out")
            .expect("a datagram (not eof)")
            .expect("a datagram");
    assert_eq!(&echoed[..], b"fresh");
    let reaped = fresh.shutdown().await.expect("shutdown");
    assert!(reaped);
    no_leaked_channels(&topo).await;
}

/// Drop teardown: dropping a live association mid-stream reaps the
/// adopted channel entry (no leak), the producer's stage sees the
/// structural EOF (the mux's implicit-EOF sentinel on the dropped
/// write half), and the connection stays healthy — a fresh associate
/// succeeds immediately after.
#[tokio::test]
async fn drop_teardown_reaps_and_connection_stays_healthy() {
    let topo = wire_consumer_side_with(echo_dial(), Some(echo_egress()), None, None).await;

    let mut assoc = AssociateSession::associate(topo.consumer.as_ref())
        .await
        .expect("the association associates");
    let channel_id = assoc.channel_id();
    assoc
        .send_datagram(TargetAddr::Ip("10.0.0.5:53".parse().unwrap()), b"mid")
        .await
        .expect("send before the drop");
    let (_, echoed) =
        tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
            .await
            .expect("recv timed out")
            .expect("a datagram (not eof)")
            .expect("a datagram");
    assert_eq!(&echoed[..], b"mid");

    // Drop mid-stream (the teardown-ownership posture).
    drop(assoc);
    assert!(
        !topo.consumer.manager().has_channel(channel_id),
        "the adopted channel entry was reaped on Drop"
    );
    no_leaked_channels(&topo).await;

    // The connection is healthy: a fresh associate round-trips.
    let mut fresh = AssociateSession::associate(topo.consumer.as_ref())
        .await
        .expect("fresh associate after the drop");
    let target = TargetAddr::Ip("10.0.0.5:53".parse().unwrap());
    fresh
        .send_datagram(&target, b"again")
        .await
        .expect("send after the drop");
    let (_, echoed) =
        tokio::time::timeout(std::time::Duration::from_secs(10), fresh.recv_datagram())
            .await
            .expect("recv timed out")
            .expect("a datagram (not eof)")
            .expect("a datagram");
    assert_eq!(&echoed[..], b"again");
    fresh.shutdown().await.expect("shutdown");
    no_leaked_channels(&topo).await;
}

/// A missing egress seam: the ASSOCIATE refusal surfaces the RFC
/// reply code through the session type (`CommandNotSupported` 0x07 —
/// ADR-008's missing-callback class), and no phantom channel remains
/// on either side (teardown-at-birth).
#[tokio::test]
async fn missing_egress_refusal_surfaces_rfc_reply_code_no_phantom_channel() {
    let topo = wire_consumer_side_with(echo_dial(), None, None, None).await;

    let err = AssociateSession::associate(topo.consumer.as_ref())
        .await
        .expect_err("no egress wired — the ASSOCIATE must refuse");

    match &err {
        SocksSessionError::Rfc(SocksErrorKind::ReplyError(code)) => {
            assert_eq!(code.as_u8(), consts::SOCKS5_REPLY_COMMAND_NOT_SUPPORTED);
        }
        other => panic!("expected an Rfc reply-code refusal, got {other:?}"),
    }

    no_leaked_channels(&topo).await;

    // The connection is healthy: wiring the egress serves ASSOCIATE.
    let topo2 = wire_consumer_side_with(echo_dial(), Some(echo_egress()), None, None).await;
    let mut assoc = AssociateSession::associate(topo2.consumer.as_ref())
        .await
        .expect("fresh associate after the refusal");
    assoc
        .send_datagram(TargetAddr::Ip("10.0.0.6:53".parse().unwrap()), b"fresh")
        .await
        .expect("send");
    let (_, echoed) =
        tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
            .await
            .expect("recv timed out")
            .expect("a datagram (not eof)")
            .expect("a datagram");
    assert_eq!(&echoed[..], b"fresh");
    assoc.shutdown().await.expect("shutdown");
    no_leaked_channels(&topo2).await;
}

/// An egress-factory failure: the mapped RFC reply code surfaces from
/// `associate` (reply#1's code — the factory refused before the
/// reply), and no phantom channel remains.
#[tokio::test]
async fn egress_factory_refusal_surfaces_the_mapped_reply_code() {
    let topo = wire_consumer_side_with(
        echo_dial(),
        Some(refusing_egress(EgressError::Unreachable(
            "no egress".into(),
        ))),
        None,
        None,
    )
    .await;

    let err = AssociateSession::associate(topo.consumer.as_ref())
        .await
        .expect_err("the egress factory refused");
    match &err {
        SocksSessionError::Rfc(SocksErrorKind::ReplyError(code)) => {
            assert_eq!(code.as_u8(), consts::SOCKS5_REPLY_NETWORK_UNREACHABLE);
        }
        other => panic!("expected an Rfc reply-code refusal, got {other:?}"),
    }
    no_leaked_channels(&topo).await;
}

/// CONNECT + ASSOCIATE on the same registration (the regression
/// guard extended): one ALPN, both commands, interleaved — an
/// association and a CONNECT session live concurrently on the same
/// registration, both serving.
#[tokio::test]
async fn connect_and_associate_served_on_the_same_registration() {
    let topo = wire_consumer_side_with(echo_dial(), Some(echo_egress()), None, None).await;

    let mut assoc = AssociateSession::associate(topo.consumer.as_ref())
        .await
        .expect("the association associates");
    let mut connect = Socks5ChannelSession::connect(
        topo.consumer.as_ref(),
        ("echo.test", 443u16).to_target_addr().unwrap(),
    )
    .await
    .expect("the CONNECT session connects");

    // The association round-trips while the CONNECT pipe is open.
    let target = TargetAddr::Ip("10.0.0.7:53".parse().unwrap());
    assoc
        .send_datagram(&target, b"udp")
        .await
        .expect("send datagram");
    let (_, echoed) =
        tokio::time::timeout(std::time::Duration::from_secs(10), assoc.recv_datagram())
            .await
            .expect("recv timed out")
            .expect("a datagram (not eof)")
            .expect("a datagram");
    assert_eq!(&echoed[..], b"udp");

    // The CONNECT pipe round-trips while the association is open.
    connect
        .split()
        .expect("pipe")
        .1
        .write_all(b"tcp")
        .await
        .expect("write");
    let (mut read, _) = connect.split().expect("pipe");
    let mut buf = [0u8; 3];
    tokio::time::timeout(
        std::time::Duration::from_secs(10),
        read.read_exact(&mut buf),
    )
    .await
    .expect("echo timed out")
    .expect("read the echo");
    assert_eq!(&buf, b"tcp");

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