peerspeak/tests/file_transfer_loopback.rs

//! End-to-end loopback test for the chat file-transfer plane (`FILES_ALPN`).
//!
//! Spins up two real iroh endpoints on localhost (relay disabled, addresses
//! exchanged directly), registers the production [`FileRouter`] on each, serves a
//! multi-megabyte blob on one side, and fetches it from the other through the
//! real `serve_attachment`/`fetch_attachment` path.
//!
//! This is the regression guard for the "file fetch: read failed: connection
//! lost" bug: the serve handler used to return (and drop the connection) the
//! instant it called `finish()`, so the CONNECTION_CLOSE raced ahead of the
//! still-in-flight stream data and the fetcher's `read_to_end` aborted. A blob
//! large enough to span many packets makes that race deterministic — the fix
//! (waiting on `connection.closed()` before returning) keeps the link up until
//! the fetcher has the bytes.

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

use iroh::address_lookup::memory::MemoryLookup;
use iroh::endpoint::presets;
use iroh::protocol::Router;
use iroh::{Endpoint, RelayMode};

use peerspeak::files::{ChatAttachment, AttachmentKind};
use peerspeak::network::NetworkTransport;
use peerspeak::network::iroh_impl::{FileRouter, IrohTransport};
use peerspeak::protocol::FILES_ALPN;

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

async fn spawn_node() -> 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()));
    // Mirror production: a persistent FileRouter bound to this session's transport
    // is what the router accepts inbound file fetches on.
    let file_router = FileRouter::new();
    file_router.bind(&transport);
    let router = Router::builder(endpoint.clone())
        .accept(FILES_ALPN, file_router)
        .spawn();

    Node { endpoint, transport, _router: router, lookup }
}

/// A pseudo-random-ish multi-megabyte payload spanning many QUIC packets, so a
/// premature connection close on the serve side reliably corrupts/aborts the read.
fn big_blob() -> Vec<u8> {
    (0..(2 * 1024 * 1024u32))
        .map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)
        .collect()
}

#[tokio::test]
async fn loopback_attachment_round_trips_intact() {
    let server = spawn_node().await;
    let client = spawn_node().await;

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

    let server_id = server.endpoint.id();
    let client_id = client.endpoint.id();

    // The serve handler gates on room membership (the audio admission roster), so
    // the server must admit the client before it will answer the fetch.
    server.transport.admit_audio_sender(client_id);
    client.transport.admit_audio_sender(server_id);

    // The fetcher dials the retained full address; seed it so fetch_attachment
    // doesn't have to fall back to a bare-id lookup.
    client.transport.connect_peer(server.endpoint.addr()).await;

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

    let att = ChatAttachment {
        name: "exterior-landscape.jpg".to_string(),
        size: blob.len() as u64,
        kind: AttachmentKind::Image,
        id,
    };

    let fetched = tokio::time::timeout(
        Duration::from_secs(30),
        client.transport.fetch_attachment(server_id, &att),
    )
    .await
    .expect("fetch did not time out")
    .expect("fetch succeeded");

    assert_eq!(fetched.len(), blob.len(), "fetched the full blob");
    assert_eq!(fetched, blob, "fetched bytes match served bytes exactly");
}

#[tokio::test]
async fn loopback_unknown_id_reports_gone() {
    let server = spawn_node().await;
    let client = spawn_node().await;

    server.lookup.add_endpoint_info(client.endpoint.addr());
    client.lookup.add_endpoint_info(server.endpoint.addr());

    let server_id = server.endpoint.id();
    let client_id = client.endpoint.id();
    server.transport.admit_audio_sender(client_id);
    client.transport.admit_audio_sender(server_id);
    client.transport.connect_peer(server.endpoint.addr()).await;

    // Never served — the handler closes with an empty body and the fetcher must
    // surface that as an error, not hang or return empty bytes.
    let att = ChatAttachment {
        name: "missing.bin".to_string(),
        size: 4096,
        kind: AttachmentKind::File,
        id: [7u8; 32],
    };

    let result = tokio::time::timeout(
        Duration::from_secs(30),
        client.transport.fetch_attachment(server_id, &att),
    )
    .await
    .expect("fetch did not time out");

    assert!(result.is_err(), "unknown id should error, got {result:?}");
}