commonware_p2p/authenticated/lookup/

mod.rs

//! Communicate with a fixed set of authenticated peers with known addresses over encrypted connections.
//!
//! `lookup` provides multiplexed communication between fully-connected peers
//! identified by a developer-specified cryptographic identity (i.e. BLS, ed25519, etc.).
//! Unlike `discovery`, peers in `lookup` don't use a discovery mechanism to find each other;
//! each peer's address is supplied by the application layer.
//!
//! # Features
//!
//! - Configurable Cryptography Scheme for Peer Identities (BLS, ed25519, etc.)
//! - Multiplexing With Configurable Rate Limiting Per Channel and Send Prioritization
//!
//! # Design
//!
//! ## Discovery
//!
//! This module operates under the assumption that all peers are aware of and synchronized on
//! the composition of peer sets at specific, user-provided indices (`u64`). Each index maps to a
//! list of peer `PublicKey`/`SocketAddr` pairs (`(u64, Vec<(PublicKey, SocketAddr)>)`).
//!
//! On startup, the application supplies the initial set of peers. The [`Oracle`] implements
//! [`AddressableManager`](crate::AddressableManager) which provides two ways to update peer addresses:
//!
//! - [`AddressableManager::track`](crate::AddressableManager::track): Track a new peer set at a
//!   monotonically increasing index. Use this when the peer set composition changes (peers added/removed).
//!   This accepts either a list of primary peers or an
//!   [`AddressableTrackedPeers`](crate::AddressableTrackedPeers) value containing both primary and
//!   secondary peers.
//! - [`AddressableManager::overwrite`](crate::AddressableManager::overwrite): Update multiple
//!   peers' addresses in-place without creating a new peer set. Use this when only peer IPs change but
//!   the peer set composition stays the same. Peers not in the directory (or unchanged) are silently skipped (so the application doesn't
//!   need to remember what their last submitted peer set was).
//!
//! Secondary peers remain visible in [`PeerSetUpdate`](crate::PeerSetUpdate)
//! notifications, are accepted for inbound connections, and may receive
//! `Recipients::All` traffic on established connections, but outbound dialing
//! is restricted to primary peers.
//!
//! Any inbound connection attempts from an IP address that is not in the union of all registered
//! primary or secondary peers will be rejected.
//!
//! ## Messages
//!
//! Application-level data is exchanged using the `Data` message type. This structure contains:
//! - `channel`: A `u32` identifier used to route the message to the correct application handler.
//! - `message`: The arbitrary application payload as `IoBuf`.
//!
//! The size of the `message` bytes must not exceed the configured
//! `max_message_size`. If it does, the sending operation will panic. Messages can be sent with `priority`, allowing certain
//! communications to potentially bypass lower-priority messages waiting in send queues across all
//! channels. Each registered channel ([Sender], [Receiver]) handles its own message queuing
//! and rate limiting.
//!
//! ## Compression
//!
//! Stream compression is not provided at the transport layer to avoid inadvertently
//! enabling known attacks such as BREACH and CRIME. These attacks exploit the interaction
//! between compression and encryption by analyzing patterns in the resulting data.
//! By compressing secrets alongside attacker-controlled content, these attacks can infer
//! sensitive information through compression ratio analysis. Applications that choose
//! to compress data should do so with full awareness of these risks and implement
//! appropriate mitigations (such as ensuring no attacker-controlled data is compressed
//! alongside sensitive information).
//!
//! ## Batching
//!
//! Applications seeking higher performance should prefer batching messages
//! above `p2p`. Larger application-level batches amortize per-message
//! encryption overhead and, if the application also compresses its payloads,
//! can improve compression ratio.
//!
//! ## Rate Limiting
//!
//! There are five primary rate limits:
//!
//! - `max_concurrent_handshakes`: The maximum number of concurrent handshake attempts allowed.
//! - `allowed_handshake_rate_per_ip`: The rate limit for handshake attempts originating from a single IP address.
//! - `allowed_handshake_rate_per_subnet`: The rate limit for handshake attempts originating from a single IP subnet.
//! - `peer_connection_cooldown`: The per-peer rate limit for inbound and outbound connection reservations, expressed as a minimum cooldown between attempts.
//! - `rate` (per channel): The rate limit for messages sent on a single channel.
//!
//! _Users should consider these rate limits as best-effort protection against moderate abuse. Targeted abuse (e.g. DDoS)
//! must be mitigated with an external proxy (that limits inbound connection attempts to authorized IPs)._
//!
//! ## IP Poisoning
//!
//! A malicious peer can claim an ingress [std::net::SocketAddr] that collides with an honest
//! peer, drawing invalid dial attempts to
//! the honest peer (where we expect the malicious public key rather than the honest public key).
//!
//! Because we rate limit inbound connection attempts per IP/subnet, this poisoning can lead to us dropping legitimate
//! dial attempts (if quota was already exhausted on useless dial attempts). Recall, an honest dialer doesn't know which public
//! key actually resides at an address and must try all that collide.
//!
//! To mitigate this issue, we shuffle peer dial order on each dial queue refresh. This ensures we eventually dial a poisoned
//! IP with the correct public key before hitting the rate limit imposed by the listener at said IP.
//!
//! _Applications that wish to entirely prevent this class of attack can assert uniqueness of
//! ingress [std::net::SocketAddr] during
//! peer registration._
//!
//! ## Message Delivery
//!
//! Outgoing message submissions can be rejected when a peer's send buffer is full, preventing slow
//! peers from blocking sends to other peers. Incoming messages are dropped when the application's
//! receive buffer is full, ensuring ping messages continue to flow and connections remain healthy.
//!
//! # Example
//!
//! ```rust
//! use commonware_p2p::{authenticated::lookup::{self, Network}, Address, AddressableManager, Sender, Recipients};
//! use commonware_cryptography::{ed25519, Signer, PrivateKey as _, PublicKey as _, };
//! use commonware_runtime::{deterministic, IoBuf, Metrics, Quota, Runner, Spawner, Supervisor};
//! use commonware_utils::{NZU32, ordered::Map};
//! use std::net::{IpAddr, Ipv4Addr, SocketAddr};
//!
//! // Configure context
//! let runtime_cfg = deterministic::Config::default();
//! let runner = deterministic::Runner::new(runtime_cfg);
//!
//! // Generate identity
//! //
//! // In production, the signer should be generated from a secure source of entropy.
//! let my_sk = ed25519::PrivateKey::from_seed(0);
//! let my_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0);
//!
//! // Generate peers
//! //
//! // In production, peer identities will be provided by some external source of truth
//! // (like the staking set of a blockchain).
//! let peer1 = ed25519::PrivateKey::from_seed(1).public_key();
//! let peer1_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 3001);
//! let peer2 = ed25519::PrivateKey::from_seed(2).public_key();
//! let peer2_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 3002);
//! let peer3 = ed25519::PrivateKey::from_seed(3).public_key();
//! let peer3_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 3003);
//!
//! // Configure namespace
//! //
//! // In production, use a unique application namespace to prevent cryptographic replay attacks.
//! let application_namespace = b"my-app-namespace";
//!
//! // Configure network
//! //
//! // In production, use a more conservative configuration like `Config::recommended`.
//! const MAX_MESSAGE_SIZE: u32 = 1_024; // 1KB
//! let p2p_cfg = lookup::Config::local(
//!     my_sk.clone(),
//!     application_namespace,
//!     my_addr,
//!     MAX_MESSAGE_SIZE,
//! );
//!
//! // Start context
//! runner.start(|context| async move {
//!     // Initialize network
//!     let (mut network, mut oracle) = Network::new(context.child("network"), p2p_cfg);
//!
//!     // Register authorized peers
//!     //
//!     // In production, this would be updated as new peer sets are created (like when
//!     // the composition of a validator set changes).
//!     let peers: Map<_, Address> = [
//!         (my_sk.public_key(), my_addr.into()),
//!         (peer1, peer1_addr.into()),
//!         (peer2, peer2_addr.into()),
//!         (peer3, peer3_addr.into()),
//!     ].try_into().unwrap();
//!     oracle.track(0, peers);
//!
//!     // Register some channel
//!     const MAX_MESSAGE_BACKLOG: usize = 128;
//!     let (mut sender, receiver) = network.register(
//!         0,
//!         Quota::per_second(NZU32!(1)),
//!         MAX_MESSAGE_BACKLOG,
//!     );
//!
//!     // Run network
//!     network.start();
//!
//!     // Example: Use sender
//!     let _ = sender.send(Recipients::All, IoBuf::from(b"hello"), false);
//!
//!     // Graceful shutdown (stops all spawned tasks)
//!     context.stop(0, None).await.unwrap();
//! });
//! ```

mod actors;
mod channels;
mod config;
mod metrics;
mod network;
mod types;

use thiserror::Error;

/// Errors that can occur when interacting with the network.
#[derive(Error, Debug)]
pub enum Error {
    #[error("network closed")]
    NetworkClosed,
}

pub use actors::tracker::Oracle;
pub use channels::{Receiver, Sender};
pub use config::Config;
pub use network::Network;

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        authenticated::{
            lookup::actors::router::{Actor as RouterActor, Config as RouterConfig},
            relay::Relay,
        },
        Address, AddressableManager, CheckedSender as _, Ingress, LimitedSender as _, Provider,
        Receiver, Recipients, Sender,
    };
    use commonware_actor::{Feedback, Unreliable};
    use commonware_cryptography::{ed25519, Signer as _};
    use commonware_macros::{select, test_group, test_traced};
    use commonware_runtime::{
        deterministic, telemetry::metrics::count_running_tasks, tokio, BufferPooler, Clock, IoBuf,
        Metrics, Network as RNetwork, Quota, Resolver, Runner, Spawner, Supervisor as _,
    };
    use commonware_utils::{
        channel::mpsc,
        hostname,
        ordered::{Map, Set},
        Hostname, NZUsize, TryCollect, NZU32,
    };
    use rand_core::{CryptoRngCore, RngCore};
    use std::{
        collections::HashSet,
        net::{IpAddr, Ipv4Addr, SocketAddr},
        time::Duration,
    };

    #[derive(Copy, Clone)]
    enum Mode {
        All,
        Some,
        One,
    }

    const MAX_MESSAGE_SIZE: u32 = 1_024 * 1_024; // 1MB
    const DEFAULT_MESSAGE_BACKLOG: usize = 128;

    /// Ensure no message rate limiting occurred.
    ///
    /// If a message is rate limited, it would be formatted as:
    ///
    /// ```text
    /// peer-9_network_spawner_messages_rate_limited_total{peer="e2e8aa145e1ec5cb01ebfaa40e10e12f0230c832fd8135470c001cb86d77de00",message="data_0"} 1
    /// peer-9_network_spawner_messages_rate_limited_total{peer="e2e8aa145e1ec5cb01ebfaa40e10e12f0230c832fd8135470c001cb86d77de00",message="ping"} 1
    /// ```
    fn assert_no_rate_limiting(metrics: &str) {
        assert!(
            !metrics.contains("messages_rate_limited_total{"),
            "no messages should be rate limited: {metrics}"
        );
    }

    /// Test connectivity between `n` peers.
    ///
    /// We set a unique `base_port` for each test to avoid "address already in use"
    /// errors when tests are run immediately after each other.
    async fn run_network(
        context: impl Spawner + BufferPooler + Clock + CryptoRngCore + RNetwork + Resolver + Metrics,
        max_message_size: u32,
        base_port: u16,
        n: usize,
        mode: Mode,
    ) {
        // Create peers
        let mut peers_and_sks = Vec::new();
        for i in 0..n {
            let private_key = ed25519::PrivateKey::from_seed(i as u64);
            let public_key = private_key.public_key();
            let address = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + i as u16);
            peers_and_sks.push((private_key, public_key, address));
        }
        let peers: Vec<(ed25519::PublicKey, Address)> = peers_and_sks
            .iter()
            .map(|(_, pub_key, addr)| (pub_key.clone(), (*addr).into()))
            .collect::<Vec<_>>();

        // Create networks
        let (complete_sender, mut complete_receiver) = mpsc::channel(peers.len());
        for (i, (private_key, public_key, address)) in peers_and_sks.iter().enumerate() {
            let public_key = public_key.clone();

            // Create peer context
            let context = context.child("peer").with_attribute("index", i);

            // Create network
            let config = Config::test(private_key.clone(), *address, max_message_size);
            let (mut network, mut oracle) = Network::new(context.child("network"), config);

            // Register peers
            oracle.track(0, Map::try_from(peers.clone()).unwrap());

            // Register basic application
            let (mut sender, mut receiver) =
                network.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);

            // Wait to connect to all peers, and then send messages to everyone
            network.start();

            // Send/Receive messages
            context.child("agent").spawn({
                let complete_sender = complete_sender.clone();
                let peers = peers.clone();
                move |context| async move {
                    // Wait for all peers to send their identity
                    let receiver = context.child("receiver").spawn(move |_| async move {
                        // Wait for all peers to send their identity
                        let mut received = HashSet::new();
                        while received.len() < n - 1 {
                            // Ensure message equals sender identity
                            let (sender, message) = receiver.recv().await.unwrap();
                            assert_eq!(message, sender.as_ref());

                            // Add to received set
                            received.insert(sender);
                        }
                        complete_sender.send(()).await.unwrap();

                        // Process messages until all finished (or else sender loops could get stuck as a peer may drop)
                        loop {
                            receiver.recv().await.unwrap();
                        }
                    });

                    // Send identity to all peers
                    let sender = context.child("sender").spawn(move |context| async move {
                        // Get all peers not including self
                        let mut recipients: Vec<_> = peers
                            .iter()
                            .enumerate()
                            .filter(|(j, _)| i != *j)
                            .map(|(_, (pk, _))| pk.clone())
                            .collect();
                        recipients.sort();

                        // Loop forever to account for unexpected message drops
                        loop {
                            match mode {
                                Mode::One => {
                                    for pub_key in &recipients {
                                        // Loop until success
                                        loop {
                                            let sent = sender.send(
                                                Recipients::One(pub_key.clone()),
                                                public_key.as_ref().to_vec(),
                                                true,
                                            );
                                            if sent.len() != 1 {
                                                context.sleep(Duration::from_millis(100)).await;
                                                continue;
                                            }
                                            break;
                                        }
                                    }
                                }
                                Mode::Some | Mode::All => {
                                    // Loop until all peer sends successful
                                    loop {
                                        let sent = sender.send(
                                            match mode {
                                                Mode::Some => Recipients::Some(recipients.clone()),
                                                Mode::All => Recipients::All,
                                                _ => unreachable!(),
                                            },
                                            public_key.as_ref().to_vec(),
                                            true,
                                        );
                                        if sent.len() != recipients.len() {
                                            context.sleep(Duration::from_millis(100)).await;
                                            continue;
                                        }
                                        break;
                                    }
                                }
                            };

                            // Sleep to avoid busy loop
                            context.sleep(Duration::from_secs(10)).await;
                        }
                    });

                    // Neither task should exit
                    select! {
                        receiver = receiver => {
                            panic!("receiver exited: {receiver:?}");
                        },
                        sender = sender => {
                            panic!("sender exited: {sender:?}");
                        },
                    }
                }
            });
        }

        // Wait for all peers to finish
        for _ in 0..n {
            complete_receiver.recv().await.unwrap();
        }

        // Ensure no message rate limiting occurred
        assert_no_rate_limiting(&context.encode());
    }

    fn run_deterministic_test(seed: u64, mode: Mode) {
        // Configure test
        const NUM_PEERS: usize = 25;
        const BASE_PORT: u16 = 3000;

        // Run first instance
        let executor = deterministic::Runner::seeded(seed);
        let state = executor.start(|context| async move {
            run_network(
                context.child("network"),
                MAX_MESSAGE_SIZE,
                BASE_PORT,
                NUM_PEERS,
                mode,
            )
            .await;
            context.auditor().state()
        });

        // Compare result to second instance
        let executor = deterministic::Runner::seeded(seed);
        let state2 = executor.start(|context| async move {
            run_network(
                context.child("network"),
                MAX_MESSAGE_SIZE,
                BASE_PORT,
                NUM_PEERS,
                mode,
            )
            .await;
            context.auditor().state()
        });
        assert_eq!(state, state2);
    }

    #[test_group("slow")]
    #[test_traced]
    fn test_determinism_one() {
        for i in 0..10 {
            run_deterministic_test(i, Mode::One);
        }
    }

    #[test_group("slow")]
    #[test_traced]
    fn test_determinism_some() {
        for i in 0..10 {
            run_deterministic_test(i, Mode::Some);
        }
    }

    #[test_group("slow")]
    #[test_traced]
    fn test_determinism_all() {
        for i in 0..10 {
            run_deterministic_test(i, Mode::All);
        }
    }

    #[test_traced]
    fn test_tokio_connectivity() {
        let executor = tokio::Runner::default();
        executor.start(|context| async move {
            let base_port = 4000;
            let n = 10;
            run_network(context, MAX_MESSAGE_SIZE, base_port, n, Mode::One).await;
        });
    }

    #[test_traced]
    fn test_multi_index_oracle() {
        // Configure test
        let base_port = 3000;
        let n: usize = 10;

        // Initialize context
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            // Create peers
            let mut peers_and_sks = Vec::new();
            for i in 0..n {
                let sk = ed25519::PrivateKey::from_seed(i as u64);
                let pk = sk.public_key();
                let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + i as u16);
                peers_and_sks.push((sk, pk, addr));
            }
            let peers = peers_and_sks
                .iter()
                .map(|(_, pk, addr)| (pk.clone(), (*addr).into()))
                .collect::<Vec<_>>();

            // Create networks
            let mut waiters = Vec::new();
            for (i, (peer_sk, peer_pk, peer_addr)) in peers_and_sks.iter().enumerate() {
                // Create peer context
                let context = context.child("peer").with_attribute("index", i);

                // Create network
                let config = Config::test(
                    peer_sk.clone(),
                    *peer_addr,
                    1_024 * 1_024, // 1MB
                );
                let (mut network, mut oracle) = Network::new(context.child("network"), config);

                // Register peers at separate indices
                oracle.track(0, Map::try_from([peers[0].clone()]).unwrap());
                oracle.track(
                    1,
                    Map::try_from([peers[1].clone(), peers[2].clone()]).unwrap(),
                );
                oracle.track(
                    2,
                    peers
                        .iter()
                        .skip(2)
                        .cloned()
                        .try_collect::<Map<_, _>>()
                        .unwrap(),
                );

                // Register basic application
                let (mut sender, mut receiver) =
                    network.register(0, Quota::per_second(NZU32!(10)), DEFAULT_MESSAGE_BACKLOG);

                // Wait to connect to all peers, and then send messages to everyone
                network.start();

                // Send/Receive messages
                let msg = peer_pk.clone();
                let handler = context.child("agent").spawn(move |context| async move {
                    if i == 0 {
                        // Loop until success
                        loop {
                            if sender
                                .send(Recipients::All, msg.as_ref().to_vec(), true)
                                .len()
                                == n - 1
                            {
                                break;
                            }

                            // Sleep and try again (avoid busy loop)
                            context.sleep(Duration::from_millis(100)).await;
                        }
                    } else {
                        // Ensure message equals sender identity
                        let (sender, message) = receiver.recv().await.unwrap();
                        assert_eq!(message, sender.as_ref());
                    }
                });

                // Add to waiters
                waiters.push(handler);
            }

            // Wait for waiters to finish (receiver before sender)
            for waiter in waiters.into_iter().rev() {
                waiter.await.unwrap();
            }

            // Ensure no message rate limiting occurred
            assert_no_rate_limiting(&context.encode());
        });
    }

    #[test_traced]
    #[should_panic(expected = "message too large")]
    fn test_message_too_large() {
        // Configure test
        let base_port = 3000;
        let n: usize = 2;

        // Initialize context
        let executor = deterministic::Runner::seeded(0);
        executor.start(|mut context| async move {
            // Create peers
            let mut peers_and_sks = Vec::new();
            for i in 0..n {
                let peer_sk = ed25519::PrivateKey::from_seed(i as u64);
                let peer_pk = peer_sk.public_key();
                let peer_addr =
                    SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + i as u16);
                peers_and_sks.push((peer_sk, peer_pk, peer_addr));
            }
            let peers: Map<_, _> = peers_and_sks
                .iter()
                .map(|(_, pk, addr)| (pk.clone(), (*addr).into()))
                .try_collect()
                .unwrap();

            // Create network
            let (sk, _, addr) = peers_and_sks[0].clone();
            let config = Config::test(
                sk,
                addr,
                1_024 * 1_024, // 1MB
            );
            let (mut network, mut oracle) = Network::new(context.child("network"), config);

            // Register peers
            oracle.track(0, peers.clone());

            // Register basic application
            let (mut sender, _) =
                network.register(0, Quota::per_second(NZU32!(10)), DEFAULT_MESSAGE_BACKLOG);

            // Wait to connect to all peers, and then send messages to everyone
            network.start();

            // Crate random message
            let mut msg = vec![0u8; 10 * 1024 * 1024]; // 10MB (greater than frame capacity)
            context.fill_bytes(&mut msg[..]);

            // Send message
            let recipient = Recipients::One(peers[1].clone());
            sender.send(recipient, msg, true);
        });
    }

    #[test_traced]
    fn test_rate_limiting() {
        // Configure test
        let base_port = 3000;
        let n: usize = 2;

        // Initialize context
        let executor = deterministic::Runner::seeded(0);
        executor.start(|context| async move {
            // Create peers
            let mut peers_and_sks = Vec::new();
            for i in 0..n {
                let sk = ed25519::PrivateKey::from_seed(i as u64);
                let pk = sk.public_key();
                let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + i as u16);
                peers_and_sks.push((sk, pk, addr));
            }
            let peers: Map<_, _> = peers_and_sks
                .iter()
                .map(|(_, pk, addr)| (pk.clone(), (*addr).into()))
                .try_collect()
                .unwrap();
            let (sk0, _, addr0) = peers_and_sks[0].clone();
            let (sk1, pk1, addr1) = peers_and_sks[1].clone();

            // Create network for peer 0
            let config0 = Config::test(sk0, addr0, 1_024 * 1_024); // 1MB
            let (mut network0, mut oracle0) =
                Network::new(context.child("peer").with_attribute("index", 0), config0);
            oracle0.track(0, peers.clone());
            let (mut sender0, _receiver0) =
                network0.register(0, Quota::per_minute(NZU32!(1)), DEFAULT_MESSAGE_BACKLOG);
            network0.start();

            // Create network for peer 1
            let config1 = Config::test(sk1, addr1, 1_024 * 1_024); // 1MB
            let (mut network1, mut oracle1) =
                Network::new(context.child("peer").with_attribute("index", 1), config1);
            oracle1.track(0, peers.clone());
            let (_sender1, _receiver1) =
                network1.register(0, Quota::per_minute(NZU32!(1)), DEFAULT_MESSAGE_BACKLOG);
            network1.start();

            // Send first message, which should be allowed and consume the quota.
            let msg = vec![0u8; 1024]; // 1KB
            loop {
                // Confirm message is sent to peer
                let checked = sender0.check(Recipients::All).unwrap();
                if !checked.recipients().is_empty() {
                    checked.send(msg.clone(), true);
                    break;
                }

                // Ensure we don't rate limit outbound sends while
                // waiting for peers to connect
                context.sleep(Duration::from_mins(1)).await
            }

            // Immediately send the second message to trigger the rate limit.
            let sent = sender0.send(Recipients::One(pk1), msg, true);
            assert!(sent.is_empty());

            // Give the metrics time to reflect the rate-limited message.
            for _ in 0..10 {
                assert_no_rate_limiting(&context.encode());
                context.sleep(Duration::from_millis(100)).await;
            }
        });
    }

    #[test_traced]
    fn test_unordered_peer_sets() {
        let (n, base_port) = (10, 3000);
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            // Create peers
            let mut peers_and_sks = Vec::new();
            for i in 0..n {
                let sk = ed25519::PrivateKey::from_seed(i as u64);
                let pk = sk.public_key();
                let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + i as u16);
                peers_and_sks.push((sk, pk, addr));
            }
            let peer0 = peers_and_sks[0].clone();
            let config = Config::test(peer0.0, peer0.2, 1_024 * 1_024);
            let (network, mut oracle) = Network::new(context.child("network"), config);
            network.start();

            // Subscribe to peer sets
            let mut subscription = oracle.subscribe().await;

            // Register initial peer set
            let set10: Map<_, _> = peers_and_sks
                .iter()
                .take(2)
                .map(|(_, pk, addr)| (pk.clone(), (*addr).into()))
                .try_collect()
                .unwrap();
            oracle.track(10, set10.clone());
            let update = subscription.recv().await.unwrap();
            assert_eq!(update.index, 10);
            assert_eq!(&update.latest.primary, set10.keys());
            assert!(update.latest.secondary.is_empty());
            assert_eq!(&update.all.primary, set10.keys());
            assert!(update.all.secondary.is_empty());

            // Register old peer sets (ignored)
            let set9: Map<_, _> = peers_and_sks
                .iter()
                .skip(2)
                .map(|(_, pk, addr)| (pk.clone(), (*addr).into()))
                .try_collect()
                .unwrap();
            oracle.track(9, set9.clone());

            // Add new peer set
            let set11: Map<_, _> = peers_and_sks
                .iter()
                .skip(4)
                .map(|(_, pk, addr)| (pk.clone(), (*addr).into()))
                .try_collect()
                .unwrap();
            oracle.track(11, set11.clone());
            let update = subscription.recv().await.unwrap();
            assert_eq!(update.index, 11);
            assert_eq!(&update.latest.primary, set11.keys());
            assert!(update.latest.secondary.is_empty());
            let all_keys: Set<_> = set10
                .into_keys()
                .into_iter()
                .chain(set11.into_keys().into_iter())
                .try_collect()
                .unwrap();
            assert_eq!(update.all.primary, all_keys);
            assert!(update.all.secondary.is_empty());
        });
    }

    #[test_traced]
    fn test_graceful_shutdown() {
        let base_port = 3000;
        let n: usize = 5;

        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            // Create peers
            let mut peers_and_sks = Vec::new();
            for i in 0..n {
                let sk = ed25519::PrivateKey::from_seed(i as u64);
                let pk = sk.public_key();
                let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + i as u16);
                peers_and_sks.push((sk, pk, addr));
            }
            let peers: Map<_, _> = peers_and_sks
                .iter()
                .map(|(_, pk, addr)| (pk.clone(), (*addr).into()))
                .try_collect()
                .unwrap();

            // Create networks for all peers
            let (complete_sender, mut complete_receiver) = mpsc::channel(n);
            for (i, (sk, pk, addr)) in peers_and_sks.iter().enumerate() {
                let peer_context = context.child("peer").with_attribute("index", i);
                let config = Config::test(sk.clone(), *addr, 1_024 * 1_024);
                let (mut network, mut oracle) = Network::new(peer_context.child("network"), config);

                // Register peer set
                oracle.track(0, peers.clone());

                let (mut sender, mut receiver) =
                    network.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
                network.start();

                peer_context.child("agent").spawn({
                    let complete_sender = complete_sender.clone();
                    let pk = pk.clone();
                    move |context| async move {
                        // Wait to connect to at least one other peer
                        let expected_connections = if i == 0 { n - 1 } else { 1 };

                        // Send a message
                        loop {
                            let sent = sender.send(Recipients::All, pk.as_ref().to_vec(), true);
                            if sent.len() >= expected_connections {
                                break;
                            }
                            context.sleep(Duration::from_millis(100)).await;
                        }

                        // Signal that this peer is connected
                        complete_sender.send(()).await.unwrap();

                        // Keep receiving messages until shutdown
                        loop {
                            select! {
                                result = receiver.recv() => {
                                    if result.is_err() {
                                        // Channel closed due to shutdown
                                        break;
                                    }
                                },
                                _ = context.stopped() => {
                                    // Graceful shutdown signal received
                                    break;
                                },
                            }
                        }
                    }
                });
            }

            // Wait for all peers to establish connectivity
            for _ in 0..n {
                complete_receiver.recv().await.unwrap();
            }

            // Verify that network actors started for all peers
            let metrics_before = context.encode();
            let tasks_running = |metrics: &str, name: &str| -> Option<u64> {
                metrics.lines().find_map(|line| {
                    (line.starts_with("runtime_tasks_running{")
                        && line.contains(&format!("name=\"{name}\""))
                        && line.contains("kind=\"Task\""))
                    .then(|| {
                        line.split_whitespace()
                            .next_back()
                            .expect("metric line must have a value")
                            .parse::<u64>()
                            .expect("running task count must be an integer")
                    })
                })
            };

            for actor in ["tracker", "router", "spawner", "listener", "dialer"] {
                let name = format!("peer_network_{actor}");
                assert_eq!(
                    tasks_running(&metrics_before, &name),
                    Some(n as u64),
                    "{name} should have {n} running tasks before shutdown"
                );
            }

            // All peers are connected - now trigger graceful shutdown
            context.child("shutdown").spawn(move |context| async move {
                // Trigger graceful shutdown
                let result = context.stop(0, Some(Duration::from_secs(5))).await;

                // Shutdown should complete successfully without timeout
                assert!(
                    result.is_ok(),
                    "graceful shutdown should complete: {result:?}"
                );
            });

            // Wait for shutdown to complete
            context.stopped().await.unwrap();

            // Give the runtime a tick to process task completions and update metrics
            context.sleep(Duration::from_millis(100)).await;

            // Verify that all network actors stopped
            let metrics_after = context.encode();
            for actor in ["tracker", "router", "spawner", "listener", "dialer"] {
                let name = format!("peer_network_{actor}");
                assert_eq!(
                    tasks_running(&metrics_after, &name),
                    Some(0),
                    "{name} should be stopped"
                );
            }
        });
    }

    #[test_traced]
    fn test_subscription_includes_self_when_registered() {
        let base_port = 3000;
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            // Create self (peer0) and other peers
            let self_sk = ed25519::PrivateKey::from_seed(0);
            let self_pk = self_sk.public_key();
            let self_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port);

            let other_pk = ed25519::PrivateKey::from_seed(1).public_key();
            let other_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + 1);

            // Create network for peer0 (self)
            let config = Config::test(self_sk, self_addr, 1_024 * 1_024);
            let (network, mut oracle) = Network::new(context.child("network"), config);
            network.start();

            // Subscribe to peer sets
            let mut subscription = oracle.subscribe().await;

            // Register a peer set that does NOT include self
            let peer_set: Map<_, _> = [(other_pk.clone(), other_addr.into())].try_into().unwrap();
            oracle.track(1, peer_set.clone());

            // Receive subscription notification
            let update = subscription.recv().await.unwrap();
            assert_eq!(update.index, 1);
            assert_eq!(update.latest.primary.len(), 1);
            assert_eq!(update.all.primary.len(), 1);
            assert!(update.all.secondary.is_empty());

            // Self should NOT be in the new set
            assert!(
                update.latest.primary.position(&self_pk).is_none(),
                "new set should not include self"
            );
            assert!(
                update.latest.primary.position(&other_pk).is_some(),
                "new set should include other"
            );
            assert!(
                update.latest.secondary.is_empty(),
                "new secondary set should be empty"
            );

            // Self should NOT be in the peer set (not tracked)
            assert!(
                update.all.primary.position(&self_pk).is_none(),
                "peer set should not include self"
            );
            assert!(
                update.all.primary.position(&other_pk).is_some(),
                "peer set should include other"
            );

            // Now register a peer set that DOES include self
            let peer_set: Map<_, _> = [
                (self_pk.clone(), self_addr.into()),
                (other_pk.clone(), other_addr.into()),
            ]
            .try_into()
            .unwrap();
            oracle.track(2, peer_set.clone());

            // Receive subscription notification
            let update = subscription.recv().await.unwrap();
            assert_eq!(update.index, 2);
            assert_eq!(update.latest.primary.len(), 2);
            assert_eq!(update.all.primary.len(), 2);
            assert!(update.all.secondary.is_empty());

            // Both peers should be in the new set
            assert!(
                update.latest.primary.position(&self_pk).is_some(),
                "new set should include self"
            );
            assert!(
                update.latest.primary.position(&other_pk).is_some(),
                "new set should include other"
            );
            assert!(
                update.latest.secondary.is_empty(),
                "new secondary set should be empty"
            );

            // Both peers should be in the peer set
            assert!(
                update.all.primary.position(&self_pk).is_some(),
                "peer set should include self"
            );
            assert!(
                update.all.primary.position(&other_pk).is_some(),
                "peer set should include other"
            );
        });
    }

    #[test_traced]
    fn test_dns_peer_addresses() {
        let base_port = 3200;
        let n: usize = 3;

        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            // Create peers
            let mut peers_and_sks = Vec::new();
            for i in 0..n {
                let private_key = ed25519::PrivateKey::from_seed(i as u64);
                let public_key = private_key.public_key();
                let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + i as u16);
                let host_str = format!("peer-{i}.local");
                let host = Hostname::new(&host_str).unwrap();
                peers_and_sks.push((private_key, public_key, socket, host_str, host));
            }

            // Register DNS mappings for all peers
            for (_, _, socket, host_str, _) in &peers_and_sks {
                context.resolver_register(host_str.clone(), Some(vec![socket.ip()]));
            }

            // Create peer addresses with DNS ingress
            let peers: Vec<(_, Address)> = peers_and_sks
                .iter()
                .map(|(_, pk, socket, _, host)| {
                    (
                        pk.clone(),
                        Address::Asymmetric {
                            ingress: Ingress::Dns {
                                host: host.clone(),
                                port: socket.port(),
                            },
                            egress: *socket,
                        },
                    )
                })
                .collect();

            // Create networks
            let (complete_sender, mut complete_receiver) = mpsc::channel(n);
            for (i, (private_key, public_key, socket, _, _)) in peers_and_sks.iter().enumerate() {
                let context = context.child("peer").with_attribute("index", i);

                // Create network
                let config = Config::test(private_key.clone(), *socket, 1_024 * 1_024);
                let (mut network, mut oracle) = Network::new(context.child("network"), config);

                // Register peers with DNS addresses
                oracle.track(0, Map::try_from(peers.clone()).unwrap());

                // Register channel
                let (mut sender, mut receiver) =
                    network.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);

                network.start();

                // Send/Receive messages
                let pk = public_key.clone();
                context.child("agent").spawn({
                    let complete_sender = complete_sender.clone();
                    let peers = peers.clone();
                    move |context| async move {
                        // Wait for messages from other peers
                        let receiver = context.child("receiver").spawn(move |_| async move {
                            let mut received = HashSet::new();
                            while received.len() < n - 1 {
                                let (sender, message) = receiver.recv().await.unwrap();
                                assert_eq!(message, sender.as_ref());
                                received.insert(sender);
                            }
                            complete_sender.send(()).await.unwrap();

                            loop {
                                receiver.recv().await.unwrap();
                            }
                        });

                        // Send identity to all peers
                        let sender_task =
                            context.child("sender").spawn(move |context| async move {
                                loop {
                                    let mut recipients: Vec<_> = peers
                                        .iter()
                                        .filter(|(p, _)| p != &pk)
                                        .map(|(p, _)| p.clone())
                                        .collect();
                                    recipients.sort();

                                    loop {
                                        let mut sent = sender.send(
                                            Recipients::All,
                                            pk.as_ref().to_vec(),
                                            true,
                                        );
                                        if sent.len() != n - 1 {
                                            context.sleep(Duration::from_millis(100)).await;
                                            continue;
                                        }
                                        sent.sort();
                                        assert_eq!(sent, recipients);
                                        break;
                                    }

                                    context.sleep(Duration::from_secs(10)).await;
                                }
                            });

                        select! {
                            receiver = receiver => {
                                panic!("receiver exited: {receiver:?}")
                            },
                            sender = sender_task => {
                                panic!("sender exited: {sender:?}")
                            },
                        }
                    }
                });
            }

            // Wait for all peers to exchange messages
            for _ in 0..n {
                complete_receiver.recv().await.unwrap();
            }

            assert_no_rate_limiting(&context.encode());
        });
    }

    #[test_traced]
    fn test_mixed_socket_and_dns_addresses() {
        let base_port = 3300;
        let n: usize = 4;

        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            // Create peers
            let mut peers_and_sks = Vec::new();
            for i in 0..n {
                let private_key = ed25519::PrivateKey::from_seed(i as u64);
                let public_key = private_key.public_key();
                let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + i as u16);
                peers_and_sks.push((private_key, public_key, socket));
            }

            // Register DNS mappings for peers 2 and 3 only
            for (i, (_, _, socket)) in peers_and_sks.iter().enumerate().skip(2) {
                context.resolver_register(format!("peer-{i}.local"), Some(vec![socket.ip()]));
            }

            // Create peer addresses - peers 0,1 use Symmetric, peers 2,3 use DNS Asymmetric
            let peers: Vec<(_, Address)> = peers_and_sks
                .iter()
                .enumerate()
                .map(|(i, (_, pk, socket))| {
                    let addr = if i < 2 {
                        // Peers 0, 1: Symmetric socket address
                        Address::Symmetric(*socket)
                    } else {
                        // Peers 2, 3: Asymmetric with DNS ingress
                        Address::Asymmetric {
                            ingress: Ingress::Dns {
                                host: hostname!(&format!("peer-{i}.local")),
                                port: socket.port(),
                            },
                            egress: *socket,
                        }
                    };
                    (pk.clone(), addr)
                })
                .collect();

            // Create networks
            let (complete_sender, mut complete_receiver) = mpsc::channel(n);
            for (i, (private_key, public_key, socket)) in peers_and_sks.iter().enumerate() {
                let context = context.child("peer").with_attribute("index", i);

                // Create network
                let config = Config::test(private_key.clone(), *socket, 1_024 * 1_024);
                let (mut network, mut oracle) = Network::new(context.child("network"), config);

                // Register peers with mixed addresses
                oracle.track(0, Map::try_from(peers.clone()).unwrap());

                // Register channel
                let (mut sender, mut receiver) =
                    network.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);

                network.start();

                // Send/Receive messages
                let pk = public_key.clone();
                context.child("agent").spawn({
                    let complete_sender = complete_sender.clone();
                    let peers = peers.clone();
                    move |context| async move {
                        // Wait for messages from other peers
                        let receiver = context.child("receiver").spawn(move |_| async move {
                            let mut received = HashSet::new();
                            while received.len() < n - 1 {
                                let (sender, message) = receiver.recv().await.unwrap();
                                assert_eq!(message, sender.as_ref());
                                received.insert(sender);
                            }
                            complete_sender.send(()).await.unwrap();

                            loop {
                                receiver.recv().await.unwrap();
                            }
                        });

                        // Send identity to all peers
                        let sender_task =
                            context.child("sender").spawn(move |context| async move {
                                loop {
                                    let mut recipients: Vec<_> = peers
                                        .iter()
                                        .filter(|(p, _)| p != &pk)
                                        .map(|(p, _)| p.clone())
                                        .collect();
                                    recipients.sort();

                                    loop {
                                        let mut sent = sender.send(
                                            Recipients::All,
                                            pk.as_ref().to_vec(),
                                            true,
                                        );
                                        if sent.len() != n - 1 {
                                            context.sleep(Duration::from_millis(100)).await;
                                            continue;
                                        }
                                        sent.sort();
                                        assert_eq!(sent, recipients);
                                        break;
                                    }

                                    context.sleep(Duration::from_secs(10)).await;
                                }
                            });

                        select! {
                            receiver = receiver => {
                                panic!("receiver exited: {receiver:?}")
                            },
                            sender = sender_task => {
                                panic!("sender exited: {sender:?}")
                            },
                        }
                    }
                });
            }

            // Wait for all peers to exchange messages
            for _ in 0..n {
                complete_receiver.recv().await.unwrap();
            }

            assert_no_rate_limiting(&context.encode());
        });
    }

    #[test_traced]
    fn test_dns_resolving_to_private_ip_not_dialed() {
        // Test that when allow_private_ips=false, DNS addresses that resolve
        // to private IPs are not dialed.
        let base_port = 4400;
        let executor = deterministic::Runner::timed(Duration::from_secs(10));
        executor.start(|context| async move {
            let peer0 = ed25519::PrivateKey::from_seed(0);
            let peer1 = ed25519::PrivateKey::from_seed(1);

            let socket0 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port);
            let socket1 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + 1);

            // Register DNS mapping that resolves to localhost (private IP)
            context.resolver_register("peer-0.local".to_string(), Some(vec![socket0.ip()]));

            // Create peer 0 with allow_private_ips=true
            let mut config0 = Config::test(peer0.clone(), socket0, 1_024 * 1_024);
            config0.allow_private_ips = true;
            let (mut network0, mut oracle0) =
                Network::new(context.child("peer").with_attribute("index", 0), config0);

            // Peer 0 knows about peer 1 with a socket address
            let peers0: Vec<(_, Address)> = vec![
                (peer0.public_key(), Address::Symmetric(socket0)),
                (peer1.public_key(), Address::Symmetric(socket1)),
            ];
            oracle0.track(0, Map::try_from(peers0).unwrap());

            let (_sender0, mut receiver0) =
                network0.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
            network0.start();

            // Create peer 1 with allow_private_ips=false
            let mut config1 = Config::test(peer1.clone(), socket1, 1_024 * 1_024);
            config1.allow_private_ips = false; // This should prevent dialing the private IP
            let (mut network1, mut oracle1) =
                Network::new(context.child("peer").with_attribute("index", 1), config1);

            // Peer 1 knows about peer 0 with a DNS address that resolves to private IP
            let peers1: Vec<(_, Address)> = vec![
                (
                    peer0.public_key(),
                    Address::Asymmetric {
                        ingress: Ingress::Dns {
                            host: hostname!("peer-0.local"),
                            port: socket0.port(),
                        },
                        egress: socket0,
                    },
                ),
                (peer1.public_key(), Address::Symmetric(socket1)),
            ];
            oracle1.track(0, Map::try_from(peers1).unwrap());

            let (mut sender1, _receiver1) =
                network1.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
            network1.start();

            // Wait for a period during which peer 1 would normally connect
            context.sleep(Duration::from_secs(5)).await;

            // Try to send from peer 1 - should not reach anyone since private IPs are blocked
            let sent = sender1.send(Recipients::All, peer1.public_key().as_ref().to_vec(), true);
            assert!(
                sent.is_empty(),
                "peer 1 should not have connected to peer 0 (private IP)"
            );

            // Verify peer 0 received nothing from peer 1
            select! {
                msg = receiver0.recv() => {
                    panic!("peer 0 should not have received any message, got: {msg:?}");
                },
                _ = context.sleep(Duration::from_secs(1)) => {
                    // Expected: timeout with no message
                },
            }
        });
    }

    #[test_traced]
    fn test_dns_mixed_ips_connectivity() {
        // Test that peers can connect even when DNS resolves to multiple IPs
        // where most are unreachable. The dialer randomly picks an IP, so
        // eventually it should pick the reachable one.
        //
        // Run over 25 seeds to ensure we exercise the random IP selection.
        for seed in 0..25 {
            let base_port = 3500;

            let cfg = deterministic::Config::default()
                .with_seed(seed)
                .with_timeout(Some(Duration::from_secs(120)));
            let executor = deterministic::Runner::new(cfg);
            executor.start(|context| async move {
                let peer0 = ed25519::PrivateKey::from_seed(0);
                let peer1 = ed25519::PrivateKey::from_seed(1);

                let good_ip = IpAddr::V4(Ipv4Addr::LOCALHOST);
                let socket0 = SocketAddr::new(good_ip, base_port);
                let socket1 = SocketAddr::new(good_ip, base_port + 1);

                // Register DNS mappings with 3 bad IPs and 1 good IP for both peers
                let mut all_ips0: Vec<IpAddr> = (1..=3)
                    .map(|i| IpAddr::V4(Ipv4Addr::new(127, 0, 0, 100 + i)))
                    .collect();
                all_ips0.push(good_ip);
                context.resolver_register("peer-0.local", Some(all_ips0));

                let mut all_ips1: Vec<IpAddr> = (1..=3)
                    .map(|i| IpAddr::V4(Ipv4Addr::new(127, 0, 0, 110 + i)))
                    .collect();
                all_ips1.push(good_ip);
                context.resolver_register("peer-1.local", Some(all_ips1));

                // Create peer addresses - both peers use DNS with mixed IPs
                let peers: Vec<(_, Address)> = vec![
                    (
                        peer0.public_key(),
                        Address::Asymmetric {
                            ingress: Ingress::Dns {
                                host: hostname!("peer-0.local"),
                                port: base_port,
                            },
                            egress: socket0,
                        },
                    ),
                    (
                        peer1.public_key(),
                        Address::Asymmetric {
                            ingress: Ingress::Dns {
                                host: hostname!("peer-1.local"),
                                port: base_port + 1,
                            },
                            egress: socket1,
                        },
                    ),
                ];

                // Create peer 0
                let config0 = Config::test(peer0.clone(), socket0, 1_024 * 1_024);
                let (mut network0, mut oracle0) =
                    Network::new(context.child("peer").with_attribute("index", 0), config0);
                oracle0.track(0, Map::try_from(peers.clone()).unwrap());
                let (_sender0, mut receiver0) =
                    network0.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
                network0.start();

                // Create peer 1
                let config1 = Config::test(peer1.clone(), socket1, 1_024 * 1_024);
                let (mut network1, mut oracle1) =
                    Network::new(context.child("peer").with_attribute("index", 1), config1);
                oracle1.track(0, Map::try_from(peers.clone()).unwrap());
                let (mut sender1, _receiver1) =
                    network1.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
                network1.start();

                // Send until peers connect (may take multiple attempts due to random IP selection).
                loop {
                    let checked = sender1.check(Recipients::All).unwrap();
                    if !checked.recipients().is_empty() {
                        checked.send(peer1.public_key().as_ref().to_vec(), true);
                    }

                    select! {
                        result = receiver0.recv() => {
                            let (sender, msg) = result.unwrap();
                            assert_eq!(sender, peer1.public_key());
                            assert_eq!(msg, peer1.public_key().as_ref());
                            break;
                        },
                        _ = context.sleep(Duration::from_millis(100)) => {},
                    }
                }
            });
        }
    }

    #[test_traced]
    fn test_many_peer_restart_with_new_address() {
        let base_port = 9500;
        let n = 5;

        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            // Create peers
            let peers: Vec<_> = (0..n)
                .map(|i| ed25519::PrivateKey::from_seed(i as u64))
                .collect();
            let addresses: Vec<_> = peers.iter().map(|p| p.public_key()).collect();

            // Track senders/receivers/oracles/handles across restarts
            let mut senders: Vec<Option<channels::Sender<_, _>>> = (0..n).map(|_| None).collect();
            let mut receivers: Vec<Option<channels::Receiver<_>>> = (0..n).map(|_| None).collect();
            let mut oracles: Vec<Option<Oracle<_>>> = (0..n).map(|_| None).collect();
            let mut handles: Vec<Option<commonware_runtime::Handle<()>>> =
                (0..n).map(|_| None).collect();

            // Track port allocations (updated on restart)
            let mut ports: Vec<u16> = (0..n).map(|i| base_port + i as u16).collect();

            // Create initial peer set with addresses
            let peer_set: Vec<(_, Address)> = addresses
                .iter()
                .enumerate()
                .map(|(i, pk)| {
                    (
                        pk.clone(),
                        Address::Symmetric(SocketAddr::new(
                            IpAddr::V4(Ipv4Addr::LOCALHOST),
                            ports[i],
                        )),
                    )
                })
                .collect();

            // Create networks for all peers
            for (i, peer) in peers.iter().enumerate() {
                let peer_context = context.child("peer").with_attribute("index", i);

                let config = Config::test(
                    peer.clone(),
                    SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), ports[i]),
                    MAX_MESSAGE_SIZE,
                );
                let (mut network, mut oracle) = Network::new(peer_context.child("network"), config);

                // Register peer set
                oracle.track(0, Map::try_from(peer_set.clone()).unwrap());

                let (sender, receiver) =
                    network.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
                senders[i] = Some(sender);
                receivers[i] = Some(receiver);
                oracles[i] = Some(oracle);

                let handle = network.start();
                handles[i] = Some(handle);
            }

            // Wait for full connectivity (each peer can send to all others)
            for (i, sender) in senders.iter_mut().enumerate() {
                let sender = sender.as_mut().unwrap();
                loop {
                    let sent = sender.send(
                        Recipients::All,
                        peers[i].public_key().as_ref().to_vec(),
                        true,
                    );
                    if sent.len() == n - 1 {
                        break;
                    }
                    context.sleep(Duration::from_millis(100)).await;
                }
            }

            // Verify each peer can receive from all others
            for receiver in receivers.iter_mut() {
                let receiver = receiver.as_mut().unwrap();
                let mut received = HashSet::new();
                while received.len() < n - 1 {
                    let (sender, message): (ed25519::PublicKey, _) = receiver.recv().await.unwrap();
                    assert_eq!(message, sender.as_ref());
                    received.insert(sender);
                }
            }

            // Restart each non-first peer with a new port, multiple rounds
            let mut restart_counter = 0u16;
            for round in 0..3 {
                for restart_peer_idx in 1..n {
                    // Allocate a new unique port
                    restart_counter += 1;
                    let new_port = base_port + 100 + restart_counter;
                    ports[restart_peer_idx] = new_port;

                    // Abort the peer's network
                    if let Some(handle) = handles[restart_peer_idx].take() {
                        handle.abort();
                    }
                    senders[restart_peer_idx] = None;
                    receivers[restart_peer_idx] = None;
                    oracles[restart_peer_idx] = None;

                    // Create updated peer set with new port
                    let updated_peer_set: Vec<(_, Address)> = addresses
                        .iter()
                        .enumerate()
                        .map(|(i, pk)| {
                            (
                                pk.clone(),
                                Address::Symmetric(SocketAddr::new(
                                    IpAddr::V4(Ipv4Addr::LOCALHOST),
                                    ports[i],
                                )),
                            )
                        })
                        .collect();

                    // Update oracle on all running peers
                    for oracle in oracles.iter_mut().flatten() {
                        oracle.track(
                            (round * (n - 1) + restart_peer_idx) as u64,
                            Map::try_from(updated_peer_set.clone()).unwrap(),
                        );
                    }

                    // Restart the peer with new port
                    let peer_context = context
                        .child("peer_round")
                        .with_attribute("index", restart_peer_idx)
                        .with_attribute("round", round);
                    let config = Config::test(
                        peers[restart_peer_idx].clone(),
                        SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), new_port),
                        MAX_MESSAGE_SIZE,
                    );
                    let (mut network, mut oracle) =
                        Network::new(peer_context.child("network"), config);

                    oracle.track(
                        (round * (n - 1) + restart_peer_idx) as u64,
                        Map::try_from(updated_peer_set.clone()).unwrap(),
                    );

                    let (sender, receiver) = network.register(
                        0,
                        Quota::per_second(NZU32!(100)),
                        DEFAULT_MESSAGE_BACKLOG,
                    );
                    senders[restart_peer_idx] = Some(sender);
                    receivers[restart_peer_idx] = Some(receiver);
                    oracles[restart_peer_idx] = Some(oracle);

                    let handle = network.start();
                    handles[restart_peer_idx] = Some(handle);

                    // Wait for the restarted peer to reconnect to all others
                    let restarted_sender = senders[restart_peer_idx].as_mut().unwrap();
                    loop {
                        let sent = restarted_sender.send(
                            Recipients::All,
                            peers[restart_peer_idx].public_key().as_ref().to_vec(),
                            true,
                        );
                        if sent.len() == n - 1 {
                            break;
                        }
                        context.sleep(Duration::from_millis(100)).await;
                    }

                    // Verify all other peers can send to the restarted peer
                    for i in 0..n {
                        if i == restart_peer_idx {
                            continue;
                        }
                        let sender = senders[i].as_mut().unwrap();
                        loop {
                            let sent = sender.send(
                                Recipients::One(addresses[restart_peer_idx].clone()),
                                peers[i].public_key().as_ref().to_vec(),
                                true,
                            );
                            if sent.len() == 1 {
                                break;
                            }
                            context.sleep(Duration::from_millis(100)).await;
                        }
                    }

                    // Verify the restarted peer receives from all others
                    let restarted_receiver = receivers[restart_peer_idx].as_mut().unwrap();
                    let mut received = HashSet::new();
                    while received.len() < n - 1 {
                        let (sender, message): (ed25519::PublicKey, _) =
                            restarted_receiver.recv().await.unwrap();
                        assert_eq!(message, sender.as_ref());
                        received.insert(sender);
                    }
                }
            }

            assert_no_rate_limiting(&context.encode());
        });
    }

    #[test_traced]
    fn test_simultaneous_peer_restart() {
        let base_port = 9700;
        let n = 5;

        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            // Create peers
            let peers: Vec<_> = (0..n)
                .map(|i| ed25519::PrivateKey::from_seed(i as u64))
                .collect();
            let addresses: Vec<_> = peers.iter().map(|p| p.public_key()).collect();

            // Track port allocations (updated on restart)
            let mut ports: Vec<u16> = (0..n).map(|i| base_port + i as u16).collect();

            // Create initial peer set with addresses
            let peer_set: Vec<(_, Address)> = addresses
                .iter()
                .enumerate()
                .map(|(i, pk)| {
                    (
                        pk.clone(),
                        Address::Symmetric(SocketAddr::new(
                            IpAddr::V4(Ipv4Addr::LOCALHOST),
                            ports[i],
                        )),
                    )
                })
                .collect();

            // Track senders/receivers/oracles/handles across restarts
            let mut senders: Vec<Option<channels::Sender<_, _>>> = (0..n).map(|_| None).collect();
            let mut receivers: Vec<Option<channels::Receiver<_>>> = (0..n).map(|_| None).collect();
            let mut oracles: Vec<Option<Oracle<_>>> = (0..n).map(|_| None).collect();
            let mut handles: Vec<Option<commonware_runtime::Handle<()>>> =
                (0..n).map(|_| None).collect();

            // Create networks for all peers
            for (i, peer) in peers.iter().enumerate() {
                let peer_context = context.child("peer").with_attribute("index", i);

                let config = Config::test(
                    peer.clone(),
                    SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), ports[i]),
                    MAX_MESSAGE_SIZE,
                );
                let (mut network, mut oracle) = Network::new(peer_context.child("network"), config);

                // Register peer set
                oracle.track(0, Map::try_from(peer_set.clone()).unwrap());

                let (sender, receiver) =
                    network.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
                senders[i] = Some(sender);
                receivers[i] = Some(receiver);
                oracles[i] = Some(oracle);

                let handle = network.start();
                handles[i] = Some(handle);
            }

            // Wait for full connectivity
            for (i, sender) in senders.iter_mut().enumerate() {
                let sender = sender.as_mut().unwrap();
                loop {
                    let sent = sender.send(
                        Recipients::All,
                        peers[i].public_key().as_ref().to_vec(),
                        true,
                    );
                    if sent.len() == n - 1 {
                        break;
                    }
                    context.sleep(Duration::from_millis(100)).await;
                }
            }

            // Verify each peer can receive from all others
            for receiver in receivers.iter_mut() {
                let receiver = receiver.as_mut().unwrap();
                let mut received = HashSet::new();
                while received.len() < n - 1 {
                    let (sender, message): (ed25519::PublicKey, _) = receiver.recv().await.unwrap();
                    assert_eq!(message, sender.as_ref());
                    received.insert(sender);
                }
            }

            // Shutdown all peers except peer 0 simultaneously.
            //
            // We keep peer 0 alive to exercise the case where multiple
            // peers churn at once.
            let restart_peers: Vec<usize> = (1..n).collect();
            for &idx in &restart_peers {
                if let Some(handle) = handles[idx].take() {
                    handle.abort();
                }
                senders[idx] = None;
                receivers[idx] = None;
                oracles[idx] = None;
                ports[idx] = base_port + 100 + idx as u16;
            }

            // Wait for connections to be detected as closed
            context.sleep(Duration::from_secs(2)).await;

            // Create updated peer set with new ports
            let updated_peer_set: Vec<(_, Address)> = addresses
                .iter()
                .enumerate()
                .map(|(i, pk)| {
                    (
                        pk.clone(),
                        Address::Symmetric(SocketAddr::new(
                            IpAddr::V4(Ipv4Addr::LOCALHOST),
                            ports[i],
                        )),
                    )
                })
                .collect();

            // Update oracle on peer 0 (the only running peer)
            oracles[0]
                .as_mut()
                .unwrap()
                .track(1, Map::try_from(updated_peer_set.clone()).unwrap());

            // Restart all shutdown peers with new ports
            for &idx in &restart_peers {
                let peer_context = context.child("peer_restarted").with_attribute("index", idx);
                let config = Config::test(
                    peers[idx].clone(),
                    SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), ports[idx]),
                    MAX_MESSAGE_SIZE,
                );
                let (mut network, mut oracle) = Network::new(peer_context.child("network"), config);

                oracle.track(1, Map::try_from(updated_peer_set.clone()).unwrap());

                let (sender, receiver) =
                    network.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
                senders[idx] = Some(sender);
                receivers[idx] = Some(receiver);
                oracles[idx] = Some(oracle);

                let handle = network.start();
                handles[idx] = Some(handle);
            }

            // Wait for full connectivity after restart
            for (i, sender) in senders.iter_mut().enumerate() {
                let sender = sender.as_mut().unwrap();
                loop {
                    let sent = sender.send(
                        Recipients::All,
                        peers[i].public_key().as_ref().to_vec(),
                        true,
                    );
                    if sent.len() == n - 1 {
                        break;
                    }
                    context.sleep(Duration::from_millis(100)).await;
                }
            }

            // Verify all peers can receive from all others after restart
            for receiver in receivers.iter_mut() {
                let receiver = receiver.as_mut().unwrap();
                let mut received = HashSet::new();
                while received.len() < n - 1 {
                    let (sender, message): (ed25519::PublicKey, _) = receiver.recv().await.unwrap();
                    assert_eq!(message, sender.as_ref());
                    received.insert(sender);
                }
            }

            assert_no_rate_limiting(&context.encode());
        });
    }
    #[test_traced]
    fn test_operations_after_shutdown_do_not_panic() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let peer = ed25519::PrivateKey::from_seed(0);
            let address = peer.public_key();

            let peer_context = context.child("peer");
            let config = Config::test(
                peer.clone(),
                SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 5200),
                MAX_MESSAGE_SIZE,
            );
            let (mut network, mut oracle) = Network::new(peer_context.child("network"), config);

            // Register channel and peer set
            let (mut sender, _receiver) =
                network.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
            let peer_addr =
                Address::Symmetric(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 5200));
            let peers: Map<ed25519::PublicKey, Address> =
                vec![(address.clone(), peer_addr)].try_into().unwrap();
            oracle.track(0, peers.clone());

            // Start and immediately abort the network
            let handle = network.start();
            handle.abort();

            // Wait for shutdown to propagate
            context.sleep(Duration::from_millis(100)).await;

            // Oracle operations should not panic even after shutdown
            oracle.track(1, peers.clone());
            let _ = oracle.peer_set(0).await;
            let _ = oracle.subscribe().await;
            crate::block_peer(&mut oracle, address.clone());

            // Sender operations should not panic even after shutdown
            let sent = sender.send(Recipients::All, address.as_ref().to_vec(), true);
            assert!(sent.is_empty());
        });
    }

    fn clean_shutdown(seed: u64) {
        let cfg = deterministic::Config::default()
            .with_seed(seed)
            .with_timeout(Some(Duration::from_secs(30)));
        let executor = deterministic::Runner::new(cfg);
        executor.start(|context| async move {
            let peer = ed25519::PrivateKey::from_seed(0);

            let peer_context = context.child("peer");
            let config = Config::test(
                peer.clone(),
                SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 5200),
                MAX_MESSAGE_SIZE,
            );
            let (mut network, mut oracle) = Network::new(peer_context.child("network"), config);

            // Register channel and peer set
            let (_, _) =
                network.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
            let peer_addr =
                Address::Symmetric(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 5200));
            let peers: Map<ed25519::PublicKey, Address> =
                vec![(peer.public_key(), peer_addr)].try_into().unwrap();
            oracle.track(0, peers);

            // Start the network
            let handle = network.start();

            // Allow tasks to start
            context.sleep(Duration::from_millis(100)).await;

            // Count running tasks under the network prefix
            let running_before = count_running_tasks(&context, "peer_network");
            assert!(
                running_before > 0,
                "at least one network task should be running"
            );

            // Abort the network
            handle.abort();
            let _ = handle.await;

            // Give the runtime a tick to process aborts
            context.sleep(Duration::from_millis(100)).await;

            // Verify all network tasks are stopped
            let running_after = count_running_tasks(&context, "peer_network");
            assert_eq!(
                running_after, 0,
                "all network tasks should be stopped, but {running_after} still running"
            );
        });
    }

    #[test_traced]
    fn test_clean_shutdown() {
        for seed in 0..25 {
            clean_shutdown(seed);
        }
    }

    fn duplicate_addresses_disconnected(seed: u64) {
        let base_port = 6000;
        let executor = deterministic::Runner::seeded(seed);
        executor.start(|context| async move {
            let peer0 = ed25519::PrivateKey::from_seed(0);
            let socket0 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port);
            let wrong_socket0 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + 100);
            let peer1 = ed25519::PrivateKey::from_seed(1);
            let socket1 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + 1);
            let peer2 = ed25519::PrivateKey::from_seed(2);

            // Start peer 0
            let config0 = Config::test(peer0.clone(), socket0, MAX_MESSAGE_SIZE);
            let (mut network0, mut oracle0) =
                Network::new(context.child("peer").with_attribute("index", 0), config0);
            let (mut sender0, _receiver0) =
                network0.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
            network0.start();

            // Create peer set with an incorrect address for peer 2
            let peer_set0: Vec<(_, Address)> = vec![
                (peer0.public_key(), Address::Symmetric(socket0)),
                (peer1.public_key(), Address::Symmetric(socket1)),
                (peer2.public_key(), Address::Symmetric(socket1)),
            ];
            oracle0.track(0, Map::try_from(peer_set0).unwrap());

            // Wait for connections to be attempted
            context.sleep(Duration::from_secs(30)).await;

            // Peer 0 can't send to anyone
            let sent = sender0.send(Recipients::All, peer1.public_key().as_ref().to_vec(), true);
            assert!(sent.is_empty());

            // Start peer 1 (has duplicate but correct address)
            let config1 = Config::test(peer1.clone(), socket1, MAX_MESSAGE_SIZE);
            let (mut network1, mut oracle1) =
                Network::new(context.child("peer").with_attribute("index", 1), config1);
            let (_sender1, mut receiver1) =
                network1.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
            network1.start();

            // Create peer set where peer 0 must dial peer 1 for connection
            let peer_set1: Vec<(_, Address)> = vec![
                (peer0.public_key(), Address::Symmetric(wrong_socket0)),
                (peer1.public_key(), Address::Symmetric(socket1)),
                (peer2.public_key(), Address::Symmetric(socket1)),
            ];
            oracle1.track(0, Map::try_from(peer_set1).unwrap());

            // Wait for connections to be made
            context.sleep(Duration::from_secs(30)).await;

            // Now peer 0 should connect to peer 1 at correct address
            loop {
                let sent =
                    sender0.send(Recipients::All, peer0.public_key().as_ref().to_vec(), true);
                if sent.len() == 1 {
                    assert_eq!(sent[0], peer1.public_key());
                    break;
                }
                context.sleep(Duration::from_millis(100)).await;
            }
            let (sender, _) = receiver1.recv().await.unwrap();
            assert_eq!(sender, peer0.public_key());
        });
    }

    #[test_traced]
    fn test_duplicate_addresses_disconnected() {
        // Ensure different dial orders explored by running with different seeds
        for seed in 0..25 {
            duplicate_addresses_disconnected(seed);
        }
    }

    #[test_traced]
    fn test_duplicate_addresses_connected() {
        let base_port = 6000;
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let peer0 = ed25519::PrivateKey::from_seed(0);
            let socket0 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port);
            let wrong_socket0 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + 100);
            let peer1 = ed25519::PrivateKey::from_seed(1);
            let socket1 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + 1);
            let peer2 = ed25519::PrivateKey::from_seed(2);
            let socket2 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), base_port + 2);

            // Start peer 0
            let config0 = Config::test(peer0.clone(), socket0, MAX_MESSAGE_SIZE);
            let (mut network0, mut oracle0) =
                Network::new(context.child("peer").with_attribute("index", 0), config0);
            let (mut sender0, mut receiver0) =
                network0.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
            network0.start();

            // Start peer 2
            let config2 = Config::test(peer2.clone(), socket2, MAX_MESSAGE_SIZE);
            let (mut network2, mut oracle2) =
                Network::new(context.child("peer").with_attribute("index", 2), config2);
            let (_sender2, mut receiver2) =
                network2.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
            network2.start();

            // Create peer set with an incorrect address for peer 2
            let peer_set: Vec<(_, Address)> = vec![
                (peer0.public_key(), Address::Symmetric(socket0)),
                (peer1.public_key(), Address::Symmetric(socket1)),
                (peer2.public_key(), Address::Symmetric(socket1)),
            ];
            oracle0.track(0, Map::try_from(peer_set.clone()).unwrap());
            oracle2.track(0, Map::try_from(peer_set).unwrap());

            // Wait for connections to be made
            context.sleep(Duration::from_secs(30)).await;

            // Peer 0 can send to peer 2
            loop {
                let sent =
                    sender0.send(Recipients::All, peer2.public_key().as_ref().to_vec(), true);
                if sent.len() == 1 {
                    assert_eq!(sent[0], peer2.public_key());
                    break;
                }
                context.sleep(Duration::from_millis(100)).await;
            }
            let (sender, _) = receiver2.recv().await.unwrap();
            assert_eq!(sender, peer0.public_key());

            // Start peer 1 (has duplicate but correct address)
            let config1 = Config::test(peer1.clone(), socket1, MAX_MESSAGE_SIZE);
            let (mut network1, mut oracle1) =
                Network::new(context.child("peer").with_attribute("index", 1), config1);
            let (mut sender1, _receiver1) =
                network1.register(0, Quota::per_second(NZU32!(100)), DEFAULT_MESSAGE_BACKLOG);
            network1.start();

            // Create peer set where peer 0 must dial peer 1 for connection
            let peer_set1: Vec<(_, Address)> = vec![
                (peer0.public_key(), Address::Symmetric(wrong_socket0)),
                (peer1.public_key(), Address::Symmetric(socket1)),
                (peer2.public_key(), Address::Symmetric(socket1)),
            ];
            oracle1.track(0, Map::try_from(peer_set1).unwrap());

            // Wait for connections to be made
            context.sleep(Duration::from_secs(30)).await;

            // Now peer 0 should connect to peer 1 at correct address and peer 2 should dial peer 1
            loop {
                let sent =
                    sender1.send(Recipients::All, peer1.public_key().as_ref().to_vec(), true);
                if sent.len() == 2 {
                    assert!(sent.contains(&peer0.public_key()));
                    assert!(sent.contains(&peer2.public_key()));
                    break;
                }
                context.sleep(Duration::from_millis(100)).await;
            }
            let mut received0 = false;
            while let Ok((sender, _)) = receiver0.recv().await {
                // May have some items around from the initial send
                if sender == peer1.public_key() {
                    received0 = true;
                    break;
                }
            }
            assert!(received0);
            let (sender, _) = receiver2.recv().await.unwrap();
            assert_eq!(sender, peer1.public_key());
        });
    }

    #[test]
    fn test_broadcast_slow_peer_no_blocking() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let cfg = RouterConfig {
                mailbox_size: NZUsize!(1),
            };
            let (router, mailbox, messenger) =
                RouterActor::<_, ed25519::PublicKey>::new(context.child("router"), cfg);

            let channels = channels::Channels::new(messenger.clone(), MAX_MESSAGE_SIZE);
            let _handle = router.start(channels);

            let slow_peer = ed25519::PrivateKey::from_seed(0).public_key();
            let (slow_relay, _slow_receivers) =
                Relay::new(context.child("slow_relay"), NZUsize!(10));
            assert!(
                mailbox.ready(slow_peer, slow_relay).await.is_some(),
                "Failed to register slow peer"
            );

            let fast_peer = ed25519::PrivateKey::from_seed(1).public_key();
            let (fast_relay, mut fast_receivers) =
                Relay::new(context.child("fast_relay"), NZUsize!(100));
            assert!(
                mailbox.ready(fast_peer, fast_relay).await.is_some(),
                "Failed to register fast peer"
            );

            let message = IoBuf::from(vec![0u8; 100]);
            for i in 0..11 {
                let sent = messenger.content(Recipients::All, 0, message.clone().into(), false);
                assert_ne!(
                    sent,
                    Unreliable::new(Feedback::Closed),
                    "Broadcast {i} should be accepted"
                );

                assert!(fast_receivers.low.recv().await.is_some());
            }
        });
    }
}