diff --git a/crates/papyrus_network/src/network_manager/swarm_trait.rs b/crates/papyrus_network/src/network_manager/swarm_trait.rs index d7599ca51f..270154f962 100644 --- a/crates/papyrus_network/src/network_manager/swarm_trait.rs +++ b/crates/papyrus_network/src/network_manager/swarm_trait.rs @@ -9,7 +9,7 @@ use crate::gossipsub_impl::Topic; use crate::mixed_behaviour; use crate::peer_manager::ReputationModifier; use crate::sqmr::behaviour::{PeerNotConnected, SessionIdNotFoundError}; -use crate::sqmr::{Bytes, InboundSessionId, OutboundSessionId}; +use crate::sqmr::{Bytes, InboundSessionId, OutboundSessionId, SessionId}; pub type Event = SwarmEvent<::ToSwarm>; @@ -38,6 +38,11 @@ pub trait SwarmTrait: Stream + Unpin { fn behaviour_mut(&mut self) -> &mut mixed_behaviour::MixedBehaviour; + fn get_peer_id_from_session_id( + &self, + session_id: SessionId, + ) -> Result; + fn add_external_address(&mut self, address: Multiaddr); fn subscribe_to_topic(&mut self, topic: &Topic) -> Result<(), SubscriptionError>; @@ -86,6 +91,16 @@ impl SwarmTrait for Swarm { self.behaviour_mut() } + fn get_peer_id_from_session_id( + &self, + session_id: SessionId, + ) -> Result { + self.behaviour() + .sqmr + .get_peer_id_and_connection_id_from_session_id(session_id) + .map(|(peer_id, _)| peer_id) + } + fn add_external_address(&mut self, address: Multiaddr) { self.add_external_address(address); } diff --git a/crates/papyrus_network/src/network_manager/test.rs b/crates/papyrus_network/src/network_manager/test.rs index b5cb5a4b9f..44ea54c4e0 100644 --- a/crates/papyrus_network/src/network_manager/test.rs +++ b/crates/papyrus_network/src/network_manager/test.rs @@ -194,6 +194,13 @@ impl SwarmTrait for MockSwarm { sender.unbounded_send(protocol_name.clone()).unwrap(); } } + + fn get_peer_id_from_session_id( + &self, + _session_id: crate::sqmr::SessionId, + ) -> Result { + Ok(PeerId::random()) + } } const BUFFER_SIZE: usize = 100; diff --git a/crates/papyrus_network/src/sqmr/behaviour.rs b/crates/papyrus_network/src/sqmr/behaviour.rs index cf80e0d0a3..8997ea09fb 100644 --- a/crates/papyrus_network/src/sqmr/behaviour.rs +++ b/crates/papyrus_network/src/sqmr/behaviour.rs @@ -234,7 +234,7 @@ impl Behaviour { Ok(()) } - fn get_peer_id_and_connection_id_from_session_id( + pub(crate) fn get_peer_id_and_connection_id_from_session_id( &self, session_id: SessionId, ) -> Result<(PeerId, ConnectionId), SessionIdNotFoundError> {