feat: introduce MessageDispatcher in network manager (#383)

* feat: subscription-based message routing

This PR:
- introduces `MessageRouter` to `PeerNetworkManager` which routes messages based on previously subscribed topics to the related mpsc channels. This allows for different tasks to listen to the messages they are interested in without locking/polling the single channel that we currently have.
- intoduces `network::Message` which encapsulates actual `NetworkMessages` with a `SocketAddress` of the sending peer to allow more streamlined reputation updates in message handlers.
- drops all the "last message peer id" stuff and uses the peer address from the new message format instead.

* some refactoring:

- refactor `subscribers: Vec<MessageSubscriber>` to `senders: HashMap<MessageType, Vec<UnboundedSender<Message>>>`
- introduce `define_message_types` macro to generate `MessageType` enum with `Message` -> `MessageType` conversion for the hashmap lookups
- rename `MessageRouter` -> `MessageDispatcher` and accordingly the `route()` to `dispatch()`
- rename `subscribe` -> `message_receiver` and `new_subscriber` -> `message_receiver`

* fix: deduplicate message types in MessageDispatcher registration

Prevents duplicate message deliveries when message_receiver is called
with duplicate MessageType entries in the input slice.

* fix and test macro

* prune dead senders in `dispatch()`

* lock `message_dispatcher` with a mutex since all access is mutable

Author: xdustinface

dash-spv/src/client/chainlock.rs

@@ -6,6 +6,7 @@
 //! - ChainLock validation updates
 //! - Pending ChainLock validation
 
+use std::net::SocketAddr;
 use std::sync::Arc;
 
 use crate::error::{Result, SpvError};
@@ -21,6 +22,7 @@ impl<W: WalletInterface, N: NetworkManager, S: StorageManager> DashSpvClient<W,
     /// Process and validate a ChainLock.
     pub async fn process_chainlock(
         &mut self,
+        peer_address: SocketAddr,
         chainlock: dashcore::ephemerealdata::chain_lock::ChainLock,
     ) -> Result<()> {
         tracing::info!(
@@ -40,7 +42,7 @@ impl<W: WalletInterface, N: NetworkManager, S: StorageManager> DashSpvClient<W,
             {
                 // Penalize the peer that relayed the invalid ChainLock
                 let reason = format!("Invalid ChainLock: {}", e);
-                let _ = self.network.penalize_last_message_peer_invalid_chainlock(&reason).await;
+                self.network.penalize_peer_invalid_chainlock(peer_address, &reason).await;
                 return Err(SpvError::Validation(e));
             }
         }
@@ -87,6 +89,7 @@ impl<W: WalletInterface, N: NetworkManager, S: StorageManager> DashSpvClient<W,
     /// Process and validate an InstantSendLock.
     pub(super) async fn process_instantsendlock(
         &mut self,
+        peer_address: SocketAddr,
         islock: dashcore::ephemerealdata::instant_lock::InstantLock,
     ) -> Result<()> {
         tracing::info!("Processing InstantSendLock for tx {}", islock.txid);
@@ -107,7 +110,7 @@ impl<W: WalletInterface, N: NetworkManager, S: StorageManager> DashSpvClient<W,
             tracing::warn!("{}", reason);
 
             // Ban the peer using the reputation system
-            let _ = self.network.penalize_last_message_peer_invalid_instantlock(&reason).await;
+            self.network.penalize_peer_invalid_instantlock(peer_address, &reason).await;
 
             return Err(SpvError::Validation(e));
         }

dash-spv/src/client/message_handler.rs

@@ -3,7 +3,7 @@
 use crate::client::ClientConfig;
 use crate::error::{Result, SpvError};
 use crate::mempool_filter::MempoolFilter;
-use crate::network::NetworkManager;
+use crate::network::{Message, NetworkManager};
 use crate::storage::StorageManager;
 use crate::sync::legacy::SyncManager;
 use crate::types::{MempoolState, SpvEvent};
@@ -47,17 +47,17 @@ impl<'a, S: StorageManager, N: NetworkManager, W: WalletInterface> MessageHandle
     }
 
     /// Handle incoming network messages during monitoring.
-    pub async fn handle_network_message(
-        &mut self,
-        message: &dashcore::network::message::NetworkMessage,
-    ) -> Result<()> {
+    pub async fn handle_network_message(&mut self, message: &Message) -> Result<()> {
         use dashcore::network::message::NetworkMessage;
 
-        tracing::debug!("Client handling network message: {:?}", std::mem::discriminant(message));
+        tracing::debug!(
+            "Client handling network message: {:?}",
+            std::mem::discriminant(message.inner())
+        );
 
         // First check if this is a message that ONLY the sync manager handles
         // These messages can be moved to the sync manager without cloning
-        match message {
+        match message.inner() {
             NetworkMessage::Headers2(ref headers2) => {
                 tracing::info!(
                     "📋 Received Headers2 message with {} compressed headers",
@@ -89,22 +89,11 @@ impl<'a, S: StorageManager, N: NetworkManager, W: WalletInterface> MessageHandle
             }
             NetworkMessage::CFHeaders(ref cf_headers) => {
                 // Try to include the peer address for better diagnostics
-                let peer_addr = self.network.get_last_message_peer_addr().await;
-                match peer_addr {
-                    Some(addr) => {
-                        tracing::info!(
-                            "📨 Client received CFHeaders message with {} filter headers from {}",
-                            cf_headers.filter_hashes.len(),
-                            addr
-                        );
-                    }
-                    None => {
-                        tracing::info!(
-                            "📨 Client received CFHeaders message with {} filter headers (peer unknown)",
-                            cf_headers.filter_hashes.len()
-                        );
-                    }
-                }
+                tracing::info!(
+                    "📨 Client received CFHeaders message with {} filter headers from {}",
+                    cf_headers.filter_hashes.len(),
+                    message.peer_address()
+                );
                 // Move to sync manager without cloning
                 return self
                     .sync_manager
@@ -164,23 +153,15 @@ impl<'a, S: StorageManager, N: NetworkManager, W: WalletInterface> MessageHandle
         }
 
         // Then handle client-specific message processing
-        match message {
+        match message.inner() {
             NetworkMessage::Headers(headers) => {
                 // For post-sync headers, we need special handling
                 if self.sync_manager.is_synced() && !headers.is_empty() {
-                    let peer_addr = self.network.get_last_message_peer_addr().await;
-                    if let Some(addr) = peer_addr {
-                        tracing::info!(
+                    tracing::info!(
                             "📋 Post-sync headers received from {} ({} headers), additional processing may be needed",
-                            addr,
-                            headers.len()
-                        );
-                    } else {
-                        tracing::info!(
-                            "📋 Post-sync headers received ({} headers), additional processing may be needed",
+                            message.peer_address(),
                             headers.len()
                         );
-                    }
                 }
             }
             NetworkMessage::Block(block) => {
@@ -194,7 +175,10 @@ impl<'a, S: StorageManager, N: NetworkManager, W: WalletInterface> MessageHandle
 
                 // 1) Ensure header processing and chain tip update for this block
                 //    Route the header through the sequential sync manager as a Headers message
-                let headers_msg = NetworkMessage::Headers(vec![block.header]);
+                let headers_msg = Message::new(
+                    message.peer_address(),
+                    NetworkMessage::Headers(vec![block.header]),
+                );
                 if let Err(e) = self
                     .sync_manager
                     .handle_message(&headers_msg, &mut *self.network, &mut *self.storage)
@@ -297,7 +281,10 @@ impl<'a, S: StorageManager, N: NetworkManager, W: WalletInterface> MessageHandle
             }
             _ => {
                 // Ignore other message types for now
-                tracing::debug!("Received network message: {:?}", std::mem::discriminant(message));
+                tracing::debug!(
+                    "Received network message: {:?}",
+                    std::mem::discriminant(message.inner())
+                );
             }
         }
 

dash-spv/src/client/message_handler_test.rs

@@ -3,9 +3,10 @@
 #[cfg(test)]
 mod tests {
     use crate::client::{ClientConfig, MessageHandler};
+    use crate::network::Message;
     use crate::storage::DiskStorageManager;
     use crate::sync::legacy::SyncManager;
-    use crate::test_utils::MockNetworkManager;
+    use crate::test_utils::{test_socket_address, MockNetworkManager};
     use crate::types::{MempoolState, SpvEvent, SpvStats};
     use crate::ChainState;
     use dashcore::block::Header as BlockHeader;
@@ -69,7 +70,7 @@ mod tests {
         let headers2 = dashcore::network::message_headers2::Headers2Message {
             headers: vec![],
         };
-        let message = NetworkMessage::Headers2(headers2);
+        let message = Message::new(test_socket_address(1), NetworkMessage::Headers2(headers2));
 
         // Handle the message
         let result = handler.handle_network_message(&message).await;
@@ -115,7 +116,7 @@ mod tests {
             new_quorums: vec![],
             quorums_chainlock_signatures: vec![],
         };
-        let message = NetworkMessage::MnListDiff(mnlistdiff);
+        let message = Message::new(test_socket_address(1), NetworkMessage::MnListDiff(mnlistdiff));
 
         // Handle the message
         let result = handler.handle_network_message(&message).await;
@@ -144,7 +145,7 @@ mod tests {
             previous_filter_header: dashcore::hash_types::FilterHeader::from([0u8; 32]),
             filter_hashes: vec![],
         };
-        let message = NetworkMessage::CFHeaders(cfheaders);
+        let message = Message::new(test_socket_address(1), NetworkMessage::CFHeaders(cfheaders));
 
         // Handle the message
         let result = handler.handle_network_message(&message).await;
@@ -172,7 +173,7 @@ mod tests {
             block_hash: BlockHash::from([0u8; 32]),
             filter: vec![],
         };
-        let message = NetworkMessage::CFilter(cfilter);
+        let message = Message::new(test_socket_address(1), NetworkMessage::CFilter(cfilter));
 
         // Handle the message - should be passed to sync manager
         let result = handler.handle_network_message(&message).await;
@@ -206,7 +207,7 @@ mod tests {
             },
             txdata: vec![],
         };
-        let message = NetworkMessage::Block(block.clone());
+        let message = Message::new(test_socket_address(1), NetworkMessage::Block(block.clone()));
 
         // Handle the message
         let result = handler.handle_network_message(&message).await;
@@ -234,7 +235,7 @@ mod tests {
 
         // Create an Inv message with transaction
         let inv = vec![Inventory::Transaction(dashcore::Txid::all_zeros())];
-        let message = NetworkMessage::Inv(inv);
+        let message = Message::new(test_socket_address(1), NetworkMessage::Inv(inv));
 
         // Handle the message
         let result = handler.handle_network_message(&message).await;
@@ -267,7 +268,7 @@ mod tests {
             output: vec![],
             special_transaction_payload: None,
         };
-        let message = NetworkMessage::Tx(tx.clone());
+        let message = Message::new(test_socket_address(1), NetworkMessage::Tx(tx.clone()));
 
         // Handle the message
         let result = handler.handle_network_message(&message).await;
@@ -300,7 +301,7 @@ mod tests {
             block_hash: BlockHash::from([0u8; 32]),
             signature: dashcore::bls_sig_utils::BLSSignature::from([0u8; 96]),
         };
-        let message = NetworkMessage::CLSig(chainlock);
+        let message = Message::new(test_socket_address(1), NetworkMessage::CLSig(chainlock));
 
         // Handle the message
         let result = handler.handle_network_message(&message).await;
@@ -323,7 +324,7 @@ mod tests {
         );
 
         // Create a Ping message
-        let message = NetworkMessage::Ping(12345);
+        let message = Message::new(test_socket_address(1), NetworkMessage::Ping(12345));
 
         // Handle the message
         let result = handler.handle_network_message(&message).await;
@@ -352,7 +353,7 @@ mod tests {
         let headers2 = dashcore::network::message_headers2::Headers2Message {
             headers: vec![], // Empty headers might cause validation error
         };
-        let message = NetworkMessage::Headers2(headers2);
+        let message = Message::new(test_socket_address(1), NetworkMessage::Headers2(headers2));
 
         // Handle the message - error should be propagated
         let result = handler.handle_network_message(&message).await;

dash-spv/src/client/sync_coordinator.rs

@@ -14,7 +14,7 @@ use super::{DashSpvClient, MessageHandler};
 use crate::client::interface::DashSpvClientCommand;
 use crate::error::{Result, SpvError};
 use crate::network::constants::MESSAGE_RECEIVE_TIMEOUT;
-use crate::network::NetworkManager;
+use crate::network::{Message, MessageType, NetworkManager};
 use crate::storage::StorageManager;
 use crate::types::{DetailedSyncProgress, SyncProgress};
 use key_wallet_manager::wallet_interface::WalletInterface;
@@ -78,6 +78,22 @@ impl<W: WalletInterface, N: NetworkManager, S: StorageManager> DashSpvClient<W,
         let mut last_emitted_filters_downloaded: u64 = 0;
         let mut last_emitted_phase_name: Option<String> = None;
 
+        let mut message_receiver = self
+            .network
+            .message_receiver(&[
+                MessageType::Headers,
+                MessageType::Headers2,
+                MessageType::CFHeaders,
+                MessageType::CFilter,
+                MessageType::Block,
+                MessageType::MnListDiff,
+                MessageType::QRInfo,
+                MessageType::CLSig,
+                MessageType::ISLock,
+                MessageType::Inv,
+            ])
+            .await;
+
         loop {
             // Check if we should stop
             let running = self.running.read().await;
@@ -415,12 +431,13 @@ impl<W: WalletInterface, N: NetworkManager, S: StorageManager> DashSpvClient<W,
                         }
                     }
                 }
-                received = self.network.receive_message() => {
+                received = message_receiver.recv() => {
                     match received {
-                        Ok(None) => {
-                            continue;
+                        None => {
+                            tracing::info!("Network message subscription channel closed.");
+                            break;
                         }
-                        Ok(Some(message)) => {
+                        Some(message) => {
                             // Wrap message handling in comprehensive error handling
                             match self.handle_network_message(message).await {
                                 Ok(_) => {
@@ -452,36 +469,6 @@ impl<W: WalletInterface, N: NetworkManager, S: StorageManager> DashSpvClient<W,
                                 }
                             }
                         },
-                        Err(err) => {
-                            // Handle specific network error types
-                            if let crate::error::NetworkError::ConnectionFailed(msg) = &err {
-                                if msg.contains("No connected peers") || self.network.peer_count() == 0 {
-                                    tracing::warn!("All peers disconnected during monitoring, checking connection health");
-
-                                    // Wait for potential reconnection
-                                    let mut wait_count = 0;
-                                    while wait_count < 10 && self.network.peer_count() == 0 {
-                                        tokio::time::sleep(Duration::from_millis(500)).await;
-                                        wait_count += 1;
-                                    }
-
-                                    if self.network.peer_count() > 0 {
-                                        tracing::info!(
-                                            "✅ Reconnected to {} peer(s), resuming monitoring",
-                                            self.network.peer_count()
-                                        );
-                                        continue
-                                    } else {
-                                        tracing::warn!(
-                                            "No peers available after waiting, will retry monitoring"
-                                        );
-                                    }
-                                }
-                            }
-
-                            tracing::error!("Network error during monitoring: {}", err);
-                            tokio::time::sleep(Duration::from_secs(5)).await;
-                        }
                     }
                 }
                 _ = tokio::time::sleep(MESSAGE_RECEIVE_TIMEOUT) => {}
@@ -546,13 +533,10 @@ impl<W: WalletInterface, N: NetworkManager, S: StorageManager> DashSpvClient<W,
     }
 
     /// Handle incoming network messages during monitoring.
-    pub(super) async fn handle_network_message(
-        &mut self,
-        message: dashcore::network::message::NetworkMessage,
-    ) -> Result<()> {
+    pub(super) async fn handle_network_message(&mut self, message: Message) -> Result<()> {
         // Check if this is a special message that needs client-level processing
         let needs_special_processing = matches!(
-            &message,
+            &message.inner(),
             dashcore::network::message::NetworkMessage::CLSig(_)
                 | dashcore::network::message::NetworkMessage::ISLock(_)
         );
@@ -582,17 +566,21 @@ impl<W: WalletInterface, N: NetworkManager, S: StorageManager> DashSpvClient<W,
                 if needs_special_processing {
                     // Special handling for messages that need client-level processing
                     use dashcore::network::message::NetworkMessage;
-                    match &message {
+                    match message.inner() {
                         NetworkMessage::CLSig(clsig) => {
                             // Additional client-level ChainLock processing
-                            self.process_chainlock(clsig.clone()).await?;
+                            self.process_chainlock(message.peer_address(), clsig.clone()).await?;
                         }
                         NetworkMessage::ISLock(islock_msg) => {
                             // Only process InstantLocks when fully synced and masternode engine is available
                             if self.sync_manager.is_synced()
                                 && self.sync_manager.get_masternode_engine().is_some()
                             {
-                                self.process_instantsendlock(islock_msg.clone()).await?;
+                                self.process_instantsendlock(
+                                    message.peer_address(),
+                                    islock_msg.clone(),
+                                )
+                                .await?;
                             } else {
                                 tracing::debug!(
                                     "Skipping InstantLock processing - not fully synced or masternode engine unavailable"