peerspeak/tests/gossip_rebootstrap.rs

//! Phase-0 spike for post-grace gossip recovery.
//!
//! These tests prove that `GossipSender::join_peers` can restore an existing
//! topic subscription after the other peer drops and rejoins without its own
//! bootstrap target. The second case disables relays, clears the surviving
//! node's lookup, and moves the peer to a fresh endpoint address so only the
//! retained full address passed to `rebootstrap_peers` can drive recovery.

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, EndpointAddr, EndpointId, RelayMode, SecretKey};
use iroh_gossip::net::Gossip;
use tokio::sync::mpsc;

use peerspeak::network::gossip::IrohGossipState;
use peerspeak::network::{PeerSpeakTicket, PeerState, RoomEvent, RoomState};

const EVENT_TIMEOUT: Duration = Duration::from_secs(10);

struct GossipNode {
    endpoint: Endpoint,
    lookup: MemoryLookup,
    room: Arc<IrohGossipState>,
    _router: Router,
}

async fn spawn_node(secret: SecretKey) -> GossipNode {
    let lookup = MemoryLookup::new();
    let endpoint = Endpoint::builder(presets::Minimal)
        .secret_key(secret.clone())
        .relay_mode(RelayMode::Disabled)
        .address_lookup(lookup.clone())
        .bind()
        .await
        .expect("bind gossip endpoint");
    let gossip = Gossip::builder().spawn(endpoint.clone());
    let router = Router::builder(endpoint.clone())
        .accept(iroh_gossip::net::GOSSIP_ALPN, gossip.clone())
        .spawn();
    let room = Arc::new(IrohGossipState::new(
        endpoint.clone(),
        gossip,
        lookup.clone(),
        secret,
    ));

    GossipNode {
        endpoint,
        lookup,
        room,
        _router: router,
    }
}

fn state(name: &str, addr: EndpointAddr) -> PeerState {
    PeerState {
        name: name.to_string(),
        is_muted: false,
        addr,
        sharing: None,
        avatar: Default::default(),
        game: None,
        music: None,
    }
}

async fn await_joined(rx: &mut mpsc::Receiver<RoomEvent>, peer_id: EndpointId) -> PeerState {
    let deadline = tokio::time::Instant::now() + EVENT_TIMEOUT;
    loop {
        match tokio::time::timeout_at(deadline, rx.recv()).await {
            Ok(Some(RoomEvent::PeerJoined(id, peer_state))) if id == peer_id => return peer_state,
            Ok(Some(_)) => continue,
            Ok(None) => panic!("room event channel closed while waiting for PeerJoined"),
            Err(_) => panic!("timed out waiting for PeerJoined({peer_id:?})"),
        }
    }
}

async fn await_joined_all(rx: &mut mpsc::Receiver<RoomEvent>, peer_ids: &[EndpointId]) {
    let deadline = tokio::time::Instant::now() + EVENT_TIMEOUT;
    let mut remaining = peer_ids.to_vec();
    while !remaining.is_empty() {
        match tokio::time::timeout_at(deadline, rx.recv()).await {
            Ok(Some(RoomEvent::PeerJoined(id, _))) => remaining.retain(|wanted| *wanted != id),
            Ok(Some(_)) => {}
            Ok(None) => panic!("room event channel closed while waiting for PeerJoined set"),
            Err(_) => panic!("timed out waiting for PeerJoined set: {remaining:?}"),
        }
    }
}

async fn await_left(rx: &mut mpsc::Receiver<RoomEvent>, peer_id: EndpointId) {
    let deadline = tokio::time::Instant::now() + EVENT_TIMEOUT;
    loop {
        match tokio::time::timeout_at(deadline, rx.recv()).await {
            Ok(Some(RoomEvent::PeerLeft(id))) if id == peer_id => return,
            Ok(Some(_)) => continue,
            Ok(None) => panic!("room event channel closed while waiting for PeerLeft"),
            Err(_) => panic!("timed out waiting for PeerLeft({peer_id:?})"),
        }
    }
}

async fn await_absent(room: &IrohGossipState, peer_id: EndpointId) {
    let deadline = tokio::time::Instant::now() + EVENT_TIMEOUT;
    loop {
        if !room.active_peers().iter().any(|(id, _)| *id == peer_id) {
            return;
        }
        assert!(
            tokio::time::Instant::now() < deadline,
            "timed out waiting for peer to leave the roster"
        );
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
}

fn ticket(host_addr: EndpointAddr) -> String {
    PeerSpeakTicket {
        host_addr,
        topic_id: rand::random(),
        name: "rebootstrap-spike".to_string(),
    }
    .to_string()
}

async fn establish_room(
    a: &GossipNode,
    b: &GossipNode,
    ticket: &str,
    events_a: &mut mpsc::Receiver<RoomEvent>,
) {
    // B is the ticket host. Its own bootstrap set is empty; A is the only side
    // that initially dials, which is also how the recovery setup is controlled.
    b.room
        .join(ticket, state("Bob", b.endpoint.addr()), vec![])
        .await
        .expect("host joins topic");
    a.room
        .join(ticket, state("Alice", a.endpoint.addr()), vec![])
        .await
        .expect("client joins topic");
    let joined = await_joined(events_a, b.endpoint.id()).await;
    assert_eq!(joined.name, "Bob");
}

#[tokio::test]
async fn rebootstrap_restores_roster_on_existing_subscription() {
    let a = spawn_node(SecretKey::generate()).await;
    let b = spawn_node(SecretKey::generate()).await;
    let ticket = ticket(b.endpoint.addr());
    let mut events_a = a.room.subscribe_events().await.expect("subscribe A events");

    establish_room(&a, &b, &ticket, &mut events_a).await;

    // Drop only B's topic subscription. A stays subscribed. B then rejoins as
    // the ticket host, so compute_bootstrap removes self and B has nobody to dial.
    b.room.leave().await.expect("B leaves topic");
    await_absent(&a.room, b.endpoint.id()).await;
    b.room
        .join(&ticket, state("Bob recovered", b.endpoint.addr()), vec![])
        .await
        .expect("B rejoins without bootstrap peers");

    tokio::time::sleep(Duration::from_millis(250)).await;
    assert!(
        !a.room
            .active_peers()
            .iter()
            .any(|(id, _)| *id == b.endpoint.id()),
        "B must not recover before A explicitly re-bootstraps it"
    );

    a.room
        .rebootstrap_peers(vec![b.endpoint.addr()])
        .await
        .expect("targeted gossip re-bootstrap");

    let recovered = await_joined(&mut events_a, b.endpoint.id()).await;
    assert_eq!(recovered.name, "Bob recovered");
}

#[tokio::test]
async fn rebootstrap_uses_retained_full_address_with_empty_lookup() {
    let a = spawn_node(SecretKey::generate()).await;
    let b_secret = SecretKey::generate();
    let b = spawn_node(b_secret.clone()).await;
    let b_id = b.endpoint.id();
    let old_b_addr = b.endpoint.addr();
    let ticket = ticket(old_b_addr.clone());
    let mut events_a = a.room.subscribe_events().await.expect("subscribe A events");

    establish_room(&a, &b, &ticket, &mut events_a).await;
    b.room.leave().await.expect("old B leaves topic");
    await_absent(&a.room, b_id).await;

    // Move the same authenticated identity to a newly-bound direct-only endpoint.
    // The old cached path is now dead; the new full address is the only valid one.
    b.endpoint.close().await;
    drop(b);
    tokio::time::sleep(Duration::from_millis(250)).await;
    let b_rebound = spawn_node(b_secret).await;
    let new_b_addr = b_rebound.endpoint.addr();
    assert_eq!(new_b_addr.id, b_id, "identity must survive the rebind");
    assert_ne!(
        new_b_addr, old_b_addr,
        "rebound peer must have a fresh address"
    );

    b_rebound
        .room
        .join(&ticket, state("Bob rebound", new_b_addr.clone()), vec![])
        .await
        .expect("rebound host joins without bootstrap peers");

    // Remove the stale lookup entry. `rebootstrap_peers` must seed the retained
    // new full address before asking gossip to join the peer by id.
    a.lookup.remove_endpoint_info(b_id);
    assert!(
        a.lookup.get_endpoint_info(b_id).is_none(),
        "A lookup starts empty for B"
    );
    tokio::time::sleep(Duration::from_millis(250)).await;
    assert!(
        !a.room.active_peers().iter().any(|(id, _)| *id == b_id),
        "rebound B must not be rediscovered without the retained address"
    );

    a.room
        .rebootstrap_peers(vec![new_b_addr])
        .await
        .expect("retained-address gossip re-bootstrap");
    assert!(
        a.lookup.get_endpoint_info(b_id).is_some(),
        "re-bootstrap must restore B's address to the lookup"
    );

    let recovered = await_joined(&mut events_a, b_id).await;
    assert_eq!(recovered.name, "Bob rebound");
}

#[tokio::test]
async fn demoted_peer_requires_a_fresh_signed_announce_to_rejoin() {
    let a = spawn_node(SecretKey::generate()).await;
    let b = spawn_node(SecretKey::generate()).await;
    let ticket = ticket(b.endpoint.addr());
    let mut events_a = a.room.subscribe_events().await.expect("subscribe A events");

    establish_room(&a, &b, &ticket, &mut events_a).await;
    a.room.mark_peer_disconnected(b.endpoint.id());
    assert!(
        !a.room
            .active_peers()
            .iter()
            .any(|(id, _)| *id == b.endpoint.id()),
        "demotion must revoke live roster membership"
    );

    b.room
        .update_self_state(state("Bob authenticated again", b.endpoint.addr()))
        .await
        .expect("broadcast fresh signed announce");

    let recovered = await_joined(&mut events_a, b.endpoint.id()).await;
    assert_eq!(recovered.name, "Bob authenticated again");
}

#[tokio::test]
async fn signed_leave_after_demotion_still_emits_peer_left() {
    let a = spawn_node(SecretKey::generate()).await;
    let b = spawn_node(SecretKey::generate()).await;
    let ticket = ticket(b.endpoint.addr());
    let mut events_a = a.room.subscribe_events().await.expect("subscribe A events");

    establish_room(&a, &b, &ticket, &mut events_a).await;
    a.room.mark_peer_disconnected(b.endpoint.id());

    b.room
        .leave()
        .await
        .expect("broadcast signed Leave after demotion");
    await_left(&mut events_a, b.endpoint.id()).await;
}

#[tokio::test]
async fn targeted_rebootstrap_preserves_healthy_peer_in_three_peer_room() {
    let a = spawn_node(SecretKey::generate()).await;
    let b = spawn_node(SecretKey::generate()).await;
    let c_secret = SecretKey::generate();
    let c = spawn_node(c_secret.clone()).await;
    let c_id = c.endpoint.id();
    let ticket = ticket(c.endpoint.addr());
    let mut events_a = a.room.subscribe_events().await.expect("subscribe A events");
    let mut events_b = b.room.subscribe_events().await.expect("subscribe B events");

    c.room
        .join(&ticket, state("Carol", c.endpoint.addr()), vec![])
        .await
        .expect("C hosts topic");
    b.room
        .join(&ticket, state("Bob", b.endpoint.addr()), vec![])
        .await
        .expect("B joins C");
    await_joined(&mut events_b, c_id).await;
    a.room
        .join(&ticket, state("Alice", a.endpoint.addr()), vec![])
        .await
        .expect("A joins C");
    await_joined_all(&mut events_a, &[b.endpoint.id(), c_id]).await;
    await_joined(&mut events_b, a.endpoint.id()).await;

    // Remove only C. A and B keep their existing topic subscriptions and remain
    // mutually present while C is rebound to a fresh address.
    c.endpoint.close().await;
    drop(c);
    a.room.mark_peer_disconnected(c_id);
    b.room.mark_peer_disconnected(c_id);
    let c_rebound = spawn_node(c_secret).await;
    let rebound_addr = c_rebound.endpoint.addr();
    c_rebound
        .room
        .join(
            &ticket,
            state("Carol recovered", rebound_addr.clone()),
            vec![],
        )
        .await
        .expect("rebound C rejoins as host without bootstrap");

    a.room
        .rebootstrap_peers(vec![rebound_addr])
        .await
        .expect("A targets only C for recovery");
    let recovered = await_joined(&mut events_a, c_id).await;
    assert_eq!(recovered.name, "Carol recovered");
    await_joined(&mut events_b, c_id).await;

    assert!(
        a.room
            .active_peers()
            .iter()
            .any(|(id, _)| *id == b.endpoint.id()),
        "healthy B must remain present at A throughout C recovery"
    );
    assert!(
        b.room
            .active_peers()
            .iter()
            .any(|(id, _)| *id == a.endpoint.id()),
        "healthy A must remain present at B throughout C recovery"
    );
}