From cae0c15d10235bf0ec0bd6f8b20814dc7b63dfd5 Mon Sep 17 00:00:00 2001
From: hozan23 <hozan23@karyontech.net>
Date: Tue, 16 Jul 2024 08:16:57 +0200
Subject: p2p: check for the endpoints before listen/connect to them

---
 p2p/src/discovery/lookup.rs | 27 +++++++++++++++++++--------
 p2p/src/discovery/mod.rs    |  1 -
 2 files changed, 19 insertions(+), 9 deletions(-)

(limited to 'p2p/src/discovery')

diff --git a/p2p/src/discovery/lookup.rs b/p2p/src/discovery/lookup.rs
index 47a1d09..ba15da3 100644
--- a/p2p/src/discovery/lookup.rs
+++ b/p2p/src/discovery/lookup.rs
@@ -41,6 +41,9 @@ pub struct LookupService {
     /// Resolved listen endpoint
     listen_endpoint: RwLock<Option<Endpoint>>,
 
+    /// Resolved discovery endpoint
+    discovery_endpoint: RwLock<Option<Endpoint>>,
+
     /// Holds the configuration for the P2P network.
     config: Arc<Config>,
 
@@ -52,7 +55,6 @@ impl LookupService {
     /// Creates a new lookup service
     pub fn new(
         key_pair: &KeyPair,
-        id: &PeerID,
         table: Arc<RoutingTable>,
         config: Arc<Config>,
         monitor: Arc<Monitor>,
@@ -78,13 +80,18 @@ impl LookupService {
             ex,
         );
 
+        let id = key_pair
+            .public()
+            .try_into()
+            .expect("Get PeerID from KeyPair");
         Self {
-            id: id.clone(),
+            id,
             table,
             listener,
             connector,
             outbound_slots,
             listen_endpoint: RwLock::new(None),
+            discovery_endpoint: RwLock::new(None),
             config,
             monitor,
         }
@@ -98,19 +105,23 @@ impl LookupService {
 
     /// Set the resolved listen endpoint.
     pub fn set_listen_endpoint(&self, resolved_endpoint: &Endpoint) -> Result<()> {
-        let resolved_endpoint = Endpoint::Tcp(
+        let discovery_endpoint = Endpoint::Tcp(
             resolved_endpoint.addr()?.clone(),
             self.config.discovery_port,
         );
-        *self.listen_endpoint.write() = Some(resolved_endpoint);
+        *self.listen_endpoint.write() = Some(resolved_endpoint.clone());
+        *self.discovery_endpoint.write() = Some(discovery_endpoint.clone());
         Ok(())
     }
 
-    /// Get the listening endpoint.
     pub fn listen_endpoint(&self) -> Option<Endpoint> {
         self.listen_endpoint.read().clone()
     }
 
+    pub fn discovery_endpoint(&self) -> Option<Endpoint> {
+        self.discovery_endpoint.read().clone()
+    }
+
     /// Shuts down the lookup service.
     pub async fn shutdown(&self) {
         self.connector.shutdown().await;
@@ -278,7 +289,7 @@ impl LookupService {
 
         trace!("Send Peer msg");
         if let Some(endpoint) = self.listen_endpoint() {
-            self.send_peer_msg(&conn, endpoint.clone()).await?;
+            self.send_peer_msg(&conn, endpoint).await?;
         }
 
         trace!("Send Shutdown msg");
@@ -289,8 +300,8 @@ impl LookupService {
 
     /// Start a listener.
     async fn start_listener(self: &Arc<Self>) -> Result<()> {
-        let endpoint: Endpoint = match self.listen_endpoint() {
-            Some(e) => e.clone(),
+        let endpoint: Endpoint = match self.discovery_endpoint() {
+            Some(e) => e,
             None => return Ok(()),
         };
 
diff --git a/p2p/src/discovery/mod.rs b/p2p/src/discovery/mod.rs
index a81a817..b4c20bf 100644
--- a/p2p/src/discovery/mod.rs
+++ b/p2p/src/discovery/mod.rs
@@ -80,7 +80,6 @@ impl Discovery {
 
         let lookup_service = Arc::new(LookupService::new(
             key_pair,
-            peer_id,
             table.clone(),
             config.clone(),
             monitor.clone(),
-- 
cgit v1.2.3