irtt-rs/crates/irtt-client/tests/support/mod.rs

#![allow(dead_code)]

mod in_tree_server;
mod real_irtt;

use std::{
    net::{SocketAddr, UdpSocket},
    process::Command,
    sync::mpsc,
    thread::{self, JoinHandle},
    time::{Duration, SystemTime},
};

use irtt_client::{Client, ClientConfig, ClientEvent, NegotiatedParams, OpenOutcome, SocketConfig};
use irtt_proto::{
    compute_hmac_in_place, decode_request, echo_packet_len, encode_echo_reply, encode_open_reply,
    flags, verify_packet_hmac, Clock, DecodedRequest, DecodedRequestKind, EchoReply, OpenReply,
    PacketLayout, Params, ProtoError, ReceivedStats, ServerFill, StampAt, TimestampFields,
    HMAC_SIZE, MAGIC, PROTOCOL_VERSION,
};
use irtt_server::ServerConfig;

pub use in_tree_server::InTreeServer;
pub use real_irtt::RealIrtServer;

const HMAC_OFFSET: usize = 4;
pub const TOKEN: u64 = 0x1234_5678_90ab_cdef;
pub const RECV_COUNT: u32 = 42;
pub const RECV_WINDOW: u64 = 0x07;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TestBackendKind {
    InTree,
    Real,
}

pub fn selected_backend() -> TestBackendKind {
    match std::env::var("IRTT_TEST_BACKEND").as_deref() {
        Ok("real") => {
            let irtt_bin = std::env::var("IRTT_BIN").unwrap_or_else(|_| "irtt".to_string());
            match Command::new(&irtt_bin).arg("version").output() {
                Ok(output) if output.status.success() => {
                    debug_backend("[backend] selected backend=real");
                    TestBackendKind::Real
                }
                _ => panic!("IRTT_TEST_BACKEND=real but irtt binary not found at '{irtt_bin}'"),
            }
        }
        Ok("in-tree") | Err(_) => {
            debug_backend("[backend] selected backend=in-tree");
            TestBackendKind::InTree
        }
        Ok(other) => panic!("unknown IRTT_TEST_BACKEND value: {other}"),
    }
}

fn debug_backend(message: &str) {
    if std::env::var("IRTT_TEST_BACKEND_DEBUG").as_deref() == Ok("1") {
        eprintln!("{message}");
    }
}

pub enum BackendPeer {
    InTree(InTreeServer),
    Real(RealIrtServer),
}

impl BackendPeer {
    pub fn addr(&self) -> SocketAddr {
        match self {
            BackendPeer::InTree(s) => s.addr,
            BackendPeer::Real(s) => s.addr(),
        }
    }

    pub fn start_open_echo(hmac_key: Option<Vec<u8>>) -> Self {
        match selected_backend() {
            TestBackendKind::InTree => BackendPeer::InTree(InTreeServer::start(
                hmac_key.map_or_else(ServerConfig::default, |key| {
                    ServerConfig::default().with_hmac_key(key)
                }),
            )),
            TestBackendKind::Real => {
                BackendPeer::Real(RealIrtServer::start(hmac_key.as_deref()).unwrap())
            }
        }
    }

    pub fn start_hmac_required(key: Vec<u8>) -> Self {
        match selected_backend() {
            TestBackendKind::InTree => BackendPeer::InTree(InTreeServer::start(
                ServerConfig::default().with_hmac_key(key),
            )),
            TestBackendKind::Real => BackendPeer::Real(RealIrtServer::start(Some(&key)).unwrap()),
        }
    }
}

pub struct FakeServer {
    pub addr: SocketAddr,
    rx: mpsc::Receiver<ServerObservation>,
    done: JoinHandle<()>,
}

impl FakeServer {
    pub fn observations(&self, count: usize) -> Vec<ServerObservation> {
        self.rx.iter().take(count).collect()
    }

    pub fn join(self) {
        self.done.join().unwrap();
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ServerObservation {
    Open {
        params: Params,
        hmac: bool,
    },
    RejectedHmac {
        hmac: bool,
        bad_hmac: bool,
    },
    Echo {
        len: usize,
        hmac: bool,
        token: u64,
        sequence: u32,
    },
    Close {
        hmac: bool,
        token: u64,
    },
}

pub struct OneProbeRun {
    pub negotiated: NegotiatedParams,
    pub sent: ClientEvent,
    pub reply: ClientEvent,
    pub observations: Vec<ServerObservation>,
}

pub fn default_params() -> Params {
    Params {
        protocol_version: PROTOCOL_VERSION,
        duration_ns: 3_000_000_000,
        interval_ns: 1_000_000_000,
        received_stats: ReceivedStats::Both,
        stamp_at: StampAt::Both,
        clock: Clock::Both,
        ..Params::default()
    }
}

pub fn params_for_modes(received_stats: ReceivedStats, stamp_at: StampAt, clock: Clock) -> Params {
    Params {
        received_stats,
        stamp_at,
        clock,
        ..default_params()
    }
}

/// Builds a [`ClientConfig`] whose negotiated `Params` will match `params`.
///
/// `params.dscp` is the raw IP TOS / Traffic Class byte, while
/// [`ClientConfig::dscp`] is the six-bit codepoint that produces it; shift
/// right by two to recover the codepoint the resulting config must request.
/// Panics if `params.dscp` has any ECN bits set, since those cannot be
/// recovered from a codepoint and this helper's contract would silently stop
/// holding otherwise.
pub fn config_for_params(addr: SocketAddr, params: &Params) -> ClientConfig {
    let traffic_class =
        u8::try_from(params.dscp).expect("test DSCP must fit the raw traffic-class byte");
    assert_eq!(
        traffic_class & 0b11,
        0,
        "test Params::dscp must be representable by ClientConfig's DSCP codepoint"
    );
    ClientConfig {
        server_addr: addr.to_string(),
        duration: if params.duration_ns == 0 {
            None
        } else {
            Some(Duration::from_nanos(
                u64::try_from(params.duration_ns).expect("test duration must be non-negative"),
            ))
        },
        interval: Duration::from_nanos(
            u64::try_from(params.interval_ns).expect("test interval must be non-negative"),
        ),
        length: u32::try_from(params.length).unwrap(),
        received_stats: params.received_stats,
        stamp_at: params.stamp_at,
        clock: params.clock,
        dscp: traffic_class >> 2,
        server_fill: params.server_fill.as_ref().map(|fill| fill.value.clone()),
        open_timeouts: vec![Duration::from_millis(200)],
        socket_config: SocketConfig {
            recv_timeout: Some(Duration::from_millis(500)),
            ..Default::default()
        },
        ..ClientConfig::default()
    }
}

pub fn standard_timestamps() -> TimestampFields {
    TimestampFields {
        recv_wall: Some(0),
        recv_mono: Some(5_000_000),
        midpoint_wall: Some(0),
        midpoint_mono: Some(5_000_000),
        send_wall: Some(1),
        send_mono: Some(5_000_001),
    }
}

pub fn run_one_probe(params: Params, timestamps: TimestampFields) -> OneProbeRun {
    let config_params = params.clone();
    run_one_probe_with_config(params, timestamps, None, |addr| {
        config_for_params(addr, &config_params)
    })
}

pub fn run_one_probe_with_config<F>(
    server_params: Params,
    timestamps: TimestampFields,
    hmac_key: Option<Vec<u8>>,
    build_config: F,
) -> OneProbeRun
where
    F: FnOnce(SocketAddr) -> ClientConfig,
{
    let server = start_one_probe_server(server_params, timestamps, hmac_key);
    let mut client = Client::connect(build_config(server.addr)).unwrap();
    let negotiated = assert_started(client.open().unwrap());

    let sent_events = client.send_probe().unwrap();
    assert_eq!(sent_events.len(), 1);

    let reply_events = client.recv_once().unwrap();
    assert_eq!(reply_events.len(), 1);

    let observations = server.observations(2);
    server.join();

    OneProbeRun {
        negotiated,
        sent: sent_events.into_iter().next().unwrap(),
        reply: reply_events.into_iter().next().unwrap(),
        observations,
    }
}

pub fn start_open_server(params: Params, hmac_key: Option<Vec<u8>>) -> FakeServer {
    start_fake_server(move |socket, tx| {
        let (request, peer) = recv_request(&socket);
        let (open_params, hmac) = open_request_params(&request, hmac_key.as_deref());
        tx.send(ServerObservation::Open {
            params: open_params,
            hmac,
        })
        .unwrap();

        let reply = open_reply(
            flags::FLAG_OPEN | flags::FLAG_REPLY,
            TOKEN,
            &params,
            hmac_key.as_deref(),
        );
        socket.send_to(&reply, peer).unwrap();
    })
}

pub fn start_hmac_required_open_drop_server(key: Vec<u8>, wait: Duration) -> FakeServer {
    start_fake_server(move |socket, tx| {
        socket.set_read_timeout(Some(wait)).unwrap();
        while let Some(request) = recv_request_timeout(&socket) {
            let hmac = decode_request(&request).is_ok_and(|request| request.hmac_present);
            // A server with a key drops any request that fails authentication,
            // whether it carried no HMAC field at all or a field signed with
            // the wrong key.
            match verify_packet_hmac(&key, &request) {
                Ok(()) => {}
                Err(error) => {
                    tx.send(ServerObservation::RejectedHmac {
                        hmac,
                        bad_hmac: error == ProtoError::BadHmac,
                    })
                    .unwrap();
                }
            }
        }
    })
}

pub fn start_hmac_close_server(params: Params, key: Vec<u8>) -> FakeServer {
    start_fake_server(move |socket, tx| {
        let (request, peer) = recv_request(&socket);
        let (open_params, hmac) = open_request_params(&request, Some(&key));
        tx.send(ServerObservation::Open {
            params: open_params,
            hmac,
        })
        .unwrap();

        let reply = open_reply(
            flags::FLAG_OPEN | flags::FLAG_REPLY,
            TOKEN,
            &params,
            Some(&key),
        );
        socket.send_to(&reply, peer).unwrap();

        socket
            .set_read_timeout(Some(Duration::from_secs(2)))
            .unwrap();
        let (request, peer) = recv_request(&socket);
        let close = observe_close_request(&request, &key);
        tx.send(close).unwrap();

        let reply = close_reply(TOKEN, &key);
        socket.send_to(&reply, peer).unwrap();
    })
}

pub fn start_bad_hmac_echo_reply_server(params: Params, key: Vec<u8>) -> FakeServer {
    start_fake_server(move |socket, tx| {
        let (request, peer) = recv_request(&socket);
        let (open_params, hmac) = open_request_params(&request, Some(&key));
        tx.send(ServerObservation::Open {
            params: open_params,
            hmac,
        })
        .unwrap();

        let reply = open_reply(
            flags::FLAG_OPEN | flags::FLAG_REPLY,
            TOKEN,
            &params,
            Some(&key),
        );
        socket.send_to(&reply, peer).unwrap();

        socket
            .set_read_timeout(Some(Duration::from_secs(2)))
            .unwrap();
        let (request, peer) = recv_request(&socket);
        let echo = observe_echo_request(&request, Some(&key));
        let sequence = match echo {
            ServerObservation::Echo { sequence, .. } => sequence,
            _ => unreachable!(),
        };
        tx.send(echo).unwrap();

        let mut reply_packet = echo_reply_packet(
            TOKEN,
            sequence,
            &params,
            &TimestampFields::default(),
            Some(&key),
        );
        reply_packet[HMAC_OFFSET] ^= 0xff;
        socket.send_to(&reply_packet, peer).unwrap();
    })
}

pub fn start_bad_hmac_open_reply_server(
    params: Params,
    request_key: Vec<u8>,
    reply_key: Vec<u8>,
) -> FakeServer {
    start_fake_server(move |socket, tx| {
        let (request, peer) = recv_request(&socket);
        let (open_params, hmac) = open_request_params(&request, Some(&request_key));
        tx.send(ServerObservation::Open {
            params: open_params,
            hmac,
        })
        .unwrap();

        let reply = open_reply(
            flags::FLAG_OPEN | flags::FLAG_REPLY,
            TOKEN,
            &params,
            Some(&reply_key),
        );
        socket.send_to(&reply, peer).unwrap();
    })
}

pub fn server_fill(value: &str) -> Option<ServerFill> {
    Some(ServerFill {
        value: value.to_owned(),
    })
}

fn start_one_probe_server(
    params: Params,
    timestamps: TimestampFields,
    hmac_key: Option<Vec<u8>>,
) -> FakeServer {
    start_fake_server(move |socket, tx| {
        let (request, peer) = recv_request(&socket);
        let (open_params, hmac) = open_request_params(&request, hmac_key.as_deref());
        tx.send(ServerObservation::Open {
            params: open_params,
            hmac,
        })
        .unwrap();

        let reply = open_reply(
            flags::FLAG_OPEN | flags::FLAG_REPLY,
            TOKEN,
            &params,
            hmac_key.as_deref(),
        );
        socket.send_to(&reply, peer).unwrap();

        socket
            .set_read_timeout(Some(Duration::from_secs(2)))
            .unwrap();
        let (request, peer) = recv_request(&socket);
        let echo = observe_echo_request(&request, hmac_key.as_deref());
        let sequence = match echo {
            ServerObservation::Echo { sequence, .. } => sequence,
            _ => unreachable!(),
        };
        tx.send(echo).unwrap();

        let timestamps = materialize_wall_timestamps(timestamps);
        let reply_packet =
            echo_reply_packet(TOKEN, sequence, &params, &timestamps, hmac_key.as_deref());
        socket.send_to(&reply_packet, peer).unwrap();
    })
}

fn start_fake_server<F>(handler: F) -> FakeServer
where
    F: FnOnce(UdpSocket, mpsc::Sender<ServerObservation>) + Send + 'static,
{
    let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
    let addr = socket.local_addr().unwrap();
    let (tx, rx) = mpsc::channel();
    let done = thread::spawn(move || handler(socket, tx));
    FakeServer { addr, rx, done }
}

fn recv_request(socket: &UdpSocket) -> (Vec<u8>, SocketAddr) {
    let mut buf = [0_u8; 8192];
    let (size, peer) = socket.recv_from(&mut buf).unwrap();
    (buf[..size].to_vec(), peer)
}

fn recv_request_timeout(socket: &UdpSocket) -> Option<Vec<u8>> {
    let mut buf = [0_u8; 8192];
    match socket.recv_from(&mut buf) {
        Ok((size, _)) => Some(buf[..size].to_vec()),
        Err(_) => None,
    }
}

fn assert_started(outcome: OpenOutcome) -> NegotiatedParams {
    match outcome {
        OpenOutcome::Started { negotiated, .. } => negotiated,
        OpenOutcome::NoTestCompleted { .. } => panic!("expected started outcome"),
    }
}

/// Admits a request the way a compliant server does: structural decode first,
/// then HMAC presence policy, then authentication.
fn admit_request<'a>(packet: &'a [u8], hmac_key: Option<&[u8]>) -> DecodedRequest<'a> {
    let request = decode_request(packet).expect("fake peers only receive well-formed requests");
    assert_eq!(request.hmac_present, hmac_key.is_some());
    if let Some(key) = hmac_key {
        verify_packet_hmac(key, packet).unwrap();
    }
    request
}

fn open_request_params(packet: &[u8], hmac_key: Option<&[u8]>) -> (Params, bool) {
    let request = admit_request(packet, hmac_key);
    match request.kind {
        DecodedRequestKind::Open { params, .. } => {
            (Params::decode(params).unwrap(), request.hmac_present)
        }
        other => panic!("expected an open request, got {other:?}"),
    }
}

fn observe_echo_request(packet: &[u8], hmac_key: Option<&[u8]>) -> ServerObservation {
    let request = admit_request(packet, hmac_key);
    match request.kind {
        DecodedRequestKind::Echo {
            token, sequence, ..
        } => ServerObservation::Echo {
            len: packet.len(),
            hmac: request.hmac_present,
            token,
            sequence,
        },
        other => panic!("expected an echo request, got {other:?}"),
    }
}

fn observe_close_request(packet: &[u8], hmac_key: &[u8]) -> ServerObservation {
    let request = admit_request(packet, Some(hmac_key));
    match request.kind {
        DecodedRequestKind::Close { token } => ServerObservation::Close { hmac: true, token },
        other => panic!("expected a close request, got {other:?}"),
    }
}

// irtt-proto has no encode_close_reply/CloseReply API yet, so this stays a
// manual encoder. Migrate it once server-side close encoding is added.
fn close_reply(token: u64, hmac_key: &[u8]) -> Vec<u8> {
    let mut packet = Vec::new();
    packet.extend_from_slice(&MAGIC);
    packet.push(flags::FLAG_REPLY | flags::FLAG_CLOSE | flags::FLAG_HMAC);
    packet.extend_from_slice(&[0_u8; HMAC_SIZE]);
    packet.extend_from_slice(&token.to_le_bytes());
    compute_hmac_in_place(hmac_key, &mut packet, HMAC_OFFSET).unwrap();
    packet
}

fn open_reply(flags: u8, token: u64, params: &Params, hmac_key: Option<&[u8]>) -> Vec<u8> {
    encode_open_reply(
        &OpenReply {
            flags,
            token,
            params: params.clone(),
        },
        hmac_key,
    )
    .unwrap()
}

pub(super) fn test_echo_packet_len(hmac: bool, params: &Params) -> usize {
    echo_packet_len(hmac, params).expect("test params must have a representable length")
}

fn echo_reply_packet(
    token: u64,
    seq: u32,
    params: &Params,
    timestamps: &TimestampFields,
    hmac_key: Option<&[u8]>,
) -> Vec<u8> {
    let layout = PacketLayout::echo(hmac_key.is_some(), params);
    let reply = EchoReply {
        flags: flags::FLAG_REPLY,
        token,
        sequence: seq,
        recv_count: layout.recv_count.then_some(RECV_COUNT),
        recv_window: layout.recv_window.then_some(RECV_WINDOW),
        timestamps: TimestampFields {
            recv_wall: layout.recv_wall.then(|| timestamps.recv_wall.unwrap_or(0)),
            recv_mono: layout.recv_mono.then(|| timestamps.recv_mono.unwrap_or(0)),
            midpoint_wall: layout
                .midpoint_wall
                .then(|| timestamps.midpoint_wall.unwrap_or(0)),
            midpoint_mono: layout
                .midpoint_mono
                .then(|| timestamps.midpoint_mono.unwrap_or(0)),
            send_wall: layout.send_wall.then(|| timestamps.send_wall.unwrap_or(0)),
            send_mono: layout.send_mono.then(|| timestamps.send_mono.unwrap_or(0)),
        },
        payload: Vec::new(),
    };
    encode_echo_reply(&reply, params, hmac_key).unwrap()
}

fn materialize_wall_timestamps(mut timestamps: TimestampFields) -> TimestampFields {
    let now_ns: i64 = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_nanos()
        .try_into()
        .expect("current wall-clock time fits i64 nanoseconds");
    if let Some(offset) = timestamps.recv_wall {
        timestamps.recv_wall = Some(now_ns + offset);
    }
    if let Some(offset) = timestamps.midpoint_wall {
        timestamps.midpoint_wall = Some(now_ns + offset);
    }
    if let Some(offset) = timestamps.send_wall {
        timestamps.send_wall = Some(now_ns + offset);
    }
    timestamps
}