track message via message_id

Author: borngraced

mm2src/kdf_walletconnect/src/connection_handler.rs

@@ -82,7 +82,7 @@ pub(crate) async fn handle_disconnections(
         loop {
             match this.reconnect_and_subscribe().await {
                 Ok(_) => {
-                    error!("Reconnection process complete.");
+                    info!("Reconnection process complete.");
                     backoff = 1;
                     break;
                 },

mm2src/kdf_walletconnect/src/error.rs

@@ -90,6 +90,8 @@ pub enum WalletConnectError {
     InvalidChainId(String),
     #[error("ChainId not supported: {0}")]
     ChainIdNotSupported(String),
+    #[error("Request timeout error")]
+    TimeoutError,
 }
 
 impl From<Error<PublishError>> for WalletConnectError {

mm2src/kdf_walletconnect/src/lib.rs

@@ -344,9 +344,9 @@ impl WalletConnectCtxImpl {
         &self,
         topic: &Topic,
         param: RequestParams,
+        message_id: MessageId,
     ) -> MmResult<(), WalletConnectError> {
         let irn_metadata = param.irn_metadata();
-        let message_id = MessageIdGenerator::new().next();
         let request = Request::new(message_id, param.into());
 
         self.publish_payload(topic, irn_metadata, Payload::Request(request))
@@ -463,6 +463,7 @@ impl WalletConnectCtxImpl {
         session: &Session,
         chain_id: &WcChainId,
     ) -> MmResult<(), WalletConnectError> {
+        println!("{:?}", session.namespaces.get(chain_id.chain.as_ref()));
         if let Some(Namespace {
             chains: Some(chains), ..
         }) = session.namespaces.get(chain_id.chain.as_ref())
@@ -581,21 +582,41 @@ impl WalletConnectCtxImpl {
         };
         let request = RequestParams::SessionRequest(request);
         let ttl = request.irn_metadata().ttl;
-        self.publish_request(&active_topic, request).await?;
-
-        if let Ok(Some(resp)) = timeout(Duration::from_secs(ttl), async {
-            self.message_rx.lock().await.next().await
+        let message_id = MessageIdGenerator::new().next();
+        self.publish_request(&active_topic, request, message_id).await?;
+
+        match timeout(Duration::from_secs(ttl), async {
+            // Check if the message exists and matches the expected message ID and
+            // wait till we get a message with expected id or timeout.
+            loop {
+                let next_message = {
+                    let mut lock = self.message_rx.lock().await;
+                    lock.next().await
+                };
+
+                if let Some(Ok(message)) = &next_message {
+                    if message.message_id == message_id {
+                        return next_message;
+                    }
+                }
+            }
         })
         .await
         {
-            let result = resp.mm_err(WalletConnectError::InternalError)?;
-            if let ResponseParamsSuccess::Arbitrary(data) = result.data {
-                let data = serde_json::from_value::<T>(data)?;
-                return callback(data);
-            }
+            Ok(Some(result)) => {
+                let result = result.mm_err(WalletConnectError::InternalError);
+                match result?.data {
+                    ResponseParamsSuccess::Arbitrary(data) => {
+                        let data = serde_json::from_value::<T>(data)
+                            .map_err(|e| WalletConnectError::SerdeError(e.to_string()))?;
+                        callback(data)
+                    },
+                    _ => MmError::err(WalletConnectError::PayloadError("Unexpected response type".to_string())),
+                }
+            },
+            Ok(None) => MmError::err(WalletConnectError::NoWalletFeedback),
+            Err(_) => MmError::err(WalletConnectError::TimeoutError),
         }
-
-        MmError::err(WalletConnectError::NoWalletFeedback)
     }
 
     pub async fn drop_session(&self, topic: &Topic) -> MmResult<(), WalletConnectError> {

mm2src/kdf_walletconnect/src/session/rpc/delete.rs

@@ -4,6 +4,7 @@ use crate::{error::{WalletConnectError, USER_REQUESTED},
 
 use common::log::debug;
 use mm2_err_handle::prelude::{MapMmError, MmResult};
+use relay_client::MessageIdGenerator;
 use relay_rpc::domain::{MessageId, Topic};
 use relay_rpc::rpc::params::{session_delete::SessionDeleteRequest, RequestParams, ResponseParamsSuccess};
 
@@ -28,8 +29,9 @@ pub(crate) async fn send_session_delete_request(
         message: "User Disconnected".to_owned(),
     };
     let param = RequestParams::SessionDelete(delete_request);
+    let message_id = MessageIdGenerator::new().next();
 
-    ctx.publish_request(session_topic, param).await?;
+    ctx.publish_request(session_topic, param, message_id).await?;
 
     session_delete_cleanup(ctx, session_topic).await
 }

mm2src/kdf_walletconnect/src/session/rpc/ping.rs

@@ -5,6 +5,7 @@ use crate::{error::WalletConnectError, WalletConnectCtxImpl};
 use common::custom_futures::timeout::FutureTimerExt;
 use futures::StreamExt;
 use mm2_err_handle::prelude::*;
+use relay_client::MessageIdGenerator;
 use relay_rpc::{domain::{MessageId, Topic},
                 rpc::params::{RelayProtocolMetadata, RequestParams, ResponseParamsSuccess}};
 
@@ -22,7 +23,8 @@ pub(crate) async fn reply_session_ping_request(
 pub async fn send_session_ping_request(ctx: &WalletConnectCtxImpl, topic: &Topic) -> MmResult<(), WalletConnectError> {
     let param = RequestParams::SessionPing(());
     let ttl = param.irn_metadata().ttl;
-    ctx.publish_request(topic, param).await?;
+    let message_id = MessageIdGenerator::new().next();
+    ctx.publish_request(topic, param, message_id).await?;
 
     let wait_duration = Duration::from_secs(ttl);
     if let Ok(Some(resp)) = ctx.message_rx.lock().await.next().timeout(wait_duration).await {

mm2src/kdf_walletconnect/src/session/rpc/propose.rs

@@ -8,6 +8,7 @@ use crate::{error::WalletConnectError,
 use chrono::Utc;
 use mm2_err_handle::map_to_mm::MapToMmResult;
 use mm2_err_handle::prelude::*;
+use relay_client::MessageIdGenerator;
 use relay_rpc::rpc::params::session::ProposeNamespaces;
 use relay_rpc::{domain::{MessageId, Topic},
                 rpc::params::{session_propose::{Proposer, SessionProposeRequest, SessionProposeResponse},
@@ -30,7 +31,8 @@ pub(crate) async fn send_proposal_request(
         required_namespaces,
         optional_namespaces,
     });
-    ctx.publish_request(topic, session_proposal).await?;
+    let message_id = MessageIdGenerator::new().next();
+    ctx.publish_request(topic, session_proposal, message_id).await?;
 
     Ok(())
 }