From 44bb722861e23ae5011f6f52fdc11011069da19c Mon Sep 17 00:00:00 2001 From: Andronik Ordian Date: Fri, 9 Oct 2020 16:12:57 +0200 Subject: [PATCH] =?UTF-8?q?sc-network:=20expose=20`add=5Fto=5Fpriority=5Fg?= =?UTF-8?q?roup`=20and=20`remove=5Ffrom=5Fpriority=5F=E2=80=A6=20(#7247)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * sc-network: expose `add_to_priority_group` and `remove_from_priority_group` in `NetworkService` * sc-network: fix a typo * Update client/network/src/service.rs Co-authored-by: Max Inden * s/parse_multiaddr/split_multiaddr_and_peer_id/g * sc-network: mark new functions as async and add comments * Apply suggestions from code review Co-authored-by: Max Inden Co-authored-by: Max Inden --- substrate/client/network/src/service.rs | 83 ++++++++++++++++++++----- 1 file changed, 66 insertions(+), 17 deletions(-) diff --git a/substrate/client/network/src/service.rs b/substrate/client/network/src/service.rs index 5b675186e9..af123e94fa 100644 --- a/substrate/client/network/src/service.rs +++ b/substrate/client/network/src/service.rs @@ -973,23 +973,7 @@ impl NetworkService { /// Returns an `Err` if one of the given addresses is invalid or contains an /// invalid peer ID (which includes the local peer ID). pub fn set_priority_group(&self, group_id: String, peers: HashSet) -> Result<(), String> { - let peers = peers.into_iter() - .map(|mut addr| { - let peer = match addr.pop() { - Some(multiaddr::Protocol::P2p(key)) => PeerId::from_multihash(key) - .map_err(|_| "Invalid PeerId format".to_string())?, - _ => return Err("Missing PeerId from address".to_string()), - }; - - // 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::, String>>()?; + let peers = self.split_multiaddr_and_peer_id(peers)?; let peer_ids = peers.iter().map(|(peer_id, _addr)| peer_id.clone()).collect(); @@ -1004,6 +988,47 @@ impl NetworkService { Ok(()) } + /// Add peers to a peerset priority group. + /// + /// Each `Multiaddr` must end with a `/p2p/` component containing the `PeerId`. + /// + /// Returns an `Err` if one of the given addresses is invalid or contains an + /// invalid peer ID (which includes the local peer ID). + // + // NOTE: even though this function is currently sync, it's marked as async for + // future-proofing, see https://github.com/paritytech/substrate/pull/7247#discussion_r502263451. + pub async fn add_to_priority_group(&self, group_id: String, peers: HashSet) -> Result<(), String> { + let peers = self.split_multiaddr_and_peer_id(peers)?; + + for (peer_id, addr) in peers.into_iter() { + self.peerset.add_to_priority_group(group_id.clone(), peer_id.clone()); + + let _ = self + .to_worker + .unbounded_send(ServiceToWorkerMsg::AddKnownAddress(peer_id, addr)); + } + + Ok(()) + } + + /// Remove peers from a peerset priority group. + /// + /// Each `Multiaddr` must end with a `/p2p/` component containing the `PeerId`. + /// + /// Returns an `Err` if one of the given addresses is invalid or contains an + /// invalid peer ID (which includes the local peer ID). + // + // NOTE: even though this function is currently sync, it's marked as async for + // future-proofing, see https://github.com/paritytech/substrate/pull/7247#discussion_r502263451. + // NOTE: technically, this function only needs `Vec`, but we use `Multiaddr` here for convenience. + pub async fn remove_from_priority_group(&self, group_id: String, peers: HashSet) -> Result<(), String> { + let peers = self.split_multiaddr_and_peer_id(peers)?; + for (peer_id, _) in peers.into_iter() { + self.peerset.remove_from_priority_group(group_id.clone(), peer_id); + } + Ok(()) + } + /// Returns the number of peers we're connected to. pub fn num_connected(&self) -> usize { self.num_connected.load(Ordering::Relaxed) @@ -1025,6 +1050,30 @@ impl NetworkService { .to_worker .unbounded_send(ServiceToWorkerMsg::OwnBlockImported(hash, number)); } + + /// Utility function to extract `PeerId` from each `Multiaddr` for priority group updates. + /// + /// Returns an `Err` if one of the given addresses is invalid or contains an + /// invalid peer ID (which includes the local peer ID). + fn split_multiaddr_and_peer_id(&self, peers: HashSet) -> Result, String> { + peers.into_iter() + .map(|mut addr| { + let peer = match addr.pop() { + Some(multiaddr::Protocol::P2p(key)) => PeerId::from_multihash(key) + .map_err(|_| "Invalid PeerId format".to_string())?, + _ => return Err("Missing PeerId from address".to_string()), + }; + + // 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::, String>>() + } } impl sp_consensus::SyncOracle