network/discovery: Add to DHT only peers that support genesis-based protocol

substrate/client/network/src/discovery.rs

@@ -105,7 +105,8 @@ pub struct DiscoveryConfig {
 	discovery_only_if_under_num: u64,
 	enable_mdns: bool,
 	kademlia_disjoint_query_paths: bool,
-	kademlia_protocols: Vec<Vec<u8>>,
+	kademlia_protocol: Vec<u8>,
+	kademlia_legacy_protocol: Vec<u8>,
 	kademlia_replication_factor: NonZeroUsize,
 }
 
@@ -121,7 +122,8 @@ impl DiscoveryConfig {
 			discovery_only_if_under_num: std::u64::MAX,
 			enable_mdns: false,
 			kademlia_disjoint_query_paths: false,
-			kademlia_protocols: Vec::new(),
+			kademlia_protocol: Vec::new(),
+			kademlia_legacy_protocol: Vec::new(),
 			kademlia_replication_factor: NonZeroUsize::new(DEFAULT_KADEMLIA_REPLICATION_FACTOR)
 				.expect("value is a constant; constant is non-zero; qed."),
 		}
@@ -177,9 +179,8 @@ impl DiscoveryConfig {
 		fork_id: Option<&str>,
 		protocol_id: &ProtocolId,
 	) -> &mut Self {
-		self.kademlia_protocols = Vec::new();
-		self.kademlia_protocols.push(kademlia_protocol_name(genesis_hash, fork_id));
-		self.kademlia_protocols.push(legacy_kademlia_protocol_name(protocol_id));
+		self.kademlia_protocol = kademlia_protocol_name(genesis_hash, fork_id);
+		self.kademlia_legacy_protocol = legacy_kademlia_protocol_name(protocol_id);
 		self
 	}
 
@@ -207,14 +208,19 @@ impl DiscoveryConfig {
 			discovery_only_if_under_num,
 			enable_mdns,
 			kademlia_disjoint_query_paths,
-			kademlia_protocols,
+			kademlia_protocol,
+			kademlia_legacy_protocol,
 			kademlia_replication_factor,
 		} = self;
 
-		let kademlia = if !kademlia_protocols.is_empty() {
+		let kademlia = if !kademlia_protocol.is_empty() {
 			let mut config = KademliaConfig::default();
 
 			config.set_replication_factor(kademlia_replication_factor);
+			// Populate kad with both the legacy and the new protocol names.
+			// Remove the legacy protocol:
+			// https://github.com/paritytech/polkadot-sdk/issues/504
+			let kademlia_protocols = [kademlia_protocol.clone(), kademlia_legacy_protocol];
 			config.set_protocol_names(kademlia_protocols.into_iter().map(Into::into).collect());
 			// By default Kademlia attempts to insert all peers into its routing table once a
 			// dialing attempt succeeds. In order to control which peer is added, disable the
@@ -266,6 +272,7 @@ impl DiscoveryConfig {
 					.expect("value is a constant; constant is non-zero; qed."),
 			),
 			records_to_publish: Default::default(),
+			kademlia_protocol,
 		}
 	}
 }
@@ -309,6 +316,11 @@ pub struct DiscoveryBehaviour {
 	/// did not return the record(in `FinishedWithNoAdditionalRecord`). We will then put the record
 	/// to these peers.
 	records_to_publish: HashMap<QueryId, Record>,
+	/// The chain based kademlia protocol name (including genesis hash and fork id).
+	///
+	/// Remove when all nodes are upgraded to genesis hash and fork ID-based Kademlia:
+	/// https://github.com/paritytech/polkadot-sdk/issues/504.
+	kademlia_protocol: Vec<u8>,
 }
 
 impl DiscoveryBehaviour {
@@ -366,23 +378,29 @@ impl DiscoveryBehaviour {
 				return
 			}
 
-			if let Some(matching_protocol) = supported_protocols
+			// The supported protocols must include the chain-based Kademlia protocol.
+			//
+			// Extract the chain-based Kademlia protocol from `kademlia.protocol_name()`
+			// when all nodes are upgraded to genesis hash and fork ID-based Kademlia:
+			// https://github.com/paritytech/polkadot-sdk/issues/504.
+			if !supported_protocols
 				.iter()
-				.find(|p| kademlia.protocol_names().iter().any(|k| k.as_ref() == p.as_ref()))
+				.any(|p| p.as_ref() == self.kademlia_protocol.as_slice())
 			{
-				trace!(
-					target: "sub-libp2p",
-					"Adding self-reported address {} from {} to Kademlia DHT {}.",
-					addr, peer_id, String::from_utf8_lossy(matching_protocol.as_ref()),
-				);
-				kademlia.add_address(peer_id, addr.clone());
-			} else {
 				trace!(
 					target: "sub-libp2p",
 					"Ignoring self-reported address {} from {} as remote node is not part of the \
 					 Kademlia DHT supported by the local node.", addr, peer_id,
 				);
+				return
 			}
+
+			trace!(
+				target: "sub-libp2p",
+				"Adding self-reported address {} from {} to Kademlia DHT.",
+				addr, peer_id
+			);
+			kademlia.add_address(peer_id, addr.clone());
 		}
 	}
 
@@ -1075,17 +1093,20 @@ mod tests {
 												.unwrap();
 											// Test both genesis hash-based and legacy
 											// protocol names.
-											let protocol_name = if swarm_n % 2 == 0 {
-												kademlia_protocol_name(genesis_hash, fork_id)
+											let protocol_names = if swarm_n % 2 == 0 {
+												vec![kademlia_protocol_name(genesis_hash, fork_id)]
 											} else {
-												legacy_kademlia_protocol_name(&protocol_id)
+												vec![
+													legacy_kademlia_protocol_name(&protocol_id),
+													kademlia_protocol_name(genesis_hash, fork_id),
+												]
 											};
 											swarms[swarm_n]
 												.0
 												.behaviour_mut()
 												.add_self_reported_address(
 													&other,
-													&[protocol_name],
+													protocol_names.as_slice(),
 													addr,
 												);
 
@@ -1181,12 +1202,33 @@ mod tests {
 			&[kademlia_protocol_name(supported_genesis_hash, None)],
 			remote_addr.clone(),
 		);
+		// Note: legacy protocol is not supported without genesis hash and fork ID,
+		// if the legacy is the only protocol supported, then the peer will not be added.
 		discovery.add_self_reported_address(
 			&another_peer_id,
 			&[legacy_kademlia_protocol_name(&supported_protocol_id)],
 			another_addr.clone(),
 		);
 
+		{
+			let kademlia = discovery.kademlia.as_mut().unwrap();
+			assert_eq!(
+				1,
+				kademlia.kbuckets().fold(0, |acc, bucket| acc + bucket.num_entries()),
+				"Expect peers with supported protocol to be added."
+			);
+		}
+
+		// Supported legacy and genesis based protocols are allowed to be added.
+		discovery.add_self_reported_address(
+			&another_peer_id,
+			&[
+				legacy_kademlia_protocol_name(&supported_protocol_id),
+				kademlia_protocol_name(supported_genesis_hash, None),
+			],
+			another_addr.clone(),
+		);
+
 		{
 			let kademlia = discovery.kademlia.as_mut().unwrap();
 			assert_eq!(