chore(network): get peer_id from session_id in swarmtrait (#2205)

starkware-libs · Jul 11, 2024 · cb1c4b7 · cb1c4b7
1 parent 3fa9b2f
commit cb1c4b7
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 2 deletions.
diff --git 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<<mixed_behaviour::MixedBehaviour as NetworkBehaviour>::ToSwarm>;
 
@@ -38,6 +38,11 @@ pub trait SwarmTrait: Stream<Item = Event> + Unpin {
 
     fn behaviour_mut(&mut self) -> &mut mixed_behaviour::MixedBehaviour;
 
+    fn get_peer_id_from_session_id(
+        &self,
+        session_id: SessionId,
+    ) -> Result<PeerId, SessionIdNotFoundError>;
+
     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<mixed_behaviour::MixedBehaviour> {
         self.behaviour_mut()
     }
 
+    fn get_peer_id_from_session_id(
+        &self,
+        session_id: SessionId,
+    ) -> Result<PeerId, SessionIdNotFoundError> {
+        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
@@ -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<PeerId, SessionIdNotFoundError> {
+        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
@@ -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> {