peerspeak/tests/transport_loopback.rs

//! End-to-end loopback test for the redesigned audio transport.
//!
//! Spins up two real iroh endpoints on localhost (relay disabled, addresses
//! exchanged directly) and drives the actual production path: reactive
//! `connect_peer`, sequenced `broadcast`, `receive_datagrams`, and the
//! per-peer `JitterBuffer` decode. No microphone, speakers, or GUI required.

use std::sync::Arc;
use std::time::Duration;

use tokio::sync::mpsc;

use bytes::Bytes;
use iroh::address_lookup::memory::MemoryLookup;
use iroh::endpoint::presets;
use iroh::protocol::Router;
use iroh::{Endpoint, RelayMode};
use opus::{Application, Channels};

use iroh::endpoint::Connection;
use iroh::protocol::{AcceptError, ProtocolHandler};

use peerspeak::codec::AudioEncoder;
use peerspeak::codec::opus_impl::OpusEncoder;
use peerspeak::core::jitter::{FRAME_SAMPLES, JitterBuffer};
use peerspeak::network::{ConnEvent, NetworkTransport};
use peerspeak::network::iroh_impl::IrohTransport;

const AUDIO_ALPN: &[u8] = b"peerspeak-audio";

struct Node {
    endpoint: Endpoint,
    transport: Arc<IrohTransport>,
    _router: Router,
    lookup: MemoryLookup,
}

async fn spawn_node() -> Node {
    spawn_node_with_key(iroh::SecretKey::generate()).await
}

/// A minimal protocol handler for the reconnect test's "peer" side: it hands
/// every accepted audio connection out through a channel so the test can hold
/// it, read datagrams from it, and explicitly `close()` it to drive a drop —
/// all on a single, stable endpoint address that the dialer can re-reach.
#[derive(Clone)]
struct CaptureProtocol {
    conns_tx: mpsc::UnboundedSender<Connection>,
}

impl std::fmt::Debug for CaptureProtocol {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CaptureProtocol").finish_non_exhaustive()
    }
}

impl ProtocolHandler for CaptureProtocol {
    fn accept(
        &self,
        connection: Connection,
    ) -> impl std::future::Future<Output = Result<(), AcceptError>> + Send {
        let conns_tx = self.conns_tx.clone();
        async move {
            // The channel retains the connection, which keeps it open after this
            // future resolves (iroh closes accepted connections with no retained
            // handle).
            let _ = conns_tx.send(connection);
            Ok(())
        }
    }
}

struct CapturePeer {
    endpoint: Endpoint,
    _router: Router,
    lookup: MemoryLookup,
    conns_rx: mpsc::UnboundedReceiver<Connection>,
}

async fn spawn_capture_peer(secret: iroh::SecretKey) -> CapturePeer {
    let lookup = MemoryLookup::new();
    let endpoint = Endpoint::builder(presets::Minimal)
        .secret_key(secret)
        .relay_mode(RelayMode::Disabled)
        .address_lookup(lookup.clone())
        .bind()
        .await
        .expect("bind capture peer");

    let (conns_tx, conns_rx) = mpsc::unbounded_channel();
    let router = Router::builder(endpoint.clone())
        .accept(AUDIO_ALPN, CaptureProtocol { conns_tx })
        .spawn();

    CapturePeer { endpoint, _router: router, lookup, conns_rx }
}

/// Spawn a node with a specific secret key. Reusing a key gives the respawned
/// node the same `EndpointId`, which is how we simulate a peer dropping off the
/// network and coming back (new address, same identity) for the reconnect test.
async fn spawn_node_with_key(secret: iroh::SecretKey) -> Node {
    let lookup = MemoryLookup::new();
    let endpoint = Endpoint::builder(presets::Minimal)
        .secret_key(secret)
        // Direct-only: two endpoints on the same host reach each other via the
        // addresses we seed into each other's lookup, with no external relay.
        .relay_mode(RelayMode::Disabled)
        .address_lookup(lookup.clone())
        .bind()
        .await
        .expect("bind endpoint");

    let (transport, audio_proto) = IrohTransport::new(endpoint.clone());
    let router = Router::builder(endpoint.clone())
        .accept(AUDIO_ALPN, audio_proto)
        .spawn();

    Node {
        endpoint,
        transport: Arc::new(transport),
        _router: router,
        lookup,
    }
}

/// One real, decodable Opus packet for a 20ms mono frame, prefixed with the
/// 4-byte little-endian sequence header the transport/jitter buffer expect.
fn packet(enc: &mut OpusEncoder, seq: u32) -> Bytes {
    let pcm: Vec<i16> = (0..FRAME_SAMPLES)
        .map(|i| if i % 2 == 0 { 2000 } else { -2000 })
        .collect();
    let encoded = enc.encode(&pcm).unwrap();
    let mut buf = Vec::with_capacity(4 + encoded.len());
    buf.extend_from_slice(&seq.to_le_bytes());
    buf.extend_from_slice(&encoded);
    Bytes::from(buf)
}

/// Drain link-state events until a `Connected` arrives, returning whether a
/// `Connecting` was seen first (the down→up transition). Panics on timeout so a
/// stuck reconnect fails loudly rather than hanging the suite.
async fn await_reconnect(
    rx: &mut mpsc::Receiver<ConnEvent>,
    deadline: tokio::time::Instant,
) -> bool {
    let mut saw_connecting = false;
    loop {
        match tokio::time::timeout_at(deadline, rx.recv()).await {
            Ok(Some(ConnEvent::Connecting(_))) => saw_connecting = true,
            Ok(Some(ConnEvent::Connected(_))) => return saw_connecting,
            Ok(None) => panic!("conn events channel closed unexpectedly"),
            Err(_) => panic!("timed out waiting for Connected"),
        }
    }
}

#[tokio::test]
async fn loopback_sequenced_audio_reaches_peer_and_decodes() {
    let a = spawn_node().await;
    let b = spawn_node().await;

    // Seed each side with the other's full address so direct dialing works.
    a.lookup.add_endpoint_info(b.endpoint.addr());
    b.lookup.add_endpoint_info(a.endpoint.addr());

    let a_id = a.endpoint.id();
    let b_id = b.endpoint.id();

    // Subscribe to incoming datagrams on B before any are sent.
    let mut b_rx = b.transport.receive_datagrams().await.expect("subscribe B");

    // Reactive connection setup, exactly as core does on peer-join. Calling on
    // both sides is fine: the lower EndpointId dials, the higher accepts, and a
    // single shared connection forms.
    a.transport.connect_peer(b_id).await;
    b.transport.connect_peer(a_id).await;

    // Let the dial + accept registration settle.
    tokio::time::sleep(Duration::from_millis(500)).await;

    // A sends 50 sequenced frames.
    const N: u32 = 50;
    let mut enc = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap();
    for seq in 0..N {
        a.transport.broadcast(packet(&mut enc, seq));
        tokio::time::sleep(Duration::from_millis(5)).await;
    }

    // Collect what B receives and feed it through a real jitter buffer,
    // popping a frame per arrival to mirror the mixer's steady 20ms cadence
    // (so the buffer stays shallow rather than overflowing its cap).
    let mut jitter = JitterBuffer::new().unwrap();
    let mut received = 0u32;
    let mut decoded_frames = 0u32;
    let deadline = tokio::time::Instant::now() + Duration::from_secs(2);
    // Loop ends when the channel closes or the deadline is hit (pattern stops matching).
    while let Ok(Some((from, bytes))) = tokio::time::timeout_at(deadline, b_rx.recv()).await {
        assert_eq!(from, a_id, "datagram should be attributed to sender A");
        assert!(bytes.len() >= 4, "packet carries a sequence header");
        let seq = u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
        jitter.insert(seq, bytes[4..].to_vec());
        received += 1;

        if let Some(frame) = jitter.pop_frame() {
            assert_eq!(frame.len(), FRAME_SAMPLES, "decoded frame is one 20ms frame");
            decoded_frames += 1;
        }
        if received >= N {
            break;
        }
    }

    // Drain whatever remains buffered behind the playout delay.
    while let Some(frame) = jitter.pop_frame() {
        assert_eq!(frame.len(), FRAME_SAMPLES);
        decoded_frames += 1;
    }

    // On localhost essentially nothing should be lost over a real QUIC datagram path.
    assert!(
        received >= N - 2,
        "expected to receive ~{N} datagrams, got {received}"
    );
    // And nearly all received packets should decode to PCM (a few absorbed by
    // the initial priming delay).
    assert!(
        decoded_frames >= N - 5,
        "expected to decode ~{N} frames, got {decoded_frames} (received {received})"
    );
}

/// Read datagrams off a raw connection until `target` arrive or the deadline
/// passes, asserting each carries the 4-byte sequence header.
async fn count_audio(conn: &Connection, target: u32, deadline: tokio::time::Instant) -> u32 {
    let mut received = 0;
    while let Ok(Ok(bytes)) = tokio::time::timeout_at(deadline, conn.read_datagram()).await {
        assert!(bytes.len() >= 4, "packet carries a sequence header");
        received += 1;
        if received >= target {
            break;
        }
    }
    received
}

#[tokio::test]
async fn dialer_reconnects_after_link_drops() {
    // The unit under test is the IrohTransport dialer. The peer is a hand-rolled
    // endpoint we fully control, so we can explicitly close the live connection
    // (which the dialer observes promptly, unlike a silent handle drop) and then
    // accept the dialer's re-dial — all on one stable address it can re-reach.
    //
    // The dialer must hold the lower id so the transport's deterministic rule
    // makes it dial, so order the two identities up front.
    let s1 = iroh::SecretKey::generate();
    let s2 = iroh::SecretKey::generate();
    let (dialer_secret, peer_secret) = if s1.public().to_string() < s2.public().to_string() {
        (s1, s2)
    } else {
        (s2, s1)
    };

    let dialer = spawn_node_with_key(dialer_secret).await;
    let mut peer = spawn_capture_peer(peer_secret).await;

    let dialer_id = dialer.endpoint.id();
    let peer_id = peer.endpoint.id();
    assert!(dialer_id.to_string() < peer_id.to_string());

    // Seed addresses both ways and form the initial link.
    dialer.lookup.add_endpoint_info(peer.endpoint.addr());
    peer.lookup.add_endpoint_info(dialer.endpoint.addr());

    let mut conn_events = dialer
        .transport
        .subscribe_conn_events()
        .await
        .expect("subscribe dialer conn events");
    dialer.transport.connect_peer(peer_id).await;

    // The peer accepts the dialer's initial connection.
    let conn1 = tokio::time::timeout(Duration::from_secs(10), peer.conns_rx.recv())
        .await
        .expect("timed out awaiting initial connection")
        .expect("connection channel closed");
    assert!(
        await_reconnect(&mut conn_events, tokio::time::Instant::now() + Duration::from_secs(10)).await,
        "initial link should report Connecting then Connected"
    );

    // Drop the live link by closing it from the peer side. An explicit close
    // sends a CONNECTION_CLOSE the dialer sees right away.
    conn1.close(0u32.into(), b"transient drop");
    drop(conn1);

    // The dialer's supervisor should re-dial the (unchanged) peer address, and
    // the peer accepts the fresh connection.
    let conn2 = tokio::time::timeout(Duration::from_secs(15), peer.conns_rx.recv())
        .await
        .expect("timed out awaiting reconnect")
        .expect("connection channel closed");
    assert!(
        await_reconnect(&mut conn_events, tokio::time::Instant::now() + Duration::from_secs(15)).await,
        "dropped link should report Connecting (down) then Connected (recovered)"
    );

    // Audio must flow over the rebuilt link. Give the freshly-registered send
    // path a beat, then stream frames and confirm they reach the peer on conn2.
    tokio::time::sleep(Duration::from_millis(200)).await;
    let mut enc = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap();
    for seq in 0..30u32 {
        dialer.transport.broadcast(packet(&mut enc, seq));
        tokio::time::sleep(Duration::from_millis(5)).await;
    }

    let received = count_audio(&conn2, 25, tokio::time::Instant::now() + Duration::from_secs(3)).await;
    assert!(
        received >= 20,
        "audio should resume after reconnect; got {received} frames"
    );
}