diff --git a/client/network/src/service.rs b/client/network/src/service.rs index 5b675186e9065..af123e94fac4b 100644 --- a/client/network/src/service.rs +++ b/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