refactor: move Headers2 processing into network layer (#369)

This PR:
- drops the `Headers2StateManager` and instead creates the state per peer in the peer run loop.
- moves `Headers2StateManager.process_headers()` to `CompressionState.process_headers()`
- upgrades `GetHeaders` to `GetHeaders2` in `send_message` when peer supports it
- decompresses `Headers2` to `Headers` before forwarding to sync layer

This simplifies the sync code and centralizes compression handling.

Author: xdustinface

dash-spv/src/network/manager.rs

@@ -10,12 +10,6 @@ use tokio::sync::{mpsc, Mutex};
 use tokio::task::JoinSet;
 use tokio::time;
 
-use async_trait::async_trait;
-use dashcore::network::constants::ServiceFlags;
-use dashcore::network::message::NetworkMessage;
-use dashcore::Network;
-use tokio_util::sync::CancellationToken;
-
 use crate::client::config::MempoolStrategy;
 use crate::client::ClientConfig;
 use crate::error::{NetworkError, NetworkResult, SpvError as Error};
@@ -29,6 +23,12 @@ use crate::network::reputation::{
 };
 use crate::network::{HandshakeManager, NetworkManager, Peer};
 use crate::types::PeerInfo;
+use async_trait::async_trait;
+use dashcore::network::constants::ServiceFlags;
+use dashcore::network::message::NetworkMessage;
+use dashcore::network::message_headers2::CompressionState;
+use dashcore::Network;
+use tokio_util::sync::CancellationToken;
 
 /// Peer network manager
 pub struct PeerNetworkManager {
@@ -71,6 +71,8 @@ pub struct PeerNetworkManager {
     exclusive_mode: bool,
     /// Cached count of currently connected peers for fast, non-blocking queries
     connected_peer_count: Arc<AtomicUsize>,
+    /// Disable headers2 after decompression failure
+    headers2_disabled: Arc<Mutex<HashSet<SocketAddr>>>,
 }
 
 impl PeerNetworkManager {
@@ -124,6 +126,7 @@ impl PeerNetworkManager {
             user_agent: config.user_agent.clone(),
             exclusive_mode,
             connected_peer_count: Arc::new(AtomicUsize::new(0)),
+            headers2_disabled: Arc::new(Mutex::new(HashSet::new())),
         })
     }
 
@@ -210,6 +213,7 @@ impl PeerNetworkManager {
         let mempool_strategy = self.mempool_strategy;
         let user_agent = self.user_agent.clone();
         let connected_peer_count = self.connected_peer_count.clone();
+        let headers2_disabled = self.headers2_disabled.clone();
 
         // Spawn connection task
         let mut tasks = self.tasks.lock().await;
@@ -249,6 +253,7 @@ impl PeerNetworkManager {
                                 shutdown_token,
                                 reputation_manager.clone(),
                                 connected_peer_count.clone(),
+                                headers2_disabled.clone(),
                             )
                             .await;
                         }
@@ -283,6 +288,7 @@ impl PeerNetworkManager {
     }
 
     /// Start reading messages from a peer
+    #[allow(clippy::too_many_arguments)] // TODO: refactor to reduce arguments
     async fn start_peer_reader(
         addr: SocketAddr,
         pool: Arc<PeerPool>,
@@ -291,10 +297,12 @@ impl PeerNetworkManager {
         shutdown_token: CancellationToken,
         reputation_manager: Arc<PeerReputationManager>,
         connected_peer_count: Arc<AtomicUsize>,
+        headers2_disabled: Arc<Mutex<HashSet<SocketAddr>>>,
     ) {
         tokio::spawn(async move {
             log::debug!("Starting peer reader loop for {}", addr);
             let mut loop_iteration = 0;
+            let mut headers2_state = CompressionState::default();
 
             loop {
                 loop_iteration += 1;
@@ -440,9 +448,49 @@ impl PeerNetworkManager {
                                 // Forward to client
                             }
                             NetworkMessage::Headers2(headers2) => {
-                                // Log compressed headers messages specifically
-                                log::info!("📨 Received Headers2 message from {} with {} compressed headers!", addr, headers2.headers.len());
-                                // Forward to client (decompression handled by sync manager)
+                                // Decompress headers in network layer and forward as regular Headers
+                                log::info!(
+                                    "Received Headers2 from {} with {} compressed headers - decompressing",
+                                    addr,
+                                    headers2.headers.len()
+                                );
+
+                                match headers2_state.process_headers(&headers2.headers) {
+                                    Ok(headers) => {
+                                        log::info!(
+                                            "Decompressed {} headers from {} - forwarding as regular Headers",
+                                            headers.len(),
+                                            addr
+                                        );
+                                        // Forward as regular Headers message
+                                        let headers_msg = NetworkMessage::Headers(headers);
+                                        if message_tx.send((addr, headers_msg)).await.is_err() {
+                                            log::warn!(
+                                                "Breaking peer reader loop for {} - failed to send decompressed headers",
+                                                addr
+                                            );
+                                            break;
+                                        }
+                                        continue; // Already sent, don't forward the original Headers2
+                                    }
+                                    Err(e) => {
+                                        log::error!(
+                                            "Headers2 decompression failed from {}: {} - disabling headers2",
+                                            addr,
+                                            e
+                                        );
+                                        headers2_disabled.lock().await.insert(addr);
+                                        // Apply reputation penalty
+                                        reputation_manager
+                                            .update_reputation(
+                                                addr,
+                                                misbehavior_scores::INVALID_MESSAGE,
+                                                "Headers2 decompression failed",
+                                            )
+                                            .await;
+                                        continue; // Don't forward corrupted message
+                                    }
+                                }
                             }
                             NetworkMessage::GetHeaders(_) => {
                                 // SPV clients don't serve headers to peers
@@ -573,6 +621,8 @@ impl PeerNetworkManager {
                 connected_peer_count.fetch_sub(1, Ordering::Relaxed);
             }
 
+            headers2_disabled.lock().await.remove(&addr);
+
             // Give small positive reputation if peer maintained long connection
             let conn_duration = Duration::from_secs(60 * loop_iteration); // Rough estimate
             if conn_duration > Duration::from_secs(3600) {
@@ -780,7 +830,8 @@ impl PeerNetworkManager {
         // For filter-related messages, we need a peer that supports compact filters
         let requires_compact_filters =
             matches!(&message, NetworkMessage::GetCFHeaders(_) | NetworkMessage::GetCFilters(_));
-        let requires_headers2 = matches!(&message, NetworkMessage::GetHeaders2(_));
+        let check_headers2 =
+            matches!(&message, NetworkMessage::GetHeaders(_) | NetworkMessage::GetHeaders2(_));
 
         let selected_peer = if requires_compact_filters {
             // Find a peer that supports compact filters
@@ -808,7 +859,7 @@ impl PeerNetworkManager {
                     ));
                 }
             }
-        } else if requires_headers2 {
+        } else if check_headers2 {
             // Prefer a peer that advertises headers2 support
             let mut current_sync_peer = self.current_sync_peer.lock().await;
             let mut selected: Option<SocketAddr> = None;
@@ -875,6 +926,25 @@ impl PeerNetworkManager {
             .find(|(a, _)| *a == selected_peer)
             .ok_or_else(|| NetworkError::ConnectionFailed("Selected peer not found".to_string()))?;
 
+        // Upgrade GetHeaders to GetHeaders2 if this specific peer supports it and not disabled
+        let peer_supports_headers2 = {
+            let peer_guard = peer.read().await;
+            peer_guard.can_request_headers2()
+        };
+        let message = match message {
+            NetworkMessage::GetHeaders(get_headers)
+                if !self.headers2_disabled.lock().await.contains(addr)
+                    && peer_supports_headers2 =>
+            {
+                log::debug!(
+                    "Upgrading GetHeaders to GetHeaders2 for peer {}: {:?}",
+                    addr,
+                    get_headers
+                );
+                NetworkMessage::GetHeaders2(get_headers)
+            }
+            other => other,
+        };
         // Reduce verbosity for common sync messages
         match &message {
             NetworkMessage::GetHeaders(_)
@@ -1069,6 +1139,7 @@ impl Clone for PeerNetworkManager {
             user_agent: self.user_agent.clone(),
             exclusive_mode: self.exclusive_mode,
             connected_peer_count: self.connected_peer_count.clone(),
+            headers2_disabled: self.headers2_disabled.clone(),
         }
     }
 }

dash-spv/src/sync/headers/manager.rs

@@ -13,7 +13,6 @@ use crate::error::{SyncError, SyncResult};
 use crate::network::NetworkManager;
 use crate::storage::StorageManager;
 use crate::sync::headers::validate_headers;
-use crate::sync::headers2::Headers2StateManager;
 use crate::types::{ChainState, HashedBlockHeader};
 use crate::ValidationMode;
 use std::sync::Arc;
@@ -51,11 +50,8 @@ pub struct HeaderSyncManager<S: StorageManager, N: NetworkManager> {
     checkpoint_manager: CheckpointManager,
     reorg_config: ReorgConfig,
     chain_state: Arc<RwLock<ChainState>>,
-    // WalletState removed - wallet functionality is now handled externally
-    headers2_state: Headers2StateManager,
     syncing_headers: bool,
     last_sync_progress: std::time::Instant,
-    headers2_failed: bool,
     // Cached flag for quick access without locking
     cached_sync_base_height: u32,
 }
@@ -83,11 +79,8 @@ impl<S: StorageManager, N: NetworkManager> HeaderSyncManager<S, N> {
             checkpoint_manager,
             reorg_config,
             chain_state,
-            // WalletState removed
-            headers2_state: Headers2StateManager::new(),
             syncing_headers: false,
             last_sync_progress: std::time::Instant::now(),
-            headers2_failed: false,
             cached_sync_base_height: 0,
             _phantom_s: std::marker::PhantomData,
             _phantom_n: std::marker::PhantomData,
@@ -324,144 +317,22 @@ impl<S: StorageManager, N: NetworkManager> HeaderSyncManager<S, N> {
             getheaders_msg.stop_hash
         );
 
-        // Use headers2 if peer supports it and we haven't had failures
-        let use_headers2 = network.has_headers2_peer().await && !self.headers2_failed;
-
         // Log details about the request
         tracing::info!(
-            "Preparing headers request - height: {}, base_hash: {:?}, headers2_supported: {}",
+            "Preparing headers request - height: {}, base_hash: {:?}",
             storage.get_tip_height().await.unwrap_or(0),
-            base_hash,
-            use_headers2
+            base_hash
         );
 
-        // Try GetHeaders2 first if peer supports it, with fallback to regular GetHeaders
-        if use_headers2 {
-            tracing::info!("📤 Sending GetHeaders2 message (compressed headers)");
-            tracing::debug!(
-                "GetHeaders2 details: version={}, locator_hashes={:?}, stop_hash={}",
-                getheaders_msg.version,
-                getheaders_msg.locator_hashes,
-                getheaders_msg.stop_hash
-            );
-
-            // Log the raw message bytes for debugging
-            let msg_bytes = dashcore::consensus::encode::serialize(&getheaders_msg);
-            tracing::debug!(
-                "GetHeaders2 raw bytes ({}): {:02x?}",
-                msg_bytes.len(),
-                &msg_bytes[..std::cmp::min(100, msg_bytes.len())]
-            );
-
-            // Send GetHeaders2 message for compressed headers
-            let result =
-                network.send_message(NetworkMessage::GetHeaders2(getheaders_msg.clone())).await;
-
-            match result {
-                Ok(_) => {
-                    // TODO: Implement timeout and fallback mechanism
-                    // For now, we rely on the network layer's timeout handling
-                    // In the future, we should:
-                    // 1. Track the request with a unique ID
-                    // 2. Set a specific timeout for GetHeaders2 response
-                    // 3. Fall back to GetHeaders if no response within timeout
-                    // 4. Mark peers that don't respond to GetHeaders2 properly
-                }
-                Err(e) => {
-                    tracing::warn!("Failed to send GetHeaders2, falling back to GetHeaders: {}", e);
-                    // Fall back to regular GetHeaders
-                    network
-                        .send_message(NetworkMessage::GetHeaders(getheaders_msg))
-                        .await
-                        .map_err(|e| {
-                            SyncError::Network(format!("Failed to send GetHeaders: {}", e))
-                        })?;
-                }
-            }
-        } else {
-            tracing::info!("📤 Sending GetHeaders message (uncompressed headers)");
-            // Send regular GetHeaders message
-            network
-                .send_message(NetworkMessage::GetHeaders(getheaders_msg))
-                .await
-                .map_err(|e| SyncError::Network(format!("Failed to send GetHeaders: {}", e)))?;
-        }
+        tracing::debug!("Sending GetHeaders message");
+        network
+            .send_message(NetworkMessage::GetHeaders(getheaders_msg))
+            .await
+            .map_err(|e| SyncError::Network(format!("Failed to send GetHeaders: {}", e)))?;
 
         Ok(())
     }
 
-    /// Handle a Headers2 message with compressed headers.
-    /// Returns true if the message was processed and sync should continue, false if sync is complete.
-    pub async fn handle_headers2_message(
-        &mut self,
-        headers2: &dashcore::network::message_headers2::Headers2Message,
-        peer_id: crate::types::PeerId,
-        storage: &mut S,
-        network: &mut N,
-    ) -> SyncResult<(bool, usize)> {
-        tracing::info!(
-            "📦 Received {} compressed headers from peer {}",
-            headers2.headers.len(),
-            peer_id
-        );
-
-        // If this is the first headers2 message, and we need to initialize compression state
-        if !headers2.headers.is_empty() {
-            // Check if we need to initialize the compression state
-            let state = self.headers2_state.get_state(peer_id);
-            if state.prev_header.is_none() {
-                // Initialize with header at current tip height (works for both genesis and later)
-                let tip_height = storage.get_tip_height().await.unwrap_or(0);
-                if let Some(tip_header) = storage.get_header(tip_height).await.map_err(|e| {
-                    SyncError::Storage(format!(
-                        "Error trying to get block at height {} from storage: {}",
-                        tip_height, e
-                    ))
-                })? {
-                    tracing::info!(
-                        "Initializing headers2 compression state for peer {} with header at height {}",
-                        peer_id,
-                        tip_height
-                    );
-                    self.headers2_state.init_peer_state(peer_id, tip_header);
-                }
-            }
-        }
-
-        // Decompress headers using the peer's compression state
-        let headers = match self.headers2_state.process_headers(peer_id, &headers2.headers) {
-            Ok(headers) => headers,
-            Err(e) => {
-                tracing::error!(
-                    "Failed to decompress headers2 from peer {}: {}. Headers count: {}, first header compressed: {}, chain height: {}",
-                    peer_id,
-                    e,
-                    headers2.headers.len(),
-                    if headers2.headers.is_empty() {
-                        "N/A (empty)".to_string()
-                    } else {
-                        (!headers2.headers[0].is_full()).to_string()
-                    },
-                    storage.get_tip_height().await.unwrap_or(0)
-                );
-
-                // Mark that headers2 failed for this sync session to trigger fallback to regular headers
-                self.headers2_failed = true;
-                return Err(SyncError::Headers2DecompressionFailed(format!(
-                    "Failed to decompress headers: {}",
-                    e
-                )));
-            }
-        };
-
-        let headers_count = headers.len();
-
-        // Process decompressed headers through the normal flow
-        let continue_sync = self.handle_headers_message(&headers, storage, network).await?;
-
-        Ok((continue_sync, headers_count))
-    }
-
     /// Prepare sync state without sending network requests.
     /// This allows monitoring to be set up before requests are sent.
     pub async fn prepare_sync(&mut self, storage: &mut S) -> SyncResult<Option<BlockHash>> {