diff --git a/substrate/client/network/src/protocol.rs b/substrate/client/network/src/protocol.rs index 70a17dc44b..a0c52d14fa 100644 --- a/substrate/client/network/src/protocol.rs +++ b/substrate/client/network/src/protocol.rs @@ -1102,7 +1102,7 @@ impl Protocol { /// Removes a `PeerId` from the list of reserved peers for syncing purposes. pub fn remove_reserved_peer(&self, peer: PeerId) { - self.peerset_handle.remove_reserved_peer(HARDCODED_PEERSETS_SYNC, peer.clone()); + self.peerset_handle.remove_reserved_peer(HARDCODED_PEERSETS_SYNC, peer); } /// Returns the list of reserved peers. @@ -1112,12 +1112,26 @@ impl Protocol { /// Adds a `PeerId` to the list of reserved peers for syncing purposes. pub fn add_reserved_peer(&self, peer: PeerId) { - self.peerset_handle.add_reserved_peer(HARDCODED_PEERSETS_SYNC, peer.clone()); + self.peerset_handle.add_reserved_peer(HARDCODED_PEERSETS_SYNC, peer); } /// Sets the list of reserved peers for syncing purposes. pub fn set_reserved_peers(&self, peers: HashSet) { - self.peerset_handle.set_reserved_peers(HARDCODED_PEERSETS_SYNC, peers.clone()); + self.peerset_handle.set_reserved_peers(HARDCODED_PEERSETS_SYNC, peers); + } + + /// Sets the list of reserved peers for the given protocol/peerset. + pub fn set_reserved_peerset_peers(&self, protocol: Cow<'static, str>, peers: HashSet) { + if let Some(index) = self.notification_protocols.iter().position(|p| *p == protocol) { + self.peerset_handle + .set_reserved_peers(sc_peerset::SetId::from(index + NUM_HARDCODED_PEERSETS), peers); + } else { + error!( + target: "sub-libp2p", + "set_reserved_peerset_peers with unknown protocol: {}", + protocol + ); + } } /// Removes a `PeerId` from the list of reserved peers. diff --git a/substrate/client/network/src/service.rs b/substrate/client/network/src/service.rs index 23f9c614d9..72b382c990 100644 --- a/substrate/client/network/src/service.rs +++ b/substrate/client/network/src/service.rs @@ -1092,6 +1092,44 @@ impl NetworkService { let _ = self.to_worker.unbounded_send(ServiceToWorkerMsg::RemoveReserved(peer_id)); } + /// Sets the reserved set of a protocol to the given set of peers. + /// + /// Each `Multiaddr` must end with a `/p2p/` component containing the `PeerId`. It can also + /// consist of only `/p2p/`. + /// + /// 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_reserved_peers( + &self, + protocol: Cow<'static, str>, + peers: HashSet, + ) -> Result<(), String> { + let peers_addrs = self.split_multiaddr_and_peer_id(peers)?; + + let mut peers: HashSet = HashSet::with_capacity(peers_addrs.len()); + + for (peer_id, addr) in peers_addrs.into_iter() { + // 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()) + } + + peers.insert(peer_id); + + if !addr.is_empty() { + let _ = self + .to_worker + .unbounded_send(ServiceToWorkerMsg::AddKnownAddress(peer_id, addr)); + } + } + + let _ = self + .to_worker + .unbounded_send(ServiceToWorkerMsg::SetPeersetReserved(protocol, peers)); + + Ok(()) + } + /// Add peers to a peer set. /// /// Each `Multiaddr` must end with a `/p2p/` component containing the `PeerId`. It can also @@ -1400,6 +1438,7 @@ enum ServiceToWorkerMsg { AddReserved(PeerId), RemoveReserved(PeerId), SetReserved(HashSet), + SetPeersetReserved(Cow<'static, str>, HashSet), AddSetReserved(Cow<'static, str>, PeerId), RemoveSetReserved(Cow<'static, str>, PeerId), AddToPeersSet(Cow<'static, str>, PeerId), @@ -1541,6 +1580,11 @@ impl Future for NetworkWorker { .behaviour_mut() .user_protocol_mut() .set_reserved_peers(peers), + ServiceToWorkerMsg::SetPeersetReserved(protocol, peers) => this + .network_service + .behaviour_mut() + .user_protocol_mut() + .set_reserved_peerset_peers(protocol, peers), ServiceToWorkerMsg::AddReserved(peer_id) => this .network_service .behaviour_mut()