tsunagi/tests/discovery.rs

//! Discovery backends: static bootstrap candidates, composition, and the fact
//! that discovery only ever supplies *candidates*.

#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]

mod common;

use std::sync::Arc;

use common::{TestAgent, local_config, network, settle, wait_for_peers};
use tsunagi::Agent;
use tsunagi::discovery::{
    CandidateSource, CompositeDiscovery, NetworkDiscovery, SharedMemoryDiscovery, StaticBootstrap,
};
use tsunagi::identity::NetworkKeys;

#[tokio::test]
async fn a_static_bootstrap_candidate_is_enough_to_join() {
    let (name, secret) = network("bootstrap-only");

    // The listener runs with no discovery at all: it only accepts.
    let listener_dir = tempfile::TempDir::new().unwrap();
    let listener = Agent::spawn(local_config(listener_dir.path()))
        .await
        .unwrap();
    let network_id = listener.join_network(&name, &secret).await.unwrap();

    // The joiner is handed the listener's iroh id and addresses up front, which
    // is exactly what a static bootstrap entry is.
    let bootstrap = Arc::new(StaticBootstrap::new([listener.local_addr()]));
    let joiner_dir = tempfile::TempDir::new().unwrap();
    let joiner = Agent::spawn(local_config(joiner_dir.path()).with_discovery(bootstrap))
        .await
        .unwrap();
    joiner.join_network(&name, &secret).await.unwrap();

    wait_for_peers(&joiner, network_id, 1).await;
    wait_for_peers(&listener, network_id, 1).await;

    let status = joiner.network_status(network_id).await.unwrap();
    assert_eq!(
        status
            .candidates
            .iter()
            .find(|candidate| candidate.endpoint_id == listener.endpoint_id())
            .map(|candidate| candidate.source),
        Some(CandidateSource::Bootstrap)
    );

    joiner.shutdown().await;
    listener.shutdown().await;
    drop(joiner);
    drop(listener);
    drop(joiner_dir);
    drop(listener_dir);
}

#[tokio::test]
async fn a_composite_backend_merges_its_sources() {
    let (name, secret) = network("composite");
    let keys = NetworkKeys::derive(&name, &secret);

    let shared = SharedMemoryDiscovery::new();
    let via_shared = TestAgent::spawn(&shared).await.unwrap();
    let network_id = via_shared.agent.join_network(&name, &secret).await.unwrap();

    let bootstrap_dir = tempfile::TempDir::new().unwrap();
    let via_bootstrap = Agent::spawn(local_config(bootstrap_dir.path()))
        .await
        .unwrap();
    via_bootstrap.join_network(&name, &secret).await.unwrap();

    // One backend knows the bootstrap peer, the other knows the shared-table
    // peer. Composed, the joiner reaches both.
    let composite = Arc::new(CompositeDiscovery::new([
        Arc::new(StaticBootstrap::new([via_bootstrap.local_addr()])) as Arc<dyn NetworkDiscovery>,
        Arc::new(shared.clone()) as Arc<dyn NetworkDiscovery>,
    ]));
    let joiner_dir = tempfile::TempDir::new().unwrap();
    let joiner = Agent::spawn(local_config(joiner_dir.path()).with_discovery(composite))
        .await
        .unwrap();
    joiner.join_network(&name, &secret).await.unwrap();

    wait_for_peers(&joiner, network_id, 2).await;

    // Publishing went to the shared backend, so the other agent finds us too.
    assert!(shared.len(&keys.discovery_key()) >= 2);

    joiner.shutdown().await;
    via_bootstrap.shutdown().await;
    via_shared.agent.shutdown().await;
    drop(joiner);
    drop(via_bootstrap);
    drop(joiner_dir);
    drop(bootstrap_dir);
}

#[tokio::test]
async fn discovery_entries_are_withdrawn_when_a_network_stops() {
    let discovery = SharedMemoryDiscovery::new();
    let (name, secret) = network("withdrawn");
    let keys = NetworkKeys::derive(&name, &secret);

    let agent = TestAgent::spawn(&discovery).await.unwrap();
    let network_id = agent.agent.join_network(&name, &secret).await.unwrap();

    settle().await;
    assert_eq!(discovery.len(&keys.discovery_key()), 1);

    agent.agent.deactivate_network(network_id).await.unwrap();
    assert!(discovery.is_empty(&keys.discovery_key()));

    agent.agent.shutdown().await;
}

#[tokio::test]
async fn forgetting_a_network_removes_it_from_the_state_store() {
    let discovery = SharedMemoryDiscovery::new();
    let (name, secret) = network("forgettable");

    let agent = TestAgent::spawn(&discovery).await.unwrap();
    let network_id = agent.agent.join_network(&name, &secret).await.unwrap();
    assert_eq!(agent.agent.list_networks().await.unwrap().len(), 1);

    agent.agent.forget_network(network_id).await.unwrap();
    assert!(agent.agent.list_networks().await.unwrap().is_empty());
    assert!(!agent.agent.is_active(network_id).await);

    let dir = agent.stop().await;
    let restarted =
        Agent::spawn(local_config(dir.path()).with_discovery(Arc::new(discovery.clone())))
            .await
            .unwrap();
    assert!(restarted.list_networks().await.unwrap().is_empty());
    assert_eq!(restarted.status().await.unwrap().networks.len(), 0);

    restarted.shutdown().await;
    drop(restarted);
    drop(dir);
}