ghost-walker/src/walker.rs

use std::{
    collections::{hash_map::Entry, HashMap, HashSet},
    time::Duration,
    error::Error,
};

use futures::StreamExt;
use codec::Decode;
use ip_network::IpNetwork;
use primitive_types::H256;
use libp2p::{
    identify::Info, identity, kad::{
        Event as KademliaEvent, GetClosestPeersError, GetClosestPeersOk, 
        QueryId, QueryResult,
    }, 
    multiaddr::Protocol, 
    swarm::{Config as SwarmConfig, SwarmEvent}, 
    Multiaddr, PeerId, Swarm 
};

use crate::{
    locator::Locator,
    p2p::{notifications::behaviour::NotificationsToSwarm, BehaviourEvent},
};

use super::p2p::{
    discovery::DiscoveryBuilder,
    peer_behaviour::{PeerInfoEvent, PeerBehaviour},
    transport::{TransportBuilder, MIB},
    notifications::{
        behaviour::{Notifications, ProtocolsData},
        messages::ProtocolRole,
    },
    Behaviour,
};

pub struct Walker {
    genesis: String,
    timeout: Duration,
    swarm: Swarm<Behaviour>,
    queries: HashSet<QueryId>,
    discovered_with_address: HashMap<PeerId, HashSet<Multiaddr>>,
    peer_details: HashMap<PeerId, Info>,
    peer_role: HashMap<PeerId, ProtocolRole>,
    dialed_peers: HashMap<PeerId, usize>,
}

impl Walker {
    pub fn new() -> WalkerBuilder {
        WalkerBuilder::default()
    }

    fn insert_queries(&mut self, num: usize) {
        for _ in 0..num {
            self.queries.insert(
                self.swarm
                    .behaviour_mut()
                    .discovery
                    .get_closest_peers(PeerId::random()));
        }
    }

    fn dialed_peer(&mut self, peer_id: Option<PeerId>) {
        let Some(peer_id) = peer_id else { return };

        self.dialed_peers
            .entry(peer_id)
            .and_modify(|num| *num += 1)
            .or_insert(0);
    }

    async fn inner_walk(&mut self) {
        self.insert_queries(128);
        let mut previous_tracing_time = std::time::Instant::now();

        loop {
            let event = self.swarm.select_next_some().await;

            match event {
                SwarmEvent::Dialing { peer_id, .. } => self.dialed_peer(peer_id),
                SwarmEvent::Behaviour(BehaviourEvent::Discovery(event)) => match event {
                    KademliaEvent::OutboundQueryProgressed { 
                        id, 
                        result: QueryResult::GetClosestPeers(result),
                        ..
                    } => {
                        self.queries.remove(&id);

                        let peers = match result {
                            Ok(GetClosestPeersOk { peers, ..}) => peers,
                            Err(GetClosestPeersError::Timeout { peers, .. }) => peers
                        };
                        let num_discovered = peers.len();

                        let now = std::time::Instant::now();
                        if now.duration_since(previous_tracing_time) > Duration::from_secs(10) {
                            previous_tracing_time = now;
                            tracing::info!("...Discovery in progress last query #{num_discovered}");
                        }

                        if self.queries.is_empty() {
                            self.insert_queries(128);
                        }
                    }
                    KademliaEvent::RoutingUpdated { peer, addresses, .. } => {
                        match self.discovered_with_address.entry(peer) {
                            Entry::Occupied(mut occupied) => {
                                occupied.get_mut().extend(addresses.into_vec());
                            },
                            Entry::Vacant(vacant) => {
                                vacant.insert(addresses.iter().cloned().collect());
                            },
                        };
                    }
                    KademliaEvent::RoutablePeer { peer, address } | KademliaEvent::PendingRoutablePeer { peer, address } => {
                        match self.discovered_with_address.entry(peer) {
                            Entry::Occupied(mut occupied) => {
                                occupied.get_mut().insert(address);
                            },
                            Entry::Vacant(vacant) => {
                                let mut set = HashSet::new();
                                set.insert(address);
                                vacant.insert(set);
                            },
                        };
                    }
                    _ => (),
                },
                SwarmEvent::Behaviour(BehaviourEvent::PeerInfo(info_event)) => match info_event {
                    PeerInfoEvent::Identified { peer_id, info } => {
                        tracing::debug!("Identified peer_id={:?} info={:?}", peer_id, info);
                        self.peer_details.insert(peer_id, info);
                    }
                },
                SwarmEvent::Behaviour(BehaviourEvent::Notifications(
                        NotificationsToSwarm::CustomProtocolOpen { 
                            peer_id, 
                            index, 
                            received_handshake, 
                            inbound, 
                            .. 
                        }
                )) => {
                    if let Ok(role) = ProtocolRole::decode(&mut &received_handshake[..]) {
                        tracing::debug!("Identified peer_id={:?} role={:?}", peer_id, role);
                        self.peer_role.insert(peer_id, role);
                    }

                    tracing::debug!(
                        "Protocol open peer={:?} index={:?} handshake={:?} inbound={:?}",
                        peer_id,
                        index,
                        received_handshake,
                        inbound,
                    );
                }
                _ => (),
            }
        }
    }

    fn is_public_adress(&self, addr: &Multiaddr) -> bool {
        let ip = match addr.iter().next() {
            Some(Protocol::Ip4(ip)) => IpNetwork::from(ip),
            Some(Protocol::Ip6(ip)) => IpNetwork::from(ip),
            Some(Protocol::Dns(_)) | Some(Protocol::Dns4(_)) | Some(Protocol::Dns6(_)) => return true,
            _ => return false,
        };
        ip.is_global()
    }

    pub async fn walk_around(&mut self) -> Result<(), Box<dyn Error>> {
        let _ = tokio::time::timeout(self.timeout, self.inner_walk()).await;

        println!("Number of dialed peers: {}", self.dialed_peers.len());
        println!("Total peers discovered: {}", self.discovered_with_address.len());
        println!("Peers with local iden.: {}", self.peer_details.len());

        let infos: HashMap<_, _> = self
            .peer_details
            .iter()
            .filter(|(_, info)| {
                info.protocols
                    .iter()
                    .any(|stream_proto| stream_proto.as_ref().contains(&self.genesis))
            })
            .collect();
        println!("Support correct genesis {}: {}", &self.genesis, infos.len());

        let peers_with_public: HashMap<_, _> = infos
            .iter()
            .filter(|(_, info)| info.listen_addrs
                .iter()
                .any(|a| self.is_public_adress(&a)))
            .collect();
        println!("Peers with public address: {}", peers_with_public.len());
        println!("Peers with private address: {}", infos.len() - peers_with_public.len());
        println!("Peers with associated role: {}", self.peer_role.len());

        Locator::new(&infos).locate_peers().await;

        Ok(())
    }
}

pub struct WalkerBuilder {
    genesis: String,
    bootnodes: Vec<String>,
    timeout: Duration,
}

impl Default for WalkerBuilder {
    fn default() -> Self {
        Self {
            genesis: Default::default(), 
            bootnodes: Default::default(),
            timeout: Default::default(),
        }
    }
}

impl WalkerBuilder {
    fn create_swarm(&self) -> Result<Swarm<Behaviour>, Box<dyn Error>> {
        let local_key = identity::Keypair::generate_ed25519();
        let local_peer_id = PeerId::from(local_key.public());
        tracing::info!("Local peer Id {:?}", local_peer_id);

        let bootnodes = self.bootnodes
            .iter()
            .map(|bootnode| {
                let parts: Vec<_> = bootnode.split('/').collect();
                let peer = parts.last().expect("Bootnode peer should be valid; qed");
                let multiaddress: Multiaddr = bootnode
                    .parse()
                    .expect("Bootnode multiaddress should be valid; qed");
                let peer_id: PeerId = peer
                    .parse()
                    .expect("Bootnode should have valid peer ID; qed");

                tracing::info!("Bootnode peer={:?}", peer_id);
                (peer_id, multiaddress)
            })
            .collect::<Vec<_>>();

        let protocol_data = ProtocolsData {
            genesis_hash: H256::from_slice(hex::decode(self.genesis.clone())?.as_slice()),
            node_role: ProtocolRole::FullNode,
        };

        let mut swarm: Swarm<Behaviour> = {
            let transport = TransportBuilder::new()
                .yamux_maximum_buffer_size(256 * MIB)
                .build(local_key.clone());

            let discovery = DiscoveryBuilder::new()
                .record_ttl(Some(Duration::from_secs(0)))
                .provider_ttl(Some(Duration::from_secs(0)))
                .query_timeout(Duration::from_secs(5 * 60))
                .build(local_peer_id, &self.genesis);

            let peer_info = PeerBehaviour::new(local_key.public());
            let notifications = Notifications::new(protocol_data);

            let behavior = Behaviour {
                notifications,
                peer_info,
                discovery,
            };

            let config = SwarmConfig::with_tokio_executor();
            Swarm::new(transport, behavior, local_peer_id, config)
        };

        for (peer, multiaddress) in &bootnodes {
            swarm
                .behaviour_mut()
                .discovery
                .add_address(peer, multiaddress.clone());
        }

        Ok(swarm)
    }

    pub fn with_genesis(mut self, genesis: String) -> Self {
        self.genesis = genesis.trim_start_matches("0x").to_string();
        self
    }

    pub fn with_bootnodes(mut self, bootnodes: Vec<String>) -> Self {
        self.bootnodes = bootnodes;
        self
    }

    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }

    pub fn build(&self) -> Result<Walker, Box<dyn Error>> {
        let swarm = self.create_swarm()?;
        Ok(Walker {
            swarm,
            genesis: self.genesis.clone(),
            timeout: self.timeout,
            queries: HashSet::with_capacity(1024),
            discovered_with_address: HashMap::with_capacity(1024),
            peer_details: HashMap::with_capacity(1024),
            peer_role: HashMap::with_capacity(1024),
            dialed_peers: HashMap::with_capacity(1024),
        })
    }
}