From 36755d314fa7e7de5d58045af68d47b950bf4119 Mon Sep 17 00:00:00 2001
From: rustaceanrob <rob.netzke@gmail.com>
Date: Sun, 13 Oct 2024 09:05:56 -0700
Subject: [PATCH] feat(core): set connection timeout with client message

---
 src/core/client.rs   | 15 +++++++++++++++
 src/core/messages.rs |  4 +++-
 src/core/node.rs     |  4 ++++
 src/core/peer_map.rs |  5 +++++
 4 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/core/client.rs b/src/core/client.rs
index 3f7ddba..0fe4bdd 100644
--- a/src/core/client.rs
+++ b/src/core/client.rs
@@ -1,6 +1,7 @@
 #[cfg(feature = "silent-payments")]
 use bitcoin::BlockHash;
 use bitcoin::ScriptBuf;
+use std::time::Duration;
 use tokio::sync::broadcast;
 pub use tokio::sync::broadcast::Receiver;
 use tokio::sync::mpsc::Sender;
@@ -185,6 +186,20 @@ macro_rules! impl_core_client {
                     .await
                     .map_err(|_| ClientError::SendError)
             }
+            /// Set a new connection timeout for peers to respond to messages.
+            ///
+            /// # Errors
+            ///
+            /// If the node has stopped running.
+            pub async fn set_response_timeout(
+                &self,
+                duration: Duration,
+            ) -> Result<(), ClientError> {
+                self.ntx
+                    .send(ClientMessage::SetDuration(duration))
+                    .await
+                    .map_err(|_| ClientError::SendError)
+            }
 
             /// Explicitly start the block filter syncing process. Note that the node will automatically download and check
             /// filters unless the policy is to explicitly halt.
diff --git a/src/core/messages.rs b/src/core/messages.rs
index e9cc997..2c7bf3e 100644
--- a/src/core/messages.rs
+++ b/src/core/messages.rs
@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, time::Duration};
 
 #[cfg(feature = "silent-payments")]
 use bitcoin::BlockHash;
@@ -111,6 +111,8 @@ pub(crate) enum ClientMessage {
     /// Explicitly request a block from the node.
     #[cfg(feature = "silent-payments")]
     GetBlock(BlockHash),
+    /// Set a new connection timeout.
+    SetDuration(Duration),
 }
 
 /// Warnings a node may issue while running.
diff --git a/src/core/node.rs b/src/core/node.rs
index 8ed1e29..c2a4653 100644
--- a/src/core/node.rs
+++ b/src/core/node.rs
@@ -309,6 +309,10 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
                                 let mut chain = self.chain.lock().await;
                                 chain.get_block(hash);
                             },
+                            ClientMessage::SetDuration(duration) => {
+                                let mut peer_map = self.peer_map.lock().await;
+                                peer_map.set_duration(duration);
+                            }
                         }
                     }
                 }
diff --git a/src/core/peer_map.rs b/src/core/peer_map.rs
index ffe37ec..79f119b 100644
--- a/src/core/peer_map.rs
+++ b/src/core/peer_map.rs
@@ -153,6 +153,11 @@ impl<P: PeerStore> PeerMap<P> {
         }
     }
 
+    // Set a new timeout duration
+    pub fn set_duration(&mut self, duration: Duration) {
+        self.response_timeout = duration;
+    }
+
     // Send out a TCP connection to a new peer and begin tracking the task
     pub async fn dispatch(&mut self, loaded_peer: PersistedPeer) -> Result<(), PeerError> {
         let (ptx, prx) = mpsc::channel::<MainThreadMessage>(32);