ckb_network/

network.rs

//! Global state struct and start function
use crate::errors::Error;
#[cfg(not(target_family = "wasm"))]
use crate::errors::P2PError;
use crate::peer_registry::{ConnectionStatus, PeerRegistry};
use crate::peer_store::{
    PeerStore,
    types::{AddrInfo, BannedAddr},
};
use crate::protocols::{
    disconnect_message::DisconnectMessageProtocol,
    discovery::{DiscoveryAddressManager, DiscoveryProtocol},
    feeler::Feeler,
    identify::{Flags, IdentifyCallback, IdentifyProtocol},
    ping::PingHandler,
    support_protocols::SupportProtocols,
};
use crate::services::{
    dump_peer_store::DumpPeerStoreService, outbound_peer::OutboundPeerService,
    protocol_type_checker::ProtocolTypeCheckerService,
};
use crate::{Behaviour, CKBProtocol, Peer, PeerIndex, ProtocolId, ServiceControl};
use ckb_app_config::{NetworkConfig, SupportProtocol, default_support_all_protocols};
use ckb_logger::{debug, error, info, trace, warn};
use ckb_spawn::Spawn;
use ckb_stop_handler::{CancellationToken, broadcast_exit_signals, new_tokio_exit_rx};
use ckb_systemtime::{Duration, Instant};
use ckb_util::{Condvar, Mutex, RwLock};
use futures::{Future, channel::mpsc::Sender};
use ipnetwork::IpNetwork;
use p2p::{
    SessionId, async_trait,
    builder::ServiceBuilder,
    bytes::Bytes,
    context::{ServiceContext, SessionContext},
    error::{DialerErrorKind, HandshakeErrorKind, ProtocolHandleErrorKind, SendErrorKind},
    multiaddr::{Multiaddr, Protocol},
    secio::{self, PeerId, SecioKeyPair, error::SecioError},
    service::{
        ProtocolHandle, Service, ServiceAsyncControl, ServiceError, ServiceEvent, TargetProtocol,
        TargetSession,
    },
    traits::ServiceHandle,
    utils::{extract_peer_id, is_reachable, multiaddr_to_socketaddr},
    yamux::config::Config as YamuxConfig,
};
use rand::prelude::IteratorRandom;
#[cfg(feature = "with_sentry")]
use sentry::{Level, capture_message, with_scope};
#[cfg(not(target_family = "wasm"))]
use std::sync::mpsc;
use std::{
    borrow::Cow,
    cmp::max,
    collections::{HashMap, HashSet},
    pin::Pin,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    thread,
};
use tokio::{self, sync::oneshot};

const P2P_SEND_TIMEOUT: Duration = Duration::from_secs(6);
const P2P_TRY_SEND_INTERVAL: Duration = Duration::from_millis(100);
// After 5 minutes we consider this dial hang
const DIAL_HANG_TIMEOUT: Duration = Duration::from_secs(300);

/// The global shared state of the network module
pub struct NetworkState {
    pub(crate) peer_registry: RwLock<PeerRegistry>,
    pub(crate) peer_store: Mutex<PeerStore>,
    /// Node listened addresses
    pub(crate) listened_addrs: RwLock<Vec<Multiaddr>>,
    dialing_addrs: RwLock<HashMap<PeerId, Instant>>,
    /// Node public addresses by config
    public_addrs: HashSet<Multiaddr>,
    observed_addrs: RwLock<HashMap<PeerIndex, Multiaddr>>,
    local_private_key: secio::SecioKeyPair,
    local_peer_id: PeerId,
    pub(crate) bootnodes: Vec<Multiaddr>,
    pub(crate) config: NetworkConfig,
    pub(crate) active: AtomicBool,
    /// Node supported protocols
    /// fields: ProtocolId, Protocol Name, Supported Versions
    pub(crate) protocols: RwLock<Vec<(ProtocolId, String, Vec<String>)>>,
    pub(crate) required_flags: Flags,
}

impl NetworkState {
    /// Init from config
    #[cfg(not(target_family = "wasm"))]
    pub fn from_config(config: NetworkConfig) -> Result<NetworkState, Error> {
        config.create_dir_if_not_exists()?;
        let local_private_key = config.fetch_private_key()?;
        let local_peer_id = local_private_key.peer_id();
        // set max score to public addresses
        let public_addrs: HashSet<Multiaddr> = config
            .listen_addresses
            .iter()
            .chain(config.public_addresses.iter())
            .cloned()
            .filter_map(|mut addr| match multiaddr_to_socketaddr(&addr) {
                Some(socket_addr) if !is_reachable(socket_addr.ip()) => None,
                _ => {
                    match extract_peer_id(&addr) {
                        Some(peer_id) if peer_id != local_peer_id => {
                            error!("Don't include addresses that not associated with this node in the public_addresses list: {:?}", addr);
                            std::process::exit(1);
                        }
                        Some(_) => (),
                        None => addr.push(Protocol::P2P(Cow::Borrowed(local_peer_id.as_bytes()))),
                    }
                    Some(addr)
                }
            })
            .collect();
        info!("Loading the peer store. This process may take a few seconds to complete.");

        let peer_store = Mutex::new(PeerStore::load_from_dir_or_default(
            config.peer_store_path(),
        ));
        let bootnodes = config.bootnodes();

        let peer_registry = PeerRegistry::new(
            config.max_inbound_peers(),
            config.max_outbound_peers(),
            config.whitelist_only,
            config.whitelist_peers(),
            config.disable_block_relay_only_connection,
        );

        Ok(NetworkState {
            peer_store,
            config,
            bootnodes,
            peer_registry: RwLock::new(peer_registry),
            dialing_addrs: RwLock::new(HashMap::default()),
            public_addrs,
            listened_addrs: RwLock::new(Vec::new()),
            observed_addrs: RwLock::new(HashMap::default()),
            local_private_key,
            local_peer_id,
            active: AtomicBool::new(true),
            protocols: RwLock::new(Vec::new()),
            required_flags: Flags::SYNC | Flags::DISCOVERY | Flags::RELAY,
        })
    }

    #[cfg(target_family = "wasm")]
    pub async fn from_config(config: NetworkConfig) -> Result<NetworkState, Error> {
        let local_private_key = config.fetch_private_key()?;
        let local_peer_id = local_private_key.peer_id();
        // set max score to public addresses
        let public_addrs: HashSet<Multiaddr> = config
            .listen_addresses
            .iter()
            .chain(config.public_addresses.iter())
            .cloned()
            .filter_map(|mut addr| match multiaddr_to_socketaddr(&addr) {
                Some(socket_addr) if !is_reachable(socket_addr.ip()) => None,
                _ => {
                    if extract_peer_id(&addr).is_none() {
                        addr.push(Protocol::P2P(Cow::Borrowed(local_peer_id.as_bytes())));
                    }
                    Some(addr)
                }
            })
            .collect();
        info!("Loading the peer store. This process may take a few seconds to complete.");
        let peer_store = Mutex::new(PeerStore::load_from_idb(config.peer_store_path()).await);
        let bootnodes = config.bootnodes();

        let peer_registry = PeerRegistry::new(
            config.max_inbound_peers(),
            config.max_outbound_peers(),
            config.whitelist_only,
            config.whitelist_peers(),
            config.disable_block_relay_only_connection,
        );
        Ok(NetworkState {
            peer_store,
            config,
            bootnodes,
            peer_registry: RwLock::new(peer_registry),
            dialing_addrs: RwLock::new(HashMap::default()),
            public_addrs,
            listened_addrs: RwLock::new(Vec::new()),
            observed_addrs: RwLock::new(HashMap::default()),
            local_private_key,
            local_peer_id,
            active: AtomicBool::new(true),
            protocols: RwLock::new(Vec::new()),
            required_flags: Flags::SYNC | Flags::DISCOVERY | Flags::RELAY,
        })
    }

    /// use to discovery get nodes message to announce what kind of node information need from the other peer
    /// default with `Flags::SYNC | Flags::DISCOVERY | Flags::RELAY`
    pub fn required_flags(mut self, flags: Flags) -> Self {
        self.required_flags = flags;
        self
    }

    pub(crate) fn report_session(
        &self,
        p2p_control: &ServiceControl,
        session_id: SessionId,
        behaviour: Behaviour,
    ) {
        if let Some(addr) = self.with_peer_registry(|reg| {
            reg.get_peer(session_id)
                .filter(|peer| !peer.is_whitelist)
                .map(|peer| peer.connected_addr.clone())
        }) {
            trace!("Report {:?} because {:?}", addr, behaviour);
            let report_result = self.peer_store.lock().report(&addr, behaviour);
            if report_result.is_banned() {
                if let Err(err) = disconnect_with_message(p2p_control, session_id, "banned") {
                    debug!("Disconnect failed {:?}, error: {:?}", session_id, err);
                }
            }
        } else {
            debug!(
                "Report {} failure: not found in peer registry or it is on the whitelist",
                session_id
            );
        }
    }

    pub(crate) fn ban_session(
        &self,
        p2p_control: &ServiceControl,
        session_id: SessionId,
        duration: Duration,
        reason: String,
    ) {
        if let Some(addr) = self.with_peer_registry(|reg| {
            reg.get_peer(session_id)
                .filter(|peer| !peer.is_whitelist)
                .map(|peer| peer.connected_addr.clone())
        }) {
            info!(
                "Ban peer {:?} for {} seconds, reason: {}",
                addr,
                duration.as_secs(),
                reason
            );
            if let Some(metrics) = ckb_metrics::handle() {
                metrics.ckb_network_ban_peer.inc();
            }
            if let Some(peer) = self.with_peer_registry_mut(|reg| reg.remove_peer(session_id)) {
                let message = format!("Ban for {} seconds, reason: {}", duration.as_secs(), reason);
                self.peer_store.lock().ban_addr(
                    &peer.connected_addr,
                    duration.as_millis() as u64,
                    reason,
                );
                if let Err(err) =
                    disconnect_with_message(p2p_control, peer.session_id, message.as_str())
                {
                    debug!("Disconnect failed {:?}, error: {:?}", peer.session_id, err);
                }
            }
        } else {
            debug!(
                "Ban session({}) failed: not found in peer registry or it is on the whitelist",
                session_id
            );
        }
    }

    pub(crate) fn accept_peer(
        &self,
        session_context: &SessionContext,
    ) -> Result<Option<Peer>, Error> {
        // NOTE: be careful, here easy cause a deadlock,
        //    because peer_store's lock scope across peer_registry's lock scope
        let mut peer_store = self.peer_store.lock();
        let accept_peer_result = {
            self.peer_registry.write().accept_peer(
                session_context.address.clone(),
                session_context.id,
                session_context.ty,
                &mut peer_store,
            )
        };
        accept_peer_result
    }

    /// For restrict lock in inner scope
    pub fn with_peer_registry<F, T>(&self, callback: F) -> T
    where
        F: FnOnce(&PeerRegistry) -> T,
    {
        callback(&self.peer_registry.read())
    }

    // For restrict lock in inner scope
    pub(crate) fn with_peer_registry_mut<F, T>(&self, callback: F) -> T
    where
        F: FnOnce(&mut PeerRegistry) -> T,
    {
        callback(&mut self.peer_registry.write())
    }

    // For restrict lock in inner scope
    pub(crate) fn with_peer_store_mut<F, T>(&self, callback: F) -> T
    where
        F: FnOnce(&mut PeerStore) -> T,
    {
        callback(&mut self.peer_store.lock())
    }

    /// Get peer id of local node
    pub fn local_peer_id(&self) -> &PeerId {
        &self.local_peer_id
    }

    /// Use on test
    pub fn local_private_key(&self) -> &secio::SecioKeyPair {
        &self.local_private_key
    }

    /// Get local node's peer id in base58 format string
    pub fn node_id(&self) -> String {
        self.local_peer_id().to_base58()
    }

    pub(crate) fn public_addrs(&self, count: usize) -> Vec<Multiaddr> {
        if self.public_addrs.len() <= count {
            return self.public_addrs.iter().cloned().collect();
        } else {
            self.public_addrs
                .iter()
                .cloned()
                .choose_multiple(&mut rand::thread_rng(), count)
        }
    }

    pub(crate) fn connection_status(&self) -> ConnectionStatus {
        self.peer_registry.read().connection_status()
    }

    /// Get local node's listen address list
    pub fn public_urls(&self, max_urls: usize) -> Vec<(String, u8)> {
        let listened_addrs = self.listened_addrs.read();
        self.public_addrs(max_urls.saturating_sub(listened_addrs.len()))
            .into_iter()
            .filter_map(|addr| {
                if !listened_addrs.contains(&addr) {
                    Some((addr, 1))
                } else {
                    None
                }
            })
            .chain(listened_addrs.iter().map(|addr| (addr.to_owned(), 1)))
            .map(|(addr, score)| (addr.to_string(), score))
            .collect()
    }

    pub(crate) fn add_node(&self, p2p_control: &ServiceControl, address: Multiaddr) {
        self.dial_identify(p2p_control, address);
    }

    /// use a filter to get protocol id list
    pub fn get_protocol_ids<F: Fn(ProtocolId) -> bool>(&self, filter: F) -> Vec<ProtocolId> {
        self.protocols
            .read()
            .iter()
            .filter_map(|&(id, _, _)| if filter(id) { Some(id) } else { None })
            .collect::<Vec<_>>()
    }

    pub(crate) fn can_dial(&self, addr: &Multiaddr) -> bool {
        let peer_id = extract_peer_id(addr);
        if peer_id.is_none() {
            error!("Do not dial addr without peer id, addr: {}", addr);
            return false;
        }
        let peer_id = peer_id.as_ref().unwrap();

        if self.local_peer_id() == peer_id {
            trace!("Do not dial self: {:?}, {}", peer_id, addr);
            return false;
        }
        if self.public_addrs.contains(addr) {
            trace!(
                "Do not dial listened address(self): {:?}, {}",
                peer_id, addr
            );
            return false;
        }

        let peer_in_registry = self.with_peer_registry(|reg| {
            reg.get_key_by_peer_id(peer_id).is_some() || reg.is_feeler(addr)
        });
        if peer_in_registry {
            trace!("Do not dial peer in registry: {:?}, {}", peer_id, addr);
            return false;
        }

        if let Some(dial_started) = self.dialing_addrs.read().get(peer_id) {
            trace!(
                "Do not send repeated dial commands to network service: {:?}, {}",
                peer_id, addr
            );
            if Instant::now().saturating_duration_since(*dial_started) > DIAL_HANG_TIMEOUT {
                #[cfg(feature = "with_sentry")]
                with_scope(
                    |scope| scope.set_fingerprint(Some(&["ckb-network", "dialing-timeout"])),
                    || {
                        capture_message(
                            &format!(
                                "Dialing {:?}, {:?} for more than {} seconds, \
                                 something is wrong in network service",
                                peer_id,
                                addr,
                                DIAL_HANG_TIMEOUT.as_secs(),
                            ),
                            Level::Warning,
                        )
                    },
                );
            }
            return false;
        }

        true
    }

    pub(crate) fn dial_success(&self, addr: &Multiaddr) {
        if let Some(peer_id) = extract_peer_id(addr) {
            self.dialing_addrs.write().remove(&peer_id);
        }
    }

    pub(crate) fn dial_failed(&self, addr: &Multiaddr) {
        self.with_peer_registry_mut(|reg| {
            reg.remove_feeler(addr);
        });

        if let Some(peer_id) = extract_peer_id(addr) {
            self.dialing_addrs.write().remove(&peer_id);
        }
    }

    /// Dial
    /// return value indicates the dialing is actually sent or denied.
    fn dial_inner(
        &self,
        p2p_control: &ServiceControl,
        addr: Multiaddr,
        target: TargetProtocol,
    ) -> Result<(), Error> {
        if !self.can_dial(&addr) {
            return Err(Error::Dial(format!("ignore dialing addr {addr}")));
        }

        debug!("Dialing {addr}");
        p2p_control.dial(addr.clone(), target)?;
        self.dialing_addrs.write().insert(
            extract_peer_id(&addr).expect("verified addr"),
            Instant::now(),
        );
        Ok(())
    }

    /// Dial just identify protocol
    pub fn dial_identify(&self, p2p_control: &ServiceControl, addr: Multiaddr) {
        if let Err(err) = self.dial_inner(
            p2p_control,
            addr,
            TargetProtocol::Single(SupportProtocols::Identify.protocol_id()),
        ) {
            debug!("dial_identify error: {err}");
        }
    }

    /// Dial just feeler protocol
    pub fn dial_feeler(&self, p2p_control: &ServiceControl, addr: Multiaddr) {
        if let Err(err) = self.dial_inner(
            p2p_control,
            addr.clone(),
            TargetProtocol::Single(SupportProtocols::Identify.protocol_id()),
        ) {
            debug!("dial_feeler error {err}");
        } else {
            self.with_peer_registry_mut(|reg| {
                reg.add_feeler(&addr);
            });
        }
    }

    /// add observed address for identify protocol
    pub(crate) fn add_observed_addr(&self, session_id: SessionId, addr: Multiaddr) {
        let mut pending_observed_addrs = self.observed_addrs.write();
        pending_observed_addrs.insert(session_id, addr);
    }

    // randomly select count addresses from observed_addrs
    pub(crate) fn observed_addrs(&self, count: usize) -> Vec<Multiaddr> {
        let observed_addrs = self
            .observed_addrs
            .read()
            .values()
            .cloned()
            .collect::<HashSet<_>>();
        if observed_addrs.len() <= count {
            return observed_addrs.into_iter().collect();
        } else {
            observed_addrs
                .into_iter()
                .choose_multiple(&mut rand::thread_rng(), count)
        }
    }

    /// Network message processing controller, default is true, if false, discard any received messages
    pub fn is_active(&self) -> bool {
        self.active.load(Ordering::Acquire)
    }
}

/// Used to handle global events of tentacle, such as session open/close
pub struct EventHandler {
    pub(crate) network_state: Arc<NetworkState>,
}

impl EventHandler {
    /// init an event handler
    pub fn new(network_state: Arc<NetworkState>) -> Self {
        Self { network_state }
    }
}

/// Exit trait used to notify all other module to exit
pub trait ExitHandler: Send + Unpin + 'static {
    /// notify other module to exit
    fn notify_exit(&self);
}

/// Default exit handle
#[derive(Clone, Default)]
pub struct DefaultExitHandler {
    lock: Arc<Mutex<()>>,
    exit: Arc<Condvar>,
}

impl DefaultExitHandler {
    /// Block on current thread util exit notify
    pub fn wait_for_exit(&self) {
        self.exit.wait(&mut self.lock.lock());
    }
}

impl ExitHandler for DefaultExitHandler {
    fn notify_exit(&self) {
        self.exit.notify_all();
    }
}

impl EventHandler {
    fn inbound_eviction(&self) -> Vec<PeerIndex> {
        if self.network_state.config.bootnode_mode {
            let status = self.network_state.connection_status();

            if status.max_inbound <= status.non_whitelist_inbound.saturating_add(10) {
                self.network_state
                    .with_peer_registry(|registry| {
                        registry
                            .peers()
                            .values()
                            .filter(|peer| peer.is_inbound() && !peer.is_whitelist)
                            .map(|peer| peer.session_id)
                            .collect::<Vec<SessionId>>()
                    })
                    .into_iter()
                    .enumerate()
                    .filter_map(|(index, peer)| if index & 0x1 != 0 { Some(peer) } else { None })
                    .collect()
            } else {
                Vec::new()
            }
        } else {
            Vec::new()
        }
    }
}

#[async_trait]
impl ServiceHandle for EventHandler {
    async fn handle_error(&mut self, context: &mut ServiceContext, error: ServiceError) {
        match error {
            ServiceError::DialerError { address, error } => {
                match error {
                    DialerErrorKind::HandshakeError(HandshakeErrorKind::SecioError(
                        SecioError::ConnectSelf,
                    )) => {
                        debug!("dial observed address success: {:?}", address);
                    }
                    DialerErrorKind::IoError(e)
                        if e.kind() == std::io::ErrorKind::AddrNotAvailable =>
                    {
                        warn!("DialerError({}) {}", address, e);
                    }
                    _ => {
                        debug!("DialerError({}) {}", address, error);
                    }
                }
                self.network_state.dial_failed(&address);
            }
            ServiceError::ProtocolError {
                id,
                proto_id,
                error,
            } => {
                debug!("ProtocolError({}, {}) {}", id, proto_id, error);
                let message = format!("ProtocolError id={proto_id}");
                // Ban because misbehave of remote peer
                self.network_state.ban_session(
                    &context.control().clone().into(),
                    id,
                    Duration::from_secs(300),
                    message,
                );
            }
            ServiceError::SessionTimeout { session_context } => {
                debug!(
                    "SessionTimeout({}, {})",
                    session_context.id, session_context.address,
                );
            }
            ServiceError::MuxerError {
                session_context,
                error,
            } => {
                debug!(
                    "MuxerError({}, {}), substream error {}, disconnect it",
                    session_context.id, session_context.address, error,
                );
            }
            ServiceError::ListenError { address, error } => {
                debug!("ListenError: address={:?}, error={:?}", address, error);
            }
            ServiceError::ProtocolSelectError {
                proto_name,
                session_context,
            } => {
                debug!(
                    "ProtocolSelectError: proto_name={:?}, session_id={}",
                    proto_name, session_context.id,
                );
            }
            ServiceError::SessionBlocked { session_context } => {
                debug!("SessionBlocked: {}", session_context.id);
            }
            ServiceError::ProtocolHandleError { proto_id, error } => {
                debug!("ProtocolHandleError: {:?}, proto_id: {}", error, proto_id);

                let ProtocolHandleErrorKind::AbnormallyClosed(opt_session_id) = error;
                {
                    if let Some(id) = opt_session_id {
                        self.network_state.ban_session(
                            &context.control().clone().into(),
                            id,
                            Duration::from_secs(300),
                            format!("protocol {proto_id} panic when process peer message"),
                        );
                    }
                    #[cfg(feature = "with_sentry")]
                    with_scope(
                        |scope| scope.set_fingerprint(Some(&["ckb-network", "p2p-service-error"])),
                        || {
                            capture_message(
                                &format!(
                                    "ProtocolHandleError: AbnormallyClosed, proto_id: {opt_session_id:?}, session id: {opt_session_id:?}"
                                ),
                                Level::Warning,
                            )
                        },
                    );
                    error!(
                        "ProtocolHandleError: AbnormallyClosed, proto_id: {opt_session_id:?}, session id: {opt_session_id:?}"
                    );

                    broadcast_exit_signals();
                }
            }
        }
    }

    async fn handle_event(&mut self, context: &mut ServiceContext, event: ServiceEvent) {
        // When session disconnect update status anyway
        match event {
            ServiceEvent::SessionOpen { session_context } => {
                debug!(
                    "SessionOpen({}, {})",
                    session_context.id, session_context.address,
                );

                self.network_state.dial_success(&session_context.address);

                let iter = self.inbound_eviction();

                let control = context.control().clone().into();

                for peer in iter {
                    if let Err(err) =
                        disconnect_with_message(&control, peer, "bootnode random eviction")
                    {
                        debug!("Inbound eviction failed {:?}, error: {:?}", peer, err);
                    }
                }

                if self
                    .network_state
                    .with_peer_registry(|reg| reg.is_feeler(&session_context.address))
                {
                    debug!(
                        "Feeler connected {} => {}",
                        session_context.id, session_context.address,
                    );
                } else {
                    match self.network_state.accept_peer(&session_context) {
                        Ok(Some(evicted_peer)) => {
                            debug!(
                                "Disconnect peer, {} => {}",
                                evicted_peer.session_id, evicted_peer.connected_addr,
                            );
                            if let Err(err) = disconnect_with_message(
                                &control,
                                evicted_peer.session_id,
                                "evict because accepted better peer",
                            ) {
                                debug!(
                                    "Disconnect failed {:?}, error: {:?}",
                                    evicted_peer.session_id, err
                                );
                            }
                        }
                        Ok(None) => debug!(
                            "{} open, registry {} success",
                            session_context.id, session_context.address,
                        ),
                        Err(err) => {
                            debug!(
                                "Peer registry failed {:?}. Disconnect {} => {}",
                                err, session_context.id, session_context.address,
                            );
                            if let Err(err) = disconnect_with_message(
                                &control,
                                session_context.id,
                                "reject peer connection",
                            ) {
                                debug!(
                                    "Disconnect failed {:?}, error: {:?}",
                                    session_context.id, err
                                );
                            }
                        }
                    }
                }
            }
            ServiceEvent::SessionClose { session_context } => {
                debug!(
                    "SessionClose({}, {})",
                    session_context.id, session_context.address,
                );
                let peer_exists = self.network_state.with_peer_registry_mut(|reg| {
                    // should make sure feelers is clean
                    reg.remove_feeler(&session_context.address);
                    reg.remove_peer(session_context.id).is_some()
                });
                if peer_exists {
                    debug!(
                        "{} closed. Remove {} from peer_registry",
                        session_context.id, session_context.address,
                    );
                    self.network_state.with_peer_store_mut(|peer_store| {
                        peer_store.remove_disconnected_peer(&session_context.address);
                    });
                }
                self.network_state
                    .observed_addrs
                    .write()
                    .remove(&session_context.id);
            }
            _ => {
                info!("p2p service event: {:?}", event);
            }
        }
    }
}

/// Ckb network service, use to start p2p network
pub struct NetworkService {
    p2p_service: Service<EventHandler, SecioKeyPair>,
    network_state: Arc<NetworkState>,
    ping_controller: Option<Sender<()>>,
    // Background services
    bg_services: Vec<Pin<Box<dyn Future<Output = ()> + 'static + Send>>>,
    version: String,
}

impl NetworkService {
    /// init with all config
    pub fn new(
        network_state: Arc<NetworkState>,
        protocols: Vec<CKBProtocol>,
        required_protocol_ids: Vec<ProtocolId>,
        // name, version, flags
        identify_announce: (String, String, Flags),
        transport_type: TransportType,
    ) -> Self {
        let config = &network_state.config;

        if config.support_protocols.iter().collect::<HashSet<_>>()
            != default_support_all_protocols()
                .iter()
                .collect::<HashSet<_>>()
        {
            warn!(
                "Customized supported protocols: {:?}",
                config.support_protocols
            );
        }

        // == Build p2p service struct
        let mut protocol_metas = protocols
            .into_iter()
            .map(CKBProtocol::build)
            .collect::<Vec<_>>();

        // == Build special protocols

        // Identify is a core protocol, user cannot disable it via config
        let identify_callback = IdentifyCallback::new(
            Arc::clone(&network_state),
            identify_announce.0,
            identify_announce.1.clone(),
            identify_announce.2,
        );
        let identify_meta = SupportProtocols::Identify.build_meta_with_service_handle(move || {
            ProtocolHandle::Callback(Box::new(IdentifyProtocol::new(identify_callback)))
        });
        protocol_metas.push(identify_meta);

        // Ping protocol
        let ping_controller = if config.support_protocols.contains(&SupportProtocol::Ping) {
            let ping_interval = Duration::from_secs(config.ping_interval_secs);
            let ping_timeout = Duration::from_secs(config.ping_timeout_secs);

            let ping_network_state = Arc::clone(&network_state);
            let (ping_handler, ping_controller) =
                PingHandler::new(ping_interval, ping_timeout, ping_network_state);
            let ping_meta = SupportProtocols::Ping.build_meta_with_service_handle(move || {
                ProtocolHandle::Callback(Box::new(ping_handler))
            });
            protocol_metas.push(ping_meta);
            Some(ping_controller)
        } else {
            None
        };

        // Discovery protocol
        if config
            .support_protocols
            .contains(&SupportProtocol::Discovery)
        {
            let addr_mgr = DiscoveryAddressManager {
                network_state: Arc::clone(&network_state),
                discovery_local_address: config.discovery_local_address,
            };
            let disc_meta = SupportProtocols::Discovery.build_meta_with_service_handle(move || {
                ProtocolHandle::Callback(Box::new(DiscoveryProtocol::new(
                    addr_mgr,
                    config
                        .discovery_announce_check_interval_secs
                        .map(Duration::from_secs),
                )))
            });
            protocol_metas.push(disc_meta);
        }

        // Feeler protocol
        if config.support_protocols.contains(&SupportProtocol::Feeler) {
            let feeler_meta = SupportProtocols::Feeler.build_meta_with_service_handle({
                let network_state = Arc::clone(&network_state);
                move || ProtocolHandle::Callback(Box::new(Feeler::new(Arc::clone(&network_state))))
            });
            protocol_metas.push(feeler_meta);
        }

        // DisconnectMessage protocol
        if config
            .support_protocols
            .contains(&SupportProtocol::DisconnectMessage)
        {
            let disconnect_message_state = Arc::clone(&network_state);
            let disconnect_message_meta = SupportProtocols::DisconnectMessage
                .build_meta_with_service_handle(move || {
                    ProtocolHandle::Callback(Box::new(DisconnectMessageProtocol::new(
                        disconnect_message_state,
                    )))
                });
            protocol_metas.push(disconnect_message_meta);
        }

        // HolePunching protocol
        #[cfg(not(target_family = "wasm"))]
        if config
            .support_protocols
            .contains(&SupportProtocol::HolePunching)
        {
            let hole_punching_state = Arc::clone(&network_state);
            let hole_punching_meta =
                SupportProtocols::HolePunching.build_meta_with_service_handle(move || {
                    ProtocolHandle::Callback(Box::new(
                        crate::protocols::hole_punching::HolePunching::new(hole_punching_state),
                    ))
                });
            protocol_metas.push(hole_punching_meta);
        }

        let mut service_builder = ServiceBuilder::default();
        let yamux_config = YamuxConfig {
            max_stream_count: protocol_metas.len(),
            max_stream_window_size: 1024 * 1024,
            ..Default::default()
        };
        for meta in protocol_metas.into_iter() {
            network_state
                .protocols
                .write()
                .push((meta.id(), meta.name(), meta.support_versions()));
            service_builder = service_builder.insert_protocol(meta);
        }
        let event_handler = EventHandler {
            network_state: Arc::clone(&network_state),
        };
        service_builder = service_builder
            .handshake_type(network_state.local_private_key.clone().into())
            .yamux_config(yamux_config)
            .forever(true)
            .max_connection_number(1024)
            .set_send_buffer_size(config.max_send_buffer())
            .set_channel_size(config.channel_size())
            .timeout(Duration::from_secs(5));

        #[cfg(not(target_family = "wasm"))]
        {
            service_builder = service_builder.upnp(config.upnp);
        }

        #[cfg(target_os = "linux")]
        let p2p_service = {
            if config.reuse_port_on_linux {
                let iter = config.listen_addresses.iter();

                #[derive(Clone, Copy, Debug, Eq, PartialEq)]
                enum BindType {
                    None,
                    Ws,
                    Tcp,
                    Both,
                }
                impl BindType {
                    fn transform(&mut self, other: TransportType) {
                        match (&self, other) {
                            (BindType::None, TransportType::Ws) => *self = BindType::Ws,
                            (BindType::None, TransportType::Tcp) => *self = BindType::Tcp,
                            (BindType::Ws, TransportType::Tcp) => *self = BindType::Both,
                            (BindType::Tcp, TransportType::Ws) => *self = BindType::Both,
                            _ => (),
                        }
                    }

                    fn is_ready(&self) -> bool {
                        // should change to Both if ckb enable ws
                        matches!(self, BindType::Both)
                    }
                }

                let mut init = BindType::None;

                for multi_addr in iter {
                    if init.is_ready() {
                        break;
                    }
                    match find_type(multi_addr) {
                        TransportType::Tcp => {
                            // only bind once
                            if matches!(init, BindType::Tcp) {
                                continue;
                            }
                            if let Some(addr) = multiaddr_to_socketaddr(multi_addr) {
                                let domain = socket2::Domain::for_address(addr);
                                let bind_fn = move |socket: p2p::service::TcpSocket, ctxt: p2p::service::TransformerContext| {
                                    let socket_ref = socket2::SockRef::from(&socket);
                                    #[cfg(all(
                                        unix,
                                        not(target_os = "solaris"),
                                        not(target_os = "illumos")
                                    ))]
                                    socket_ref.set_reuse_port(true)?;
                                    socket_ref.set_reuse_address(true)?;
                                    match ctxt.state {
                                        p2p::service::SocketState::Listen => {
                                            Ok(socket)
                                        }
                                        p2p::service::SocketState::Dial => {
                                            if socket_ref.domain()? == domain {
                                                socket_ref.bind(&addr.into())?;
                                            }
                                            Ok(socket)
                                        }
                                    }
                                };
                                init.transform(TransportType::Tcp);
                                service_builder = service_builder.tcp_config(bind_fn);
                            }
                        }
                        TransportType::Ws | TransportType::Wss => {
                            // only bind once
                            if matches!(init, BindType::Ws) {
                                continue;
                            }
                            if let Some(addr) = multiaddr_to_socketaddr(multi_addr) {
                                let domain = socket2::Domain::for_address(addr);
                                let bind_fn = move |socket: p2p::service::TcpSocket, ctxt: p2p::service::TransformerContext| {
                                    let socket_ref = socket2::SockRef::from(&socket);
                                    #[cfg(all(
                                        unix,
                                        not(target_os = "solaris"),
                                        not(target_os = "illumos")
                                    ))]
                                    socket_ref.set_reuse_port(true)?;
                                    socket_ref.set_reuse_address(true)?;
                                    match ctxt.state {
                                        p2p::service::SocketState::Listen => {
                                            Ok(socket)
                                        }
                                        p2p::service::SocketState::Dial => {
                                            if socket_ref.domain()? == domain {
                                                socket_ref.bind(&addr.into())?;
                                            }
                                            Ok(socket)
                                        }
                                    }
                                };
                                init.transform(TransportType::Ws);
                                service_builder = service_builder.tcp_config_on_ws(bind_fn);
                            }
                        }
                    }
                }
            }

            service_builder.build(event_handler)
        };

        #[cfg(not(target_os = "linux"))]
        // The default permissions of Windows are not enough to enable this function,
        // and the administrator permissions of group permissions must be turned on.
        // This operation is very burdensome for windows users, so it is turned off by default
        //
        // The integration test fails after MacOS is turned on, the behavior is different from linux.
        // Decision to turn off it
        let p2p_service = service_builder.build(event_handler);

        // == Build background service tasks
        let dump_peer_store_service = DumpPeerStoreService::new(Arc::clone(&network_state));
        let protocol_type_checker_service = ProtocolTypeCheckerService::new(
            Arc::clone(&network_state),
            p2p_service.control().to_owned().into(),
            required_protocol_ids,
        );
        let mut bg_services = vec![
            Box::pin(dump_peer_store_service) as Pin<Box<_>>,
            Box::pin(protocol_type_checker_service) as Pin<Box<_>>,
        ];
        if config.outbound_peer_service_enabled() {
            let outbound_peer_service = OutboundPeerService::new(
                Arc::clone(&network_state),
                p2p_service.control().to_owned().into(),
                Duration::from_secs(config.connect_outbound_interval_secs),
                transport_type,
            );
            bg_services.push(Box::pin(outbound_peer_service) as Pin<Box<_>>);
        };

        #[cfg(feature = "with_dns_seeding")]
        if config.dns_seeding_service_enabled() {
            let dns_seeding_service = crate::services::dns_seeding::DnsSeedingService::new(
                Arc::clone(&network_state),
                config.dns_seeds.clone(),
            );
            bg_services.push(Box::pin(dns_seeding_service.start()) as Pin<Box<_>>);
        };

        NetworkService {
            p2p_service,
            network_state,
            ping_controller,
            bg_services,
            version: identify_announce.1,
        }
    }

    /// Start the network in the background and return a controller
    pub fn start<S: Spawn>(self, handle: &S) -> Result<NetworkController, Error> {
        let config = self.network_state.config.clone();

        let p2p_control: ServiceControl = self.p2p_service.control().to_owned().into();

        // dial whitelist_nodes
        for addr in self.network_state.config.whitelist_peers() {
            debug!("Dial whitelist_peers {:?}", addr);
            self.network_state.dial_identify(&p2p_control, addr);
        }

        let target = &self.network_state.required_flags;

        // get bootnodes
        // try get addrs from peer_store, if peer_store have no enough addrs then use bootnodes
        let bootnodes = self.network_state.with_peer_store_mut(|peer_store| {
            let count = max((config.max_outbound_peers >> 1) as usize, 1);
            let mut addrs: Vec<_> = peer_store
                .fetch_addrs_to_attempt(count, *target, |_| true)
                .into_iter()
                .map(|paddr| paddr.addr)
                .collect();
            // tried to re-connect to anchors on startup
            let anchors: Vec<_> = peer_store.mut_anchors().drain().collect();
            addrs.extend(anchors);
            // Get bootnodes randomly
            let bootnodes = self
                .network_state
                .bootnodes
                .iter()
                .choose_multiple(&mut rand::thread_rng(), count.saturating_sub(addrs.len()))
                .into_iter()
                .cloned();
            addrs.extend(bootnodes);
            addrs
        });

        // dial half bootnodes
        for addr in bootnodes {
            debug!("Dial bootnode {:?}", addr);
            self.network_state.dial_identify(&p2p_control, addr);
        }

        let Self {
            mut p2p_service,
            network_state,
            ping_controller,
            bg_services,
            version,
        } = self;

        // NOTE: for ensure background task finished
        let (bg_signals, bg_receivers): (Vec<_>, Vec<_>) = bg_services
            .into_iter()
            .map(|bg_service| {
                let (signal_sender, signal_receiver) = oneshot::channel::<()>();
                (signal_sender, (bg_service, signal_receiver))
            })
            .unzip();

        let receiver: CancellationToken = new_tokio_exit_rx();
        #[cfg(not(target_family = "wasm"))]
        let (start_sender, start_receiver) = mpsc::channel();
        {
            #[cfg(not(target_family = "wasm"))]
            let network_state = Arc::clone(&network_state);
            let p2p_control: ServiceAsyncControl = p2p_control.clone().into();
            handle.spawn_task(async move {
                #[cfg(not(target_family = "wasm"))]
                {
                    let listen_addresses = {
                        let mut addresses = config.listen_addresses.clone();
                        if config.reuse_tcp_with_ws {
                            let ws_listens = addresses
                                .iter()
                                .cloned()
                                .filter_map(|mut addr| {
                                    if matches!(find_type(&addr), TransportType::Tcp) {
                                        addr.push(Protocol::Ws);
                                        Some(addr)
                                    } else {
                                        None
                                    }
                                })
                                .collect::<Vec<_>>();

                            addresses.extend(ws_listens);
                        }
                        let mut addresses = addresses
                            .into_iter()
                            .collect::<HashSet<_>>()
                            .into_iter()
                            .collect::<Vec<_>>();
                        addresses.sort_by(|a, b| {
                            let ty_a = find_type(a);
                            let ty_b = find_type(b);

                            ty_a.cmp(&ty_b)
                        });

                        addresses
                    };

                    for addr in &listen_addresses {
                        match p2p_service.listen(addr.to_owned()).await {
                            Ok(listen_address) => {
                                info!("Listen on address: {}", listen_address);
                                network_state
                                    .listened_addrs
                                    .write()
                                    .push(listen_address.clone());
                            }
                            Err(err) => {
                                warn!(
                                    "Listen on address {} failed, due to error: {}",
                                    addr.clone(),
                                    err
                                );
                                start_sender
                                    .send(Err(Error::P2P(P2PError::Transport(err))))
                                    .expect("channel abnormal shutdown");
                                return;
                            }
                        };
                    }
                    start_sender.send(Ok(())).unwrap();
                }

                p2p::runtime::spawn(async move { p2p_service.run().await });
                tokio::select! {
                    _ = receiver.cancelled() => {
                        info!("NetworkService receive exit signal, start shutdown...");
                        let _ = p2p_control.shutdown().await;
                        // Drop senders to stop all corresponding background task
                        drop(bg_signals);

                        info!("NetworkService shutdown now");
                    },
                    else => {
                        let _ = p2p_control.shutdown().await;
                        // Drop senders to stop all corresponding background task
                        drop(bg_signals);
                    },
                }
            });
        }
        for (mut service, mut receiver) in bg_receivers {
            handle.spawn_task(async move {
                loop {
                    tokio::select! {
                        _ = &mut service => {},
                        _ = &mut receiver => break
                    }
                }
            });
        }
        #[cfg(not(target_family = "wasm"))]
        if let Ok(Err(e)) = start_receiver.recv() {
            return Err(e);
        }

        Ok(NetworkController {
            version,
            network_state,
            p2p_control,
            ping_controller,
        })
    }
}

/// Network controller
#[derive(Clone)]
pub struct NetworkController {
    version: String,
    network_state: Arc<NetworkState>,
    p2p_control: ServiceControl,
    ping_controller: Option<Sender<()>>,
}

impl NetworkController {
    /// Node listen address list
    pub fn public_urls(&self, max_urls: usize) -> Vec<(String, u8)> {
        self.network_state.public_urls(max_urls)
    }

    /// ckb version
    pub fn version(&self) -> &String {
        &self.version
    }

    /// Node peer id's base58 format string
    pub fn node_id(&self) -> String {
        self.network_state.node_id()
    }

    /// p2p service control
    pub fn p2p_control(&self) -> &ServiceControl {
        &self.p2p_control
    }

    /// Dial remote node
    pub fn add_node(&self, address: Multiaddr) {
        self.network_state.add_node(&self.p2p_control, address)
    }

    /// Disconnect session with peer id
    pub fn remove_node(&self, peer_id: &PeerId) {
        if let Some(session_id) = self
            .network_state
            .peer_registry
            .read()
            .get_key_by_peer_id(peer_id)
        {
            if let Err(err) =
                disconnect_with_message(&self.p2p_control, session_id, "disconnect manually")
            {
                debug!("Disconnect failed {:?}, error: {:?}", session_id, err);
            }
        } else {
            error!("Cannot find peer {:?}", peer_id);
        }
    }

    /// Get banned peer list
    pub fn get_banned_addrs(&self) -> Vec<BannedAddr> {
        self.network_state
            .peer_store
            .lock()
            .ban_list()
            .get_banned_addrs()
    }

    /// Clear banned list
    pub fn clear_banned_addrs(&self) {
        self.network_state.peer_store.lock().clear_ban_list();
    }

    /// Get address info from peer store
    pub fn addr_info(&self, addr: &Multiaddr) -> Option<AddrInfo> {
        self.network_state
            .peer_store
            .lock()
            .addr_manager()
            .get(addr)
            .cloned()
    }

    /// Ban an ip
    pub fn ban(&self, address: IpNetwork, ban_until: u64, ban_reason: String) {
        self.disconnect_peers_in_ip_range(address, &ban_reason);
        self.network_state
            .peer_store
            .lock()
            .ban_network(address, ban_until, ban_reason)
    }

    /// Unban an ip
    pub fn unban(&self, address: &IpNetwork) {
        self.network_state
            .peer_store
            .lock()
            .mut_ban_list()
            .unban_network(address);
    }

    /// Return all connected peers' information
    pub fn connected_peers(&self) -> Vec<(PeerIndex, Peer)> {
        self.network_state.with_peer_registry(|reg| {
            reg.peers()
                .iter()
                .map(|(peer_index, peer)| (*peer_index, peer.clone()))
                .collect::<Vec<_>>()
        })
    }

    /// Ban an peer through peer index
    pub fn ban_peer(&self, peer_index: PeerIndex, duration: Duration, reason: String) {
        self.network_state
            .ban_session(&self.p2p_control, peer_index, duration, reason);
    }

    /// disconnect peers with matched peer_ip or peer_ip_network, eg: 192.168.0.2 or 192.168.0.0/24
    fn disconnect_peers_in_ip_range(&self, address: IpNetwork, reason: &str) {
        self.network_state.with_peer_registry(|reg| {
            reg.peers().iter().for_each(|(peer_index, peer)| {
                if let Some(addr) = multiaddr_to_socketaddr(&peer.connected_addr) {
                    if address.contains(addr.ip()) {
                        let _ = disconnect_with_message(
                            &self.p2p_control,
                            *peer_index,
                            &format!("Ban peer {}, reason: {}", addr.ip(), reason),
                        );
                    }
                }
            })
        });
    }

    fn try_broadcast(
        &self,
        quick: bool,
        target: Option<SessionId>,
        proto_id: ProtocolId,
        data: Bytes,
    ) -> Result<(), SendErrorKind> {
        let now = Instant::now();
        loop {
            let target = target
                .map(TargetSession::Single)
                .unwrap_or(TargetSession::All);
            let result = if quick {
                self.p2p_control
                    .quick_filter_broadcast(target, proto_id, data.clone())
            } else {
                self.p2p_control
                    .filter_broadcast(target, proto_id, data.clone())
            };
            match result {
                Ok(()) => {
                    return Ok(());
                }
                Err(SendErrorKind::WouldBlock) => {
                    if Instant::now().saturating_duration_since(now) > P2P_SEND_TIMEOUT {
                        warn!("Broadcast message to {} timeout", proto_id);
                        return Err(SendErrorKind::WouldBlock);
                    }
                    thread::sleep(P2P_TRY_SEND_INTERVAL);
                }
                Err(err) => {
                    warn!("Broadcast message to {} failed: {:?}", proto_id, err);
                    return Err(err);
                }
            }
        }
    }

    /// Broadcast a message to all connected peers
    pub fn broadcast(&self, proto_id: ProtocolId, data: Bytes) -> Result<(), SendErrorKind> {
        self.try_broadcast(false, None, proto_id, data)
    }

    /// Broadcast a message to all connected peers through quick queue
    pub fn quick_broadcast(&self, proto_id: ProtocolId, data: Bytes) -> Result<(), SendErrorKind> {
        self.try_broadcast(true, None, proto_id, data)
    }

    /// Send message to one connected peer
    pub fn send_message_to(
        &self,
        session_id: SessionId,
        proto_id: ProtocolId,
        data: Bytes,
    ) -> Result<(), SendErrorKind> {
        self.try_broadcast(false, Some(session_id), proto_id, data)
    }

    /// network message processing controller, always true, if false, discard any received messages
    pub fn is_active(&self) -> bool {
        self.network_state.is_active()
    }

    /// Change active status, if set false discard any received messages
    pub fn set_active(&self, active: bool) {
        self.network_state.active.store(active, Ordering::Release);
    }

    /// Return all connected peers' protocols info
    pub fn protocols(&self) -> Vec<(ProtocolId, String, Vec<String>)> {
        self.network_state.protocols.read().clone()
    }

    /// Try ping all connected peers
    pub fn ping_peers(&self) {
        if let Some(mut ping_controller) = self.ping_controller.clone() {
            let _ignore = ping_controller.try_send(());
        }
    }
}

// Send an optional message before disconnect a peer
pub(crate) fn disconnect_with_message(
    control: &ServiceControl,
    peer_index: SessionId,
    message: &str,
) -> Result<(), SendErrorKind> {
    if !message.is_empty() {
        let data = Bytes::from(message.as_bytes().to_vec());
        // Must quick send, otherwise this message will be dropped.
        control.quick_send_message_to(
            peer_index,
            SupportProtocols::DisconnectMessage.protocol_id(),
            data,
        )?;
    }
    control.disconnect(peer_index)
}

pub(crate) async fn async_disconnect_with_message(
    control: &ServiceAsyncControl,
    peer_index: SessionId,
    message: &str,
) -> Result<(), SendErrorKind> {
    if !message.is_empty() {
        let data = Bytes::from(message.as_bytes().to_vec());
        // Must quick send, otherwise this message will be dropped.
        control
            .quick_send_message_to(
                peer_index,
                SupportProtocols::DisconnectMessage.protocol_id(),
                data,
            )
            .await?;
    }
    control.disconnect(peer_index).await
}

/// Transport type on ckb
#[derive(Clone, Copy, Debug, Eq, PartialEq, PartialOrd, Ord)]
pub enum TransportType {
    /// Tcp
    Tcp,
    /// Ws
    Ws,
    /// Wss only on wasm
    Wss,
}

#[allow(dead_code)]
pub(crate) fn find_type(addr: &Multiaddr) -> TransportType {
    let mut iter = addr.iter();

    iter.find_map(|proto| match proto {
        Protocol::Ws => Some(TransportType::Ws),
        Protocol::Wss => Some(TransportType::Wss),
        _ => None,
    })
    .unwrap_or(TransportType::Tcp)
}