From 654d0a89f381bba2199a7b1d683fd5718c58ba55 Mon Sep 17 00:00:00 2001
From: Roland Sherwin <RolandSherwin@protonmail.com>
Date: Fri, 17 Jan 2025 16:56:36 +0530
Subject: [PATCH] fix(relay): push identify to all peers

---
 ant-networking/src/event/swarm.rs   |  1 -
 ant-networking/src/relay_manager.rs | 16 +++++++---------
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/ant-networking/src/event/swarm.rs b/ant-networking/src/event/swarm.rs
index e997e576d2..73f67547bc 100644
--- a/ant-networking/src/event/swarm.rs
+++ b/ant-networking/src/event/swarm.rs
@@ -69,7 +69,6 @@ impl SwarmDriver {
                             relay_manager.on_successful_reservation_by_client(
                                 &relay_peer_id,
                                 &mut self.swarm,
-                                &self.live_connected_peers,
                             );
                         }
                     } else {
diff --git a/ant-networking/src/relay_manager.rs b/ant-networking/src/relay_manager.rs
index 6c491c37e3..175d5d22c7 100644
--- a/ant-networking/src/relay_manager.rs
+++ b/ant-networking/src/relay_manager.rs
@@ -7,7 +7,6 @@
 // permissions and limitations relating to use of the SAFE Network Software.
 
 use crate::driver::{BadNodes, NodeBehaviour};
-use itertools::Itertools;
 use libp2p::{
     core::transport::ListenerId, multiaddr::Protocol, swarm::ConnectionId, Multiaddr, PeerId,
     StreamProtocol, Swarm,
@@ -19,7 +18,7 @@ use rand::Rng;
 use std::sync::atomic::AtomicU64;
 use std::{
     collections::{btree_map::Entry, BTreeMap, HashMap, HashSet, VecDeque},
-    time::{Instant, SystemTime},
+    time::SystemTime,
 };
 
 const MAX_CONCURRENT_RELAY_CONNECTIONS: usize = 4;
@@ -220,7 +219,6 @@ impl RelayManager {
         &mut self,
         peer_id: &PeerId,
         swarm: &mut Swarm<NodeBehaviour>,
-        live_connected_peers: &BTreeMap<ConnectionId, (PeerId, Multiaddr, Instant)>,
     ) {
         match self.waiting_for_reservation.remove(peer_id) {
             Some(addr) => {
@@ -235,12 +233,12 @@ impl RelayManager {
 
         if self.connected_relay_servers.len() == MAX_CONCURRENT_RELAY_CONNECTIONS {
             debug!("We have reached the maximum number of relay connections. Push new identify info to all connected peers");
-            swarm.behaviour_mut().identify.push(
-                live_connected_peers
-                    .values()
-                    .map(|(peer_id, _, _)| *peer_id)
-                    .unique(),
-            );
+            // send identify to all peers
+            let mut all_peers = HashSet::new();
+            for bucket in swarm.behaviour_mut().kademlia.kbuckets() {
+                all_peers.extend(bucket.iter().map(|entry| entry.node.key.preimage()));
+            }
+            swarm.behaviour_mut().identify.push(all_peers);
         }
     }