From 931b89998ee5b50037aa790c6c1f3d29cb357ea8 Mon Sep 17 00:00:00 2001
From: magine <magine233@gmail.com>
Date: Tue, 16 Jan 2024 09:51:16 +0800
Subject: [PATCH] fix: check if destination is already connected when infer the
 next hop of message

---
 crates/core/src/message/payload.rs | 7 +++++++
 crates/core/src/swarm/mod.rs       | 8 ++++++++
 2 files changed, 15 insertions(+)

diff --git a/crates/core/src/message/payload.rs b/crates/core/src/message/payload.rs
index 1fb60fade..b17198ea3 100644
--- a/crates/core/src/message/payload.rs
+++ b/crates/core/src/message/payload.rs
@@ -240,6 +240,9 @@ pub trait PayloadSender {
     /// Get access to DHT.
     fn dht(&self) -> Arc<PeerRing>;
 
+    /// Used to check if destination is already connected when `infer_next_hop`
+    fn is_connected(&self, did: Did) -> bool;
+
     /// Send a message payload to a specified DID.
     async fn do_send_payload(&self, did: Did, payload: MessagePayload) -> Result<()>;
 
@@ -249,6 +252,10 @@ pub trait PayloadSender {
             return Ok(next_hop);
         }
 
+        if self.is_connected(destination) {
+            return Ok(destination);
+        }
+
         match self.dht().find_successor(destination)? {
             PeerRingAction::Some(did) => Ok(did),
             PeerRingAction::RemoteAction(did, _) => Ok(did),
diff --git a/crates/core/src/swarm/mod.rs b/crates/core/src/swarm/mod.rs
index e88e9cd3f..a8671e6fc 100644
--- a/crates/core/src/swarm/mod.rs
+++ b/crates/core/src/swarm/mod.rs
@@ -18,6 +18,7 @@ use rings_derive::JudgeConnection;
 use rings_transport::core::transport::BoxedTransport;
 use rings_transport::core::transport::ConnectionInterface;
 use rings_transport::core::transport::TransportMessage;
+use rings_transport::core::transport::WebrtcConnectionState;
 use rings_transport::error::Error as TransportError;
 pub use types::MeasureImpl;
 pub use types::WrappedDid;
@@ -322,6 +323,13 @@ impl PayloadSender for Swarm {
         Swarm::dht(self)
     }
 
+    fn is_connected(&self, did: Did) -> bool {
+        let Some(conn) = self.get_connection(did) else {
+            return false;
+        };
+        conn.webrtc_connection_state() == WebrtcConnectionState::Connected
+    }
+
     async fn do_send_payload(&self, did: Did, payload: MessagePayload) -> Result<()> {
         #[cfg(test)]
         {