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::iroh_impl::{AudioRouter, IrohTransport};
use peerspeak::network::{ConnEvent, NetworkTransport};
use peerspeak::protocol::AUDIO_ALPN;

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 = Arc::new(IrohTransport::new(endpoint.clone()));
    // Mirror production: a persistent AudioRouter bound to this session's transport
    // is what the router accepts inbound audio links on.
    let audio_router = AudioRouter::new();
    audio_router.bind(&transport);
    let router = Router::builder(endpoint.clone())
        .accept(AUDIO_ALPN, audio_router)
        .spawn();

    Node {
        endpoint,
        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(Some(ConnEvent::Left(_))) => panic!("unexpected graceful Left during reconnect"),
            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();
    a.transport.admit_audio_sender(b_id);
    b.transport.admit_audio_sender(a_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.endpoint.addr()).await;
    b.transport.connect_peer(a.endpoint.addr()).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})"
    );
}

/// Connection transparency: over a real loopback link, `connection_stats()`
/// must report the peer's selected path as direct (relay disabled here), with
/// an IP remote address and counters that advance while audio flows — and the
/// `connstats::derive` seam must turn two such snapshots into badge info with
/// live rates.
#[tokio::test]
async fn connection_stats_report_a_direct_path_with_live_counters() {
    let a = spawn_node().await;
    let b = spawn_node().await;

    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();
    a.transport.admit_audio_sender(b_id);
    b.transport.admit_audio_sender(a_id);

    // Keep B's receive path subscribed like production (drained implicitly).
    let _b_rx = b.transport.receive_datagrams().await.expect("subscribe B");

    a.transport.connect_peer(b.endpoint.addr()).await;
    b.transport.connect_peer(a.endpoint.addr()).await;
    tokio::time::sleep(Duration::from_millis(500)).await;

    let snap = |stats: Vec<(iroh::EndpointId, peerspeak::network::PathSnapshot)>| {
        stats
            .into_iter()
            .find(|(id, _)| *id == b_id)
            .map(|(_, s)| s)
            .expect("peer B should appear in A's connection stats")
    };
    let s1 = snap(a.transport.connection_stats());
    assert!(!s1.is_relay, "loopback with relay disabled must be direct");
    assert!(
        s1.remote_addr.parse::<std::net::SocketAddr>().is_ok(),
        "direct path address should be ip:port, got {}",
        s1.remote_addr
    );

    // Stream real audio so the path counters move.
    let mut enc = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap();
    for seq in 0..25u32 {
        a.transport.broadcast(packet(&mut enc, seq));
        tokio::time::sleep(Duration::from_millis(5)).await;
    }

    let s2 = snap(a.transport.connection_stats());
    assert!(s2.tx_bytes > s1.tx_bytes, "sent bytes should advance");
    assert!(
        s2.tx_datagrams > s1.tx_datagrams,
        "sent datagrams should advance"
    );

    // The derivation seam turns the two snapshots into live badge info.
    let info = peerspeak::core::connstats::derive(Some(&s1), &s2, Duration::from_millis(200));
    assert!(!info.relay);
    assert_eq!(info.remote_addr, s2.remote_addr);
    assert!(info.rtt_ms < 1000, "localhost RTT should be sane");
    assert!(
        info.up_kbps
            .expect("same path + positive window has a rate")
            > 0.0,
        "audio was flowing, so the upstream rate must be non-zero"
    );
}

/// 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.endpoint.addr()).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"
    );
}

/// A graceful close (our goodbye code `1`) must report `Left` and NOT reconnect —
/// the opposite of `dialer_reconnects_after_link_drops`, which closes with code 0
/// (a transient drop). This is what lets a clicked "disconnect" evict the peer
/// promptly instead of leaving it stuck "reconnecting". Mirrors the reconnect
/// test's harness but asserts the supervisor stops rather than re-dials.
#[tokio::test]
async fn dialer_reports_left_on_graceful_close() {
    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;

    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.endpoint.addr()).await;

    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"
    );

    // Close with the goodbye code (1) — an intentional leave, not a blip.
    conn1.close(1u32.into(), b"leave");
    drop(conn1);

    // The supervisor must surface `Left` (and must NOT try to reconnect: any
    // Connecting/Connected here would be the bug). The peer also must not see a
    // re-dial.
    let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
    match tokio::time::timeout_at(deadline, conn_events.recv()).await {
        Ok(Some(ConnEvent::Left(id))) => {
            assert_eq!(id, peer.endpoint.id(), "Left should name the peer");
        }
        Ok(Some(ConnEvent::Connecting(_))) | Ok(Some(ConnEvent::Connected(_))) => {
            panic!("a graceful close must not trigger a reconnect");
        }
        Ok(None) => panic!("conn events channel closed unexpectedly"),
        Err(_) => panic!("timed out waiting for Left after graceful close"),
    }

    // And no re-dial reaches the peer within a short window.
    let redial = tokio::time::timeout(Duration::from_secs(2), peer.conns_rx.recv()).await;
    assert!(
        redial.is_err(),
        "supervisor must not re-dial after a graceful leave"
    );
}

#[tokio::test]
async fn dialer_connects_and_reconnects_without_an_address_lookup() {
    // Regression for the redial-loop bug: a transient drop that also triggers a
    // gossip Leave/NeighborDown purges the peer from the address lookup, so a
    // dialer that resolves by bare id (the old behavior) has nothing to dial and
    // loops forever. The fix retains the peer's full address in the transport and
    // dials it directly. To prove the retained address — not the lookup — is what
    // carries the dial, we never seed the dialer's lookup with the peer at all
    // (and disable the relay), so a bare-id dial could not resolve an address for
    // either the initial connect or the reconnect; only `connect_peer`'s retained
    // EndpointAddr makes them possible.
    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());

    // Deliberately leave the dialer's lookup empty for the peer. (The peer side
    // only ever accepts, so it needs no addressing for the dialer.)
    let mut conn_events = dialer
        .transport
        .subscribe_conn_events()
        .await
        .expect("subscribe dialer conn events");

    // The full address handed here is the dialer's only way to reach the peer.
    dialer.transport.connect_peer(peer.endpoint.addr()).await;

    let conn1 = tokio::time::timeout(Duration::from_secs(10), peer.conns_rx.recv())
        .await
        .expect("timed out awaiting initial connection (retained address path)")
        .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 from the peer side so the dialer must reconnect — again
    // with no lookup to fall back on.
    conn1.close(0u32.into(), b"transient drop");
    drop(conn1);

    let conn2 = tokio::time::timeout(Duration::from_secs(15), peer.conns_rx.recv())
        .await
        .expect("timed out awaiting reconnect (retained address path)")
        .expect("connection channel closed");
    assert!(
        await_reconnect(
            &mut conn_events,
            tokio::time::Instant::now() + Duration::from_secs(15)
        )
        .await,
        "reconnect should report Connecting then Connected with no lookup at all"
    );

    // Confirm the rebuilt link actually carries audio.
    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 reconnecting via the retained address; got {received} frames"
    );
}

/// A node that also serves the file plane (`FILES_ALPN`), mirroring how core
/// registers the `FileRouter` for a session.
async fn spawn_file_server() -> Node {
    let lookup = MemoryLookup::new();
    let endpoint = Endpoint::builder(presets::Minimal)
        .secret_key(iroh::SecretKey::generate())
        .relay_mode(RelayMode::Disabled)
        .address_lookup(lookup.clone())
        .bind()
        .await
        .expect("bind endpoint");
    let transport = Arc::new(IrohTransport::new(endpoint.clone()));
    let audio_router = AudioRouter::new();
    audio_router.bind(&transport);
    let file_router = peerspeak::network::iroh_impl::FileRouter::new();
    file_router.bind(&transport);
    let router = Router::builder(endpoint.clone())
        .accept(AUDIO_ALPN, audio_router)
        .accept(peerspeak::protocol::FILES_ALPN, file_router)
        .spawn();
    Node {
        endpoint,
        transport,
        _router: router,
        lookup,
    }
}

/// Phase 3C: a file fetch must deliver EXACTLY the declared size — short,
/// overlong, and unknown-id transfers are all rejected with local errors, and
/// an exact transfer round-trips byte-identically.
#[tokio::test]
async fn file_plane_requires_exact_declared_size() {
    let fetcher = spawn_node().await;
    let server = spawn_file_server().await;
    fetcher.lookup.add_endpoint_info(server.endpoint.addr());
    server.lookup.add_endpoint_info(fetcher.endpoint.addr());

    let server_id = server.endpoint.id();
    // Member gating: the server only serves current room members.
    server.transport.admit_audio_sender(fetcher.endpoint.id());

    let blob = vec![42u8; 1000];
    let id = [7u8; 32];
    server
        .transport
        .serve_attachment(id, Arc::new(blob.clone()));

    // Exact declared size: byte-identical round trip.
    let got = fetcher
        .transport
        .fetch_blob(server_id, id, 1000)
        .await
        .expect("exact-size fetch succeeds");
    assert_eq!(got, blob);

    // Declared larger than served (short transfer): rejected, not cached as-is.
    let err = fetcher
        .transport
        .fetch_blob(server_id, id, 2000)
        .await
        .expect_err("short transfer must fail");
    assert!(
        err.to_string().contains("incomplete transfer"),
        "unexpected error: {err}"
    );

    // Declared smaller than served (overlong transfer): the bounded read
    // rejects the stream rather than truncating it into a "valid" result.
    let err = fetcher
        .transport
        .fetch_blob(server_id, id, 500)
        .await
        .expect_err("overlong transfer must fail");
    assert!(
        err.to_string().contains("read failed"),
        "unexpected error: {err}"
    );

    // Unknown id: the empty body reads as the sender no longer having it.
    let err = fetcher
        .transport
        .fetch_blob(server_id, [9u8; 32], 1000)
        .await
        .expect_err("unknown id must fail");
    assert!(
        err.to_string().contains("no longer has the file"),
        "unexpected error: {err}"
    );
}