use crate::data::error::{Error, Result};
use ant_protocol::transport::{
CoreNodeConfig, IPDiversityConfig, MultiAddr, NodeMode, P2PNode, PeerId, WitnessedCloseGroup,
};
use ant_protocol::MAX_WIRE_MESSAGE_SIZE;
use serde::{Deserialize, Serialize};
use std::net::SocketAddr;
use std::sync::Arc;
pub const REBOOTSTRAP_THRESHOLD: usize = 3;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct NetworkHealth {
pub write_ready: bool,
pub connected_peers: u32,
pub routing_table_size: u32,
pub rebootstrap_threshold: u32,
}
impl NetworkHealth {
#[must_use]
pub fn from_counts(connected_peers: usize, routing_table_size: usize) -> Self {
Self {
write_ready: routing_table_size.max(connected_peers) >= REBOOTSTRAP_THRESHOLD,
connected_peers: connected_peers.try_into().unwrap_or(u32::MAX),
routing_table_size: routing_table_size.try_into().unwrap_or(u32::MAX),
rebootstrap_threshold: REBOOTSTRAP_THRESHOLD as u32,
}
}
}
pub struct Network {
node: Arc<P2PNode>,
}
impl Network {
pub async fn new(
bootstrap_peers: &[SocketAddr],
allow_loopback: bool,
ipv6: bool,
) -> Result<Self> {
let mut core_config = CoreNodeConfig::builder()
.port(0)
.ipv6(ipv6)
.local(allow_loopback)
.mode(NodeMode::Client)
.max_message_size(MAX_WIRE_MESSAGE_SIZE)
.build()
.map_err(|e| Error::Network(format!("Failed to create core config: {e}")))?;
core_config.diversity_config = Some(IPDiversityConfig::permissive());
core_config.bootstrap_peers = bootstrap_peers
.iter()
.map(|addr| MultiAddr::quic(*addr))
.collect();
let node = P2PNode::new(core_config)
.await
.map_err(|e| Error::Network(format!("Failed to create P2P node: {e}")))?;
node.start()
.await
.map_err(|e| Error::Network(format!("Failed to start P2P node: {e}")))?;
Ok(Self {
node: Arc::new(node),
})
}
#[must_use]
pub fn from_node(node: Arc<P2PNode>) -> Self {
Self { node }
}
#[must_use]
pub fn node(&self) -> &Arc<P2PNode> {
&self.node
}
#[must_use]
pub fn peer_id(&self) -> &PeerId {
self.node.peer_id()
}
pub async fn find_closest_peers(
&self,
target: &[u8; 32],
count: usize,
) -> Result<Vec<(PeerId, Vec<MultiAddr>)>> {
let local_peer_id = self.node.peer_id();
let closest_nodes = self
.node
.dht()
.find_closest_nodes(target, count + 1)
.await
.map_err(|e| Error::Network(format!("DHT closest-nodes lookup failed: {e}")))?;
Ok(closest_nodes
.into_iter()
.filter(|n| n.peer_id != *local_peer_id)
.take(count)
.map(|n| {
let addrs = n.addresses_by_priority();
(n.peer_id, addrs)
})
.collect())
}
pub async fn find_witnessed_close_group(
&self,
target: &[u8; 32],
count: usize,
) -> Result<WitnessedCloseGroup> {
self.find_witnessed_close_group_with_view_count(target, count, count)
.await
}
pub async fn find_witnessed_close_group_with_view_count(
&self,
target: &[u8; 32],
count: usize,
view_count: usize,
) -> Result<WitnessedCloseGroup> {
self.node
.dht()
.find_witnessed_close_group_with_view_count(target, count, view_count)
.await
.map_err(|e| Error::Network(format!("DHT witnessed close-group lookup failed: {e}")))
}
pub async fn connected_peers(&self) -> Vec<PeerId> {
self.node.connected_peers().await
}
pub async fn health(&self) -> NetworkHealth {
let connected_peers = self.node.peer_count().await;
let routing_table_size = self.node.dht_manager().get_routing_table_size().await;
NetworkHealth::from_counts(connected_peers, routing_table_size)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn write_ready_false_with_no_peers() {
let h = NetworkHealth::from_counts(0, 0);
assert!(!h.write_ready);
assert_eq!(h.connected_peers, 0);
assert_eq!(h.routing_table_size, 0);
assert_eq!(h.rebootstrap_threshold, REBOOTSTRAP_THRESHOLD as u32);
}
#[test]
fn write_ready_false_below_threshold_on_both_signals() {
assert!(!NetworkHealth::from_counts(1, 0).write_ready);
assert!(!NetworkHealth::from_counts(2, 2).write_ready);
}
#[test]
fn write_ready_true_via_connections_despite_low_routing_table() {
assert!(NetworkHealth::from_counts(10, 2).write_ready);
}
#[test]
fn write_ready_true_via_routing_table_alone() {
assert!(NetworkHealth::from_counts(0, REBOOTSTRAP_THRESHOLD).write_ready);
}
#[test]
fn write_ready_true_at_exact_threshold_on_connections() {
assert!(NetworkHealth::from_counts(REBOOTSTRAP_THRESHOLD, 0).write_ready);
}
#[test]
fn counts_saturate_at_u32_max() {
let h = NetworkHealth::from_counts(usize::MAX, usize::MAX);
assert_eq!(h.connected_peers, u32::MAX);
assert_eq!(h.routing_table_size, u32::MAX);
assert!(h.write_ready);
}
}