From 35eaa4bcc45eeb83876608f7d9a460303eabe743 Mon Sep 17 00:00:00 2001
From: rustaceanrob <rob.netzke@gmail.com>
Date: Sun, 13 Oct 2024 09:34:22 -0700
Subject: [PATCH] feat: add trusted peers with client message

---
 src/core/client.rs   | 15 ++++++++++++++-
 src/core/messages.rs |  4 +++-
 src/core/node.rs     |  4 ++++
 src/core/peer_map.rs | 10 ++++++++++
 4 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/core/client.rs b/src/core/client.rs
index 0fe4bdd..34a45f0 100644
--- a/src/core/client.rs
+++ b/src/core/client.rs
@@ -6,7 +6,7 @@ use tokio::sync::broadcast;
 pub use tokio::sync::broadcast::Receiver;
 use tokio::sync::mpsc::Sender;
 
-use crate::{IndexedBlock, TxBroadcast};
+use crate::{IndexedBlock, TrustedPeer, TxBroadcast};
 
 use super::{
     error::ClientError,
@@ -186,6 +186,7 @@ macro_rules! impl_core_client {
                     .await
                     .map_err(|_| ClientError::SendError)
             }
+
             /// Set a new connection timeout for peers to respond to messages.
             ///
             /// # Errors
@@ -201,6 +202,18 @@ macro_rules! impl_core_client {
                     .map_err(|_| ClientError::SendError)
             }
 
+            /// Add another known peer to connect to.
+            ///
+            /// # Errors
+            ///
+            /// If the node has stopped running.
+            pub async fn add_peer(&self, peer: TrustedPeer) -> Result<(), ClientError> {
+                self.ntx
+                    .send(ClientMessage::AddPeer(peer))
+                    .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 2c7bf3e..f6d1219 100644
--- a/src/core/messages.rs
+++ b/src/core/messages.rs
@@ -8,7 +8,7 @@ use bitcoin::{block::Header, p2p::message_network::RejectReason, ScriptBuf, Txid
 use crate::IndexedFilter;
 use crate::{
     chain::checkpoints::HeaderCheckpoint, DisconnectedHeader, IndexedBlock, IndexedTransaction,
-    TxBroadcast,
+    TrustedPeer, TxBroadcast,
 };
 
 use super::node::NodeState;
@@ -113,6 +113,8 @@ pub(crate) enum ClientMessage {
     GetBlock(BlockHash),
     /// Set a new connection timeout.
     SetDuration(Duration),
+    /// Add another known peer to connect to.
+    AddPeer(TrustedPeer),
 }
 
 /// Warnings a node may issue while running.
diff --git a/src/core/node.rs b/src/core/node.rs
index c2a4653..ce1e95f 100644
--- a/src/core/node.rs
+++ b/src/core/node.rs
@@ -312,6 +312,10 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
                             ClientMessage::SetDuration(duration) => {
                                 let mut peer_map = self.peer_map.lock().await;
                                 peer_map.set_duration(duration);
+                            },
+                            ClientMessage::AddPeer(peer) => {
+                                let mut peer_map = self.peer_map.lock().await;
+                                peer_map.add_peer(peer);
                             }
                         }
                     }
diff --git a/src/core/peer_map.rs b/src/core/peer_map.rs
index 79f119b..eacb884 100644
--- a/src/core/peer_map.rs
+++ b/src/core/peer_map.rs
@@ -158,6 +158,16 @@ impl<P: PeerStore> PeerMap<P> {
         self.response_timeout = duration;
     }
 
+    // Add a new trusted peer to the whitelist
+    pub fn add_peer(&mut self, peer: TrustedPeer) {
+        match &mut self.whitelist {
+            Some(peers) => {
+                peers.push(peer);
+            }
+            None => self.whitelist = Some(vec![peer]),
+        }
+    }
+
     // 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);