feat(handle): is_connected command (#43)

Author: storopoli

crates/p2p/src/commands.rs

@@ -5,10 +5,11 @@ use libp2p::{identity::secp256k1, Multiaddr, PeerId};
 use musig2::{PartialSignature, PubNonce};
 use strata_p2p_types::{P2POperatorPubKey, Scope, SessionId, StakeChainId, WotsPublicKeys};
 use strata_p2p_wire::p2p::v1::{GetMessageRequest, GossipsubMsg, UnsignedGossipsubMsg};
+use tokio::sync::oneshot;
 
 /// Ask P2P implementation to distribute some data across network.
 #[expect(clippy::large_enum_variant)]
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub enum Command {
     /// Publishes message through gossip sub network of peers.
     PublishMessage(PublishMessage),
@@ -21,6 +22,9 @@ pub enum Command {
 
     /// Connects to a peer, whitelists peer, and adds peer to the gossip sub network.
     ConnectToPeer(ConnectToPeerCommand),
+
+    /// Directly queries P2P state (doesn't produce events)
+    QueryP2PState(QueryP2PStateCommand),
 }
 
 #[derive(Debug, Clone)]
@@ -192,6 +196,36 @@ pub struct ConnectToPeerCommand {
     pub peer_addr: Multiaddr,
 }
 
+impl From<ConnectToPeerCommand> for Command {
+    fn from(v: ConnectToPeerCommand) -> Self {
+        Self::ConnectToPeer(v)
+    }
+}
+
+/// Commands to directly query P2P state information
+#[derive(Debug)]
+pub enum QueryP2PStateCommand {
+    /// Query if we're connected to a specific peer
+    IsConnected {
+        /// Peer ID to check
+        peer_id: PeerId,
+        /// Channel to send the response back
+        response_sender: oneshot::Sender<bool>,
+    },
+
+    /// Get all connected peers
+    GetConnectedPeers {
+        /// Channel to send the response back
+        response_sender: oneshot::Sender<Vec<PeerId>>,
+    },
+}
+
+impl From<QueryP2PStateCommand> for Command {
+    fn from(v: QueryP2PStateCommand) -> Self {
+        Self::QueryP2PState(v)
+    }
+}
+
 /// Commands P2P to clean entries from internal key-value storage by
 /// session IDs, scopes and operator pubkeys.
 #[derive(Debug, Clone)]

crates/p2p/src/swarm/handle.rs

@@ -5,19 +5,23 @@ use std::{
     fmt::Display,
     pin::Pin,
     task::{Context, Poll},
+    time::Duration,
 };
 
 use futures::{FutureExt, Stream};
-use libp2p::identity::secp256k1::Keypair;
+use libp2p::{identity::secp256k1::Keypair, PeerId};
 use thiserror::Error;
-use tokio::sync::{
-    broadcast::{self, error::RecvError},
-    mpsc,
+use tokio::{
+    sync::{
+        broadcast::{self, error::RecvError},
+        mpsc, oneshot,
+    },
+    time::timeout,
 };
 use tracing::warn;
 
 use crate::{
-    commands::{Command, PublishMessage, UnsignedPublishMessage},
+    commands::{Command, PublishMessage, QueryP2PStateCommand, UnsignedPublishMessage},
     events::Event,
 };
 
@@ -76,6 +80,59 @@ impl P2PHandle {
     pub fn sign_message(&self, msg: UnsignedPublishMessage) -> PublishMessage {
         msg.sign_secp256k1(&self.keypair)
     }
+
+    /// Checks if the P2P node is connected to the specified peer.
+    /// Returns true if connected, false otherwise.
+    pub async fn is_connected(&self, peer_id: PeerId) -> bool {
+        let (sender, receiver) = oneshot::channel();
+
+        // Send the command to check connection
+        let cmd = Command::QueryP2PState(QueryP2PStateCommand::IsConnected {
+            peer_id,
+            response_sender: sender,
+        });
+
+        // Use a cloned sender to avoid borrow issues
+        let cmd_sender = self.commands.clone();
+
+        // Send the command
+        if cmd_sender.send(cmd).await.is_err() {
+            // If the command channel is closed, assume not connected
+            return false;
+        }
+
+        // Wait for response with timeout
+        // TODO: make this configurable
+        match timeout(Duration::from_secs(1), receiver).await {
+            Ok(Ok(is_connected)) => is_connected,
+            _ => false, // Timeout or channel closed
+        }
+    }
+
+    /// Gets the list of all currently connected peers.
+    pub async fn get_connected_peers(&self) -> Vec<PeerId> {
+        let (sender, receiver) = oneshot::channel();
+
+        // Send the command
+        let cmd = Command::QueryP2PState(QueryP2PStateCommand::GetConnectedPeers {
+            response_sender: sender,
+        });
+
+        // Use a cloned sender to avoid borrow issues
+        let cmd_sender = self.commands.clone();
+
+        // If sending fails, return empty list
+        if cmd_sender.send(cmd).await.is_err() {
+            return Vec::new();
+        }
+
+        // Wait for response with timeout
+        // TODO: make this configurable
+        match timeout(Duration::from_secs(1), receiver).await {
+            Ok(Ok(peers)) => peers,
+            _ => Vec::new(), // Timeout or channel closed
+        }
+    }
 }
 
 impl Clone for P2PHandle {

crates/p2p/src/swarm/mod.rs

@@ -38,7 +38,10 @@ use tokio::{
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info, instrument};
 
-use crate::{commands::Command, events::Event};
+use crate::{
+    commands::{Command, QueryP2PStateCommand},
+    events::Event,
+};
 
 mod behavior;
 mod codec;
@@ -522,6 +525,21 @@ impl<DB: RepositoryExt> P2P<DB> {
 
                 Ok(())
             }
+            Command::QueryP2PState(query) => match query {
+                QueryP2PStateCommand::IsConnected {
+                    peer_id,
+                    response_sender,
+                } => {
+                    let is_connected = self.swarm.is_connected(&peer_id);
+                    let _ = response_sender.send(is_connected);
+                    Ok(())
+                }
+                QueryP2PStateCommand::GetConnectedPeers { response_sender } => {
+                    let peers = self.swarm.connected_peers().cloned().collect();
+                    let _ = response_sender.send(peers);
+                    Ok(())
+                }
+            },
         }
     }
 

crates/p2p/src/tests/is_connected.rs

@@ -0,0 +1,27 @@
+//! Tests for the `is_connected` method.
+
+use crate::tests::common::Setup;
+
+#[tokio::test]
+async fn test_is_connected() -> anyhow::Result<()> {
+    // Set up two connected operators
+    let Setup {
+        operators,
+        cancel,
+        tasks,
+    } = Setup::all_to_all(2).await?;
+
+    // Verify operator 0 is connected to operator 1
+    let is_connected = operators[0].handle.is_connected(operators[1].peer_id).await;
+    assert!(is_connected);
+
+    // Also test the get_connected_peers API
+    let connected_peers = operators[0].handle.get_connected_peers().await;
+    assert!(connected_peers.contains(&operators[1].peer_id));
+
+    // Cleanup
+    cancel.cancel();
+    tasks.wait().await;
+
+    Ok(())
+}

crates/p2p/src/tests/mod.rs

@@ -2,4 +2,5 @@ pub(crate) mod common;
 pub(crate) mod connect_peer;
 pub(crate) mod db;
 pub(crate) mod gossipsub;
+pub(crate) mod is_connected;
 pub(crate) mod request;