Fix debug_assertion failing in authority discovery (#8599)

* Fix debug_assertion failing in authority discovery

* Improve test

* Change the map_or for invalid addresses

* Remove debug_assertion

substrate/client/authority-discovery/src/worker/addr_cache.rs

@@ -24,6 +24,9 @@ use sc_network::PeerId;
 
 /// Cache for [`AuthorityId`] -> [`Vec<Multiaddr>`] and [`PeerId`] -> [`AuthorityId`] mappings.
 pub(super) struct AddrCache {
+	// The addresses found in `authority_id_to_addresses` are guaranteed to always match
+	// the peerids found in `peer_id_to_authority_id`. In other words, these two hashmaps
+	// are similar to a bi-directional map.
 	authority_id_to_addresses: HashMap<AuthorityId, Vec<Multiaddr>>,
 	peer_id_to_authority_id: HashMap<PeerId, AuthorityId>,
 }
@@ -43,17 +46,44 @@ impl AddrCache {
 			return;
 		}
 
+		addresses.sort_unstable_by(|a, b| a.as_ref().cmp(b.as_ref()));
+
 		// Insert into `self.peer_id_to_authority_id`.
 		let peer_ids = addresses.iter()
 			.map(|a| peer_id_from_multiaddr(a))
 			.filter_map(|peer_id| peer_id);
-		for peer_id in  peer_ids {
-			self.peer_id_to_authority_id.insert(peer_id, authority_id.clone());
+		for peer_id in peer_ids.clone() {
+			let former_auth = match self.peer_id_to_authority_id.insert(peer_id, authority_id.clone()) {
+				Some(a) if a != authority_id => a,
+				_ => continue,
+			};
+
+			// PeerId was associated to a different authority id before.
+			// Remove corresponding authority from `self.authority_id_to_addresses`.
+			let former_auth_addrs = match self.authority_id_to_addresses.get_mut(&former_auth) {
+				Some(a) => a,
+				None => { debug_assert!(false); continue }
+			};
+			former_auth_addrs.retain(|a| peer_id_from_multiaddr(a).map_or(true, |p| p != peer_id));
 		}
 
 		// Insert into `self.authority_id_to_addresses`.
-		addresses.sort_unstable_by(|a, b| a.as_ref().cmp(b.as_ref()));
-		self.authority_id_to_addresses.insert(authority_id, addresses);
+		for former_addr in
+			self.authority_id_to_addresses.insert(authority_id.clone(), addresses.clone()).unwrap_or_default()
+		{
+			// Must remove from `self.peer_id_to_authority_id` any PeerId formerly associated
+			// to that authority but that can't be found in its new addresses.
+
+			let peer_id = match peer_id_from_multiaddr(&former_addr) {
+				Some(p) => p,
+				None => continue,
+			};
+
+			if !peer_ids.clone().any(|p| p == peer_id) {
+				let _old_auth = self.peer_id_to_authority_id.remove(&peer_id);
+				debug_assert!(_old_auth.is_some());
+			}
+		}
 	}
 
 	/// Returns the number of authority IDs in the cache.
@@ -144,6 +174,25 @@ mod tests {
 		}
 	}
 
+	#[derive(Clone, Debug)]
+	struct TestMultiaddrsSamePeerCombo(Multiaddr, Multiaddr);
+
+	impl Arbitrary for TestMultiaddrsSamePeerCombo {
+		fn arbitrary(g: &mut Gen) -> Self {
+			let seed = (0..32).map(|_| u8::arbitrary(g)).collect::<Vec<_>>();
+			let peer_id = PeerId::from_multihash(
+				Multihash::wrap(multihash::Code::Sha2_256.into(), &seed).unwrap()
+			).unwrap();
+			let multiaddr1 = "/ip6/2001:db8:0:0:0:0:0:2/tcp/30333".parse::<Multiaddr>()
+				.unwrap()
+				.with(Protocol::P2p(peer_id.clone().into()));
+			let multiaddr2 = "/ip6/2002:db8:0:0:0:0:0:2/tcp/30133".parse::<Multiaddr>()
+				.unwrap()
+				.with(Protocol::P2p(peer_id.into()));
+			TestMultiaddrsSamePeerCombo(multiaddr1, multiaddr2)
+		}
+	}
+
 	#[test]
 	fn retains_only_entries_of_provided_authority_ids() {
 		fn property(
@@ -190,4 +239,73 @@ mod tests {
 			.max_tests(10)
 			.quickcheck(property as fn(_, _, _) -> TestResult)
 	}
+
+	#[test]
+	fn keeps_consistency_between_authority_id_and_peer_id() {
+		fn property(
+			authority1: TestAuthorityId,
+			authority2: TestAuthorityId,
+			multiaddr1: TestMultiaddr,
+			multiaddr2: TestMultiaddr,
+			multiaddr3: TestMultiaddrsSamePeerCombo,
+		) -> TestResult {
+			let authority1 = authority1.0;
+			let authority2 = authority2.0;
+			let multiaddr1 = multiaddr1.0;
+			let multiaddr2 = multiaddr2.0;
+			let TestMultiaddrsSamePeerCombo(multiaddr3, multiaddr4) = multiaddr3;
+
+			let mut cache = AddrCache::new();
+
+			cache.insert(authority1.clone(), vec![multiaddr1.clone()]);
+			cache.insert(authority1.clone(), vec![multiaddr2.clone(), multiaddr3.clone(), multiaddr4.clone()]);
+
+			assert_eq!(
+				None,
+				cache.get_authority_id_by_peer_id(&peer_id_from_multiaddr(&multiaddr1).unwrap())
+			);
+			assert_eq!(
+				Some(&authority1),
+				cache.get_authority_id_by_peer_id(&peer_id_from_multiaddr(&multiaddr2).unwrap())
+			);
+			assert_eq!(
+				Some(&authority1),
+				cache.get_authority_id_by_peer_id(&peer_id_from_multiaddr(&multiaddr3).unwrap())
+			);
+			assert_eq!(
+				Some(&authority1),
+				cache.get_authority_id_by_peer_id(&peer_id_from_multiaddr(&multiaddr4).unwrap())
+			);
+
+			cache.insert(authority2.clone(), vec![multiaddr2.clone()]);
+
+			assert_eq!(
+				Some(&authority2),
+				cache.get_authority_id_by_peer_id(&peer_id_from_multiaddr(&multiaddr2).unwrap())
+			);
+			assert_eq!(
+				Some(&authority1),
+				cache.get_authority_id_by_peer_id(&peer_id_from_multiaddr(&multiaddr3).unwrap())
+			);
+			assert_eq!(cache.get_addresses_by_authority_id(&authority1).unwrap().len(), 2);
+
+			cache.insert(authority2.clone(), vec![multiaddr2.clone(), multiaddr3.clone()]);
+
+			assert_eq!(
+				Some(&authority2),
+				cache.get_authority_id_by_peer_id(&peer_id_from_multiaddr(&multiaddr2).unwrap())
+			);
+			assert_eq!(
+				Some(&authority2),
+				cache.get_authority_id_by_peer_id(&peer_id_from_multiaddr(&multiaddr3).unwrap())
+			);
+			assert!(cache.get_addresses_by_authority_id(&authority1).unwrap().is_empty());
+
+			TestResult::passed()
+		}
+
+		QuickCheck::new()
+			.max_tests(10)
+			.quickcheck(property as fn(_, _, _, _, _) -> TestResult)
+	}
 }