From 5fc77551dde64769f911fbc6c9d9008a380e4ecd Mon Sep 17 00:00:00 2001
From: Emilia Hane <elsaemiliaevahane@gmail.com>
Date: Tue, 3 Sep 2024 12:22:08 +0200
Subject: [PATCH] fix(discv5): return discv5 local node record (#10640)

---
 crates/net/network/src/discovery.rs | 5 +++++
 crates/net/network/src/manager.rs   | 2 ++
 crates/net/network/src/network.rs   | 7 +++++++
 3 files changed, 14 insertions(+)

diff --git a/crates/net/network/src/discovery.rs b/crates/net/network/src/discovery.rs
index feabeed87347..d366027d680b 100644
--- a/crates/net/network/src/discovery.rs
+++ b/crates/net/network/src/discovery.rs
@@ -196,6 +196,11 @@ impl Discovery {
         }
     }
 
+    /// Returns discv5 handle.
+    pub fn discv5(&self) -> Option<Discv5> {
+        self.discv5.clone()
+    }
+
     /// Add a node to the discv4 table.
     pub(crate) fn add_discv5_node(&self, enr: Enr<SecretKey>) -> Result<(), NetworkError> {
         if let Some(discv5) = &self.discv5 {
diff --git a/crates/net/network/src/manager.rs b/crates/net/network/src/manager.rs
index bd6c0d41c0fc..864b6791f67a 100644
--- a/crates/net/network/src/manager.rs
+++ b/crates/net/network/src/manager.rs
@@ -229,6 +229,7 @@ impl NetworkManager {
         // need to retrieve the addr here since provided port could be `0`
         let local_peer_id = discovery.local_id();
         let discv4 = discovery.discv4();
+        let discv5 = discovery.discv5();
 
         let num_active_peers = Arc::new(AtomicUsize::new(0));
 
@@ -266,6 +267,7 @@ impl NetworkManager {
             Arc::new(AtomicU64::new(chain_spec.chain.id())),
             tx_gossip_disabled,
             discv4,
+            discv5,
             event_sender.clone(),
         );
 
diff --git a/crates/net/network/src/network.rs b/crates/net/network/src/network.rs
index 8e734caaf8c3..0c20c636c3b9 100644
--- a/crates/net/network/src/network.rs
+++ b/crates/net/network/src/network.rs
@@ -9,6 +9,7 @@ use std::{
 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},
@@ -61,6 +62,7 @@ impl NetworkHandle {
         chain_id: Arc<AtomicU64>,
         tx_gossip_disabled: bool,
         discv4: Option<Discv4>,
+        discv5: Option<Discv5>,
         event_sender: EventSender<NetworkEvent>,
     ) -> Self {
         let inner = NetworkInner {
@@ -76,6 +78,7 @@ impl NetworkHandle {
             chain_id,
             tx_gossip_disabled,
             discv4,
+            discv5,
             event_sender,
         };
         Self { inner: Arc::new(inner) }
@@ -209,6 +212,8 @@ impl PeersInfo for NetworkHandle {
     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();
@@ -422,6 +427,8 @@ struct NetworkInner {
     tx_gossip_disabled: bool,
     /// The instance of the discv4 service
     discv4: Option<Discv4>,
+    /// The instance of the discv5 service
+    discv5: Option<Discv5>,
     /// Sender for high level network events.
     event_sender: EventSender<NetworkEvent>,
 }