From baa38e9d6a729b76712adee1cd21b84f3fcbe1fd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Oliveira?= <hello@jxs.pt>
Date: Mon, 8 Jan 2024 11:00:59 +0000
Subject: [PATCH] remove peer from connected_peers if connection fails

---
 protocols/gossipsub/src/behaviour.rs | 30 ++++++++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/protocols/gossipsub/src/behaviour.rs b/protocols/gossipsub/src/behaviour.rs
index 8819a117bd3e..25c9dcf1fc39 100644
--- a/protocols/gossipsub/src/behaviour.rs
+++ b/protocols/gossipsub/src/behaviour.rs
@@ -41,8 +41,8 @@ use libp2p_identity::PeerId;
 use libp2p_swarm::{
     behaviour::{AddressChange, ConnectionClosed, ConnectionEstablished, FromSwarm},
     dial_opts::DialOpts,
-    ConnectionDenied, ConnectionId, NetworkBehaviour, NotifyHandler, THandler, THandlerInEvent,
-    THandlerOutEvent, ToSwarm,
+    ConnectionDenied, ConnectionId, DialFailure, ListenFailure, NetworkBehaviour, NotifyHandler,
+    THandler, THandlerInEvent, THandlerOutEvent, ToSwarm,
 };
 
 use crate::peer_score::{PeerScore, PeerScoreParams, PeerScoreThresholds, RejectReason};
@@ -3269,6 +3269,32 @@ where
             FromSwarm::ConnectionClosed(connection_closed) => {
                 self.on_connection_closed(connection_closed)
             }
+            FromSwarm::DialFailure(DialFailure {
+                peer_id,
+                connection_id,
+                ..
+            })
+            | FromSwarm::ListenFailure(ListenFailure {
+                connection_id,
+                peer_id,
+                ..
+            }) => {
+                if let Some(peer_id) = peer_id {
+                    let mut peer = self
+                        .connected_peers
+                        .remove(&peer_id)
+                        .expect("To be connected to peer");
+                    let index = peer
+                        .connections
+                        .iter()
+                        .position(|v| v == &connection_id)
+                        .expect("Previously established connection to peer must be present");
+                    peer.connections.remove(index);
+                    if !peer.connections.is_empty() {
+                        self.connected_peers.insert(peer_id, peer);
+                    }
+                }
+            }
             FromSwarm::AddressChange(address_change) => self.on_address_change(address_change),
             _ => {}
         }