From 70cfeffe4ef63fc650a7039c0a754dd4ca9db100 Mon Sep 17 00:00:00 2001
From: Roman Borschel <romanb@users.noreply.github.com>
Date: Wed, 10 Jun 2020 18:50:37 +0200
Subject: [PATCH] Avoid self-lookups in Authority Discovery (#6317)

* Ensure authority discovery avoids self-lookups.

Thereby additionally guard the `NetworkService` against
adding the local peer to the PSM or registering a
"known address" for the local peer.

* Clarify comments.

* See if returning errors is ok.
---
 client/authority-discovery/src/lib.rs | 23 ++++++++++++++++++-----
 client/network/src/service.rs         | 27 ++++++++++++++++++++++++---
 2 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/client/authority-discovery/src/lib.rs b/client/authority-discovery/src/lib.rs
index bc76c14331403..de98e6a4a38ae 100644
--- a/client/authority-discovery/src/lib.rs
+++ b/client/authority-discovery/src/lib.rs
@@ -294,13 +294,26 @@ where
 			.authorities(&id)
 			.map_err(Error::CallingRuntime)?;
 
+		let local_keys = match &self.role {
+			Role::Authority(key_store) => {
+				key_store.read()
+					.sr25519_public_keys(key_types::AUTHORITY_DISCOVERY)
+					.into_iter()
+					.collect::<HashSet<_>>()
+			},
+			Role::Sentry => HashSet::new(),
+		};
+
 		for authority_id in authorities.iter() {
-			if let Some(metrics) = &self.metrics {
-				metrics.request.inc();
-			}
+			// Make sure we don't look up our own keys.
+			if !local_keys.contains(authority_id.as_ref()) {
+				if let Some(metrics) = &self.metrics {
+					metrics.request.inc();
+				}
 
-			self.network
-				.get_value(&hash_authority_id(authority_id.as_ref()));
+				self.network
+					.get_value(&hash_authority_id(authority_id.as_ref()));
+			}
 		}
 
 		Ok(())
diff --git a/client/network/src/service.rs b/client/network/src/service.rs
index fd58aa631d6b2..2297fe6a52f72 100644
--- a/client/network/src/service.rs
+++ b/client/network/src/service.rs
@@ -672,8 +672,15 @@ impl<B: BlockT + 'static, H: ExHashT> NetworkService<B, H> {
 
 	/// Adds a `PeerId` and its address as reserved. The string should encode the address
 	/// and peer ID of the remote node.
+	///
+	/// Returns an `Err` if the given string is not a valid multiaddress
+	/// or contains an invalid peer ID (which includes the local peer ID).
 	pub fn add_reserved_peer(&self, peer: String) -> Result<(), String> {
 		let (peer_id, addr) = parse_str_addr(&peer).map_err(|e| format!("{:?}", e))?;
+		// Make sure the local peer ID is never added to the PSM.
+		if peer_id == self.local_peer_id {
+			return Err("Local peer ID cannot be added as a reserved peer.".to_string())
+		}
 		self.peerset.add_reserved_peer(peer_id.clone());
 		let _ = self
 			.to_worker
@@ -694,12 +701,26 @@ impl<B: BlockT + 'static, H: ExHashT> NetworkService<B, H> {
 	}
 
 	/// Modify a peerset priority group.
+	///
+	/// Returns an `Err` if one of the given addresses contains an invalid
+	/// peer ID (which includes the local peer ID).
 	pub fn set_priority_group(&self, group_id: String, peers: HashSet<Multiaddr>) -> Result<(), String> {
-		let peers = peers.into_iter().map(|p| {
-			parse_addr(p).map_err(|e| format!("{:?}", e))
-		}).collect::<Result<Vec<(PeerId, Multiaddr)>, String>>()?;
+		let peers = peers.into_iter()
+			.map(|p| match parse_addr(p) {
+				Err(e) => Err(format!("{:?}", e)),
+				Ok((peer, addr)) =>
+					// Make sure the local peer ID is never added to the PSM
+					// or added as a "known address", even if given.
+					if peer == self.local_peer_id {
+						Err("Local peer ID in priority group.".to_string())
+					} else {
+						Ok((peer, addr))
+					}
+				})
+			.collect::<Result<Vec<(PeerId, Multiaddr)>, String>>()?;
 
 		let peer_ids = peers.iter().map(|(peer_id, _addr)| peer_id.clone()).collect();
+
 		self.peerset.set_priority_group(group_id, peer_ids);
 
 		for (peer_id, addr) in peers.into_iter() {