client/authority-discovery: Append PeerId to Multiaddr at most once (#6933)

* client/authority-discovery/worker: Extract address getter

* client/authority-discovery: Test for no duplicate p2p components

* client/authority-discovery: Append PeerId to Multiaddr at most once

When collecting the addresses to be published for the local node,
`addresses_to_publish` adds the local nodes `PeerId` to each
`Multiaddr`. Before doing so, ensure the `Multiaddr` does not already
contain one.

* client/authority-discovery: Remove explicit return

substrate/Cargo.lock

@@ -6250,6 +6250,7 @@ version = "0.8.0-rc6"
 dependencies = [
  "bytes 0.5.6",
  "derive_more",
+ "either",
  "env_logger",
  "futures 0.3.5",
  "futures-timer 3.0.2",

substrate/client/authority-discovery/Cargo.toml

@@ -19,6 +19,7 @@ prost-build = "0.6.1"
 bytes = "0.5.0"
 codec = { package = "parity-scale-codec", default-features = false, version = "1.3.4" }
 derive_more = "0.99.2"
+either = "1.5.3"
 futures = "0.3.4"
 futures-timer = "3.0.1"
 libp2p = { version = "0.24.0", default-features = false, features = ["kad"] }

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

@@ -30,7 +30,8 @@ use futures_timer::Delay;
 
 use addr_cache::AddrCache;
 use codec::Decode;
-use libp2p::core::multiaddr;
+use either::Either;
+use libp2p::{core::multiaddr, multihash::Multihash};
 use log::{debug, error, log_enabled};
 use prometheus_endpoint::{Counter, CounterVec, Gauge, Opts, U64, register};
 use prost::Message;
@@ -232,6 +233,26 @@ where
 		}
 	}
 
+	fn addresses_to_publish(&self) -> impl ExactSizeIterator<Item = Multiaddr> {
+		match &self.sentry_nodes {
+			Some(addrs) => Either::Left(addrs.clone().into_iter()),
+			None => {
+				let peer_id: Multihash = self.network.local_peer_id().into();
+				Either::Right(
+					self.network.external_addresses()
+						.into_iter()
+						.map(move |a| {
+							if a.iter().any(|p| matches!(p, multiaddr::Protocol::P2p(_))) {
+								a
+							} else {
+								a.with(multiaddr::Protocol::P2p(peer_id.clone()))
+							}
+						}),
+				)
+			}
+		}
+	}
+
 	/// Publish either our own or if specified the public addresses of our sentry nodes.
 	fn publish_ext_addresses(&mut self) -> Result<()> {
 		let key_store = match &self.role {
@@ -242,29 +263,15 @@ where
 			Role::Sentry => return Ok(()),
 		};
 
-		if let Some(metrics) = &self.metrics {
-			metrics.publish.inc()
-		}
-
-		let addresses: Vec<_> = match &self.sentry_nodes {
-			Some(addrs) => addrs.clone().into_iter()
-				.map(|a| a.to_vec())
-				.collect(),
-			None => self.network.external_addresses()
-				.into_iter()
-				.map(|a| a.with(multiaddr::Protocol::P2p(
-					self.network.local_peer_id().into(),
-				)))
-				.map(|a| a.to_vec())
-				.collect(),
-		};
+		let addresses = self.addresses_to_publish();
 
 		if let Some(metrics) = &self.metrics {
+			metrics.publish.inc();
 			metrics.amount_last_published.set(addresses.len() as u64);
 		}
 
 		let mut serialized_addresses = vec![];
-		schema::AuthorityAddresses { addresses }
+		schema::AuthorityAddresses { addresses: addresses.map(|a| a.to_vec()).collect() }
 			.encode(&mut serialized_addresses)
 			.map_err(Error::EncodingProto)?;
 

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

@@ -168,6 +168,7 @@ sp_api::mock_impl_runtime_apis! {
 
 pub struct TestNetwork {
 	peer_id: PeerId,
+	external_addresses: Vec<Multiaddr>,
 	// Whenever functions on `TestNetwork` are called, the function arguments are added to the
 	// vectors below.
 	pub put_value_call: Arc<Mutex<Vec<(kad::record::Key, Vec<u8>)>>>,
@@ -179,6 +180,10 @@ impl Default for TestNetwork {
 	fn default() -> Self {
 		TestNetwork {
 			peer_id: PeerId::random(),
+			external_addresses: vec![
+				"/ip6/2001:db8::/tcp/30333"
+					.parse().unwrap(),
+			],
 			put_value_call: Default::default(),
 			get_value_call: Default::default(),
 			set_priority_group_call: Default::default(),
@@ -212,7 +217,7 @@ impl NetworkStateInfo for TestNetwork {
 	}
 
 	fn external_addresses(&self) -> Vec<Multiaddr> {
-		vec!["/ip6/2001:db8::/tcp/30333".parse().unwrap()]
+		self.external_addresses.clone()
 	}
 }
 
@@ -691,3 +696,67 @@ fn do_not_cache_addresses_without_peer_id() {
 		"Expect worker to only cache `Multiaddr`s with `PeerId`s.",
 	);
 }
+
+#[test]
+fn addresses_to_publish_adds_p2p() {
+	let (_dht_event_tx, dht_event_rx) = channel(1000);
+	let network: Arc<TestNetwork> = Arc::new(Default::default());
+
+	assert!(!matches!(
+		network.external_addresses().pop().unwrap().pop().unwrap(),
+		multiaddr::Protocol::P2p(_)
+	));
+
+	let (_to_worker, from_service) = mpsc::channel(0);
+	let worker = Worker::new(
+		from_service,
+		Arc::new(TestApi {
+			authorities: vec![],
+		}),
+		network.clone(),
+		vec![],
+		dht_event_rx.boxed(),
+		Role::Authority(KeyStore::new()),
+		Some(prometheus_endpoint::Registry::new()),
+	);
+
+	assert!(
+		matches!(
+			worker.addresses_to_publish().next().unwrap().pop().unwrap(),
+			multiaddr::Protocol::P2p(_)
+		),
+		"Expect `addresses_to_publish` to append `p2p` protocol component.",
+	);
+}
+
+/// Ensure [`Worker::addresses_to_publish`] does not add an additional `p2p` protocol component in
+/// case one already exists.
+#[test]
+fn addresses_to_publish_respects_existing_p2p_protocol() {
+	let (_dht_event_tx, dht_event_rx) = channel(1000);
+	let network: Arc<TestNetwork> = Arc::new(TestNetwork {
+		external_addresses: vec![
+			"/ip6/2001:db8::/tcp/30333/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC"
+				.parse().unwrap(),
+		],
+		.. Default::default()
+	});
+
+	let (_to_worker, from_service) = mpsc::channel(0);
+	let worker = Worker::new(
+		from_service,
+		Arc::new(TestApi {
+			authorities: vec![],
+		}),
+		network.clone(),
+		vec![],
+		dht_event_rx.boxed(),
+		Role::Authority(KeyStore::new()),
+		Some(prometheus_endpoint::Registry::new()),
+	);
+
+	assert_eq!(
+		network.external_addresses, worker.addresses_to_publish().collect::<Vec<_>>(),
+		"Expected Multiaddr from `TestNetwork` to not be altered.",
+	);
+}