From f0d3a951b085601b0a06f1fe7733f3313be361e8 Mon Sep 17 00:00:00 2001
From: Max Inden <mail@max-inden.de>
Date: Mon, 24 Aug 2020 15:34:41 +0200
Subject: [PATCH 1/2] client/authority-discovery: Test addresses per authority
 limit

---
 .../authority-discovery/src/worker/tests.rs   | 66 +++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/client/authority-discovery/src/worker/tests.rs b/client/authority-discovery/src/worker/tests.rs
index 4b16b9040b8dc..2d45c93adcc22 100644
--- a/client/authority-discovery/src/worker/tests.rs
+++ b/client/authority-discovery/src/worker/tests.rs
@@ -617,6 +617,72 @@ fn never_add_own_address_to_priority_group() {
 	);
 }
 
+#[test]
+fn limit_number_of_addresses_added_to_cache_per_authority() {
+	let remote_key_store = KeyStore::new();
+	let remote_public = remote_key_store
+		.write()
+		.sr25519_generate_new(key_types::AUTHORITY_DISCOVERY, None)
+		.unwrap();
+
+	let dht_event = {
+		let addresses = (0..100).map(|_| {
+			let peer_id = PeerId::random();
+			let address: Multiaddr = "/ip6/2001:db8:0:0:0:0:0:1/tcp/30333".parse().unwrap();
+			address.with(multiaddr::Protocol::P2p(
+				peer_id.into(),
+			)).to_vec()
+		}).collect();
+
+		let mut serialized_addresses = vec![];
+		schema::AuthorityAddresses { addresses }
+		.encode(&mut serialized_addresses)
+			.map_err(Error::EncodingProto)
+			.unwrap();
+
+		let signature = remote_key_store.read()
+			.sign_with(
+				key_types::AUTHORITY_DISCOVERY,
+				&remote_public.clone().into(),
+				serialized_addresses.as_slice(),
+			)
+			.map_err(|_| Error::Signing)
+			.unwrap();
+
+		let mut signed_addresses = vec![];
+		schema::SignedAuthorityAddresses {
+			addresses: serialized_addresses.clone(),
+			signature,
+		}
+			.encode(&mut signed_addresses)
+			.map_err(Error::EncodingProto)
+			.unwrap();
+
+		let key = hash_authority_id(&remote_public.to_raw_vec());
+		let value = signed_addresses;
+		(key, value)
+	};
+
+	let (_dht_event_tx, dht_event_rx) = channel(1);
+
+	let (_to_worker, from_service) = mpsc::channel(0);
+	let mut worker = Worker::new(
+		from_service,
+		Arc::new(TestApi { authorities: vec![remote_public.into()] }),
+		Arc::new(TestNetwork::default()),
+		vec![],
+		dht_event_rx.boxed(),
+		Role::Sentry,
+		None,
+	);
+
+	worker.handle_dht_value_found_event(vec![dht_event]).unwrap();
+	assert_eq!(
+		10,
+		worker.addr_cache.get_addresses_by_authority_id(&remote_public.into()).unwrap().len(),
+	);
+}
+
 #[test]
 fn do_not_cache_addresses_without_peer_id() {
 	let remote_key_store = KeyStore::new();

From 113ae26b82a603c8f44d2e981a9378e3336d5e1c Mon Sep 17 00:00:00 2001
From: Max Inden <mail@max-inden.de>
Date: Mon, 24 Aug 2020 15:39:11 +0200
Subject: [PATCH 2/2] client/authority-discovery: Limit number of addresses per
 authority

---
 client/authority-discovery/src/worker.rs       | 4 ++++
 client/authority-discovery/src/worker/tests.rs | 2 +-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/client/authority-discovery/src/worker.rs b/client/authority-discovery/src/worker.rs
index 09cdedd93a19e..232e59d08dd78 100644
--- a/client/authority-discovery/src/worker.rs
+++ b/client/authority-discovery/src/worker.rs
@@ -67,6 +67,9 @@ const LIBP2P_KADEMLIA_BOOTSTRAP_TIME: Duration = Duration::from_secs(30);
 /// discovery module.
 const AUTHORITIES_PRIORITY_GROUP_NAME: &'static str = "authorities";
 
+/// Maximum number of addresses cached per authority. Additional addresses are discarded.
+const MAX_ADDRESSES_PER_AUTHORITY: usize = 10;
+
 /// Role an authority discovery module can run as.
 pub enum Role {
 	/// Actual authority as well as a reference to its key store.
@@ -496,6 +499,7 @@ where
 
 				false // `protocol` is not a [`Protocol::P2p`], let's keep looking.
 			}))
+			.take(MAX_ADDRESSES_PER_AUTHORITY)
 			.collect();
 
 		if !remote_addresses.is_empty() {
diff --git a/client/authority-discovery/src/worker/tests.rs b/client/authority-discovery/src/worker/tests.rs
index 2d45c93adcc22..baa6bd0fc7d62 100644
--- a/client/authority-discovery/src/worker/tests.rs
+++ b/client/authority-discovery/src/worker/tests.rs
@@ -678,7 +678,7 @@ fn limit_number_of_addresses_added_to_cache_per_authority() {
 
 	worker.handle_dht_value_found_event(vec![dht_event]).unwrap();
 	assert_eq!(
-		10,
+		MAX_ADDRESSES_PER_AUTHORITY,
 		worker.addr_cache.get_addresses_by_authority_id(&remote_public.into()).unwrap().len(),
 	);
 }