arc-malachitebft-network 0.7.0-pre

use std::error::Error;
use std::ops::ControlFlow;
use std::time::Duration;

use futures::StreamExt;
use libp2p::metrics::{Metrics, Recorder};
use libp2p::request_response::{InboundRequestId, OutboundRequestId};
use libp2p::swarm::{self, SwarmEvent};
use libp2p::{gossipsub, identify, quic, SwarmBuilder};
use libp2p_broadcast as broadcast;
use tokio::sync::{mpsc, oneshot};
use tracing::{debug, error, error_span, info, trace, warn, Instrument};

use malachitebft_discovery::{self as discovery};
use malachitebft_metrics::SharedRegistry;
use malachitebft_sync::{self as sync};

pub use malachitebft_peer::PeerId;

pub use bytes::Bytes;
pub use libp2p::gossipsub::MessageId;
pub use libp2p::identity::Keypair;
pub use libp2p::Multiaddr;

pub mod behaviour;
pub mod handle;
pub mod pubsub;

mod channel;
pub use channel::{Channel, ChannelNames};

mod metrics;
use metrics::Metrics as NetworkMetrics;

mod peer_type;
pub use peer_type::PeerType;

pub mod peer_scoring;

mod utils;

mod ip_limits;

// Re-export state types for external use (e.g., RPC)
pub use state::{LocalNodeInfo, PeerInfo, ValidatorInfo};

mod state;
pub use state::NetworkStateDump;
use state::State;

use behaviour::{Behaviour, NetworkEvent};
use handle::Handle;

const METRICS_PREFIX: &str = "malachitebft_network";
const DISCOVERY_METRICS_PREFIX: &str = "malachitebft_discovery";

#[derive(Clone, Debug, PartialEq)]
pub struct ProtocolNames {
    pub consensus: String,
    pub discovery_kad: String,
    pub discovery_regres: String,
    pub sync: String,
    pub broadcast: String,
}

impl Default for ProtocolNames {
    fn default() -> Self {
        Self {
            consensus: "/malachitebft-core-consensus/v1beta1".to_string(),
            discovery_kad: "/malachitebft-discovery/kad/v1beta1".to_string(),
            discovery_regres: "/malachitebft-discovery/reqres/v1beta1".to_string(),
            sync: "/malachitebft-sync/v1beta1".to_string(),
            broadcast: "/malachitebft-broadcast/v1beta1".to_string(),
        }
    }
}

#[derive(Copy, Clone, Debug, Default)]
pub enum PubSubProtocol {
    /// GossipSub: a pubsub protocol based on epidemic broadcast trees
    #[default]
    GossipSub,

    /// Broadcast: a simple broadcast protocol
    Broadcast,
}

impl PubSubProtocol {
    pub fn is_gossipsub(&self) -> bool {
        matches!(self, Self::GossipSub)
    }

    pub fn is_broadcast(&self) -> bool {
        matches!(self, Self::Broadcast)
    }
}

#[derive(Copy, Clone, Debug)]
pub struct GossipSubConfig {
    pub mesh_n: usize,
    pub mesh_n_high: usize,
    pub mesh_n_low: usize,
    pub mesh_outbound_min: usize,
    pub enable_peer_scoring: bool,
    pub enable_explicit_peering: bool,
    pub enable_flood_publish: bool,
}

impl Default for GossipSubConfig {
    fn default() -> Self {
        // Tests use these defaults.
        Self {
            mesh_n: 6,
            mesh_n_high: 12,
            mesh_n_low: 4,
            mesh_outbound_min: 2,
            enable_peer_scoring: false,
            enable_explicit_peering: false,
            enable_flood_publish: true,
        }
    }
}

pub type BoxError = Box<dyn Error + Send + Sync + 'static>;

pub type DiscoveryConfig = discovery::Config;
pub type BootstrapProtocol = discovery::config::BootstrapProtocol;
pub type Selector = discovery::config::Selector;

/// Node identity bundling all node-specific information.
///
/// The consensus address is derived from the keypair in the current implementation
/// where libp2p and consensus use the same key. In the future, when using separate
/// keys (e.g., cc-signer for consensus), the address will be provided separately.
///
/// If consensus_address is None, the node will not advertise a validator address
/// and cannot become a validator.
#[derive(Clone, Debug)]
pub struct NetworkIdentity {
    pub moniker: String,
    pub keypair: Keypair,
    pub consensus_address: Option<String>,
}

impl NetworkIdentity {
    /// Create a new NodeIdentity.
    ///
    /// # Arguments
    /// * `moniker` - Human-readable node identifier
    /// * `keypair` - libp2p keypair for network authentication
    /// * `consensus_address` - Optional consensus address (Some = potential validator, None = full node)
    ///
    /// In the current implementation where libp2p and consensus share the same key,
    /// the address is typically derived from the keypair before calling this method.
    /// In the future with cc-signer, the consensus address will be separate.
    pub fn new(moniker: String, keypair: Keypair, consensus_address: Option<String>) -> Self {
        Self {
            moniker,
            keypair,
            consensus_address,
        }
    }
}

#[derive(Clone, Debug)]
pub struct Config {
    pub listen_addr: Multiaddr,
    pub persistent_peers: Vec<Multiaddr>,
    pub persistent_peers_only: bool,
    pub discovery: DiscoveryConfig,
    pub idle_connection_timeout: Duration,
    pub transport: TransportProtocol,
    pub gossipsub: GossipSubConfig,
    pub pubsub_protocol: PubSubProtocol,
    pub channel_names: ChannelNames,
    pub rpc_max_size: usize,
    pub pubsub_max_size: usize,
    pub enable_consensus: bool,
    pub enable_sync: bool,
    pub protocol_names: ProtocolNames,
}

impl Config {
    fn apply_to_swarm(&self, cfg: swarm::Config) -> swarm::Config {
        cfg.with_idle_connection_timeout(self.idle_connection_timeout)
    }

    fn apply_to_quic(&self, mut cfg: quic::Config) -> quic::Config {
        // NOTE: This is set low due to quic transport not properly resetting
        // connection state when reconnecting before connection timeout.
        // See https://github.com/libp2p/rust-libp2p/issues/5097
        cfg.max_idle_timeout = 300;
        cfg.keep_alive_interval = Duration::from_millis(100);
        cfg
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum TransportProtocol {
    Tcp,
    Quic,
}

impl TransportProtocol {
    pub fn from_multiaddr(multiaddr: &Multiaddr) -> Option<TransportProtocol> {
        for protocol in multiaddr.protocol_stack() {
            match protocol {
                "tcp" => return Some(TransportProtocol::Tcp),
                "quic" | "quic-v1" => return Some(TransportProtocol::Quic),
                _ => {}
            }
        }
        None
    }
}

/// Operation to perform on a persistent peer
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PersistentPeersOp {
    /// Add a persistent peer
    Add(Multiaddr),
    /// Remove a persistent peer
    Remove(Multiaddr),
}

/// Errors that can occur during persistent peer operations
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum PersistentPeerError {
    /// Peer already exists in the persistent peers list (for Add operation)
    #[error("Persistent peer already exists")]
    AlreadyExists,
    /// Peer not found in the persistent peers list (for Remove operation)
    #[error("Persistent peer not found")]
    NotFound,
    /// Network is not started
    #[error("Network not started")]
    NetworkStopped,
    /// Internal error
    #[error("Internal error: {0}")]
    InternalError(String),
}

/// sync event details:
///
/// peer1: sync                  peer2: network                    peer2: sync              peer1: network
/// CtrlMsg::SyncRequest       --> Event::Sync      -----------> CtrlMsg::SyncReply ------> Event::Sync
/// (peer_id, height)             (RawMessage::Request           (request_id, height)       RawMessage::Response
///                           {request_id, peer_id, request}                                {request_id, response}
///
///
/// An event that can be emitted by the gossip layer
#[derive(Clone, Debug)]
pub enum Event {
    Listening(Multiaddr),
    PeerConnected(PeerId),
    PeerDisconnected(PeerId),
    ConsensusMessage(Channel, PeerId, Bytes),
    LivenessMessage(Channel, PeerId, Bytes),
    Sync(sync::RawMessage),
}

#[derive(Debug)]
pub enum CtrlMsg {
    Publish(Channel, Bytes),
    Broadcast(Channel, Bytes),
    SyncRequest(PeerId, Bytes, oneshot::Sender<OutboundRequestId>),
    SyncReply(InboundRequestId, Bytes),
    UpdateValidatorSet(Vec<ValidatorInfo>),
    DumpState(oneshot::Sender<NetworkStateDump>),
    UpdatePersistentPeers(
        PersistentPeersOp,
        oneshot::Sender<Result<(), PersistentPeerError>>,
    ),
    Shutdown,
}

pub async fn spawn(
    identity: NetworkIdentity,
    config: Config,
    registry: SharedRegistry,
) -> Result<Handle, eyre::Report> {
    let swarm = registry.with_prefix(METRICS_PREFIX, |registry| -> Result<_, eyre::Report> {
        // Pass the libp2p keypair to the behaviour, it is included in the Identify protocol
        // Required for ALL nodes
        let builder = SwarmBuilder::with_existing_identity(identity.keypair.clone()).with_tokio();
        match config.transport {
            TransportProtocol::Tcp => {
                let behaviour = Behaviour::new_with_metrics(&config, &identity, registry)?;
                Ok(builder
                    .with_tcp(
                        libp2p::tcp::Config::new().nodelay(true), // Disable Nagle's algorithm
                        libp2p::noise::Config::new,
                        libp2p::yamux::Config::default,
                    )?
                    .with_dns()?
                    .with_bandwidth_metrics(registry)
                    .with_behaviour(|_| behaviour)?
                    .with_swarm_config(|cfg| config.apply_to_swarm(cfg))
                    .build())
            }
            TransportProtocol::Quic => {
                let behaviour = Behaviour::new_with_metrics(&config, &identity, registry)?;
                Ok(builder
                    .with_quic_config(|cfg| config.apply_to_quic(cfg))
                    .with_dns()?
                    .with_bandwidth_metrics(registry)
                    .with_behaviour(|_| behaviour)?
                    .with_swarm_config(|cfg| config.apply_to_swarm(cfg))
                    .build())
            }
        }
    })?;

    let metrics = registry.with_prefix(METRICS_PREFIX, Metrics::new);

    let (tx_event, rx_event) = mpsc::channel(32);
    let (tx_ctrl, rx_ctrl) = mpsc::channel(32);

    let discovery = registry.with_prefix(DISCOVERY_METRICS_PREFIX, |reg| {
        discovery::Discovery::new(config.discovery, config.persistent_peers.clone(), reg)
    });

    let network_metrics = registry.with_prefix(METRICS_PREFIX, NetworkMetrics::new);

    let peer_id = PeerId::from_libp2p(swarm.local_peer_id());

    // Create local node info with subscribed consensus topics
    let mut subscribed_topics = std::collections::HashSet::new();
    if config.enable_consensus {
        for channel in Channel::consensus() {
            subscribed_topics.insert(channel.as_str(config.channel_names).to_string());
        }
    }

    let NetworkIdentity {
        moniker,
        consensus_address,
        ..
    } = identity;

    // Create local node info
    let local_node_info = LocalNodeInfo {
        moniker,
        peer_id: *swarm.local_peer_id(),
        listen_addr: config.listen_addr.clone(),
        subscribed_topics,
        consensus_address,
        is_validator: false, // Will be updated when validator set is received
        persistent_peers_only: config.persistent_peers_only,
    };

    // Set local node info in metrics
    network_metrics.set_local_node_info(&local_node_info);

    let state = State::new(
        discovery,
        config.persistent_peers.clone(),
        local_node_info,
        network_metrics,
    );

    let span = error_span!("network");

    info!(parent: span.clone(), %peer_id, "Starting network service");

    let task_handle =
        tokio::task::spawn(run(config, metrics, state, swarm, rx_ctrl, tx_event).instrument(span));

    Ok(Handle::new(peer_id, tx_ctrl, rx_event, task_handle))
}

async fn run(
    config: Config,
    metrics: Metrics,
    mut state: State,
    mut swarm: swarm::Swarm<Behaviour>,
    mut rx_ctrl: mpsc::Receiver<CtrlMsg>,
    tx_event: mpsc::Sender<Event>,
) {
    if let Err(e) = swarm.listen_on(config.listen_addr.clone()) {
        error!("Error listening on {}: {e}", config.listen_addr);
        return;
    }

    if config.enable_consensus {
        if let Err(e) = pubsub::subscribe(
            &mut swarm,
            config.pubsub_protocol,
            Channel::consensus(),
            config.channel_names,
        ) {
            error!("Error subscribing to consensus channels: {e}");
            return;
        };
    }

    if config.enable_sync {
        if let Err(e) = pubsub::subscribe(
            &mut swarm,
            PubSubProtocol::Broadcast,
            &[Channel::Sync],
            config.channel_names,
        ) {
            error!("Error subscribing to Sync channel: {e}");
            return;
        };
    }

    // Timer to perform periodic network operations (peer reconnection, metrics updates, etc.)
    // TODO: Using 1 second for now, for faster reconnection during testing
    // Maybe adjust via config in the future
    let mut periodic_timer = tokio::time::interval(std::time::Duration::from_secs(1));
    let mut periodic_tick_count: u32 = 0;

    loop {
        let result = tokio::select! {
            event = swarm.select_next_some() => {
                handle_swarm_event(event, &config, &metrics, &mut swarm, &mut state, &tx_event).await
            }

            Some(connection_data) = state.discovery.controller.dial.recv(), if state.discovery.can_dial() => {
                state.discovery.dial_peer(&mut swarm, connection_data);
                ControlFlow::Continue(())
            }

            Some(request_data) = state.discovery.controller.peers_request.recv(), if state.discovery.can_peers_request() => {
                state.discovery.peers_request_peer(&mut swarm, request_data);
                ControlFlow::Continue(())
            }

            Some(request_data) = state.discovery.controller.connect_request.recv(), if state.discovery.can_connect_request() => {
                state.discovery.connect_request_peer(&mut swarm, request_data);
                ControlFlow::Continue(())
            }

            Some((peer_id, connection_id)) = state.discovery.controller.close.recv(), if state.discovery.can_close() => {
                state.discovery.close_connection(&mut swarm, peer_id, connection_id);
                ControlFlow::Continue(())
            }

            Some(ctrl) = rx_ctrl.recv() => {
                handle_ctrl_msg(&mut swarm, &mut state, &config, ctrl).await
            }

            _ = periodic_timer.tick() => {
                // Attempt to dial bootstrap nodes
                state.discovery.dial_bootstrap_nodes(&swarm);

                // Update peer info in State and metrics (includes gossipsub scores and mesh membership)
                if let Some(gossipsub) = swarm.behaviour_mut().gossipsub.as_mut() {
                    state.update_peer_info(
                        gossipsub,
                        Channel::consensus(),
                        config.channel_names,
                    );
                }

                periodic_tick_count = periodic_tick_count.wrapping_add(1);
                if periodic_tick_count.is_multiple_of(5) {
                    info!("Network peer state\n{}", state.format_peer_info());
                }

                ControlFlow::Continue(())
            }
        };

        match result {
            ControlFlow::Continue(()) => continue,
            ControlFlow::Break(()) => break,
        }
    }
}

async fn handle_ctrl_msg(
    swarm: &mut swarm::Swarm<Behaviour>,
    state: &mut State,
    config: &Config,
    msg: CtrlMsg,
) -> ControlFlow<()> {
    match msg {
        CtrlMsg::Publish(channel, data) => {
            let msg_size = data.len();
            let result = pubsub::publish(
                swarm,
                config.pubsub_protocol,
                channel,
                config.channel_names,
                data,
            );

            match result {
                Ok(()) => debug!(%channel, size = %msg_size, "Published message"),
                Err(e) => error!(%channel, "Error publishing message: {e}"),
            }

            ControlFlow::Continue(())
        }

        CtrlMsg::Broadcast(channel, data) => {
            if channel == Channel::Sync && !config.enable_sync {
                trace!("Ignoring broadcast message to Sync channel: Sync not enabled");
                return ControlFlow::Continue(());
            }

            let msg_size = data.len();
            let result = pubsub::publish(
                swarm,
                PubSubProtocol::Broadcast,
                channel,
                config.channel_names,
                data,
            );

            match result {
                Ok(()) => debug!(%channel, size = %msg_size, "Broadcasted message"),
                Err(e) => error!(%channel, "Error broadcasting message: {e}"),
            }

            ControlFlow::Continue(())
        }

        CtrlMsg::SyncRequest(peer_id, request, reply_to) => {
            let Some(sync) = swarm.behaviour_mut().sync.as_mut() else {
                error!("Cannot request Sync from peer: Sync not enabled");
                return ControlFlow::Continue(());
            };

            let request_id = sync.send_request(peer_id.to_libp2p(), request);

            if let Err(e) = reply_to.send(request_id) {
                error!(%peer_id, "Error sending Sync request: {e}");
            }

            ControlFlow::Continue(())
        }

        CtrlMsg::SyncReply(request_id, data) => {
            let Some(sync) = swarm.behaviour_mut().sync.as_mut() else {
                error!("Cannot send Sync response to peer: Sync not enabled");
                return ControlFlow::Continue(());
            };

            let Some(channel) = state.sync_channels.remove(&request_id) else {
                error!(%request_id, "Received Sync reply for unknown request ID");
                return ControlFlow::Continue(());
            };

            let result = sync.send_response(channel, data);

            match result {
                Ok(()) => debug!(%request_id, "Replied to Sync request"),
                Err(e) => error!(%request_id, "Error replying to Sync request: {e}"),
            }

            ControlFlow::Continue(())
        }

        CtrlMsg::UpdateValidatorSet(validators) => {
            // Process the validator set update and get peers that need score updates
            let changed_peers = state.process_validator_set_update(validators);

            // Update GossipSub scores for peers whose type changed
            for (peer_id, new_score) in changed_peers {
                set_peer_score(swarm, peer_id, new_score);
            }

            ControlFlow::Continue(())
        }

        CtrlMsg::DumpState(reply_to) => {
            // Build a snapshot from current state
            let snapshot = NetworkStateDump {
                local_node: state.local_node.clone(),
                peers: state.peer_info.clone(),
                validator_set: state.validator_set.clone(),
                persistent_peer_ids: state.persistent_peer_ids.iter().copied().collect(),
                persistent_peer_addrs: state.persistent_peer_addrs.clone(),
            };
            if let Err(_s) = reply_to.send(snapshot) {
                error!("Error replying to DumpState");
            }
            ControlFlow::Continue(())
        }

        CtrlMsg::UpdatePersistentPeers(op, reply_to) => {
            let result = match op {
                PersistentPeersOp::Add(addr) => state.add_persistent_peer(addr, swarm),
                PersistentPeersOp::Remove(addr) => state.remove_persistent_peer(addr, swarm),
            };
            if reply_to.send(result).is_err() {
                error!("Error replying to UpdatePersistentPeers");
            }
            ControlFlow::Continue(())
        }

        CtrlMsg::Shutdown => ControlFlow::Break(()),
    }
}

/// Set a default low score for a peer immediately upon connection
/// This allows gossipsub to form an initial mesh before Identify completes
fn set_default_peer_score(swarm: &mut swarm::Swarm<Behaviour>, peer_id: libp2p::PeerId) {
    if let Some(gossipsub) = swarm.behaviour_mut().gossipsub.as_mut() {
        let score = peer_scoring::get_default_score();
        gossipsub.set_application_score(&peer_id, score);
        trace!("Set default application score {score} for peer {peer_id} before Identify");
    }
}

fn set_peer_score(swarm: &mut swarm::Swarm<Behaviour>, peer_id: libp2p::PeerId, score: f64) {
    // Set application-specific score in gossipsub if enabled
    if let Some(gossipsub) = swarm.behaviour_mut().gossipsub.as_mut() {
        gossipsub.set_application_score(&peer_id, score);
        debug!("Upgraded application score to {score} for peer {peer_id}");
    }
}

/// Add a persistent peer as an explicit peer in gossipsub (if explicit peering is enabled).
/// A node always sends and forwards messages to its explicit peers, regardless of mesh membership.
fn add_explicit_peer_to_gossipsub(
    swarm: &mut swarm::Swarm<Behaviour>,
    state: &mut State,
    peer_id: libp2p::PeerId,
) {
    let Some(peer_info) = state.peer_info.get_mut(&peer_id) else {
        return;
    };

    if peer_info.peer_type.is_persistent() {
        if let Some(gossipsub) = swarm.behaviour_mut().gossipsub.as_mut() {
            gossipsub.add_explicit_peer(&peer_id);
            state
                .metrics
                .record_explicit_peer(&peer_id, &peer_info.moniker);
            peer_info.is_explicit = true;
            info!("Added persistent peer {peer_id} as explicit peer in gossipsub");
        }
    }
}

/// Remove a persistent peer from explicit peers in gossipsub and mark the metric stale.
fn remove_explicit_peer_from_gossipsub(
    swarm: &mut swarm::Swarm<Behaviour>,
    state: &mut State,
    peer_id: &libp2p::PeerId,
) {
    let Some(peer_info) = state.peer_info.get_mut(peer_id) else {
        return;
    };

    if peer_info.peer_type.is_persistent() {
        if let Some(gossipsub) = swarm.behaviour_mut().gossipsub.as_mut() {
            gossipsub.remove_explicit_peer(peer_id);
            state
                .metrics
                .mark_explicit_peer_stale(peer_id, &peer_info.moniker);
            peer_info.is_explicit = false;
            info!("Removed persistent peer {peer_id} from explicit peers in gossipsub");
        }
    }
}

async fn handle_swarm_event(
    event: SwarmEvent<NetworkEvent>,
    config: &Config,
    metrics: &Metrics,
    swarm: &mut swarm::Swarm<Behaviour>,
    state: &mut State,
    tx_event: &mpsc::Sender<Event>,
) -> ControlFlow<()> {
    if let SwarmEvent::Behaviour(NetworkEvent::GossipSub(e)) = &event {
        metrics.record(e);
    } else if let SwarmEvent::Behaviour(NetworkEvent::Identify(e)) = &event {
        metrics.record(e.as_ref());
    }

    match event {
        SwarmEvent::NewListenAddr { address, .. } => {
            debug!(%address, "Node is listening");

            if let Err(e) = tx_event.send(Event::Listening(address)).await {
                error!("Error sending listening event to handle: {e}");
                return ControlFlow::Break(());
            }
        }

        SwarmEvent::ConnectionEstablished {
            peer_id,
            connection_id,
            endpoint,
            num_established,
            ..
        } => {
            trace!("Connected to {peer_id} with connection id {connection_id}");

            // Set a low default score immediately for gossipsub mesh formation
            // This will be upgraded later when Identify completes
            if num_established.get() == 1 {
                // Only set score on first connection to this peer
                set_default_peer_score(swarm, peer_id);
            }

            state
                .discovery
                .handle_connection(swarm, peer_id, connection_id, endpoint);
        }

        SwarmEvent::OutgoingConnectionError {
            connection_id,
            error,
            ..
        } => {
            error!("Error dialing peer: {error}");

            state
                .discovery
                .handle_failed_connection(swarm, connection_id, error);
        }

        SwarmEvent::ConnectionClosed {
            peer_id,
            connection_id,
            num_established,
            cause,
            ..
        } => {
            debug!(
                "SwarmEvent::ConnectionClosed: peer_id={}, connection_id={}, num_established={}",
                peer_id, connection_id, num_established
            );
            if let Some(cause) = cause {
                warn!("Connection closed with {peer_id}, reason: {cause}");
            } else {
                warn!("Connection closed with {peer_id}, reason: unknown");
            }

            state
                .discovery
                .handle_closed_connection(swarm, peer_id, connection_id);

            if num_established == 0 {
                // Remove explicit peer from gossipsub and mark metric stale when this peer was one
                if config.gossipsub.enable_explicit_peering {
                    remove_explicit_peer_from_gossipsub(swarm, state, &peer_id);
                }

                if let Err(e) = tx_event
                    .send(Event::PeerDisconnected(PeerId::from_libp2p(&peer_id)))
                    .await
                {
                    error!("Error sending peer disconnected event to handle: {e}");
                    return ControlFlow::Break(());
                }
            }
        }

        SwarmEvent::Behaviour(NetworkEvent::Identify(event)) => match *event {
            identify::Event::Sent { peer_id, .. } => {
                trace!("Sent identity to {peer_id}");
            }

            identify::Event::Received {
                connection_id,
                peer_id,
                info,
            } => {
                info!(
                    "Received identity from {peer_id}: protocol={:?} agent={:?}",
                    info.protocol_version, info.agent_version
                );

                if info.protocol_version == config.protocol_names.consensus {
                    trace!(
                        "Peer {peer_id} is using compatible protocol version: {:?}",
                        info.protocol_version
                    );

                    let is_already_connected = state.discovery.handle_new_peer(
                        swarm,
                        connection_id,
                        peer_id,
                        info.clone(),
                    );

                    // Update peer info in State and metrics, set peer score in gossipsub
                    let score = state.update_peer(peer_id, connection_id, &info);
                    set_peer_score(swarm, peer_id, score);

                    // If enabled, add persistent peers as explicit peers for guaranteed delivery
                    if config.gossipsub.enable_explicit_peering {
                        add_explicit_peer_to_gossipsub(swarm, state, peer_id);
                    }

                    if !is_already_connected {
                        if let Err(e) = tx_event
                            .send(Event::PeerConnected(PeerId::from_libp2p(&peer_id)))
                            .await
                        {
                            error!("Error sending peer connected event to handle: {e}");
                            return ControlFlow::Break(());
                        }
                    }
                } else {
                    trace!(
                        "Peer {peer_id} is using incompatible protocol version: {:?}",
                        info.protocol_version
                    );
                }
            }

            // Ignore other identify events
            _ => (),
        },

        SwarmEvent::Behaviour(NetworkEvent::Ping(event)) => {
            match &event.result {
                Ok(rtt) => {
                    trace!("Received pong from {} in {rtt:?}", event.peer);
                }
                Err(e) => {
                    trace!("Received pong from {} with error: {e}", event.peer);
                }
            }

            // Record metric for round-trip time sending a ping and receiving a pong
            metrics.record(&event);
        }

        SwarmEvent::Behaviour(NetworkEvent::GossipSub(event)) => {
            return handle_gossipsub_event(event, config, metrics, swarm, state, tx_event).await;
        }

        SwarmEvent::Behaviour(NetworkEvent::Broadcast(event)) => {
            return handle_broadcast_event(event, config, metrics, swarm, state, tx_event).await;
        }

        SwarmEvent::Behaviour(NetworkEvent::Sync(event)) => {
            return handle_sync_event(event, metrics, swarm, state, tx_event).await;
        }

        SwarmEvent::Behaviour(NetworkEvent::Discovery(network_event)) => {
            state.discovery.on_network_event(swarm, *network_event);
        }

        swarm_event => {
            metrics.record(&swarm_event);
        }
    }

    ControlFlow::Continue(())
}

async fn handle_gossipsub_event(
    event: gossipsub::Event,
    config: &Config,
    _metrics: &Metrics,
    _swarm: &mut swarm::Swarm<Behaviour>,
    _state: &mut State,
    tx_event: &mpsc::Sender<Event>,
) -> ControlFlow<()> {
    match event {
        gossipsub::Event::Subscribed { peer_id, topic } => {
            if !Channel::has_gossipsub_topic(&topic, config.channel_names) {
                trace!("Peer {peer_id} tried to subscribe to unknown topic: {topic}");
                return ControlFlow::Continue(());
            }

            trace!("Peer {peer_id} subscribed to {topic}");
        }

        gossipsub::Event::Unsubscribed { peer_id, topic } => {
            if !Channel::has_gossipsub_topic(&topic, config.channel_names) {
                trace!("Peer {peer_id} tried to unsubscribe from unknown topic: {topic}");
                return ControlFlow::Continue(());
            }

            trace!("Peer {peer_id} unsubscribed from {topic}");
        }

        gossipsub::Event::Message {
            message_id,
            message,
            ..
        } => {
            let Some(peer_id) = message.source else {
                return ControlFlow::Continue(());
            };

            let Some(channel) =
                Channel::from_gossipsub_topic_hash(&message.topic, config.channel_names)
            else {
                trace!(
                    "Received message {message_id} from {peer_id} on different channel: {}",
                    message.topic
                );

                return ControlFlow::Continue(());
            };

            trace!(
                "Received message {message_id} from {peer_id} on channel {channel} of {} bytes",
                message.data.len()
            );

            let peer_id = PeerId::from_libp2p(&peer_id);

            let event = if channel == Channel::Liveness {
                Event::LivenessMessage(channel, peer_id, Bytes::from(message.data))
            } else {
                Event::ConsensusMessage(channel, peer_id, Bytes::from(message.data))
            };

            if let Err(e) = tx_event.send(event).await {
                error!("Error sending message to handle: {e}");
                return ControlFlow::Break(());
            }
        }

        gossipsub::Event::SlowPeer {
            peer_id,
            failed_messages,
        } => {
            trace!(
                "Slow peer detected: {peer_id}, total failed messages: {}",
                failed_messages.total()
            );
        }

        gossipsub::Event::GossipsubNotSupported { peer_id } => {
            trace!("Peer does not support GossipSub: {peer_id}");
        }
    }

    ControlFlow::Continue(())
}

async fn handle_broadcast_event(
    event: broadcast::Event,
    config: &Config,
    _metrics: &Metrics,
    _swarm: &mut swarm::Swarm<Behaviour>,
    _state: &mut State,
    tx_event: &mpsc::Sender<Event>,
) -> ControlFlow<()> {
    match event {
        broadcast::Event::Subscribed(peer_id, topic) => {
            if !Channel::has_broadcast_topic(&topic, config.channel_names) {
                trace!("Peer {peer_id} tried to subscribe to unknown topic: {topic:?}");
                return ControlFlow::Continue(());
            }

            trace!("Peer {peer_id} subscribed to {topic:?}");
        }

        broadcast::Event::Unsubscribed(peer_id, topic) => {
            if !Channel::has_broadcast_topic(&topic, config.channel_names) {
                trace!("Peer {peer_id} tried to unsubscribe from unknown topic: {topic:?}");
                return ControlFlow::Continue(());
            }

            trace!("Peer {peer_id} unsubscribed from {topic:?}");
        }

        broadcast::Event::Received(peer_id, topic, message) => {
            let Some(channel) = Channel::from_broadcast_topic(&topic, config.channel_names) else {
                trace!("Received message from {peer_id} on different channel: {topic:?}");
                return ControlFlow::Continue(());
            };

            trace!(
                "Received message from {peer_id} on channel {channel} of {} bytes",
                message.len()
            );

            let peer_id = PeerId::from_libp2p(&peer_id);

            let event = if channel == Channel::Liveness {
                Event::LivenessMessage(channel, peer_id, message)
            } else {
                Event::ConsensusMessage(channel, peer_id, message)
            };

            if let Err(e) = tx_event.send(event).await {
                error!("Error sending message to handle: {e}");
                return ControlFlow::Break(());
            }
        }
    }

    ControlFlow::Continue(())
}

async fn handle_sync_event(
    event: sync::Event,
    _metrics: &Metrics,
    _swarm: &mut swarm::Swarm<Behaviour>,
    state: &mut State,
    tx_event: &mpsc::Sender<Event>,
) -> ControlFlow<()> {
    match event {
        sync::Event::Message { peer, message, .. } => {
            match message {
                libp2p::request_response::Message::Request {
                    request_id,
                    request,
                    channel,
                } => {
                    state.sync_channels.insert(request_id, channel);

                    let _ = tx_event
                        .send(Event::Sync(sync::RawMessage::Request {
                            request_id,
                            peer: PeerId::from_libp2p(&peer),
                            body: request.0,
                        }))
                        .await
                        .map_err(|e| {
                            error!("Error sending Sync request to handle: {e}");
                        });
                }

                libp2p::request_response::Message::Response {
                    request_id,
                    response,
                } => {
                    let _ = tx_event
                        .send(Event::Sync(sync::RawMessage::Response {
                            request_id,
                            peer: PeerId::from_libp2p(&peer),
                            body: response.0,
                        }))
                        .await
                        .map_err(|e| {
                            error!("Error sending Sync response to handle: {e}");
                        });
                }
            }

            ControlFlow::Continue(())
        }

        sync::Event::ResponseSent { .. } => ControlFlow::Continue(()),

        sync::Event::OutboundFailure { .. } => ControlFlow::Continue(()),

        sync::Event::InboundFailure { .. } => ControlFlow::Continue(()),
    }
}

pub trait PeerIdExt {
    fn to_libp2p(&self) -> libp2p::PeerId;
    fn from_libp2p(peer_id: &libp2p::PeerId) -> Self;
}

impl PeerIdExt for PeerId {
    fn to_libp2p(&self) -> libp2p::PeerId {
        libp2p::PeerId::from_bytes(&self.to_bytes()).expect("valid PeerId")
    }

    fn from_libp2p(peer_id: &libp2p::PeerId) -> Self {
        Self::from_bytes(&peer_id.to_bytes()).expect("valid PeerId")
    }
}