Merge pull request #286 from rustaceanrob/peer-id-2-6

rustaceanrob · web-flow · commit 915844176cda · 2025-02-06T19:12:52.000-10:00
Introduce `PeerId` for stricter type meaning across the crate
diff --git a/src/core/channel_messages.rs b/src/core/channel_messages.rs
@@ -11,6 +11,8 @@ use bitcoin::{
 
 use crate::core::messages::RejectPayload;
 
+use super::PeerId;
+
 #[derive(Debug, Clone)]
 pub(crate) enum MainThreadMessage {
     GetAddr,
@@ -36,7 +38,7 @@ pub struct GetBlockConfig {
 }
 
 pub(crate) struct PeerThreadMessage {
-    pub nonce: u32,
+    pub nonce: PeerId,
     pub message: PeerMessage,
 }
 
diff --git a/src/core/mod.rs b/src/core/mod.rs
@@ -7,6 +7,7 @@
 //! To build a [`Node`](node::Node) and [`Client`](client::Client), please refer to the [`NodeBuilder`](builder::NodeBuilder), which allows for node
 //! configuration.
 
+use std::hash::Hash;
 use std::time::Duration;
 
 use tokio::time::Instant;
@@ -84,3 +85,24 @@ impl PeerTimeoutConfig {
         }
     }
 }
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub(crate) struct PeerId(u32);
+
+impl PeerId {
+    pub(crate) fn increment(&mut self) {
+        self.0 = self.0.wrapping_add(1)
+    }
+}
+
+impl From<u32> for PeerId {
+    fn from(value: u32) -> Self {
+        PeerId(value)
+    }
+}
+
+impl std::fmt::Display for PeerId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Peer {}", self.0)
+    }
+}
diff --git a/src/core/node.rs b/src/core/node.rs
@@ -38,7 +38,7 @@ use super::{
     dialog::Dialog,
     error::NodeError,
     messages::{ClientMessage, Event, Log, SyncUpdate, Warning},
-    FilterSyncPolicy, LastBlockMonitor, PeerTimeoutConfig,
+    FilterSyncPolicy, LastBlockMonitor, PeerId, PeerTimeoutConfig,
 };
 
 pub(crate) const ADDR_V2_VERSION: u32 = 70015;
@@ -211,13 +211,13 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
                                     };
                                     let response = self.handle_version(peer_thread.nonce, version, best).await?;
                                     self.send_message(peer_thread.nonce, response).await;
-                                    self.dialog.send_dialog(format!("[Peer {}]: version", peer_thread.nonce))
+                                    self.dialog.send_dialog(format!("[{}]: version", peer_thread.nonce))
                                         .await;
                                 }
                                 PeerMessage::Addr(addresses) => self.handle_new_addrs(addresses).await,
                                 PeerMessage::Headers(headers) => {
                                     last_block.reset();
-                                    self.dialog.send_dialog(format!("[Peer {}]: headers", peer_thread.nonce))
+                                    self.dialog.send_dialog(format!("[{}]: headers", peer_thread.nonce))
                                         .await;
                                     match self.handle_headers(peer_thread.nonce, headers).await {
                                         Some(response) => {
@@ -227,7 +227,7 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
                                     }
                                 }
                                 PeerMessage::FilterHeaders(cf_headers) => {
-                                    self.dialog.send_dialog(format!("[Peer {}]: filter headers", peer_thread.nonce)).await;
+                                    self.dialog.send_dialog(format!("[{}]: filter headers", peer_thread.nonce)).await;
                                     match self.handle_cf_headers(peer_thread.nonce, cf_headers).await {
                                         Some(response) => {
                                             self.broadcast(response).await;
@@ -250,7 +250,7 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
                                     None => continue,
                                 },
                                 PeerMessage::NewBlocks(blocks) => {
-                                    self.dialog.send_dialog(format!("[Peer {}]: inv", peer_thread.nonce))
+                                    self.dialog.send_dialog(format!("[{}]: inv", peer_thread.nonce))
                                         .await;
                                     match self.handle_inventory_blocks(peer_thread.nonce, blocks).await {
                                         Some(response) => {
@@ -340,7 +340,7 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
     }
 
     // Send a message to a specified peer
-    async fn send_message(&self, nonce: u32, message: MainThreadMessage) {
+    async fn send_message(&self, nonce: PeerId, message: MainThreadMessage) {
         let mut peer_map = self.peer_map.lock().await;
         peer_map.send_message(nonce, message).await;
     }
@@ -524,7 +524,7 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
     // We accepted a handshake with a peer but we may disconnect if they do not support CBF
     async fn handle_version(
         &self,
-        nonce: u32,
+        nonce: PeerId,
         version_message: VersionMessage,
         best_height: u32,
     ) -> Result<MainThreadMessage, NodeError<H::Error, P::Error>> {
@@ -590,7 +590,7 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
     // We always send headers to our peers, so our next message depends on our state
     async fn handle_headers(
         &self,
-        peer_id: u32,
+        peer_id: PeerId,
         headers: Vec<Header>,
     ) -> Option<MainThreadMessage> {
         let mut chain = self.chain.lock().await;
@@ -624,11 +624,11 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
     // Compact filter headers may result in a number of outcomes, including the need to audit filters.
     async fn handle_cf_headers(
         &self,
-        peer_id: u32,
+        peer_id: PeerId,
         cf_headers: CFHeaders,
     ) -> Option<MainThreadMessage> {
         let mut chain = self.chain.lock().await;
-        match chain.sync_cf_headers(peer_id, cf_headers).await {
+        match chain.sync_cf_headers(peer_id.0, cf_headers).await {
             Ok(potential_message) => match potential_message {
                 AppendAttempt::AddedToQueue => None,
                 AppendAttempt::Extended => self.next_stateful_message(chain.deref_mut()).await,
@@ -652,7 +652,7 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
     }
 
     // Handle a new compact block filter
-    async fn handle_filter(&self, peer_id: u32, filter: CFilter) -> Option<MainThreadMessage> {
+    async fn handle_filter(&self, peer_id: PeerId, filter: CFilter) -> Option<MainThreadMessage> {
         let mut chain = self.chain.lock().await;
         match chain.sync_filter(filter).await {
             Ok(potential_message) => potential_message.map(MainThreadMessage::GetFilters),
@@ -673,7 +673,7 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
     }
 
     // Scan a block for transactions.
-    async fn handle_block(&self, peer_id: u32, block: Block) -> Option<MainThreadMessage> {
+    async fn handle_block(&self, peer_id: PeerId, block: Block) -> Option<MainThreadMessage> {
         let mut chain = self.chain.lock().await;
         if let Err(e) = chain.check_send_block(block).await {
             self.dialog.send_warning(Warning::UnexpectedSyncError {
@@ -713,7 +713,7 @@ impl<H: HeaderStore, P: PeerStore> Node<H, P> {
     // If new inventory came in, we need to download the headers and update the node state
     async fn handle_inventory_blocks(
         &self,
-        nonce: u32,
+        nonce: PeerId,
         blocks: Vec<BlockHash>,
     ) -> Option<MainThreadMessage> {
         let mut state = self.state.write().await;
diff --git a/src/core/peer_map.rs b/src/core/peer_map.rs
@@ -35,7 +35,7 @@ use super::{
     dialog::Dialog,
     error::PeerManagerError,
     messages::Warning,
-    PeerTimeoutConfig,
+    PeerId, PeerTimeoutConfig,
 };
 
 const MAX_TRIES: usize = 50;
@@ -58,11 +58,11 @@ pub(crate) struct ManagedPeer {
 // The `PeerMap` manages connections with peers, adds and bans peers, and manages the peer database
 #[derive(Debug)]
 pub(crate) struct PeerMap<P: PeerStore> {
-    num_peers: u32,
-    heights: HashMap<u32, u32>,
+    current_id: PeerId,
+    heights: HashMap<PeerId, u32>,
     network: Network,
     mtx: Sender<PeerThreadMessage>,
-    map: HashMap<u32, ManagedPeer>,
+    map: HashMap<PeerId, ManagedPeer>,
     db: Arc<Mutex<P>>,
     connector: Arc<Mutex<dyn NetworkConnector + Send + Sync>>,
     whitelist: Whitelist,
@@ -94,7 +94,7 @@ impl<P: PeerStore> PeerMap<P> {
             }
         };
         Self {
-            num_peers: 0,
+            current_id: PeerId(0),
             heights: HashMap::new(),
             network,
             mtx,
@@ -139,7 +139,7 @@ impl<P: PeerStore> PeerMap<P> {
     }
 
     // Set the time offset of a connected peer
-    pub fn set_offset(&mut self, peer: u32, time: i64) {
+    pub fn set_offset(&mut self, peer: PeerId, time: i64) {
         if let Some(peer) = self.map.get_mut(&peer) {
             let now = SystemTime::now()
                 .duration_since(UNIX_EPOCH)
@@ -162,10 +162,9 @@ impl<P: PeerStore> PeerMap<P> {
     // 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);
-        let peer_num = self.num_peers + 1;
-        self.num_peers = peer_num;
+        self.current_id.increment();
         let mut peer = Peer::new(
-            peer_num,
+            self.current_id,
             self.network,
             self.mtx.clone(),
             prx,
@@ -188,7 +187,7 @@ impl<P: PeerStore> PeerMap<P> {
             .await?;
         let handle = tokio::spawn(async move { peer.run(reader, writer).await });
         self.map.insert(
-            peer_num,
+            self.current_id,
             ManagedPeer {
                 service_flags: loaded_peer.services,
                 address: loaded_peer.addr,
@@ -203,26 +202,26 @@ impl<P: PeerStore> PeerMap<P> {
     }
 
     // Set the minimum fee rate this peer will accept
-    pub fn set_broadcast_min(&mut self, nonce: u32, fee_rate: FeeRate) {
+    pub fn set_broadcast_min(&mut self, nonce: PeerId, fee_rate: FeeRate) {
         if let Some(peer) = self.map.get_mut(&nonce) {
             peer.broadcast_min = fee_rate;
         }
     }
 
     // Set the services of a peer
-    pub fn set_services(&mut self, nonce: u32, flags: ServiceFlags) {
+    pub fn set_services(&mut self, nonce: PeerId, flags: ServiceFlags) {
         if let Some(peer) = self.map.get_mut(&nonce) {
             peer.service_flags = flags
         }
     }
 
     // Set the height of a peer upon receiving the version message
-    pub fn set_height(&mut self, nonce: u32, height: u32) {
+    pub fn set_height(&mut self, nonce: PeerId, height: u32) {
         self.heights.insert(nonce, height);
     }
 
     // Add one to the height of a peer when receiving inventory
-    pub fn add_one_height(&mut self, nonce: u32) {
+    pub fn add_one_height(&mut self, nonce: PeerId) {
         if let Some(height) = self.heights.get(&nonce) {
             self.heights.insert(nonce, height + 1);
         }
@@ -243,7 +242,7 @@ impl<P: PeerStore> PeerMap<P> {
     }
 
     // Send a message to the specified peer
-    pub async fn send_message(&mut self, nonce: u32, message: MainThreadMessage) {
+    pub async fn send_message(&mut self, nonce: PeerId, message: MainThreadMessage) {
         if let Some(peer) = self.map.get(&nonce) {
             let _ = peer.ptx.send(message).await;
         }
@@ -344,7 +343,7 @@ impl<P: PeerStore> PeerMap<P> {
     }
 
     // We tried this peer and successfully connected.
-    pub async fn tried(&mut self, nonce: u32) {
+    pub async fn tried(&mut self, nonce: PeerId) {
         if let Some(peer) = self.map.get(&nonce) {
             let mut db = self.db.lock().await;
             if let Err(e) = db
@@ -367,7 +366,7 @@ impl<P: PeerStore> PeerMap<P> {
     }
 
     // This peer misbehaved in some way.
-    pub async fn ban(&mut self, nonce: u32) {
+    pub async fn ban(&mut self, nonce: PeerId) {
         if let Some(peer) = self.map.get(&nonce) {
             let mut db = self.db.lock().await;
             if let Err(e) = db
diff --git a/src/network/peer.rs b/src/network/peer.rs
@@ -18,7 +18,7 @@ use crate::{
         channel_messages::{MainThreadMessage, PeerMessage, PeerThreadMessage},
         dialog::Dialog,
         messages::Warning,
-        PeerTimeoutConfig,
+        PeerId, PeerTimeoutConfig,
     },
     network::outbound_messages::V1OutboundMessage,
 };
@@ -42,7 +42,7 @@ const HANDSHAKE_TIMEOUT: u64 = 4;
 type MutexMessageGenerator = Mutex<Box<dyn MessageGenerator>>;
 
 pub(crate) struct Peer {
-    nonce: u32,
+    nonce: PeerId,
     main_thread_sender: Sender<PeerThreadMessage>,
     main_thread_recv: Receiver<MainThreadMessage>,
     network: Network,
@@ -54,7 +54,7 @@ pub(crate) struct Peer {
 
 impl Peer {
     pub fn new(
-        nonce: u32,
+        nonce: PeerId,
         network: Network,
         main_thread_sender: Sender<PeerThreadMessage>,
         main_thread_recv: Receiver<MainThreadMessage>,

Original file line number	Diff line number	Diff line change
`@@ -11,6 +11,8 @@ use bitcoin::{`
`11`	`11`
`12`	`12`	`use crate::core::messages::RejectPayload;`
`13`	`13`
	`14`	`+use super::PeerId;`
	`15`	`+`
`14`	`16`	`#[derive(Debug, Clone)]`
`15`	`17`	`pub(crate) enum MainThreadMessage {`
`16`	`18`	`GetAddr,`
`@@ -36,7 +38,7 @@ pub struct GetBlockConfig {`
`36`	`38`	`}`
`37`	`39`
`38`	`40`	`pub(crate) struct PeerThreadMessage {`
`39`		`- pub nonce: u32,`
	`41`	`+ pub nonce: PeerId,`
`40`	`42`	`pub message: PeerMessage,`
`41`	`43`	`}`
`42`	`44`