ave-network 0.9.1

use futures_timer::Delay;

use libp2p::{
    Multiaddr, PeerId, StreamProtocol,
    core::Endpoint,
    futures::FutureExt,
    kad::{
        Behaviour as Kademlia, BootstrapError, Config as KademliaConfig,
        Event as KademliaEvent, GetClosestPeersError, GetClosestPeersOk,
        K_VALUE, PeerInfo, QueryResult,
        store::{MemoryStore, MemoryStoreConfig},
    },
    swarm::{
        CloseConnection, ConnectionDenied, ConnectionId, FromSwarm,
        NetworkBehaviour, THandler, ToSwarm,
    },
};
use serde::{Deserialize, Serialize};

use std::{
    collections::{HashMap, VecDeque},
    task::{Poll, Waker},
    time::Duration,
};

use crate::{NodeType, utils::LimitsConfig};

#[cfg(not(any(test, feature = "test")))]
use crate::utils::{is_dns, is_global, is_loop_back, is_private, is_tcp};

/// The discovery behaviour.
pub struct Behaviour {
    /// Boolean that activates the random walk if the node has already finished the initial pre-routing phase.
    pre_routing: bool,

    /// Kademlia behavior.
    kademlia: Kademlia<MemoryStore>,

    waker: Option<Waker>,

    /// The next random walk in the Kademlia DHT. `None` if random walks are disabled.
    next_random_walk: Option<Delay>,

    /// Duration between random walks.
    duration_to_next_kad: Duration,

    /// Number of nodes we're currently connected to.
    num_connections: u64,

    /// Number of active connections over which we interrupt the discovery process.
    discovery_only_if_under_num: u64,

    #[cfg_attr(any(test, feature = "test"), allow(dead_code))]
    allow_private_address_in_dht: bool,

    #[cfg_attr(any(test, feature = "test"), allow(dead_code))]
    allow_dns_address_in_dht: bool,

    #[cfg_attr(any(test, feature = "test"), allow(dead_code))]
    allow_loop_back_address_in_dht: bool,

    /// Peers to close connection
    close_connections: VecDeque<(PeerId, Option<ConnectionId>)>,

    peer_to_remove: HashMap<PeerId, u8>,
}

impl Behaviour {
    /// Creates a new routing `Behaviour`.
    pub fn new(
        peer_id: PeerId,
        config: Config,
        protocol: StreamProtocol,
        node_type: NodeType,
        limits: LimitsConfig,
    ) -> Self {
        let Config {
            dht_random_walk,
            discovery_only_if_under_num,
            allow_dns_address_in_dht,
            allow_private_address_in_dht,
            allow_loop_back_address_in_dht,
            kademlia_disjoint_query_paths,
        } = config;

        let mut kad_config = KademliaConfig::new(protocol);

        kad_config.disjoint_query_paths(kademlia_disjoint_query_paths);
        kad_config.set_query_timeout(Duration::from_secs(
            limits.kademlia_query_timeout,
        ));
        kad_config.set_replication_interval(None);
        kad_config.set_caching(libp2p::kad::Caching::Disabled);

        // By default Kademlia attempts to insert all peers into its routing table once a
        // dialing attempt succeeds. In order to control which peer is added, disable the
        // auto-insertion and instead add peers manually.
        kad_config.set_kbucket_inserts(libp2p::kad::BucketInserts::Manual);
        kad_config.set_record_filtering(libp2p::kad::StoreInserts::FilterBoth);

        let store = MemoryStore::with_config(
            peer_id,
            MemoryStoreConfig {
                max_records: 0,
                max_value_bytes: 0,
                max_providers_per_key: K_VALUE.get(),
                max_provided_keys: 0,
            },
        );
        let mut kad = Kademlia::with_config(peer_id, store, kad_config);

        if let NodeType::Addressable | NodeType::Bootstrap = node_type {
            kad.set_mode(Some(libp2p::kad::Mode::Server));
        } else {
            kad.set_mode(Some(libp2p::kad::Mode::Client));
        }

        Self {
            kademlia: kad,
            next_random_walk: if dht_random_walk
                && node_type == NodeType::Bootstrap
            {
                Some(Delay::new(Duration::new(0, 0)))
            } else {
                None
            },
            duration_to_next_kad: Duration::from_secs(1),
            num_connections: 0,
            discovery_only_if_under_num,
            allow_dns_address_in_dht,
            allow_private_address_in_dht,
            allow_loop_back_address_in_dht,
            close_connections: VecDeque::new(),
            pre_routing: true,
            waker: None,
            peer_to_remove: HashMap::new(),
        }
    }

    pub fn new_close_connections(
        &mut self,
        peer_id: PeerId,
        connection_id: Option<ConnectionId>,
    ) {
        self.close_connections.push_back((peer_id, connection_id));

        if let Some(waker) = self.waker.take() {
            waker.wake();
        }
    }

    pub const fn finish_prerouting_state(&mut self) {
        self.pre_routing = false;
    }

    /// Returns true if the given peer is known.
    pub fn is_known_peer(&mut self, peer_id: &PeerId) -> bool {
        for b in self.kademlia.kbuckets() {
            if b.iter().any(|x| peer_id == x.node.key.preimage()) {
                return true;
            }
        }

        false
    }

    pub fn add_peer_to_remove(&mut self, peer_id: &PeerId) {
        let count = self
            .peer_to_remove
            .entry(*peer_id)
            .and_modify(|x| *x += 1)
            .or_insert(1);

        if *count >= 3 {
            self.remove_node(peer_id);
            self.clean_peer_to_remove(peer_id);
        }
    }

    pub fn clean_peer_to_remove(&mut self, peer_id: &PeerId) {
        self.peer_to_remove.remove(peer_id);
    }

    /// Add a self-reported address of a remote peer to the k-buckets of the DHT
    /// if it has compatible `supported_protocols`.
    ///
    /// **Note**: It is important that you call this method. The discovery mechanism will not
    /// automatically add connecting peers to the Kademlia k-buckets.
    pub fn add_self_reported_address(
        &mut self,
        peer_id: &PeerId,
        addr: &Multiaddr,
    ) -> bool {
        if self.is_invalid_address(addr) {
            return false;
        }

        self.kademlia.add_address(peer_id, addr.clone());
        true
    }

    #[cfg(any(test, feature = "test"))]
    pub fn is_invalid_address(&self, _addr: &Multiaddr) -> bool {
        return false;
    }

    #[cfg(not(any(test, feature = "test")))]
    pub fn is_invalid_address(&self, addr: &Multiaddr) -> bool {
        {
            // Our transport is TPC only
            if !is_tcp(addr) {
                return true;
            }

            if is_private(addr) {
                return !self.allow_private_address_in_dht;
            }

            if is_loop_back(addr) {
                return !self.allow_loop_back_address_in_dht;
            }

            if is_dns(addr) {
                return !self.allow_dns_address_in_dht;
            }

            !is_global(addr)
        }
    }

    /// Discover closet peers to the given `PeerId`.
    pub fn discover(&mut self, peer_id: &PeerId) {
        self.kademlia.get_closest_peers(*peer_id);
    }

    /// Remove node from the DHT.
    pub fn remove_node(&mut self, peer_id: &PeerId) {
        self.kademlia.remove_peer(peer_id);
    }
}

/// Event generated by the `DiscoveryBehaviour`.
#[derive(Debug)]
pub enum Event {
    /// Closest peers to a given key have been found
    ClosestPeer {
        peer_id: PeerId,
        info: Option<PeerInfo>,
    },
}

impl NetworkBehaviour for Behaviour {
    type ConnectionHandler =
        <Kademlia<MemoryStore> as NetworkBehaviour>::ConnectionHandler;
    type ToSwarm = Event;

    fn handle_established_inbound_connection(
        &mut self,
        connection_id: ConnectionId,
        peer: PeerId,
        local_addr: &Multiaddr,
        remote_addr: &Multiaddr,
    ) -> Result<THandler<Self>, ConnectionDenied> {
        self.kademlia.handle_established_inbound_connection(
            connection_id,
            peer,
            local_addr,
            remote_addr,
        )
    }

    fn handle_established_outbound_connection(
        &mut self,
        connection_id: ConnectionId,
        peer: PeerId,
        addr: &Multiaddr,
        role_override: Endpoint,
        port_use: libp2p::core::transport::PortUse,
    ) -> Result<THandler<Self>, ConnectionDenied> {
        self.kademlia.handle_established_outbound_connection(
            connection_id,
            peer,
            addr,
            role_override,
            port_use,
        )
    }

    fn on_swarm_event(&mut self, event: FromSwarm) {
        match event {
            FromSwarm::AddressChange(..)
            | FromSwarm::DialFailure(..)
            | FromSwarm::ExpiredListenAddr(..)
            | FromSwarm::ExternalAddrConfirmed(..)
            | FromSwarm::ExternalAddrExpired(..)
            | FromSwarm::ListenerClosed(..)
            | FromSwarm::ListenFailure(..)
            | FromSwarm::ListenerError(..)
            | FromSwarm::NewListener(..)
            | FromSwarm::NewListenAddr(..)
            | FromSwarm::NewExternalAddrCandidate(..)
            | FromSwarm::NewExternalAddrOfPeer(..) => {
                self.kademlia.on_swarm_event(event);
            }
            FromSwarm::ConnectionEstablished(e) => {
                self.num_connections += 1;
                self.kademlia
                    .on_swarm_event(FromSwarm::ConnectionEstablished(e));
            }
            FromSwarm::ConnectionClosed(e) => {
                self.num_connections = self.num_connections.saturating_sub(1);
                self.kademlia.on_swarm_event(FromSwarm::ConnectionClosed(e));
            }
            _ => self.kademlia.on_swarm_event(event),
        }
    }

    fn on_connection_handler_event(
        &mut self,
        peer_id: PeerId,
        connection_id: libp2p::swarm::ConnectionId,
        event: libp2p::swarm::THandlerOutEvent<Self>,
    ) {
        self.kademlia.on_connection_handler_event(
            peer_id,
            connection_id,
            event,
        );
    }

    fn poll(
        &mut self,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<
        libp2p::swarm::ToSwarm<
            Self::ToSwarm,
            libp2p::swarm::THandlerInEvent<Self>,
        >,
    > {
        if let Some((peer_id, connection_id)) =
            self.close_connections.pop_front()
        {
            if let Some(connection_id) = connection_id {
                return Poll::Ready(ToSwarm::CloseConnection {
                    peer_id,
                    connection: CloseConnection::One(connection_id),
                });
            } else {
                return Poll::Ready(ToSwarm::CloseConnection {
                    peer_id,
                    connection: CloseConnection::All,
                });
            }
        }

        while let Poll::Ready(ev) = self.kademlia.poll(cx) {
            match ev {
                ToSwarm::GenerateEvent(ev) => {
                    match ev {
                        KademliaEvent::RoutablePeer { peer, address } => {
                            self.add_self_reported_address(&peer, &address);
                        }
                        KademliaEvent::ModeChanged { .. }
                        | KademliaEvent::InboundRequest { .. }
                        | KademliaEvent::UnroutablePeer { .. }
                        | KademliaEvent::PendingRoutablePeer { .. }
                        | KademliaEvent::RoutingUpdated { .. } => {
                            // We are not interested in this event at the moment.
                        }
                        KademliaEvent::OutboundQueryProgressed {
                            result,
                            ..
                        } => {
                            match result {
                                QueryResult::Bootstrap(bootstrap_ok) => {
                                    match bootstrap_ok {
                                        Ok(ok) => {
                                            self.clean_peer_to_remove(&ok.peer);
                                        }
                                        Err(e) => {
                                            let BootstrapError::Timeout {
                                                peer,
                                                ..
                                            } = e;
                                            self.add_peer_to_remove(&peer);
                                        }
                                    };
                                }
                                QueryResult::GetClosestPeers(
                                    get_closest_peers_ok,
                                ) => {
                                    // kademlia.get_closest_peers(PeerId)
                                    match get_closest_peers_ok {
                                        Ok(GetClosestPeersOk {
                                            key,
                                            peers,
                                        }) => {
                                            if let Ok(peer_id) =
                                                PeerId::from_bytes(&key)
                                            {
                                                if let Some(info) =
                                                    peers.iter().find(|x| {
                                                        x.peer_id == peer_id
                                                    })
                                                {
                                                    return Poll::Ready(
                                                    ToSwarm::GenerateEvent(
                                                        Event::ClosestPeer {
                                                            peer_id,
                                                            info: Some(
                                                                info.clone(),
                                                            ),
                                                        },
                                                    ),
                                                );
                                                } else {
                                                    return Poll::Ready(
                                                    ToSwarm::GenerateEvent(
                                                        Event::ClosestPeer {
                                                            peer_id,
                                                            info: None,
                                                        },
                                                    ),
                                                );
                                                }
                                            };
                                        }
                                        Err(
                                            GetClosestPeersError::Timeout {
                                                key,
                                                ..
                                            },
                                        ) => {
                                            if let Ok(peer_id) =
                                                PeerId::from_bytes(&key)
                                            {
                                                return Poll::Ready(
                                                    ToSwarm::GenerateEvent(
                                                        Event::ClosestPeer {
                                                            peer_id,
                                                            info: None,
                                                        },
                                                    ),
                                                );
                                            };
                                        }
                                    }
                                }
                                QueryResult::GetProviders(..)
                                | QueryResult::StartProviding(..)
                                | QueryResult::RepublishProvider(..)
                                | QueryResult::GetRecord(..)
                                | QueryResult::PutRecord(..)
                                | QueryResult::RepublishRecord(..) => {
                                    // We are not interested in this event at the moment.
                                }
                            };
                        }
                    }
                }
                ToSwarm::Dial { opts } => {
                    return Poll::Ready(ToSwarm::Dial { opts });
                }
                ToSwarm::NotifyHandler {
                    peer_id,
                    handler,
                    event,
                } => {
                    return Poll::Ready(ToSwarm::NotifyHandler {
                        peer_id,
                        handler,
                        event,
                    });
                }
                ToSwarm::CloseConnection {
                    peer_id,
                    connection,
                } => {
                    return Poll::Ready(ToSwarm::CloseConnection {
                        peer_id,
                        connection,
                    });
                }
                ToSwarm::ExternalAddrConfirmed(e) => {
                    return Poll::Ready(ToSwarm::ExternalAddrConfirmed(e));
                }
                ToSwarm::ExternalAddrExpired(e) => {
                    return Poll::Ready(ToSwarm::ExternalAddrExpired(e));
                }
                ToSwarm::ListenOn { opts } => {
                    return Poll::Ready(ToSwarm::ListenOn { opts });
                }
                ToSwarm::NewExternalAddrCandidate(e) => {
                    return Poll::Ready(ToSwarm::NewExternalAddrCandidate(e));
                }
                ToSwarm::RemoveListener { id } => {
                    return Poll::Ready(ToSwarm::RemoveListener { id });
                }
                _ => {}
            }
        }

        // Poll the stream that fires when we need to start a random Kademlia query.
        if !self.pre_routing
            && let Some(next) = self.next_random_walk.as_mut()
            && next.poll_unpin(cx).is_ready()
        {
            if self.num_connections < self.discovery_only_if_under_num {
                self.kademlia.get_closest_peers(PeerId::random());
            }

            *next = Delay::new(self.duration_to_next_kad);
            self.duration_to_next_kad = std::cmp::min(
                self.duration_to_next_kad * 2,
                Duration::from_secs(120),
            );
        }

        self.waker = Some(cx.waker().clone());
        Poll::Pending
    }

    fn handle_pending_inbound_connection(
        &mut self,
        connection_id: libp2p::swarm::ConnectionId,
        local_addr: &Multiaddr,
        remote_addr: &Multiaddr,
    ) -> Result<(), libp2p::swarm::ConnectionDenied> {
        self.kademlia.handle_pending_inbound_connection(
            connection_id,
            local_addr,
            remote_addr,
        )
    }

    fn handle_pending_outbound_connection(
        &mut self,
        connection_id: libp2p::swarm::ConnectionId,
        maybe_peer: Option<PeerId>,
        addresses: &[Multiaddr],
        effective_role: libp2p::core::Endpoint,
    ) -> Result<Vec<Multiaddr>, libp2p::swarm::ConnectionDenied> {
        let addresses = self.kademlia.handle_pending_outbound_connection(
            connection_id,
            maybe_peer,
            addresses,
            effective_role,
        )?;

        let filter_addresses = addresses
            .iter()
            .filter(|x| !self.is_invalid_address(x))
            .cloned()
            .collect::<Vec<Multiaddr>>();

        if filter_addresses.len() != addresses.len() {
            if let Some(peer_id) = maybe_peer {
                self.kademlia.remove_peer(&peer_id);
                for addr in filter_addresses.iter() {
                    self.kademlia.add_address(&peer_id, addr.clone());
                }
            }
        } else if filter_addresses.is_empty()
            && let Some(peer_id) = maybe_peer
        {
            self.peer_to_remove.remove(&peer_id);
            self.kademlia.remove_peer(&peer_id);
        }

        Ok(filter_addresses)
    }
}

/// Configuration for the routing behaviour.
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(default)]
pub struct Config {
    /// Whether to enable random walks in the Kademlia DHT.
    dht_random_walk: bool,

    /// Number of active connections over which we interrupt the discovery process.
    discovery_only_if_under_num: u64,

    allow_private_address_in_dht: bool,

    allow_dns_address_in_dht: bool,

    allow_loop_back_address_in_dht: bool,

    /// When enabled the number of disjoint paths used equals the configured parallelism.
    kademlia_disjoint_query_paths: bool,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            dht_random_walk: true,
            discovery_only_if_under_num: 25,
            allow_private_address_in_dht: Default::default(),
            allow_dns_address_in_dht: Default::default(),
            allow_loop_back_address_in_dht: Default::default(),
            kademlia_disjoint_query_paths: true,
        }
    }
}

impl Config {
    /// Get DHT random walk.
    pub const fn get_dht_random_walk(&self) -> bool {
        self.dht_random_walk
    }

    /// Enables or disables random walks in the Kademlia DHT.
    pub const fn with_dht_random_walk(mut self, enable: bool) -> Self {
        self.dht_random_walk = enable;
        self
    }

    /// Get discovery limits.
    pub const fn get_discovery_limit(&self) -> u64 {
        self.discovery_only_if_under_num
    }

    /// Sets the number of active connections over which we interrupt the discovery process.
    pub const fn with_discovery_limit(mut self, num: u64) -> Self {
        self.discovery_only_if_under_num = num;
        self
    }

    /// Get allow_local_address_in_dht.
    pub const fn get_allow_private_address_in_dht(&self) -> bool {
        self.allow_private_address_in_dht
    }

    /// Whether to allow local addresses in the DHT.
    pub const fn with_allow_private_address_in_dht(
        mut self,
        allow: bool,
    ) -> Self {
        self.allow_private_address_in_dht = allow;
        self
    }

    /// Get allow_dns_address_in_dht.
    pub const fn get_allow_dns_address_in_dht(&self) -> bool {
        self.allow_dns_address_in_dht
    }

    /// Whether to allow non-global addresses in the DHT.
    pub const fn with_allow_dns_address_in_dht(mut self, allow: bool) -> Self {
        self.allow_dns_address_in_dht = allow;
        self
    }

    /// Get allow_non_globals_in_dht.
    pub const fn get_allow_loop_back_address_in_dht(&self) -> bool {
        self.allow_loop_back_address_in_dht
    }

    /// Whether to allow non-global addresses in the DHT.
    pub const fn with_allow_loop_back_address_in_dht(
        mut self,
        allow: bool,
    ) -> Self {
        self.allow_loop_back_address_in_dht = allow;
        self
    }

    /// Get allow kademlia disjoint query paths
    pub const fn get_kademlia_disjoint_query_paths(&self) -> bool {
        self.kademlia_disjoint_query_paths
    }

    /// When enabled the number of disjoint paths used equals the configured parallelism.
    pub const fn with_kademlia_disjoint_query_paths(
        mut self,
        enable: bool,
    ) -> Self {
        self.kademlia_disjoint_query_paths = enable;
        self
    }
}

/// A node in the routing table.
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct RoutingNode {
    /// Peer ID.
    pub peer_id: String,
    /// Address.
    pub address: Vec<String>,
}