alksocks/tests/common/mod.rs

//! Shared test helpers — the harness shapes single-sourced for both
//! the `src/` unit suites and the integration suite (review 002 N-5;
//! the standard `mod common` pattern). The lib tests include this
//! file via `#[path = "../tests/common/mod.rs"]` (unit tests in `src/`
//! cannot import `tests/harness.rs` — separate crates); the
//! integration suite reaches it through `harness`'s re-exports.
//!
//! Single-sourced here (the byte-near-identical shapes): the echo
//! dial (plain + target-witnessing), the refusing dial, the queue
//! egress + its echo factory, the raw RFC 1928 client trio
//! (`greet`/`request`/`read_reply`), the datagram frame codec
//! helpers, and the parked accept factory. Deliberately NOT collapsed
//! (distinct test shapes, stay test-local): the hot-none egress
//! variants (R-1's readiness-contract violations) and the drop-witness
//! egresses.
//!
//! `allow(dead_code)` rides the harness posture: each test binary
//! (and the lib test crate) uses a different subset of the helpers.

#![allow(dead_code)]

use std::net::SocketAddr;
use std::sync::Arc;

use alksocks::egress::{
    BindAccept, BindAcceptError, BindAcceptFactory, BindTargetHandle, DialFn, Egress, EgressError,
    EgressFactory, SocksDialError, TargetHandle,
};
use alksocks::rfc1928::consts;
use alksocks::rfc1928::util::target_addr::TargetAddr;
use alksocks::rfc1928::Socks5Command;
use bytes::Bytes;
use futures::future::BoxFuture;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::sync::Mutex;

// --- The dial seam stand-ins -------------------------------------------------

/// An echo dial closure that records every dialed target (the
/// interception-point probe: domain form arrives unresolved —
/// ADR-012; IP-string forms arrive as parsed IP targets). One duplex
/// per dial — the "target" echoes bytes until EOF. The producer-side
/// pump gets split `DuplexStream` halves; a spawned task echoes the
/// other end. No sockets, no `local` feature.
pub fn echo_dial_witnessed(targets: Arc<Mutex<Vec<TargetAddr>>>) -> DialFn {
    Arc::new(move |target: TargetAddr| {
        let targets = Arc::clone(&targets);
        Box::pin(async move {
            targets.lock().await.push(target);
            let (mut echo_end, target_end) = tokio::io::duplex(64 * 1024);
            tokio::spawn(async move {
                loop {
                    let mut buf = vec![0u8; 4096];
                    match echo_end.read(&mut buf).await {
                        Ok(0) | Err(_) => return,
                        Ok(n) => {
                            if echo_end.write_all(&buf[..n]).await.is_err() {
                                return;
                            }
                        }
                    }
                }
            });
            let (read_half, write_half) = tokio::io::split(target_end);
            Ok(TargetHandle {
                read: Box::new(read_half) as Box<dyn AsyncRead + Send + Sync + Unpin>,
                write: Box::new(write_half) as Box<dyn AsyncWrite + Send + Sync + Unpin>,
            })
        })
    })
}

/// An echo dial with a throwaway witness (the tests that never
/// inspect the dialed targets).
pub fn echo_dial() -> DialFn {
    echo_dial_witnessed(Arc::new(Mutex::new(Vec::new())))
}

/// A dial that always refuses with the given [`SocksDialError`]
/// variant (the refusal-mapping tests' dial — the refusal rides
/// in-stream as the mapped RFC reply code). The variant is
/// reconstructed per call (the seam errors are not `Clone`).
pub fn refusing_dial(err: SocksDialError) -> DialFn {
    Arc::new(move |_target: TargetAddr| {
        let err = match &err {
            SocksDialError::Refused(m) => SocksDialError::Refused(m.clone()),
            SocksDialError::NetworkUnreachable(m) => SocksDialError::NetworkUnreachable(m.clone()),
            SocksDialError::HostUnreachable(m) => SocksDialError::HostUnreachable(m.clone()),
            SocksDialError::NotAllowed(m) => SocksDialError::NotAllowed(m.clone()),
            SocksDialError::General(m) => SocksDialError::General(m.clone()),
        };
        Box::pin(async move { Err(err) })
    })
}

// --- The egress seam stand-ins ------------------------------------------------

/// The queue egress behind [`echo_egress_factory`]: `send` queues the
/// inner bytes for echo (`addr + data` — the egress speaks inner
/// bytes, no length prefix); `next_reply` parks on the queue
/// (`Pending` until a datagram arrives — the readiness contract,
/// never an instantly-ready `Ok(None)`). The receiver rides a mutex
/// so `&mut self` methods can lock it (the shape both the harness's
/// and the producer's inline copies had converged on).
pub struct QueueEgress {
    pub inbound: tokio::sync::mpsc::UnboundedSender<Bytes>,
    pub outbound: Mutex<tokio::sync::mpsc::UnboundedReceiver<Bytes>>,
}

impl Egress for QueueEgress {
    fn send(&mut self, addr: TargetAddr, payload: Bytes) -> BoxFuture<'_, Result<(), EgressError>> {
        Box::pin(async move {
            let mut inner = addr
                .to_be_bytes()
                .map_err(|e| EgressError::General(format!("addr encode: {e}")))?;
            inner.extend_from_slice(&payload);
            self.inbound
                .send(Bytes::from(inner))
                .map_err(|_| EgressError::General("egress closed".into()))
        })
    }

    fn next_reply(&mut self) -> BoxFuture<'_, Result<Option<Bytes>, EgressError>> {
        Box::pin(async move {
            // Real async park: `Pending` until a datagram arrives (the
            // readiness contract — never an instantly-ready `Ok(None)`,
            // which would starve a biased `select!` into livelock).
            match self.outbound.lock().await.recv().await {
                Some(inner) => Ok(Some(inner)),
                // All senders dropped: end the association (an `Err`,
                // never a hot `Ok(None)`).
                None => Err(EgressError::General("egress closed".into())),
            }
        })
    }
}

/// An echo egress factory: datagrams sent to the egress echo back
/// (the POC #2 shape). The readiness contract is honored (see
/// [`QueueEgress`]).
pub fn echo_egress_factory() -> EgressFactory {
    alksocks::egress::egress_factory(|| {
        let (inbound, outbound) = tokio::sync::mpsc::unbounded_channel::<Bytes>();
        Ok(Box::new(QueueEgress {
            inbound,
            outbound: Mutex::new(outbound),
        }) as Box<dyn Egress + Send>)
    })
}

/// A doubled-echo egress factory: every datagram echoes back **twice
/// back-to-back** (two frames queued as one reply burst — the
/// producer's frame writer puts both length-prefixed frames on the
/// channel; the consumer's read may batch them into one chunk). The
/// reader-state probe for the front doors' bridge (review 003 R-1):
/// a bridge that re-creates its decoder per read loses the second
/// datagram; the landed bridge queues it.
pub fn doubled_echo_egress_factory() -> EgressFactory {
    alksocks::egress::egress_factory(|| {
        let (inbound, outbound) = tokio::sync::mpsc::unbounded_channel::<Bytes>();
        Ok(Box::new(DoubledEchoEgress {
            inbound,
            outbound: Mutex::new(outbound),
        }) as Box<dyn Egress + Send>)
    })
}

/// The doubled echo behind [`doubled_echo_egress_factory`] (the
/// [`QueueEgress`] shape with the echo count two).
pub struct DoubledEchoEgress {
    pub inbound: tokio::sync::mpsc::UnboundedSender<Bytes>,
    pub outbound: Mutex<tokio::sync::mpsc::UnboundedReceiver<Bytes>>,
}

impl Egress for DoubledEchoEgress {
    fn send(&mut self, addr: TargetAddr, payload: Bytes) -> BoxFuture<'_, Result<(), EgressError>> {
        Box::pin(async move {
            let mut inner = addr
                .to_be_bytes()
                .map_err(|e| EgressError::General(format!("addr encode: {e}")))?;
            inner.extend_from_slice(&payload);
            let inner = Bytes::from(inner);
            self.inbound
                .send(inner.clone())
                .map_err(|_| EgressError::General("egress closed".into()))?;
            self.inbound
                .send(inner)
                .map_err(|_| EgressError::General("egress closed".into()))
        })
    }

    fn next_reply(&mut self) -> BoxFuture<'_, Result<Option<Bytes>, EgressError>> {
        Box::pin(async move {
            match self.outbound.lock().await.recv().await {
                Some(inner) => Ok(Some(inner)),
                None => Err(EgressError::General("egress closed".into())),
            }
        })
    }
}

/// A failing egress factory (the ASSOCIATE-refusal probe): refuses
/// with the given error, which the producer maps to reply#1's code.
/// The variant is reconstructed per call (the seam errors are not
/// `Clone`).
pub fn refusing_egress(err: EgressError) -> EgressFactory {
    Arc::new(move || {
        let err = match &err {
            EgressError::Refused(m) => EgressError::Refused(m.clone()),
            EgressError::Unreachable(m) => EgressError::Unreachable(m.clone()),
            EgressError::General(m) => EgressError::General(m.clone()),
        };
        Box::pin(async move { Err(err) })
    })
}

// --- The accept seam stand-ins -------------------------------------------------

/// A parked accept factory: the accept future never resolves (the
/// deterministic quiet-wait probe — the client-gone exit is the only
/// end this session can take). The advertised bind address matches
/// the harness's `ECHO_BIND_PORT` (2121 — the reply#1 carrier only).
pub fn parked_accept_factory() -> BindAcceptFactory {
    Arc::new(move |_target: TargetAddr| {
        Box::pin(async move {
            Ok(BindAccept::new(
                SocketAddr::from(([127, 0, 0, 1], 2121)),
                std::future::pending::<Result<BindTargetHandle, BindAcceptError>>(),
            ))
        })
    })
}

// --- The raw RFC 1928 client trio ----------------------------------------------

/// Greeting (`05 01 00` — version 5, one method, no-auth) → the
/// method-selection read. Returns the selected method byte; asserts
/// the method-selection reply's version byte is 5.
pub async fn raw_greet<S>(stream: &mut S) -> Result<u8, std::io::Error>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    stream.write_all(&[0x05, 0x01, 0x00]).await?;
    let mut chosen = [0u8; 2];
    stream.read_exact(&mut chosen).await?;
    assert_eq!(
        chosen[0],
        consts::SOCKS5_VERSION,
        "the method-selection reply is version 5"
    );
    Ok(chosen[1])
}

/// Write a command request for `target` (the request form:
/// `VER CMD RSV ATYP DST.ADDR DST.PORT`), then flush.
pub async fn raw_request<S>(
    stream: &mut S,
    cmd: &Socks5Command,
    target: &TargetAddr,
) -> Result<(), std::io::Error>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let addr = target
        .to_be_bytes()
        .map_err(|e| std::io::Error::other(format!("target encode: {e}")))?;
    let mut req = vec![consts::SOCKS5_VERSION, cmd.as_u8(), 0x00];
    req.extend_from_slice(&addr);
    stream.write_all(&req).await?;
    stream.flush().await
}

/// Read one reply (the 4-byte head + the BND.ADDR per its ATYP).
/// Returns `(reply_code, atyp, addr_bytes)` — the address bytes as
/// the RFC form carries them (v4/v6: addr + port; domain: len +
/// name + port).
pub async fn raw_read_reply<S>(stream: &mut S) -> Result<(u8, u8, Vec<u8>), std::io::Error>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let mut head = [0u8; 4];
    stream.read_exact(&mut head).await?;
    assert_eq!(head[0], consts::SOCKS5_VERSION, "the reply is version 5");
    let addr_bytes: Vec<u8> = match head[3] {
        consts::SOCKS5_ADDR_TYPE_IPV4 => {
            let mut a = [0u8; 6];
            stream.read_exact(&mut a).await?;
            a.to_vec()
        }
        consts::SOCKS5_ADDR_TYPE_IPV6 => {
            let mut a = [0u8; 18];
            stream.read_exact(&mut a).await?;
            a.to_vec()
        }
        consts::SOCKS5_ADDR_TYPE_DOMAIN_NAME => {
            let mut len = [0u8; 1];
            stream.read_exact(&mut len).await?;
            let mut a = vec![0u8; len[0] as usize + 2];
            stream.read_exact(&mut a).await?;
            let mut full = vec![len[0]];
            full.extend_from_slice(&a);
            full
        }
        other => {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("unknown address type {other}"),
            ))
        }
    };
    Ok((head[1], head[3], addr_bytes))
}

/// The producer-test trio's greet: unwraps and asserts no-auth was
/// selected.
pub async fn greet<S>(client: &mut S)
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let method = raw_greet(client).await.unwrap();
    assert_eq!(method, consts::SOCKS5_AUTH_METHOD_NONE);
}

/// The producer-test trio's request writer (unwrapping).
pub async fn write_request<S>(client: &mut S, cmd: &Socks5Command, target: &TargetAddr)
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    raw_request(client, cmd, target).await.unwrap();
}

/// The producer-test trio's reply reader: the v4-only shape every
/// producer test relies on (the BND.ADDR assertions index a
/// six-byte `addr + port`).
pub async fn read_reply_v4<S>(client: &mut S) -> (u8, [u8; 6])
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let (code, atyp, addr) = raw_read_reply(client).await.unwrap();
    assert_eq!(atyp, consts::SOCKS5_ADDR_TYPE_IPV4, "tests use v4 BND.ADDR");
    let mut out = [0u8; 6];
    out.copy_from_slice(&addr);
    (code, out)
}

// --- The datagram frame helpers -------------------------------------------------

/// Frame one datagram the way the in-band wire carries it
/// (`[len: u16 BE][addr + data]` — the wire codec's shape).
pub fn datagram_frame(addr: &TargetAddr, payload: &[u8]) -> Vec<u8> {
    alksocks::wire::frame(addr, payload)
        .expect("frame")
        .to_vec()
}

/// Decode a datagram from the wire (one frame per read batch, in
/// order).
pub fn decode_datagram(wire: &[u8]) -> (TargetAddr, Bytes) {
    let mut reader = alksocks::wire::DatagramReader::new();
    let mut got = None;
    for dg in reader.feed(wire) {
        got = Some(dg);
    }
    got.expect("one datagram decoded").expect("frame decodes")
}