use std::{
net::SocketAddr,
sync::{
atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
Arc,
},
};
use enr::Enr;
use parking_lot::Mutex;
use reth_discv4::Discv4;
use reth_discv5::Discv5;
use reth_eth_wire::{DisconnectReason, NewBlock, NewPooledTransactionHashes, SharedTransactions};
use reth_network_api::{
test_utils::{PeersHandle, PeersHandleProvider},
BlockDownloaderProvider, DiscoveryEvent, NetworkError, NetworkEvent,
NetworkEventListenerProvider, NetworkInfo, NetworkStatus, PeerInfo, PeerRequest, Peers,
PeersInfo,
};
use reth_network_p2p::{
sync::{NetworkSyncUpdater, SyncState, SyncStateProvider},
BlockClient,
};
use reth_network_peers::{NodeRecord, PeerId};
use reth_network_types::{PeerAddr, PeerKind, Reputation, ReputationChangeKind};
use reth_primitives::{Head, TransactionSigned, B256};
use reth_tokio_util::{EventSender, EventStream};
use secp256k1::SecretKey;
use tokio::sync::{
mpsc::{self, UnboundedSender},
oneshot,
};
use tokio_stream::wrappers::UnboundedReceiverStream;
use crate::{
config::NetworkMode, protocol::RlpxSubProtocol, swarm::NetworkConnectionState,
transactions::TransactionsHandle, FetchClient,
};
#[derive(Clone, Debug)]
pub struct NetworkHandle {
inner: Arc<NetworkInner>,
}
impl NetworkHandle {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
num_active_peers: Arc<AtomicUsize>,
listener_address: Arc<Mutex<SocketAddr>>,
to_manager_tx: UnboundedSender<NetworkHandleMessage>,
secret_key: SecretKey,
local_peer_id: PeerId,
peers: PeersHandle,
network_mode: NetworkMode,
chain_id: Arc<AtomicU64>,
tx_gossip_disabled: bool,
discv4: Option<Discv4>,
discv5: Option<Discv5>,
event_sender: EventSender<NetworkEvent>,
) -> Self {
let inner = NetworkInner {
num_active_peers,
to_manager_tx,
listener_address,
secret_key,
local_peer_id,
peers,
network_mode,
is_syncing: Arc::new(AtomicBool::new(false)),
initial_sync_done: Arc::new(AtomicBool::new(false)),
chain_id,
tx_gossip_disabled,
discv4,
discv5,
event_sender,
};
Self { inner: Arc::new(inner) }
}
pub fn peer_id(&self) -> &PeerId {
&self.inner.local_peer_id
}
fn manager(&self) -> &UnboundedSender<NetworkHandleMessage> {
&self.inner.to_manager_tx
}
pub fn mode(&self) -> &NetworkMode {
&self.inner.network_mode
}
pub(crate) fn send_message(&self, msg: NetworkHandleMessage) {
let _ = self.inner.to_manager_tx.send(msg);
}
pub fn update_status(&self, head: Head) {
self.send_message(NetworkHandleMessage::StatusUpdate { head });
}
pub fn announce_block(&self, block: NewBlock, hash: B256) {
self.send_message(NetworkHandleMessage::AnnounceBlock(block, hash))
}
pub fn send_request(&self, peer_id: PeerId, request: PeerRequest) {
self.send_message(NetworkHandleMessage::EthRequest { peer_id, request })
}
pub fn send_transactions_hashes(&self, peer_id: PeerId, msg: NewPooledTransactionHashes) {
self.send_message(NetworkHandleMessage::SendPooledTransactionHashes { peer_id, msg })
}
pub fn send_transactions(&self, peer_id: PeerId, msg: Vec<Arc<TransactionSigned>>) {
self.send_message(NetworkHandleMessage::SendTransaction {
peer_id,
msg: SharedTransactions(msg),
})
}
pub async fn transactions_handle(&self) -> Option<TransactionsHandle> {
let (tx, rx) = oneshot::channel();
let _ = self.manager().send(NetworkHandleMessage::GetTransactionsHandle(tx));
rx.await.unwrap()
}
pub async fn shutdown(&self) -> Result<(), oneshot::error::RecvError> {
let (tx, rx) = oneshot::channel();
self.send_message(NetworkHandleMessage::Shutdown(tx));
rx.await
}
pub fn set_network_active(&self) {
self.set_network_conn(NetworkConnectionState::Active);
}
pub fn set_network_hibernate(&self) {
self.set_network_conn(NetworkConnectionState::Hibernate);
}
fn set_network_conn(&self, network_conn: NetworkConnectionState) {
self.send_message(NetworkHandleMessage::SetNetworkState(network_conn));
}
pub fn tx_gossip_disabled(&self) -> bool {
self.inner.tx_gossip_disabled
}
pub fn secret_key(&self) -> &SecretKey {
&self.inner.secret_key
}
}
impl NetworkEventListenerProvider for NetworkHandle {
fn event_listener(&self) -> EventStream<NetworkEvent> {
self.inner.event_sender.new_listener()
}
fn discovery_listener(&self) -> UnboundedReceiverStream<DiscoveryEvent> {
let (tx, rx) = mpsc::unbounded_channel();
let _ = self.manager().send(NetworkHandleMessage::DiscoveryListener(tx));
UnboundedReceiverStream::new(rx)
}
}
impl NetworkProtocols for NetworkHandle {
fn add_rlpx_sub_protocol(&self, protocol: RlpxSubProtocol) {
self.send_message(NetworkHandleMessage::AddRlpxSubProtocol(protocol))
}
}
impl PeersInfo for NetworkHandle {
fn num_connected_peers(&self) -> usize {
self.inner.num_active_peers.load(Ordering::Relaxed)
}
fn local_node_record(&self) -> NodeRecord {
if let Some(discv4) = &self.inner.discv4 {
discv4.node_record()
} else if let Some(discv5) = &self.inner.discv5 {
discv5.node_record()
} else {
let id = *self.peer_id();
let mut socket_addr = *self.inner.listener_address.lock();
if socket_addr.ip().is_unspecified() {
if socket_addr.ip().is_ipv4() {
socket_addr.set_ip(std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST));
} else {
socket_addr.set_ip(std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST));
}
}
NodeRecord::new(socket_addr, id)
}
}
fn local_enr(&self) -> Enr<SecretKey> {
let local_node_record = self.local_node_record();
let mut builder = Enr::builder();
builder.ip(local_node_record.address);
if local_node_record.address.is_ipv4() {
builder.udp4(local_node_record.udp_port);
builder.tcp4(local_node_record.tcp_port);
} else {
builder.udp6(local_node_record.udp_port);
builder.tcp6(local_node_record.tcp_port);
}
builder.build(&self.inner.secret_key).expect("valid enr")
}
}
impl Peers for NetworkHandle {
fn add_trusted_peer_id(&self, peer: PeerId) {
self.send_message(NetworkHandleMessage::AddTrustedPeerId(peer));
}
fn add_peer_kind(
&self,
peer: PeerId,
kind: PeerKind,
tcp_addr: SocketAddr,
udp_addr: Option<SocketAddr>,
) {
let addr = PeerAddr::new(tcp_addr, udp_addr);
self.send_message(NetworkHandleMessage::AddPeerAddress(peer, kind, addr));
}
async fn get_peers_by_kind(&self, kind: PeerKind) -> Result<Vec<PeerInfo>, NetworkError> {
let (tx, rx) = oneshot::channel();
let _ = self.manager().send(NetworkHandleMessage::GetPeerInfosByPeerKind(kind, tx));
Ok(rx.await?)
}
async fn get_all_peers(&self) -> Result<Vec<PeerInfo>, NetworkError> {
let (tx, rx) = oneshot::channel();
let _ = self.manager().send(NetworkHandleMessage::GetPeerInfos(tx));
Ok(rx.await?)
}
async fn get_peer_by_id(&self, peer_id: PeerId) -> Result<Option<PeerInfo>, NetworkError> {
let (tx, rx) = oneshot::channel();
let _ = self.manager().send(NetworkHandleMessage::GetPeerInfoById(peer_id, tx));
Ok(rx.await?)
}
async fn get_peers_by_id(&self, peer_ids: Vec<PeerId>) -> Result<Vec<PeerInfo>, NetworkError> {
let (tx, rx) = oneshot::channel();
let _ = self.manager().send(NetworkHandleMessage::GetPeerInfosByIds(peer_ids, tx));
Ok(rx.await?)
}
fn remove_peer(&self, peer: PeerId, kind: PeerKind) {
self.send_message(NetworkHandleMessage::RemovePeer(peer, kind))
}
fn disconnect_peer(&self, peer: PeerId) {
self.send_message(NetworkHandleMessage::DisconnectPeer(peer, None))
}
fn disconnect_peer_with_reason(&self, peer: PeerId, reason: DisconnectReason) {
self.send_message(NetworkHandleMessage::DisconnectPeer(peer, Some(reason)))
}
fn connect_peer_kind(
&self,
peer_id: PeerId,
kind: PeerKind,
tcp_addr: SocketAddr,
udp_addr: Option<SocketAddr>,
) {
self.send_message(NetworkHandleMessage::ConnectPeer(
peer_id,
kind,
PeerAddr::new(tcp_addr, udp_addr),
))
}
fn reputation_change(&self, peer_id: PeerId, kind: ReputationChangeKind) {
self.send_message(NetworkHandleMessage::ReputationChange(peer_id, kind));
}
async fn reputation_by_id(&self, peer_id: PeerId) -> Result<Option<Reputation>, NetworkError> {
let (tx, rx) = oneshot::channel();
let _ = self.manager().send(NetworkHandleMessage::GetReputationById(peer_id, tx));
Ok(rx.await?)
}
}
impl PeersHandleProvider for NetworkHandle {
fn peers_handle(&self) -> &PeersHandle {
&self.inner.peers
}
}
impl NetworkInfo for NetworkHandle {
fn local_addr(&self) -> SocketAddr {
*self.inner.listener_address.lock()
}
async fn network_status(&self) -> Result<NetworkStatus, NetworkError> {
let (tx, rx) = oneshot::channel();
let _ = self.manager().send(NetworkHandleMessage::GetStatus(tx));
rx.await.map_err(Into::into)
}
fn chain_id(&self) -> u64 {
self.inner.chain_id.load(Ordering::Relaxed)
}
fn is_syncing(&self) -> bool {
SyncStateProvider::is_syncing(self)
}
fn is_initially_syncing(&self) -> bool {
SyncStateProvider::is_initially_syncing(self)
}
}
impl SyncStateProvider for NetworkHandle {
fn is_syncing(&self) -> bool {
self.inner.is_syncing.load(Ordering::Relaxed)
}
fn is_initially_syncing(&self) -> bool {
if self.inner.initial_sync_done.load(Ordering::Relaxed) {
return false
}
self.inner.is_syncing.load(Ordering::Relaxed)
}
}
impl NetworkSyncUpdater for NetworkHandle {
fn update_sync_state(&self, state: SyncState) {
let future_state = state.is_syncing();
let prev_state = self.inner.is_syncing.swap(future_state, Ordering::Relaxed);
let syncing_to_idle_state_transition = prev_state && !future_state;
if syncing_to_idle_state_transition {
self.inner.initial_sync_done.store(true, Ordering::Relaxed);
}
}
fn update_status(&self, head: Head) {
self.send_message(NetworkHandleMessage::StatusUpdate { head });
}
}
impl BlockDownloaderProvider for NetworkHandle {
async fn fetch_client(&self) -> Result<impl BlockClient + 'static, oneshot::error::RecvError> {
let (tx, rx) = oneshot::channel();
let _ = self.manager().send(NetworkHandleMessage::FetchClient(tx));
rx.await
}
}
#[derive(Debug)]
struct NetworkInner {
num_active_peers: Arc<AtomicUsize>,
to_manager_tx: UnboundedSender<NetworkHandleMessage>,
listener_address: Arc<Mutex<SocketAddr>>,
secret_key: SecretKey,
local_peer_id: PeerId,
peers: PeersHandle,
network_mode: NetworkMode,
is_syncing: Arc<AtomicBool>,
initial_sync_done: Arc<AtomicBool>,
chain_id: Arc<AtomicU64>,
tx_gossip_disabled: bool,
discv4: Option<Discv4>,
discv5: Option<Discv5>,
event_sender: EventSender<NetworkEvent>,
}
pub trait NetworkProtocols: Send + Sync {
fn add_rlpx_sub_protocol(&self, protocol: RlpxSubProtocol);
}
#[derive(Debug)]
pub(crate) enum NetworkHandleMessage {
AddTrustedPeerId(PeerId),
AddPeerAddress(PeerId, PeerKind, PeerAddr),
RemovePeer(PeerId, PeerKind),
DisconnectPeer(PeerId, Option<DisconnectReason>),
AnnounceBlock(NewBlock, B256),
SendTransaction {
peer_id: PeerId,
msg: SharedTransactions,
},
SendPooledTransactionHashes {
peer_id: PeerId,
msg: NewPooledTransactionHashes,
},
EthRequest {
peer_id: PeerId,
request: PeerRequest,
},
ReputationChange(PeerId, ReputationChangeKind),
FetchClient(oneshot::Sender<FetchClient>),
StatusUpdate {
head: Head,
},
GetStatus(oneshot::Sender<NetworkStatus>),
GetPeerInfosByIds(Vec<PeerId>, oneshot::Sender<Vec<PeerInfo>>),
GetPeerInfos(oneshot::Sender<Vec<PeerInfo>>),
GetPeerInfoById(PeerId, oneshot::Sender<Option<PeerInfo>>),
GetPeerInfosByPeerKind(PeerKind, oneshot::Sender<Vec<PeerInfo>>),
GetReputationById(PeerId, oneshot::Sender<Option<Reputation>>),
GetTransactionsHandle(oneshot::Sender<Option<TransactionsHandle>>),
Shutdown(oneshot::Sender<()>),
SetNetworkState(NetworkConnectionState),
DiscoveryListener(UnboundedSender<DiscoveryEvent>),
AddRlpxSubProtocol(RlpxSubProtocol),
ConnectPeer(PeerId, PeerKind, PeerAddr),
}