1use crate::{
2 config::NetworkMode, message::PeerMessage, protocol::RlpxSubProtocol,
3 swarm::NetworkConnectionState, transactions::TransactionsHandle, FetchClient,
4};
5use alloy_primitives::B256;
6use enr::Enr;
7use futures::StreamExt;
8use parking_lot::Mutex;
9use reth_discv4::{Discv4, NatResolver};
10use reth_discv5::Discv5;
11use reth_eth_wire::{
12 BlockRangeUpdate, BroadcastPoolTransactions, DisconnectReason, EthNetworkPrimitives,
13 NetworkPrimitives, NewPooledTransactionHashes, SharedTransactions,
14};
15use reth_ethereum_forks::{ForkFilter, Head};
16use reth_network_api::{
17 events::{NetworkPeersEvents, PeerEvent, PeerEventStream},
18 test_utils::{PeersHandle, PeersHandleProvider},
19 BlockDownloaderProvider, CellCustody, DiscoveryEvent, NetworkError, NetworkEvent,
20 NetworkEventListenerProvider, NetworkInfo, NetworkStatus, PeerInfo, PeerRequest, Peers,
21 PeersInfo,
22};
23use reth_network_p2p::sync::{NetworkSyncUpdater, SyncState, SyncStateProvider};
24use reth_network_peers::{NodeRecord, PeerId, TrustedPeer};
25use reth_network_types::{PeerAddr, PeerKind, Reputation, ReputationChangeKind};
26use reth_tokio_util::{EventSender, EventStream};
27use secp256k1::SecretKey;
28use std::{
29 net::SocketAddr,
30 sync::{
31 atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
32 Arc,
33 },
34};
35use tokio::sync::{
36 mpsc::{self, UnboundedSender},
37 oneshot,
38};
39use tokio_stream::wrappers::UnboundedReceiverStream;
40
41#[derive(Clone, Debug)]
45pub struct NetworkHandle<N: NetworkPrimitives = EthNetworkPrimitives> {
46 inner: Arc<NetworkInner<N>>,
48}
49
50impl<N: NetworkPrimitives> NetworkHandle<N> {
53 #[expect(clippy::too_many_arguments)]
55 pub(crate) fn new(
56 num_active_peers: Arc<AtomicUsize>,
57 listener_address: Arc<Mutex<SocketAddr>>,
58 to_manager_tx: UnboundedSender<NetworkHandleMessage<N>>,
59 secret_key: SecretKey,
60 local_peer_id: PeerId,
61 peers: PeersHandle,
62 network_mode: NetworkMode,
63 chain_id: Arc<AtomicU64>,
64 tx_gossip_disabled: bool,
65 discv4: Option<Discv4>,
66 discv5: Option<Discv5>,
67 event_sender: EventSender<NetworkEvent<PeerRequest<N>>>,
68 nat: Option<NatResolver>,
69 ) -> Self {
70 let inner = NetworkInner {
71 num_active_peers,
72 to_manager_tx,
73 listener_address,
74 secret_key,
75 local_peer_id,
76 peers,
77 network_mode,
78 is_syncing: Arc::new(AtomicBool::new(false)),
79 initial_sync_done: Arc::new(AtomicBool::new(false)),
80 chain_id,
81 cell_custody: CellCustody::default(),
82 tx_gossip_disabled,
83 discv4,
84 discv5,
85 event_sender,
86 nat,
87 };
88 Self { inner: Arc::new(inner) }
89 }
90
91 pub fn peer_id(&self) -> &PeerId {
93 &self.inner.local_peer_id
94 }
95
96 fn manager(&self) -> &UnboundedSender<NetworkHandleMessage<N>> {
97 &self.inner.to_manager_tx
98 }
99
100 pub fn mode(&self) -> &NetworkMode {
102 &self.inner.network_mode
103 }
104
105 pub(crate) fn send_message(&self, msg: NetworkHandleMessage<N>) {
107 let _ = self.inner.to_manager_tx.send(msg);
108 }
109
110 pub fn update_status(&self, head: Head) {
112 self.send_message(NetworkHandleMessage::StatusUpdate { head });
113 }
114
115 pub fn set_fork_filter(&self, fork_filter: ForkFilter) {
124 self.send_message(NetworkHandleMessage::SetForkFilter { fork_filter });
125 }
126
127 pub fn announce_block(&self, block: N::NewBlockPayload, hash: B256) {
133 self.send_message(NetworkHandleMessage::AnnounceBlock(block, hash))
134 }
135
136 pub fn send_request(&self, peer_id: PeerId, request: PeerRequest<N>) {
138 self.send_message(NetworkHandleMessage::EthRequest { peer_id, request })
139 }
140
141 pub fn send_transactions_hashes(&self, peer_id: PeerId, msg: NewPooledTransactionHashes) {
143 self.send_message(NetworkHandleMessage::SendPooledTransactionHashes { peer_id, msg })
144 }
145
146 pub fn send_transactions(&self, peer_id: PeerId, msg: Vec<Arc<N::BroadcastedTransaction>>) {
148 self.send_message(NetworkHandleMessage::SendTransaction {
149 peer_id,
150 msg: SharedTransactions(msg),
151 })
152 }
153
154 pub(crate) fn send_broadcast_pool_transactions(
156 &self,
157 peer_id: PeerId,
158 msg: BroadcastPoolTransactions,
159 ) {
160 self.send_message(NetworkHandleMessage::SendBroadcastPoolTransactions { peer_id, msg })
161 }
162
163 pub fn send_eth_message(&self, peer_id: PeerId, message: PeerMessage<N>) {
165 self.send_message(NetworkHandleMessage::EthMessage { peer_id, message })
166 }
167
168 pub async fn transactions_handle(&self) -> Option<TransactionsHandle<N>> {
172 let (tx, rx) = oneshot::channel();
173 let _ = self.manager().send(NetworkHandleMessage::GetTransactionsHandle(tx));
174 rx.await.unwrap()
175 }
176
177 pub async fn shutdown(&self) -> Result<(), oneshot::error::RecvError> {
182 let (tx, rx) = oneshot::channel();
183 self.send_message(NetworkHandleMessage::Shutdown(tx));
184 rx.await
185 }
186
187 pub fn set_network_active(&self) {
191 self.set_network_conn(NetworkConnectionState::Active);
192 }
193
194 pub fn set_network_hibernate(&self) {
198 self.set_network_conn(NetworkConnectionState::Hibernate);
199 }
200
201 fn set_network_conn(&self, network_conn: NetworkConnectionState) {
203 self.send_message(NetworkHandleMessage::SetNetworkState(network_conn));
204 }
205
206 pub fn tx_gossip_disabled(&self) -> bool {
208 self.inner.tx_gossip_disabled
209 }
210
211 pub fn secret_key(&self) -> &SecretKey {
213 &self.inner.secret_key
214 }
215
216 pub fn discv4(&self) -> Option<&Discv4> {
218 self.inner.discv4.as_ref()
219 }
220
221 pub fn discv5(&self) -> Option<&Discv5> {
223 self.inner.discv5.as_ref()
224 }
225}
226
227impl<N: NetworkPrimitives> NetworkPeersEvents for NetworkHandle<N> {
230 fn peer_events(&self) -> PeerEventStream {
232 let peer_events = self.inner.event_sender.new_listener().map(|event| match event {
233 NetworkEvent::Peer(peer_event) => peer_event,
234 NetworkEvent::ActivePeerSession { info, .. } => PeerEvent::SessionEstablished(info),
235 });
236 PeerEventStream::new(peer_events)
237 }
238}
239
240impl<N: NetworkPrimitives> NetworkEventListenerProvider for NetworkHandle<N> {
241 type Primitives = N;
242
243 fn event_listener(&self) -> EventStream<NetworkEvent<PeerRequest<Self::Primitives>>> {
244 self.inner.event_sender.new_listener()
245 }
246
247 fn discovery_listener(&self) -> UnboundedReceiverStream<DiscoveryEvent> {
248 let (tx, rx) = mpsc::unbounded_channel();
249 let _ = self.manager().send(NetworkHandleMessage::DiscoveryListener(tx));
250 UnboundedReceiverStream::new(rx)
251 }
252}
253
254impl<N: NetworkPrimitives> NetworkProtocols for NetworkHandle<N> {
255 fn add_rlpx_sub_protocol(&self, protocol: RlpxSubProtocol) {
256 self.send_message(NetworkHandleMessage::AddRlpxSubProtocol(protocol))
257 }
258}
259
260impl<N: NetworkPrimitives> PeersInfo for NetworkHandle<N> {
261 fn num_connected_peers(&self) -> usize {
262 self.inner.num_active_peers.load(Ordering::Relaxed)
263 }
264
265 fn local_node_record(&self) -> NodeRecord {
266 if let Some(discv4) = &self.inner.discv4 {
267 discv4.node_record()
270 } else if let Some(discv5) = self.inner.discv5.as_ref() {
271 if let Some(external) =
273 self.inner.nat.clone().and_then(|nat| nat.as_external_ip(discv5.local_port()))
274 {
275 NodeRecord::new((external, discv5.local_port()).into(), *self.peer_id())
276 } else {
277 self.inner.discv5.as_ref().and_then(|d| d.node_record()).unwrap_or_else(|| {
279 NodeRecord::new(
280 (std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST), discv5.local_port())
281 .into(),
282 *self.peer_id(),
283 )
284 })
285 }
286 .with_tcp_port(self.inner.listener_address.lock().port())
288 } else {
289 let mut socket_addr = *self.inner.listener_address.lock();
290
291 let external_ip =
292 self.inner.nat.clone().and_then(|nat| nat.as_external_ip(socket_addr.port()));
293
294 if let Some(ip) = external_ip {
295 socket_addr.set_ip(ip)
297 } else if socket_addr.ip().is_unspecified() {
298 if socket_addr.ip().is_ipv4() {
300 socket_addr.set_ip(std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST));
301 } else {
302 socket_addr.set_ip(std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST));
303 }
304 }
305
306 NodeRecord::new(socket_addr, *self.peer_id())
307 }
308 }
309
310 fn local_enr(&self) -> Enr<SecretKey> {
311 let local_node_record = self.local_node_record();
312 let mut builder = Enr::builder();
313 builder.ip(local_node_record.address);
314 if local_node_record.address.is_ipv4() {
315 builder.udp4(local_node_record.udp_port);
316 builder.tcp4(local_node_record.tcp_port);
317
318 if let Some(discv5) = self.inner.discv5.as_ref() {
320 let discv5_enr = discv5.local_enr();
321 if let Some(ip6) = discv5_enr.ip6() {
322 builder.ip6(ip6);
323 }
324 if let Some(udp6) = discv5_enr.udp6() {
325 builder.udp6(udp6);
326 }
327 if let Some(tcp6) = discv5_enr.tcp6() {
328 builder.tcp6(tcp6);
329 }
330 }
331 } else {
332 builder.udp6(local_node_record.udp_port);
333 builder.tcp6(local_node_record.tcp_port);
334 }
335
336 builder.build(&self.inner.secret_key).expect("valid enr")
337 }
338}
339
340impl<N: NetworkPrimitives> Peers for NetworkHandle<N> {
341 fn add_trusted_peer_id(&self, peer: PeerId) {
342 self.send_message(NetworkHandleMessage::AddTrustedPeerId(peer));
343 }
344
345 fn add_trusted_peer_node(&self, peer: TrustedPeer) {
346 self.send_message(NetworkHandleMessage::AddTrustedPeerNode(peer));
347 }
348
349 fn add_peer_kind(
352 &self,
353 peer: PeerId,
354 kind: Option<PeerKind>,
355 tcp_addr: SocketAddr,
356 udp_addr: Option<SocketAddr>,
357 ) {
358 let addr = PeerAddr::new(tcp_addr, udp_addr);
359 self.send_message(NetworkHandleMessage::AddPeerAddress(peer, kind, addr));
360 }
361
362 async fn get_peers_by_kind(&self, kind: PeerKind) -> Result<Vec<PeerInfo>, NetworkError> {
363 let (tx, rx) = oneshot::channel();
364 let _ = self.manager().send(NetworkHandleMessage::GetPeerInfosByPeerKind(kind, tx));
365 Ok(rx.await?)
366 }
367
368 async fn get_all_peers(&self) -> Result<Vec<PeerInfo>, NetworkError> {
369 let (tx, rx) = oneshot::channel();
370 let _ = self.manager().send(NetworkHandleMessage::GetPeerInfos(tx));
371 Ok(rx.await?)
372 }
373
374 async fn get_peer_by_id(&self, peer_id: PeerId) -> Result<Option<PeerInfo>, NetworkError> {
375 let (tx, rx) = oneshot::channel();
376 let _ = self.manager().send(NetworkHandleMessage::GetPeerInfoById(peer_id, tx));
377 Ok(rx.await?)
378 }
379
380 async fn get_peers_by_id(&self, peer_ids: Vec<PeerId>) -> Result<Vec<PeerInfo>, NetworkError> {
381 let (tx, rx) = oneshot::channel();
382 let _ = self.manager().send(NetworkHandleMessage::GetPeerInfosByIds(peer_ids, tx));
383 Ok(rx.await?)
384 }
385
386 fn remove_peer(&self, peer: PeerId, kind: PeerKind) {
389 self.send_message(NetworkHandleMessage::RemovePeer(peer, kind))
390 }
391
392 fn disconnect_peer(&self, peer: PeerId) {
395 self.send_message(NetworkHandleMessage::DisconnectPeer(peer, None))
396 }
397
398 fn disconnect_peer_with_reason(&self, peer: PeerId, reason: DisconnectReason) {
401 self.send_message(NetworkHandleMessage::DisconnectPeer(peer, Some(reason)))
402 }
403
404 fn ban_peer(&self, peer: PeerId) {
407 self.send_message(NetworkHandleMessage::BanPeer(peer))
408 }
409
410 fn unban_peer(&self, peer: PeerId) {
412 self.send_message(NetworkHandleMessage::UnbanPeer(peer))
413 }
414
415 fn connect_peer_kind(
421 &self,
422 peer_id: PeerId,
423 kind: PeerKind,
424 tcp_addr: SocketAddr,
425 udp_addr: Option<SocketAddr>,
426 ) {
427 self.send_message(NetworkHandleMessage::ConnectPeer(
428 peer_id,
429 kind,
430 PeerAddr::new(tcp_addr, udp_addr),
431 ))
432 }
433
434 fn reputation_change(&self, peer_id: PeerId, kind: ReputationChangeKind) {
436 self.send_message(NetworkHandleMessage::ReputationChange(peer_id, kind));
437 }
438
439 async fn reputation_by_id(&self, peer_id: PeerId) -> Result<Option<Reputation>, NetworkError> {
440 let (tx, rx) = oneshot::channel();
441 let _ = self.manager().send(NetworkHandleMessage::GetReputationById(peer_id, tx));
442 Ok(rx.await?)
443 }
444}
445
446impl<N: NetworkPrimitives> PeersHandleProvider for NetworkHandle<N> {
447 fn peers_handle(&self) -> &PeersHandle {
448 &self.inner.peers
449 }
450}
451
452impl<N: NetworkPrimitives> NetworkInfo for NetworkHandle<N> {
453 fn local_addr(&self) -> SocketAddr {
454 *self.inner.listener_address.lock()
455 }
456
457 async fn network_status(&self) -> Result<NetworkStatus, NetworkError> {
458 let (tx, rx) = oneshot::channel();
459 let _ = self.manager().send(NetworkHandleMessage::GetStatus(tx));
460 rx.await.map_err(Into::into)
461 }
462
463 fn chain_id(&self) -> u64 {
464 self.inner.chain_id.load(Ordering::Relaxed)
465 }
466
467 fn cell_custody(&self) -> &CellCustody {
468 &self.inner.cell_custody
469 }
470
471 fn is_syncing(&self) -> bool {
472 SyncStateProvider::is_syncing(self)
473 }
474
475 fn is_initially_syncing(&self) -> bool {
476 SyncStateProvider::is_initially_syncing(self)
477 }
478}
479
480impl<N: NetworkPrimitives> SyncStateProvider for NetworkHandle<N> {
481 fn is_syncing(&self) -> bool {
482 self.inner.is_syncing.load(Ordering::Relaxed)
483 }
484 fn is_initially_syncing(&self) -> bool {
486 if self.inner.initial_sync_done.load(Ordering::Relaxed) {
487 return false
488 }
489 self.inner.is_syncing.load(Ordering::Relaxed)
490 }
491}
492
493impl<N: NetworkPrimitives> NetworkSyncUpdater for NetworkHandle<N> {
494 fn update_sync_state(&self, state: SyncState) {
495 let future_state = state.is_syncing();
496 let prev_state = self.inner.is_syncing.swap(future_state, Ordering::Relaxed);
497 let syncing_to_idle_state_transition = prev_state && !future_state;
498 if syncing_to_idle_state_transition {
499 self.inner.initial_sync_done.store(true, Ordering::Relaxed);
500 }
501 }
502
503 fn update_status(&self, head: Head) {
505 self.send_message(NetworkHandleMessage::StatusUpdate { head });
506 }
507
508 fn update_block_range(&self, update: reth_eth_wire::BlockRangeUpdate) {
510 self.send_message(NetworkHandleMessage::InternalBlockRangeUpdate(update));
511 }
512
513 fn set_fork_filter(&self, fork_filter: ForkFilter) {
515 self.send_message(NetworkHandleMessage::SetForkFilter { fork_filter });
516 }
517}
518
519impl<N: NetworkPrimitives> BlockDownloaderProvider for NetworkHandle<N> {
520 type Client = FetchClient<N>;
521
522 async fn fetch_client(&self) -> Result<Self::Client, oneshot::error::RecvError> {
523 let (tx, rx) = oneshot::channel();
524 let _ = self.manager().send(NetworkHandleMessage::FetchClient(tx));
525 rx.await
526 }
527}
528
529#[derive(Debug)]
530struct NetworkInner<N: NetworkPrimitives = EthNetworkPrimitives> {
531 num_active_peers: Arc<AtomicUsize>,
533 to_manager_tx: UnboundedSender<NetworkHandleMessage<N>>,
535 listener_address: Arc<Mutex<SocketAddr>>,
537 secret_key: SecretKey,
539 local_peer_id: PeerId,
541 peers: PeersHandle,
543 network_mode: NetworkMode,
545 is_syncing: Arc<AtomicBool>,
547 initial_sync_done: Arc<AtomicBool>,
549 chain_id: Arc<AtomicU64>,
551 cell_custody: CellCustody,
553 tx_gossip_disabled: bool,
555 discv4: Option<Discv4>,
557 discv5: Option<Discv5>,
559 event_sender: EventSender<NetworkEvent<PeerRequest<N>>>,
561 nat: Option<NatResolver>,
563}
564
565pub trait NetworkProtocols: Send + Sync {
567 fn add_rlpx_sub_protocol(&self, protocol: RlpxSubProtocol);
569}
570
571#[derive(Debug)]
573pub(crate) enum NetworkHandleMessage<N: NetworkPrimitives = EthNetworkPrimitives> {
574 AddTrustedPeerId(PeerId),
576 AddTrustedPeerNode(TrustedPeer),
578 AddPeerAddress(PeerId, Option<PeerKind>, PeerAddr),
580 RemovePeer(PeerId, PeerKind),
582 DisconnectPeer(PeerId, Option<DisconnectReason>),
584 BanPeer(PeerId),
586 UnbanPeer(PeerId),
588 AnnounceBlock(N::NewBlockPayload, B256),
590 SendTransaction {
592 peer_id: PeerId,
594 msg: SharedTransactions<N::BroadcastedTransaction>,
596 },
597 SendBroadcastPoolTransactions {
599 peer_id: PeerId,
601 msg: BroadcastPoolTransactions,
603 },
604 SendPooledTransactionHashes {
606 peer_id: PeerId,
608 msg: NewPooledTransactionHashes,
610 },
611 EthRequest {
613 peer_id: PeerId,
615 request: PeerRequest<N>,
617 },
618 EthMessage {
620 peer_id: PeerId,
622 message: PeerMessage<N>,
624 },
625 ReputationChange(PeerId, ReputationChangeKind),
627 FetchClient(oneshot::Sender<FetchClient<N>>),
629 StatusUpdate {
631 head: Head,
633 },
634 SetForkFilter {
636 fork_filter: ForkFilter,
638 },
639 GetStatus(oneshot::Sender<NetworkStatus>),
641 GetPeerInfosByIds(Vec<PeerId>, oneshot::Sender<Vec<PeerInfo>>),
643 GetPeerInfos(oneshot::Sender<Vec<PeerInfo>>),
645 GetPeerInfoById(PeerId, oneshot::Sender<Option<PeerInfo>>),
647 GetPeerInfosByPeerKind(PeerKind, oneshot::Sender<Vec<PeerInfo>>),
649 GetReputationById(PeerId, oneshot::Sender<Option<Reputation>>),
651 GetTransactionsHandle(oneshot::Sender<Option<TransactionsHandle<N>>>),
653 Shutdown(oneshot::Sender<()>),
655 SetNetworkState(NetworkConnectionState),
657 DiscoveryListener(UnboundedSender<DiscoveryEvent>),
659 AddRlpxSubProtocol(RlpxSubProtocol),
661 ConnectPeer(PeerId, PeerKind, PeerAddr),
663 InternalBlockRangeUpdate(BlockRangeUpdate),
665}