peerset: fix reserved nodes (#3706)

* peerset: fix handling of reserved only peering mode

* core: add cli parameter to enable reserved nodes only

* peerset: fix tests

* peerset: add test for priority only mode

* core: fix reserved only cli flag description

* peerset: extend docs on set_priority_only

substrate/core/cli/src/lib.rs

@@ -550,7 +550,8 @@ fn fill_network_configuration(
 	);
 	config.net_config_path = config.config_path.clone();
 	config.reserved_nodes.extend(cli.reserved_nodes.into_iter());
-	if !config.reserved_nodes.is_empty() {
+
+	if cli.reserved_only {
 		config.non_reserved_mode = NonReservedPeerMode::Deny;
 	}
 

substrate/core/cli/src/params.rs

@@ -92,6 +92,13 @@ pub struct NetworkConfigurationParams {
 	#[structopt(long = "reserved-nodes", value_name = "URL")]
 	pub reserved_nodes: Vec<String>,
 
+	/// Whether to only allow connections to/from reserved nodes.
+	///
+	/// If you are a validator your node might still connect to other validator
+	/// nodes regardless of whether they are defined as reserved nodes.
+	#[structopt(long = "reserved-only")]
+	pub reserved_only: bool,
+
 	/// Listen on this multiaddress.
 	#[structopt(long = "listen-addr", value_name = "LISTEN_ADDR")]
 	pub listen_addr: Vec<String>,

substrate/core/peerset/src/lib.rs

@@ -178,7 +178,7 @@ impl Peerset {
 		};
 
 		let mut peerset = Peerset {
-			data: peersstate::PeersState::new(config.in_peers, config.out_peers),
+			data: peersstate::PeersState::new(config.in_peers, config.out_peers, config.reserved_only),
 			tx,
 			rx,
 			reserved_only: config.reserved_only,
@@ -224,9 +224,11 @@ impl Peerset {
 	}
 
 	fn on_set_reserved_only(&mut self, reserved_only: bool) {
-		// Disconnect non-reserved nodes.
 		self.reserved_only = reserved_only;
+		self.data.set_priority_only(reserved_only);
+
 		if self.reserved_only {
+			// Disconnect non-reserved nodes.
 			let reserved = self.data.get_priority_group(RESERVED_NODES).unwrap_or_default();
 			for peer_id in self.data.connected_peers().cloned().collect::<Vec<_>>().into_iter() {
 				let peer = self.data.peer(&peer_id).into_connected()

substrate/core/peerset/src/peersstate.rs

@@ -50,6 +50,9 @@ pub struct PeersState {
 
 	/// Priority groups. Each group is identified by a string ID and contains a set of peer IDs.
 	priority_nodes: HashMap<String, HashSet<PeerId>>,
+
+	/// Only allow connections to/from peers in a priority group.
+	priority_only: bool,
 }
 
 /// State of a single node that we know about.
@@ -96,7 +99,7 @@ impl ConnectionState {
 
 impl PeersState {
 	/// Builds a new empty `PeersState`.
-	pub fn new(in_peers: u32, out_peers: u32) -> Self {
+	pub fn new(in_peers: u32, out_peers: u32, priority_only: bool) -> Self {
 		PeersState {
 			nodes: HashMap::new(),
 			num_in: 0,
@@ -104,6 +107,7 @@ impl PeersState {
 			max_in: in_peers,
 			max_out: out_peers,
 			priority_nodes: HashMap::new(),
+			priority_only,
 		}
 	}
 
@@ -220,9 +224,15 @@ impl PeersState {
 
 	/// Sets the peer as connected with an outgoing connection.
 	fn try_outgoing(&mut self, peer_id: &PeerId) -> bool {
+		let is_priority = self.is_priority(peer_id);
+
+		// We are only accepting connections from priority nodes.
+		if !is_priority && self.priority_only {
+			return false;
+		}
+
 		// Note that it is possible for num_out to be strictly superior to the max, in case we were
 		// connected to reserved node then marked them as not reserved.
-		let is_priority = self.is_priority(peer_id);
 		if self.num_out >= self.max_out && !is_priority {
 			return false;
 		}
@@ -245,6 +255,12 @@ impl PeersState {
 	/// Note that reserved nodes don't count towards the number of slots.
 	fn try_accept_incoming(&mut self, peer_id: &PeerId) -> bool {
 		let is_priority = self.is_priority(peer_id);
+
+		// We are only accepting connections from priority nodes.
+		if !is_priority && self.priority_only {
+			return false;
+		}
+
 		// Note that it is possible for num_in to be strictly superior to the max, in case we were
 		// connected to reserved node then marked them as not reserved.
 		if self.num_in >= self.max_in && !is_priority {
@@ -315,6 +331,15 @@ impl PeersState {
 		self.priority_nodes.get(group_id).cloned()
 	}
 
+	/// Set whether to only allow connections to/from peers in a priority group.
+	/// Calling this method does not affect any existing connection, e.g.
+	/// enabling priority only will not disconnect from any non-priority peers
+	/// we are already connected to, only future incoming/outgoing connection
+	/// attempts will be affected.
+	pub fn set_priority_only(&mut self, priority: bool) {
+		self.priority_only = priority;
+	}
+
 	/// Check that node is any priority group.
 	fn is_priority(&self, peer_id: &PeerId) -> bool {
 		self.priority_nodes.iter().any(|(_, group)| group.contains(peer_id))
@@ -527,7 +552,7 @@ mod tests {
 
 	#[test]
 	fn full_slots_in() {
-		let mut peers_state = PeersState::new(1, 1);
+		let mut peers_state = PeersState::new(1, 1, false);
 		let id1 = PeerId::random();
 		let id2 = PeerId::random();
 
@@ -542,7 +567,7 @@ mod tests {
 
 	#[test]
 	fn priority_node_doesnt_use_slot() {
-		let mut peers_state = PeersState::new(1, 1);
+		let mut peers_state = PeersState::new(1, 1, false);
 		let id1 = PeerId::random();
 		let id2 = PeerId::random();
 
@@ -558,7 +583,7 @@ mod tests {
 
 	#[test]
 	fn disconnecting_frees_slot() {
-		let mut peers_state = PeersState::new(1, 1);
+		let mut peers_state = PeersState::new(1, 1, false);
 		let id1 = PeerId::random();
 		let id2 = PeerId::random();
 
@@ -570,7 +595,7 @@ mod tests {
 
 	#[test]
 	fn priority_not_connected_peer() {
-		let mut peers_state = PeersState::new(25, 25);
+		let mut peers_state = PeersState::new(25, 25, false);
 		let id1 = PeerId::random();
 		let id2 = PeerId::random();
 
@@ -589,7 +614,7 @@ mod tests {
 
 	#[test]
 	fn highest_not_connected_peer() {
-		let mut peers_state = PeersState::new(25, 25);
+		let mut peers_state = PeersState::new(25, 25, false);
 		let id1 = PeerId::random();
 		let id2 = PeerId::random();
 
@@ -610,7 +635,7 @@ mod tests {
 
 	#[test]
 	fn disconnect_priority_doesnt_panic() {
-		let mut peers_state = PeersState::new(1, 1);
+		let mut peers_state = PeersState::new(1, 1, false);
 		let id = PeerId::random();
 		peers_state.set_priority_group("test", vec![id.clone()].into_iter().collect());
 		let peer = peers_state.peer(&id).into_not_connected().unwrap().try_outgoing().unwrap();
@@ -619,7 +644,7 @@ mod tests {
 
 	#[test]
 	fn multiple_priority_groups_slot_count() {
-		let mut peers_state = PeersState::new(1, 1);
+		let mut peers_state = PeersState::new(1, 1, false);
 		let id = PeerId::random();
 
 		if let Peer::Unknown(p) = peers_state.peer(&id) {
@@ -636,4 +661,60 @@ mod tests {
 		peers_state.set_priority_group("test2", vec![].into_iter().collect());
 		assert_eq!(peers_state.num_in, 1);
 	}
+
+	#[test]
+	fn priority_only_mode_ignores_drops_unknown_nodes() {
+		// test whether connection to/from given peer is allowed
+		let test_connection = |peers_state: &mut PeersState, id| {
+			if let Peer::Unknown(p) = peers_state.peer(id) {
+				p.discover();
+			}
+
+			let incoming = if let Peer::NotConnected(p) = peers_state.peer(id) {
+				p.try_accept_incoming().is_ok()
+			} else {
+				panic!()
+			};
+
+			if incoming {
+				peers_state.peer(id).into_connected().map(|p| p.disconnect());
+			}
+
+			let outgoing = if let Peer::NotConnected(p) = peers_state.peer(id) {
+				p.try_outgoing().is_ok()
+			} else {
+				panic!()
+			};
+
+			if outgoing {
+				peers_state.peer(id).into_connected().map(|p| p.disconnect());
+			}
+
+			incoming || outgoing
+		};
+
+		let mut peers_state = PeersState::new(1, 1, true);
+		let id = PeerId::random();
+
+		// this is an unknown peer and our peer state is set to only allow
+		// priority peers so any connection attempt should be denied.
+		assert!(!test_connection(&mut peers_state, &id));
+
+		// disabling priority only mode should allow the connection to go
+		// through.
+		peers_state.set_priority_only(false);
+		assert!(test_connection(&mut peers_state, &id));
+
+		// re-enabling it we should again deny connections from the peer.
+		peers_state.set_priority_only(true);
+		assert!(!test_connection(&mut peers_state, &id));
+
+		// but if we add the peer to a priority group it should be accepted.
+		peers_state.set_priority_group("TEST_GROUP", vec![id.clone()].into_iter().collect());
+		assert!(test_connection(&mut peers_state, &id));
+
+		// and removing it will cause the connection to once again be denied.
+		peers_state.remove_from_priority_group("TEST_GROUP", &id);
+		assert!(!test_connection(&mut peers_state, &id));
+	}
 }