From 6319b92e1eec044600c77224e30c84b8ae6fee39 Mon Sep 17 00:00:00 2001
From: Robert Habermeier <rphmeier@gmail.com>
Date: Thu, 7 Sep 2023 18:28:18 -0500
Subject: [PATCH] properly account for peers' changing authority IDs

---
 .../node/network/gossip-support/src/lib.rs    | 43 +++++++++++++++----
 1 file changed, 34 insertions(+), 9 deletions(-)

diff --git a/polkadot/node/network/gossip-support/src/lib.rs b/polkadot/node/network/gossip-support/src/lib.rs
index c4fd9724cba1..955982a34ba9 100644
--- a/polkadot/node/network/gossip-support/src/lib.rs
+++ b/polkadot/node/network/gossip-support/src/lib.rs
@@ -104,7 +104,7 @@ pub struct GossipSupport<AD> {
 	/// By `PeerId`.
 	///
 	/// Needed for efficient handling of disconnect events.
-	connected_authorities_by_peer_id: HashMap<PeerId, HashSet<AuthorityDiscoveryId>>,
+	connected_peers: HashMap<PeerId, HashSet<AuthorityDiscoveryId>>,
 	/// Authority discovery service.
 	authority_discovery: AD,
 
@@ -130,7 +130,7 @@ where
 			failure_start: None,
 			resolved_authorities: HashMap::new(),
 			connected_authorities: HashMap::new(),
-			connected_authorities_by_peer_id: HashMap::new(),
+			connected_peers: HashMap::new(),
 			authority_discovery,
 			metrics,
 		}
@@ -407,20 +407,43 @@ where
 			}
 		}
 
-		for (peer_id, auths) in authority_ids {
+
+		// peer was authority and now isn't
+		for (peer_id, current) in self.connected_peers.iter_mut() {
+			// empty -> nonempty is handled in the next loop
+			if !current.is_empty() && !authority_ids.contains_key(peer_id) {
+				sender
+					.send_message(NetworkBridgeRxMessage::UpdatedAuthorityIds {
+						peer_id: peer_id.clone(),
+						authority_ids: HashSet::new(),
+					})
+					.await;
+
+				for a in current.drain() {
+					self.connected_authorities.remove(&a);
+				}
+			}
+		}
+
+		// peer has new authority set.
+		for (peer_id, new) in authority_ids {
 			// If the peer is connected _and_ the authority IDs have changed.
-			if self.connected_authorities_by_peer_id.get(&peer_id).map_or(false, |x| x != &auths) {
+			if let Some(prev) = self.connected_peers.get(&peer_id).filter(|x| x != &&new) {
 				sender
 					.send_message(NetworkBridgeRxMessage::UpdatedAuthorityIds {
 						peer_id,
-						authority_ids: auths.clone(),
+						authority_ids: new.clone(),
 					})
 					.await;
 
-				auths.iter().for_each(|a| {
+				prev.iter().for_each(|a| {
+					self.connected_authorities.remove(a);
+				});
+				new.iter().for_each(|a| {
 					self.connected_authorities.insert(a.clone(), peer_id);
 				});
-				self.connected_authorities_by_peer_id.insert(peer_id, auths);
+
+				self.connected_peers.insert(peer_id, new);
 			}
 		}
 	}
@@ -432,11 +455,13 @@ where
 					authority_ids.iter().for_each(|a| {
 						self.connected_authorities.insert(a.clone(), peer_id);
 					});
-					self.connected_authorities_by_peer_id.insert(peer_id, authority_ids);
+					self.connected_peers.insert(peer_id, authority_ids);
+				} else {
+					self.connected_peers.insert(peer_id, HashSet::new());
 				}
 			},
 			NetworkBridgeEvent::PeerDisconnected(peer_id) => {
-				if let Some(authority_ids) = self.connected_authorities_by_peer_id.remove(&peer_id)
+				if let Some(authority_ids) = self.connected_peers.remove(&peer_id)
 				{
 					authority_ids.into_iter().for_each(|a| {
 						self.connected_authorities.remove(&a);