alksocks/tests/composed-egress.rs

//! The composed egress integration tests (ADR-008's composed variant,
//! POC #2's `ComposedEgress` shape, end to end over channels): a
//! tunnel-producing peer (the alktunnels registry shape:
//! `register_tunnel_direct_openable` over a UDP echo dial), the
//! composed egress against it, datagram round-trips through the
//! tunnel (multiple destinations → multiple tunnels), the LRU
//! eviction at cap, the readiness contract re-pinned, the
//! closed-queue → association-end shape, and framing-once.
//!
//! The topology: the SOCKS5-style harness (consumer calls the
//! producer's open ops) carrying alktunnels' direct op; the producer
//! side registers `register_tunnel_direct_openable` (the
//! alktunnels 0.2.0 registry shape), the consumer's
//! `ChannelClient` is what the composed egress rides
//! (`TunnelSession::open_direct`).

#![cfg(feature = "composed")]

mod harness;

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

use alktunnels::params::Substrate;
use alktunnels::producer::{
    DialFn as TunnelDialFn, TargetHandle as TunnelTargetHandle, TunnelEstablishError,
};
use alktunnels::wire::DatagramReader as TunnelDatagramReader;
use bytes::Bytes;
use tokio::io::{AsyncRead, AsyncWriteExt};

use alksocks::composed::{direct_egress, ComposedEgress};
use alksocks::egress::{Egress, EgressError};
use alksocks::rfc1928::util::target_addr::TargetAddr;

/// An echo dial closure over the alktunnels direct-op `DialFn` shape:
/// an in-process framed UDP pipe per dial (the framed-adapter shape:
/// the "target" decodes the tunnel codec and echoes each datagram).
/// Counts the dialed destinations (the multiple-tunnels assertion).
fn counting_udp_echo_dial(dial_count: Arc<AtomicUsize>) -> TunnelDialFn {
    let dial_count = Arc::clone(&dial_count);
    Arc::new(move |substrate: Substrate, backing: &str| {
        let dial_count = Arc::clone(&dial_count);
        let _ = backing.to_string();
        Box::pin(async move {
            if substrate != Substrate::Udp {
                return Err(TunnelEstablishError::DialFailed(format!(
                    "non-udp dial not wired: {substrate:?}"
                )));
            }
            dial_count.fetch_add(1, Ordering::SeqCst);
            let (consumer_side, target_side) = tokio::io::duplex(64 * 1024);
            tokio::spawn(async move {
                use tokio::io::AsyncReadExt;
                let (mut t_read, mut t_write) = tokio::io::split(target_side);
                let mut reader = TunnelDatagramReader::new();
                let mut buf = vec![0u8; 16 * 1024];
                loop {
                    let n = match t_read.read(&mut buf).await {
                        Ok(0) | Err(_) => return,
                        Ok(n) => n,
                    };
                    for dg in reader.feed(&buf[..n]) {
                        let framed = match alktunnels::wire::frame_datagram(&dg) {
                            Ok(f) => f,
                            Err(_) => return,
                        };
                        if t_write.write_all(&framed).await.is_err() {
                            return;
                        }
                        let _ = t_write.flush().await;
                    }
                }
            });
            let (c_read, c_write) = tokio::io::split(consumer_side);
            Ok(TunnelTargetHandle {
                read: Box::new(c_read),
                write: Box::new(c_write),
            })
        })
    })
}

/// Wire the topology: the producer registers the alktunnels direct op
/// (over the given dial) on its serving registry; the consumer's
/// `ChannelClient` is the direct-op caller the composed egress rides.
async fn wire_direct(dial: TunnelDialFn) -> harness::TunnelDirectTopology {
    harness::wire_tunnel_direct(dial).await
}

/// The destinations the tests dial (distinct table keys).
fn dest(port: u16) -> TargetAddr {
    TargetAddr::Ip(std::net::SocketAddr::from(([10, 0, 0, 1], port)))
}

/// Build the egress against the topology (the consumer's client is
/// the direct-op caller).
async fn egress_for(
    topo: &harness::TunnelDirectTopology,
    max_tunnels: usize,
) -> Box<dyn Egress + Send> {
    let factory = direct_egress(Arc::clone(&topo.consumer), max_tunnels);
    factory().await.expect("egress builds")
}

/// A datagram round-trip through the tunnel: send to `addr`, expect
/// the echo back (the reply carries `addr + data`, i.e. inner bytes).
async fn round_trip(egress: &mut (dyn Egress + Send), addr: &TargetAddr, payload: &[u8]) {
    egress
        .send(addr.clone(), Bytes::from(payload.to_vec()))
        .await
        .expect("send egresses");
    let reply = tokio::time::timeout(std::time::Duration::from_secs(10), egress.next_reply())
        .await
        .expect("reply within timeout")
        .expect("next_reply ok")
        .expect("a reply");
    let expected = {
        let mut v = addr.to_be_bytes().expect("addr encodes");
        v.extend_from_slice(payload);
        Bytes::from(v)
    };
    assert_eq!(reply, expected, "the echo reports the destination");
}

#[tokio::test]
async fn datagrams_round_trip_per_destination_multiple_tunnels() {
    let _ = tracing_subscriber::fmt()
        .with_max_level(tracing::level_filters::LevelFilter::TRACE)
        .with_test_writer()
        .try_init();
    let dial_count = Arc::new(AtomicUsize::new(0));
    let topo = wire_direct(counting_udp_echo_dial(Arc::clone(&dial_count))).await;
    let mut egress = egress_for(&topo, 8).await;
    tracing::info!("dial_count={dial_count:?}");

    // Three destinations → three tunnels (lazy opens on first
    // datagram each), round-trips on each, in order.
    for port in [53u16, 123, 9999] {
        round_trip(egress.as_mut(), &dest(port), b"ping").await;
    }
    assert_eq!(
        dial_count.load(Ordering::SeqCst),
        3,
        "one dial per new destination"
    );

    // A repeat destination reuses its tunnel (no new dial).
    round_trip(egress.as_mut(), &dest(53), b"again").await;
    assert_eq!(
        dial_count.load(Ordering::SeqCst),
        3,
        "the table hit reuses the tunnel"
    );

    // The readiness contract: with nothing in flight, `next_reply`
    // parks (`Pending`), never a hot `Ok(None)`.
    let mut poll = egress.next_reply();
    tokio::time::timeout(std::time::Duration::from_millis(50), &mut poll)
        .await
        .expect_err("next_reply stays Pending when nothing is ready");
}

/// A fresh association (no tunnel ever opened) parks in `next_reply`:
/// the stage's select is governed by the channel branch until the
/// first datagram opens a tunnel (the assembly suite's drift finding:
/// the ended-relay `Err` fired on the never-opened table too, killing
/// every quiescent association at birth). Pinned on its own.
#[tokio::test]
async fn fresh_association_parks_before_any_tunnel_is_opened() {
    let dial_count = Arc::new(AtomicUsize::new(0));
    let topo = wire_direct(counting_udp_echo_dial(Arc::clone(&dial_count))).await;
    let mut egress = egress_for(&topo, 8).await;

    // The never-opened table: `next_reply` parks (Pending), never an
    // instantly-ready `Err` (the stage would end the association
    // before its first datagram).
    {
        let mut poll = egress.next_reply();
        tokio::time::timeout(std::time::Duration::from_millis(100), &mut poll)
            .await
            .expect_err("the fresh association's next_reply parks");
    }

    // After the first datagram opens a tunnel, the normal contract
    // holds (the round-trip works).
    round_trip(egress.as_mut(), &dest(53), b"after").await;
    assert_eq!(dial_count.load(Ordering::SeqCst), 1);
}

#[tokio::test]
async fn lru_eviction_reopens_evicted_destinations() {
    let dial_count = Arc::new(AtomicUsize::new(0));
    let topo = wire_direct(counting_udp_echo_dial(Arc::clone(&dial_count))).await;
    let mut egress = egress_for(&topo, 2).await;

    // Fill the table to the cap (2): two destinations dial.
    round_trip(egress.as_mut(), &dest(53), b"a").await;
    round_trip(egress.as_mut(), &dest(123), b"b").await;
    assert_eq!(dial_count.load(Ordering::SeqCst), 2);

    // A third destination evicts the LRU entry (dest 53) and dials.
    round_trip(egress.as_mut(), &dest(9999), b"c").await;
    assert_eq!(
        dial_count.load(Ordering::SeqCst),
        3,
        "the third destination dialed"
    );

    // The LRU (dest 53) was evicted: re-touching it re-dials (the
    // cap stayed; no fourth tunnel above it).
    round_trip(egress.as_mut(), &dest(53), b"redial").await;
    assert_eq!(
        dial_count.load(Ordering::SeqCst),
        4,
        "the evicted destination re-dialed"
    );

    // The recent entries survived: dest 9999 reuses; dest 123 was
    // the eviction victim of the redial, so it re-dials (the cap
    // stayed at 2 throughout).
    round_trip(egress.as_mut(), &dest(9999), b"reuse").await;
    assert_eq!(
        dial_count.load(Ordering::SeqCst),
        4,
        "the recency survivor reused"
    );
    round_trip(egress.as_mut(), &dest(123), b"redial too").await;
    assert_eq!(
        dial_count.load(Ordering::SeqCst),
        5,
        "the eviction victim re-dialed"
    );
}

/// The framing-once probe: the reply the stage would frame is exactly
/// `addr.to_be_bytes() + payload`, with no double prefix. Pinned through
/// the egress's own reply (the round-trip asserts equality with the
/// inner form); here the composition is asserted at the wire boundary:
/// the egress's reply, fed through `frame_inner` (the wire layer's
/// one framing application), decodes to the original (addr, payload).
#[tokio::test]
async fn framing_once_replies_decode_through_the_wire_layer_exactly_once() {
    let dial_count = Arc::new(AtomicUsize::new(0));
    let topo = wire_direct(counting_udp_echo_dial(Arc::clone(&dial_count))).await;
    let mut egress = egress_for(&topo, 8).await;

    let addr = dest(4242);
    let payload = b"frame once";
    egress
        .send(addr.clone(), Bytes::from_static(payload))
        .await
        .expect("send");
    let inner = tokio::time::timeout(std::time::Duration::from_secs(10), egress.next_reply())
        .await
        .expect("reply within timeout")
        .expect("next_reply ok")
        .expect("a reply");

    // The egress's reply carries NO length prefix (inner bytes).
    let mut expected_inner = addr.to_be_bytes().expect("addr encodes");
    expected_inner.extend_from_slice(payload);
    assert_eq!(
        inner,
        Bytes::from(expected_inner),
        "inner bytes, no double prefix"
    );

    // Feeding it through the wire layer's one framing application
    // (`frame_inner` → `DatagramReader`) decodes the original pair.
    let framed = alksocks::wire::frame_inner(&inner).expect("the wire layer frames once");
    let mut reader = alksocks::wire::DatagramReader::new();
    let mut decoded = None;
    for dg in reader.feed(&framed) {
        decoded = Some(dg.expect("frame decodes"));
    }
    let (src, payload_back) = decoded.expect("one datagram");
    assert_eq!(src, addr);
    assert_eq!(&payload_back[..], payload);
}

/// The closed-queue → association-end shape: when the tunnel's
/// producer side ends (the dial's target EOFs), the relay task exits
/// (EOF on the read half) and, with every relay ended, the egress's
/// `next_reply` returns `Err` (never a hot `Ok(None)`).
#[tokio::test]
async fn closed_queue_yields_err_not_a_hot_none() {
    // The dial's "target" echoes exactly one datagram, then EOFs the
    // tunnel (the write end drops after one reply).
    let dial: TunnelDialFn = Arc::new(move |_substrate: Substrate, _backing: &str| {
        Box::pin(async move {
            let (consumer_side, mut target_side) = tokio::io::duplex(64 * 1024);
            tokio::spawn(async move {
                use tokio::io::AsyncReadExt;
                let mut reader = TunnelDatagramReader::new();
                let mut buf = vec![0u8; 16 * 1024];
                'outer: loop {
                    let n = match target_side.read(&mut buf).await {
                        Ok(0) | Err(_) => return,
                        Ok(n) => n,
                    };
                    if let Some(dg) = reader.feed(&buf[..n]).into_iter().next() {
                        let framed = match alktunnels::wire::frame_datagram(&dg) {
                            Ok(f) => f,
                            Err(_) => return,
                        };
                        if target_side.write_all(&framed).await.is_err() {
                            return;
                        }
                        // One echo, then EOF (the channel ends
                        // producer-side; the relay task drains the
                        // EOF → exits → the queue closes).
                        break 'outer;
                    }
                }
                let _ = target_side.shutdown().await;
            });
            let (c_read, c_write) = tokio::io::split(consumer_side);
            Ok(TunnelTargetHandle {
                read: Box::new(c_read),
                write: Box::new(c_write),
            })
        })
    });
    let topo = wire_direct(dial).await;
    let mut egress = egress_for(&topo, 8).await;

    let addr = dest(53);
    egress
        .send(addr.clone(), Bytes::from_static(b"only"))
        .await
        .expect("send");
    // The one echo comes back.
    let reply = tokio::time::timeout(std::time::Duration::from_secs(10), egress.next_reply())
        .await
        .expect("reply within timeout")
        .expect("next_reply ok")
        .expect("a reply");
    let mut expected = addr.to_be_bytes().expect("addr encodes");
    expected.extend_from_slice(b"only");
    assert_eq!(reply, Bytes::from(expected));

    // The tunnel ended → the relay exited → the queue closed. The
    // next `next_reply` returns `Err` (the association-end shape);
    // awaited to prove it is an `Err`, not an instantly-ready
    // `Ok(None)` spinning the stage.
    let result = tokio::time::timeout(std::time::Duration::from_secs(10), egress.next_reply())
        .await
        .expect("the closed queue surfaces (relay task drains the EOF)")
        .expect_err("the closed queue returns Err");
    assert!(
        matches!(result, EgressError::General(ref m) if m.contains("relays ended")),
        "the association-end error: {result}"
    );
}

/// The domain form passes through unresolved (ADR-012, the ssh `-D`
/// semantic): a domain destination renders its domain into the direct
/// op's target verbatim (the dial closure sees the unresolved host).
#[tokio::test]
async fn domain_destinations_pass_through_unresolved() {
    let seen = Arc::new(tokio::sync::Mutex::new(None::<String>));
    let seen_for_dial = Arc::clone(&seen);
    let dial: TunnelDialFn = Arc::new(move |substrate: Substrate, backing: &str| {
        let backing = backing.to_string();
        let seen = Arc::clone(&seen_for_dial);
        Box::pin(async move {
            if substrate != Substrate::Udp {
                return Err(TunnelEstablishError::DialFailed("udp only".into()));
            }
            *seen.lock().await = Some(backing);
            let (consumer_side, target_side) = tokio::io::duplex(64 * 1024);
            tokio::spawn(async move {
                use tokio::io::AsyncReadExt;
                let (mut t_read, mut t_write) = tokio::io::split(target_side);
                let mut reader = TunnelDatagramReader::new();
                let mut buf = vec![0u8; 16 * 1024];
                loop {
                    let n = match t_read.read(&mut buf).await {
                        Ok(0) | Err(_) => return,
                        Ok(n) => n,
                    };
                    for dg in reader.feed(&buf[..n]) {
                        let framed = match alktunnels::wire::frame_datagram(&dg) {
                            Ok(f) => f,
                            Err(_) => return,
                        };
                        if t_write.write_all(&framed).await.is_err() {
                            return;
                        }
                        let _ = t_write.flush().await;
                    }
                }
            });
            let (c_read, c_write) = tokio::io::split(consumer_side);
            Ok(TunnelTargetHandle {
                read: Box::new(c_read),
                write: Box::new(c_write),
            })
        })
    });
    let topo = wire_direct(dial).await;
    let mut egress = egress_for(&topo, 8).await;

    let addr = TargetAddr::Domain("example.test".into(), 5353);
    egress
        .send(addr.clone(), Bytes::from_static(b"q"))
        .await
        .expect("send");
    let reply = tokio::time::timeout(std::time::Duration::from_secs(10), egress.next_reply())
        .await
        .expect("reply within timeout")
        .expect("next_reply ok")
        .expect("a reply");
    let mut expected = addr.to_be_bytes().expect("domain addr encodes");
    expected.extend_from_slice(b"q");
    assert_eq!(reply, Bytes::from(expected));
    assert_eq!(
        seen.lock().await.as_deref(),
        Some("example.test:5353"),
        "the domain passed through unresolved (producer-side resolution)"
    );
}

/// The compile-shape pin: the composed egress's session shell rides
/// the table entry (Send); the owner wrapper compiles against the
/// alktunnels surface (`TakenHalves`-shaped halves + the session).
#[test]
fn composed_shapes_compile() {
    fn assert_send<T: Send>() {}
    assert_send::<ComposedEgress>();
    assert_send::<alksocks::composed::SessionOwner<Box<dyn AsyncRead + Send + Unpin>>>();
}

/// The eviction-drives-association-end composition (review 003 D-2):
/// a cap-1 table whose only entry is evicted by a send whose
/// replacement dial FAILS: the table ends with zero entries and the
/// evicted tunnel's EOF drains its relay to `Ended`, so the parked
/// `next_reply` hits the closed-queue condition
/// (`ever_opened && live_relays == 0`) and returns the
/// association-end `Err`: the eviction path's own way to reach the
/// shape `closed_queue_yields_err_not_a_hot_none` pins by far-side
/// dial. Also pinned: the count's `saturating_sub` discipline (the
/// stage keeps serving a *live* relay beside the ended one (the
/// failing send does not poison the association).
#[tokio::test]
async fn eviction_to_a_failed_dial_ends_the_association_when_the_relay_ends() {
    // The dial's target echoes everything (a live tunnel) until the
    // consumer side EOFs; the SECOND dial fails outright (the
    // replacement the eviction can never establish).
    let dial: TunnelDialFn = {
        let inner = counting_udp_echo_dial(Arc::new(AtomicUsize::new(0)));
        let first = Arc::new(std::sync::atomic::AtomicBool::new(true));
        Arc::new(move |substrate: Substrate, backing: &str| {
            let ok = first.swap(false, Ordering::SeqCst);
            let inner = Arc::clone(&inner);
            let backing = backing.to_string();
            Box::pin(async move {
                if !ok {
                    return Err(TunnelEstablishError::DialFailed(
                        "the replacement dial fails".into(),
                    ));
                }
                inner(substrate, &backing).await
            })
        })
    };
    let topo = wire_direct(dial).await;
    let mut egress = egress_for(&topo, 1).await;

    // One live destination: the single tunnel (cap 1). The relay is
    // live; the echo works.
    round_trip(egress.as_mut(), &dest(53), b"alive").await;

    // The association still serves while the relay is live (the
    // saturating_sub discipline: a live relay keeps the stage in the
    // queue branch; a second datagram round-trips on the same
    // tunnel).
    round_trip(egress.as_mut(), &dest(53), b"again").await;

    // Evict the only entry with a send whose dial fails: the LRU
    // drop closes the tunnel (the entry's session + write half go),
    // the dial fails after the eviction, and the table is left
    // EMPTY with the evicted relay draining its EOF. The send
    // surfaces the dial failure (the datagram was not delivered).
    let err = egress
        .send(dest(9999), Bytes::from_static(b"evict"))
        .await
        .expect_err("the replacement dial fails");
    assert!(
        matches!(err, EgressError::Unreachable(ref m) if m.contains("tunnel open"),
        ),
        "the failed replacement dial: {err}"
    );

    // The evicted relay's EOF drains → `Ended` → the count hits 0 →
    // the parked `next_reply` returns the association-end `Err`
    // (awaited to prove it is the Err shape, never a hot `Ok(None)`).
    let result = tokio::time::timeout(std::time::Duration::from_secs(10), egress.next_reply())
        .await
        .expect("the association-end surfaces (the evicted relay's EOF drains)")
        .expect_err("the all-relays-ended association end");
    assert!(
        matches!(result, EgressError::General(ref m) if m.contains("relays ended")),
        "the eviction-driven association end: {result}"
    );
}