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` — 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` — 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>>>();
}