diff --git a/dash-spv/src/network/manager.rs b/dash-spv/src/network/manager.rs index b4c2ec54b..83404697e 100644 --- a/dash-spv/src/network/manager.rs +++ b/dash-spv/src/network/manager.rs @@ -1,12 +1,12 @@ //! Peer network manager for SPV client -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; use std::path::PathBuf; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::Duration; -use tokio::sync::{broadcast, Mutex, RwLock}; +use tokio::sync::{broadcast, Mutex, OwnedSemaphorePermit, RwLock, Semaphore}; use tokio::task::JoinSet; use tokio::time; @@ -29,12 +29,130 @@ use dashcore::network::constants::ServiceFlags; use dashcore::network::message::NetworkMessage; use dashcore::network::message_blockdata::Inventory; use dashcore::network::message_headers2::CompressionState; -use dashcore::Network; +use dashcore::{BlockHash, Header, Network}; use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender}; use tokio::time::Instant; use tokio_util::sync::CancellationToken; const DEFAULT_NETWORK_EVENT_CAPACITY: usize = 10000; +const MAX_CONCURRENT_HEADERS2_DECOMPRESSIONS: usize = 4; + +fn headers2_decompression_parallelism() -> usize { + std::thread::available_parallelism() + .map(|parallelism| parallelism.get().min(MAX_CONCURRENT_HEADERS2_DECOMPRESSIONS)) + .unwrap_or(1) +} + +enum HeaderDispatchOutcome { + Compressed(Vec<(Header, BlockHash)>), + Regular(Vec
), + Invalid(String), + TaskFailed(String), +} + +/// Releases header messages in wire order because an announcement can depend on the preceding +/// message from the same peer. +struct OrderedHeaderResults { + next_sequence: u64, + completed: BTreeMap, +} + +impl Default for OrderedHeaderResults { + fn default() -> Self { + Self { + next_sequence: 0, + completed: BTreeMap::new(), + } + } +} + +impl OrderedHeaderResults { + fn complete(&mut self, sequence: u64, result: T) -> Vec { + let replaced = self.completed.insert(sequence, result); + debug_assert!(replaced.is_none(), "header sequence completed twice"); + + let mut ready = Vec::new(); + while let Some(result) = self.completed.remove(&self.next_sequence) { + ready.push(result); + self.next_sequence += 1; + } + ready + } +} + +#[derive(Clone)] +struct HeaderDispatchContext { + addr: SocketAddr, + session_token: CancellationToken, + ordered_results: + Arc>>, + message_dispatcher: Arc>, + headers2_disabled: Arc>>, + reputation_manager: Arc, +} + +impl HeaderDispatchContext { + async fn complete( + &self, + sequence: u64, + outcome: HeaderDispatchOutcome, + permit: OwnedSemaphorePermit, + ) { + if self.session_token.is_cancelled() { + return; + } + + let mut ordered_results = self.ordered_results.lock().await; + // Retain each permit while its result waits for earlier messages, keeping the reorder + // buffer within the worker limit. Keep the ordering lock through dispatch so another + // completed task cannot overtake this batch after its sequence has been released. + let ready = ordered_results.complete(sequence, (outcome, permit)); + + for (outcome, _permit) in ready { + if self.session_token.is_cancelled() { + return; + } + + match outcome { + HeaderDispatchOutcome::Compressed(headers_with_hashes) => { + tracing::info!( + "Decompressed {} headers from {} - forwarding as regular Headers", + headers_with_hashes.len(), + self.addr + ); + let message = Message::new_headers(self.addr, headers_with_hashes); + let mut message_dispatcher = self.message_dispatcher.lock().await; + if self.session_token.is_cancelled() { + return; + } + message_dispatcher.dispatch(&message); + } + HeaderDispatchOutcome::Regular(headers) => { + let message = Message::new(self.addr, NetworkMessage::Headers(headers)); + let mut message_dispatcher = self.message_dispatcher.lock().await; + if self.session_token.is_cancelled() { + return; + } + message_dispatcher.dispatch(&message); + } + HeaderDispatchOutcome::Invalid(e) => { + tracing::error!( + "Headers2 decompression failed from {}: {} - disabling headers2", + self.addr, + e + ); + self.headers2_disabled.lock().await.insert(self.addr); + self.reputation_manager + .update_reputation(self.addr, ChangeReason::Headers2DecompressionFailed) + .await; + } + HeaderDispatchOutcome::TaskFailed(e) => { + tracing::error!("Headers2 decompression task failed for {}: {}", self.addr, e); + } + } + } + } +} /// Peer network manager pub struct PeerNetworkManager { @@ -72,6 +190,8 @@ pub struct PeerNetworkManager { connected_peer_count: Arc, /// Disable headers2 after decompression failure headers2_disabled: Arc>>, + /// Global bound for CPU-heavy headers2 decompression work. + headers2_decompression_semaphore: Arc, /// Dispatcher for unbounded and message-type filtered message distribution. message_dispatcher: Arc>, /// Request queue sender, cloneable handle for sending requests to the network manager. @@ -197,6 +317,9 @@ impl PeerNetworkManager { capability_rejected: Arc::new(RwLock::new(HashMap::new())), connected_peer_count: Arc::new(AtomicUsize::new(0)), headers2_disabled: Arc::new(Mutex::new(HashSet::new())), + headers2_decompression_semaphore: Arc::new(Semaphore::new( + headers2_decompression_parallelism(), + )), message_dispatcher: Arc::new(Mutex::new(MessageDispatcher::default())), request_tx, request_rx: Arc::new(Mutex::new(Some(request_rx))), @@ -321,6 +444,7 @@ impl PeerNetworkManager { let capability_rejected = self.capability_rejected.clone(); let connected_peer_count = self.connected_peer_count.clone(); let headers2_disabled = self.headers2_disabled.clone(); + let headers2_decompression_semaphore = self.headers2_decompression_semaphore.clone(); let message_dispatcher = self.message_dispatcher.clone(); let network_event_sender = self.network_event_sender.clone(); @@ -421,6 +545,7 @@ impl PeerNetworkManager { latency, connected_peer_count.clone(), headers2_disabled.clone(), + headers2_decompression_semaphore.clone(), message_dispatcher, network_event_sender.clone(), ) @@ -503,19 +628,28 @@ impl PeerNetworkManager { latency: Arc>, connected_peer_count: Arc, headers2_disabled: Arc>>, + headers2_decompression_semaphore: Arc, message_dispatcher: Arc>, network_event_sender: broadcast::Sender, ) { tokio::spawn(async move { tracing::debug!("Starting peer reader loop for {}", addr); let mut loop_iteration = 0; - let mut headers2_state = CompressionState::default(); - + let mut header_sequence = 0_u64; + let session_token = shutdown_token.child_token(); + let header_dispatch = HeaderDispatchContext { + addr, + session_token: session_token.clone(), + ordered_results: Arc::new(Mutex::new(OrderedHeaderResults::default())), + message_dispatcher: message_dispatcher.clone(), + headers2_disabled: headers2_disabled.clone(), + reputation_manager: reputation_manager.clone(), + }; loop { loop_iteration += 1; // Check shutdown signal first with detailed logging - if shutdown_token.is_cancelled() { + if session_token.is_cancelled() { tracing::info!("Breaking peer reader loop for {} - shutdown signal received (iteration {})", addr, loop_iteration); break; } @@ -549,7 +683,7 @@ impl PeerNetworkManager { _ = tokio::time::sleep(MESSAGE_POLL_INTERVAL) => { Ok(None) }, - _ = shutdown_token.cancelled() => { + _ = session_token.cancelled() => { tracing::info!("Breaking peer reader loop for {} - shutdown signal received while reading (iteration {})", addr, loop_iteration); break; } @@ -577,6 +711,77 @@ impl PeerNetworkManager { } } + let inner = msg.into_inner(); + if let NetworkMessage::Headers2(headers2) = inner { + tracing::info!( + "Received Headers2 from {} with {} compressed headers - decompressing", + addr, + headers2.headers.len() + ); + + let permit = tokio::select! { + permit = headers2_decompression_semaphore.clone().acquire_owned() => { + match permit { + Ok(permit) => permit, + Err(_) => break, + } + } + _ = session_token.cancelled() => break, + }; + let header_dispatch = header_dispatch.clone(); + let sequence = header_sequence; + header_sequence += 1; + tokio::spawn(async move { + let result = tokio::task::spawn_blocking(move || { + let mut state = CompressionState::default(); + state.process_headers_with_hashes(&headers2.headers) + }) + .await; + if header_dispatch.session_token.is_cancelled() { + return; + } + + let outcome = match result { + Ok(Ok(headers_with_hashes)) => { + HeaderDispatchOutcome::Compressed(headers_with_hashes) + } + Ok(Err(e)) => HeaderDispatchOutcome::Invalid(e.to_string()), + Err(e) => HeaderDispatchOutcome::TaskFailed(e.to_string()), + }; + header_dispatch.complete(sequence, outcome, permit).await; + }); + continue; + } + if let NetworkMessage::Headers(headers) = inner { + tracing::info!( + "📨 Received Headers message from {} with {} headers! (regular uncompressed)", + addr, + headers.len() + ); + let peer_guard = peer.read().await; + if peer_guard.supports_headers2() { + tracing::warn!("⚠️ Peer {} supports headers2 but sent regular headers - possible protocol issue", addr); + } + drop(peer_guard); + + let permit = tokio::select! { + permit = headers2_decompression_semaphore.clone().acquire_owned() => { + match permit { + Ok(permit) => permit, + Err(_) => break, + } + } + _ = session_token.cancelled() => break, + }; + let sequence = header_sequence; + header_sequence += 1; + header_dispatch + .complete(sequence, HeaderDispatchOutcome::Regular(headers), permit) + .await; + continue; + } + let msg = Message::new(addr, inner); + // Handle some messages directly match &msg.inner() { NetworkMessage::SendAddrV2 => { @@ -673,59 +878,8 @@ impl PeerNetworkManager { } continue; } - NetworkMessage::Headers(headers) => { - // Log headers messages specifically - tracing::info!( - "📨 Received Headers message from {} with {} headers! (regular uncompressed)", - addr, - headers.len() - ); - // Check if peer supports headers2 - let peer_guard = peer.read().await; - if peer_guard.supports_headers2() { - tracing::warn!("⚠️ Peer {} supports headers2 but sent regular headers - possible protocol issue", addr); - } - drop(peer_guard); - // Forward to client - } - NetworkMessage::Headers2(headers2) => { - // Decompress headers in network layer and forward as regular Headers - tracing::info!( - "Received Headers2 from {} with {} compressed headers - decompressing", - addr, - headers2.headers.len() - ); - - match headers2_state.process_headers(&headers2.headers) { - Ok(headers) => { - tracing::info!( - "Decompressed {} headers from {} - forwarding as regular Headers", - headers.len(), - addr - ); - // Forward as regular Headers message - let headers_msg = NetworkMessage::Headers(headers); - let message = Message::new(msg.peer_address(), headers_msg); - message_dispatcher.lock().await.dispatch(&message); - continue; // Already sent, don't forward the original Headers2 - } - Err(e) => { - tracing::error!( - "Headers2 decompression failed from {}: {} - disabling headers2", - addr, - e - ); - headers2_disabled.lock().await.insert(addr); - // Apply reputation penalty - reputation_manager - .update_reputation( - addr, - ChangeReason::Headers2DecompressionFailed, - ) - .await; - continue; // Don't forward corrupted message - } - } + NetworkMessage::Headers(_) | NetworkMessage::Headers2(_) => { + unreachable!() } NetworkMessage::GetHeaders(_) => { // SPV clients don't serve headers to peers @@ -843,6 +997,7 @@ impl PeerNetworkManager { } // Remove from pool and notify consumers + session_token.cancel(); tracing::warn!("Disconnecting from {} (peer reader loop ended)", addr); Self::remove_peer_and_notify( &pool, @@ -1725,6 +1880,7 @@ impl Clone for PeerNetworkManager { capability_rejected: self.capability_rejected.clone(), connected_peer_count: self.connected_peer_count.clone(), headers2_disabled: self.headers2_disabled.clone(), + headers2_decompression_semaphore: self.headers2_decompression_semaphore.clone(), message_dispatcher: self.message_dispatcher.clone(), request_tx: self.request_tx.clone(), request_rx: self.request_rx.clone(), @@ -1925,6 +2081,9 @@ impl PeerNetworkManager { capability_rejected: Arc::new(RwLock::new(HashMap::new())), connected_peer_count: Arc::new(AtomicUsize::new(0)), headers2_disabled: Arc::new(Mutex::new(HashSet::new())), + headers2_decompression_semaphore: Arc::new(Semaphore::new( + headers2_decompression_parallelism(), + )), message_dispatcher: Arc::new(Mutex::new(MessageDispatcher::default())), request_tx, request_rx: Arc::new(Mutex::new(Some(request_rx))), @@ -2028,3 +2187,128 @@ impl PeerNetworkManager { self.evict_worst_stuck_peer().await; } } + +#[cfg(test)] +mod header_ordering_tests { + use super::*; + use crate::test_utils::test_socket_address; + + fn context( + session_token: CancellationToken, + ) -> (HeaderDispatchContext, UnboundedReceiver) { + let mut dispatcher = MessageDispatcher::default(); + let receiver = dispatcher.message_receiver(&[MessageType::Headers]); + ( + HeaderDispatchContext { + addr: test_socket_address(1), + session_token, + ordered_results: Arc::new(Mutex::new(OrderedHeaderResults::default())), + message_dispatcher: Arc::new(Mutex::new(dispatcher)), + headers2_disabled: Arc::new(Mutex::new(HashSet::new())), + reputation_manager: Arc::new(PeerReputationManager::new()), + }, + receiver, + ) + } + + #[test] + fn completed_decompressions_are_released_in_receive_order() { + let mut results = OrderedHeaderResults::default(); + + assert!(results.complete(2, "third").is_empty()); + assert_eq!(results.complete(0, "first"), ["first"]); + assert_eq!(results.complete(1, "second"), ["second", "third"]); + } + + #[tokio::test] + async fn compressed_and_regular_headers_are_dispatched_in_wire_order() { + let (context, mut receiver) = context(CancellationToken::new()); + let semaphore = Arc::new(Semaphore::new(2)); + let first = Header::dummy(1); + let first_hash = first.block_hash(); + let second = Header::dummy(2); + + context + .complete( + 1, + HeaderDispatchOutcome::Regular(vec![second]), + semaphore.clone().acquire_owned().await.unwrap(), + ) + .await; + assert!(receiver.try_recv().is_err()); + + context + .complete( + 0, + HeaderDispatchOutcome::Compressed(vec![(first, first_hash)]), + semaphore.clone().acquire_owned().await.unwrap(), + ) + .await; + + let first_message = receiver.recv().await.unwrap(); + assert_eq!(first_message.header_hashes(), Some([first_hash].as_slice())); + assert_eq!(first_message.inner(), &NetworkMessage::Headers(vec![first])); + let second_message = receiver.recv().await.unwrap(); + assert_eq!(second_message.header_hashes(), None); + assert_eq!(second_message.inner(), &NetworkMessage::Headers(vec![second])); + assert_eq!(semaphore.available_permits(), 2); + } + + #[tokio::test] + async fn ended_peer_session_discards_headers_and_releases_permit() { + let session_token = CancellationToken::new(); + let (context, mut receiver) = context(session_token.clone()); + let semaphore = Arc::new(Semaphore::new(1)); + let permit = semaphore.clone().acquire_owned().await.unwrap(); + session_token.cancel(); + + context.complete(0, HeaderDispatchOutcome::Regular(vec![Header::dummy(1)]), permit).await; + + assert!(receiver.try_recv().is_err()); + assert_eq!(semaphore.available_permits(), 1); + } + + #[tokio::test] + async fn failed_decompression_does_not_block_later_headers() { + let (context, mut receiver) = context(CancellationToken::new()); + let semaphore = Arc::new(Semaphore::new(2)); + let header = Header::dummy(1); + + context + .complete( + 1, + HeaderDispatchOutcome::Regular(vec![header]), + semaphore.clone().acquire_owned().await.unwrap(), + ) + .await; + context + .complete( + 0, + HeaderDispatchOutcome::TaskFailed("cancelled".to_owned()), + semaphore.clone().acquire_owned().await.unwrap(), + ) + .await; + + assert_eq!(receiver.recv().await.unwrap().inner(), &NetworkMessage::Headers(vec![header])); + assert_eq!(semaphore.available_permits(), 2); + } + + #[tokio::test] + async fn invalid_decompression_disables_headers2_and_penalizes_peer() { + let (context, mut receiver) = context(CancellationToken::new()); + let semaphore = Arc::new(Semaphore::new(1)); + + context + .complete( + 0, + HeaderDispatchOutcome::Invalid("invalid compressed header".to_owned()), + semaphore.clone().acquire_owned().await.unwrap(), + ) + .await; + + assert!(receiver.try_recv().is_err()); + assert!(context.headers2_disabled.lock().await.contains(&context.addr)); + assert!(context.reputation_manager.scores_for([context.addr]).await[&context.addr] > 0); + assert_eq!(semaphore.available_permits(), 1); + } +} diff --git a/dash-spv/src/network/message_dispatcher.rs b/dash-spv/src/network/message_dispatcher.rs index d88540d7d..10c3de4f6 100644 --- a/dash-spv/src/network/message_dispatcher.rs +++ b/dash-spv/src/network/message_dispatcher.rs @@ -10,6 +10,7 @@ use std::collections::{HashMap, HashSet}; use std::net::SocketAddr; use dashcore::network::message::NetworkMessage; +use dashcore::BlockHash; use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender}; use crate::network::MessageType; @@ -19,6 +20,7 @@ use crate::network::MessageType; pub struct Message { peer_address: SocketAddr, inner: NetworkMessage, + header_hashes: Option>, } impl Message { @@ -27,6 +29,20 @@ impl Message { Self { peer_address, inner, + header_hashes: None, + } + } + + /// Creates a headers message carrying hashes already computed during decompression. + pub(crate) fn new_headers( + peer_address: SocketAddr, + headers_with_hashes: Vec<(dashcore::Header, BlockHash)>, + ) -> Self { + let (headers, hashes) = headers_with_hashes.into_iter().unzip(); + Self { + peer_address, + inner: NetworkMessage::Headers(headers), + header_hashes: Some(hashes), } } @@ -44,6 +60,16 @@ impl Message { pub fn inner(&self) -> &NetworkMessage { &self.inner } + + /// Consumes the wrapper and returns its network message. + pub(crate) fn into_inner(self) -> NetworkMessage { + self.inner + } + + /// Returns cached hashes for a decompressed headers2 message. + pub(crate) fn header_hashes(&self) -> Option<&[BlockHash]> { + self.header_hashes.as_deref() + } } /// Routes incoming network messages to subscribers based on message type. @@ -96,6 +122,17 @@ mod tests { assert_eq!(*msg.inner(), inner); } + #[test] + fn test_headers_message_carries_precomputed_hashes() { + let peer_address = test_socket_address(1); + let header = dashcore::Header::dummy(1); + let hash = header.block_hash(); + let msg = Message::new_headers(peer_address, vec![(header, hash)]); + + assert_eq!(msg.header_hashes(), Some([hash].as_slice())); + assert_eq!(msg.inner(), &NetworkMessage::Headers(vec![header])); + } + #[tokio::test] async fn test_dispatch_to_interested_receiver() { let mut message_dispatcher = MessageDispatcher::default(); diff --git a/dash-spv/src/sync/block_headers/manager.rs b/dash-spv/src/sync/block_headers/manager.rs index b3699e979..7c76828af 100644 --- a/dash-spv/src/sync/block_headers/manager.rs +++ b/dash-spv/src/sync/block_headers/manager.rs @@ -19,6 +19,7 @@ use crate::sync::block_headers::HeadersPipeline; use crate::sync::{BlockHeadersProgress, ProgressPercentage, SyncEvent, SyncManager, SyncState}; use crate::types::HashedBlockHeader; use crate::validation::{BlockHeaderValidator, Validator}; +#[cfg(test)] use dashcore::block::Header; use dashcore::network::message_blockdata::Inventory; use dashcore::BlockHash; @@ -122,7 +123,7 @@ impl BlockHeadersManager { /// Handle incoming headers message (used for both initial sync and post-sync). pub(super) async fn handle_headers_pipeline( &mut self, - headers: &[Header], + headers: &[HashedBlockHeader], requests: &RequestSender, ) -> SyncResult> { if !self.pipeline.is_initialized() { @@ -140,12 +141,16 @@ impl BlockHeadersManager { if matched.is_none() && !headers.is_empty() { tracing::debug!( "Headers not matched by pipeline (prev_hash: {}), may be post-sync update", - headers[0].prev_blockhash + headers[0].header().prev_blockhash ); } - // Send more requests during initial sync or active post-sync catch-up. - // Skip for unsolicited headers. + // Process ready-to-store segments + let mut events = Vec::new(); + let ready_batches = self.pipeline.take_ready_to_store(); + + // Draining can expose new segments at the end of the active window, so + // refill only after next_to_store has advanced. Skip unsolicited headers. if was_syncing || !tip_was_complete { let sent = self.pipeline.send_pending(requests)?; if sent > 0 { @@ -153,10 +158,6 @@ impl BlockHeadersManager { } } - // Process ready-to-store segments - let mut events = Vec::new(); - let ready_batches = self.pipeline.take_ready_to_store(); - for (_start_height, batch_headers) in ready_batches { if !batch_headers.is_empty() { // Validate chain continuity with current tip @@ -262,6 +263,8 @@ mod tests { use crate::storage::{ DiskStorageManager, PersistentBlockHeaderStorage, PersistentMetadataStorage, StorageManager, }; + use crate::sync::block_headers::pipeline::ACTIVE_SEGMENT_WINDOW; + use crate::sync::block_headers::segment_state::SegmentState; use crate::sync::{ManagerIdentifier, SyncManager, SyncManagerProgress}; use dashcore::network::message::NetworkMessage; use tokio::sync::mpsc::unbounded_channel; @@ -339,7 +342,7 @@ mod tests { } #[tokio::test] - async fn test_unsolicited_post_sync_header_does_not_trigger_get_headers() { + async fn test_unsolicited_post_sync_header_batch_does_not_trigger_get_headers() { let mut manager = create_test_manager().await; let tip = manager.tip().await.unwrap(); let tip_hash = *tip.hash(); @@ -351,16 +354,19 @@ mod tests { let (sender, mut rx) = create_test_request_sender(); - let header = Header::dummy_chain(1, tip_hash).remove(0); + let headers = Header::dummy_chain(2, tip_hash) + .into_iter() + .map(HashedBlockHeader::from) + .collect::>(); - let events = manager.handle_headers_pipeline(&[header], &sender).await.unwrap(); + let events = manager.handle_headers_pipeline(&headers, &sender).await.unwrap(); // Header should have been stored assert_eq!(events.len(), 1); assert!(matches!( events[0], SyncEvent::BlockHeadersStored { - tip_height: 1 + tip_height: 2 } )); @@ -371,6 +377,61 @@ mod tests { assert!(manager.pipeline.is_tip_complete()); } + #[tokio::test] + async fn test_response_cycle_refills_window_after_ordered_drain() { + let mut manager = create_test_manager().await; + let stored_tip = manager.tip().await.unwrap(); + let chain = Header::dummy_chain(ACTIVE_SEGMENT_WINDOW, *stored_tip.hash()); + + let mut segments = Vec::new(); + for (id, header) in chain.iter().enumerate() { + let start_hash = if id == 0 { + *stored_tip.hash() + } else { + chain[id - 1].block_hash() + }; + let target_hash = header.block_hash(); + let mut segment = SegmentState::new( + id, + id as u32, + start_hash, + Some(id as u32 + 1), + Some(target_hash), + ); + if id == 0 { + segment.coordinator.mark_sent(&[start_hash]); + } else { + segment.current_tip_hash = target_hash; + segment.current_height = id as u32 + 1; + segment.complete = true; + segment.buffered_headers.push((*header).into()); + } + segments.push(segment); + } + let next_locator = chain.last().expect("non-empty chain").block_hash(); + segments.push(SegmentState::new( + ACTIVE_SEGMENT_WINDOW, + ACTIVE_SEGMENT_WINDOW as u32, + next_locator, + None, + None, + )); + manager.pipeline.set_segments_for_test(segments); + manager.progress.set_state(SyncState::Syncing); + + let (requests, mut rx) = create_test_request_sender(); + let events = manager.handle_headers_pipeline(&[chain[0].into()], &requests).await.unwrap(); + + assert_eq!(events.len(), ACTIVE_SEGMENT_WINDOW); + match rx.try_recv().expect("response cycle did not refill active window") { + NetworkRequest::SendMessage(NetworkMessage::GetHeaders(request)) => { + assert_eq!(request.locator_hashes[0], next_locator); + } + other => panic!("Expected GetHeaders, got {other:?}"), + } + assert!(rx.try_recv().is_err()); + } + #[tokio::test] async fn test_peer_tip_announcement_lifecycle() { let mut manager = create_synced_manager().await; @@ -462,7 +523,7 @@ mod tests { // segment's current_tip_hash to advanced_hash. let header = Header::dummy_chain(1, initial_locator).remove(0); let advanced_hash = header.block_hash(); - manager.handle_headers_pipeline(&[header], &requests).await.unwrap(); + manager.handle_headers_pipeline(&[header.into()], &requests).await.unwrap(); // Drain the follow-up GetHeaders that send_pending issued. match rx.try_recv().expect("follow-up GetHeaders not sent") { diff --git a/dash-spv/src/sync/block_headers/pipeline.rs b/dash-spv/src/sync/block_headers/pipeline.rs index cce5baca9..286816ddf 100644 --- a/dash-spv/src/sync/block_headers/pipeline.rs +++ b/dash-spv/src/sync/block_headers/pipeline.rs @@ -6,6 +6,7 @@ use std::sync::Arc; +#[cfg(test)] use dashcore::block::Header; use dashcore::BlockHash; @@ -15,6 +16,9 @@ use crate::network::RequestSender; use crate::sync::block_headers::segment_state::SegmentState; use crate::types::HashedBlockHeader; +/// Keep enough checkpoint segments active to saturate decompression while bounding buffered data. +pub(super) const ACTIVE_SEGMENT_WINDOW: usize = 8; + /// Pipeline for parallel header downloads across checkpoint-defined segments. /// /// Divides the blockchain into segments based on checkpoints and downloads @@ -121,7 +125,8 @@ impl HeadersPipeline { /// Returns the number of requests sent. pub fn send_pending(&mut self, requests: &RequestSender) -> SyncResult { let mut sent = 0; - for segment in &mut self.segments { + let window_end = (self.next_to_store + ACTIVE_SEGMENT_WINDOW).min(self.segments.len()); + for segment in &mut self.segments[self.next_to_store..window_end] { // Skip completed segments if segment.complete { continue; @@ -137,7 +142,7 @@ impl HeadersPipeline { /// Try to match incoming headers to the correct segment. /// Returns the segment index if matched, or None if headers don't belong to any segment. /// Returns an error if checkpoint validation fails. - pub fn receive_headers(&mut self, headers: &[Header]) -> SyncResult> { + pub fn receive_headers(&mut self, headers: &[HashedBlockHeader]) -> SyncResult> { if headers.is_empty() { // Empty response means the peer has no more headers after our locator. // Route to the tip segment (target_height is None) if it has in-flight requests. @@ -159,7 +164,7 @@ impl HeadersPipeline { return Ok(None); } - let prev_hash = headers[0].prev_blockhash; + let prev_hash = headers[0].header().prev_blockhash; // Find the segment that matches for (idx, segment) in self.segments.iter_mut().enumerate() { @@ -175,7 +180,9 @@ impl HeadersPipeline { if segment.complete && segment.target_height.is_none() { segment.complete = false; self.next_to_store = idx; - // Mark as in-flight so the coordinator accepts these unsolicited headers + // A headers announcement may contain multiple consecutive headers. The + // coordinator tracks the response by its first previous hash, just like a + // requested headers batch. segment.coordinator.mark_sent(&[prev_hash]); tracing::debug!( "Tip segment {} receiving post-sync headers, reset for continued processing", @@ -189,7 +196,7 @@ impl HeadersPipeline { // Check if these are duplicate headers from another peer. The first // header's hash matches a segment's current tip, meaning we already have it. - let first_hash = headers[0].block_hash(); + let first_hash = *headers[0].hash(); if self.segments.iter().any(|s| s.current_tip_hash == first_hash) { tracing::debug!("Ignoring duplicate header {} from another peer", first_hash); return Ok(None); @@ -333,6 +340,13 @@ impl HeadersPipeline { .find(|s| s.target_height.is_none()) .is_some_and(|s| !s.complete && s.coordinator.active_count() > 0) } + + #[cfg(test)] + pub(super) fn set_segments_for_test(&mut self, segments: Vec) { + self.segments = segments; + self.next_to_store = 0; + self.initialized = true; + } } #[cfg(test)] @@ -409,8 +423,14 @@ mod tests { let sent = pipeline.send_pending(&sender).unwrap(); - // Should send at least one request per segment - assert!(sent >= pipeline.segment_count()); + assert_eq!(sent, ACTIVE_SEGMENT_WINDOW.min(pipeline.segment_count())); + + assert!(pipeline.segments[..sent] + .iter() + .all(|segment| segment.coordinator.active_count() == 1)); + assert!(pipeline.segments[sent..] + .iter() + .all(|segment| segment.coordinator.active_count() == 0)); // Verify messages were queued let mut count = 0; @@ -418,6 +438,16 @@ mod tests { count += 1; } assert_eq!(count, sent); + + if pipeline.segment_count() > ACTIVE_SEGMENT_WINDOW { + pipeline.segments[0].complete = true; + pipeline.next_to_store = 1; + + assert_eq!(pipeline.send_pending(&sender).unwrap(), 1); + assert_eq!(pipeline.segments[ACTIVE_SEGMENT_WINDOW].coordinator.active_count(), 1); + assert!(rx.try_recv().is_ok()); + assert!(rx.try_recv().is_err()); + } } #[test] @@ -464,7 +494,7 @@ mod tests { let mut header = Header::dummy(1); header.prev_blockhash = tip_hash; - let matched = pipeline.receive_headers(&[header]).unwrap(); + let matched = pipeline.receive_headers(&[header.into()]).unwrap(); assert_eq!(matched, Some(0), "Tip segment should accept unsolicited post-sync headers"); assert!(!pipeline.segments[0].complete, "Tip segment should be reset to non-complete"); @@ -472,6 +502,76 @@ mod tests { assert_eq!(pipeline.segments[0].current_height, 1001); } + #[test] + fn test_completed_tip_accepts_unsolicited_header_batch() { + let tip_hash = BlockHash::dummy(99); + let mut tip = SegmentState::new(0, 1000, tip_hash, None, None); + tip.complete = true; + + let cm = create_test_checkpoint_manager(true); + let mut pipeline = HeadersPipeline::new(cm); + pipeline.initialized = true; + pipeline.next_to_store = 1; + pipeline.segments = vec![tip]; + + let headers = Header::dummy_chain(2, tip_hash) + .into_iter() + .map(HashedBlockHeader::from) + .collect::>(); + let matched = pipeline.receive_headers(&headers).unwrap(); + + assert_eq!(matched, Some(0)); + assert!(!pipeline.segments[0].complete); + assert_eq!(pipeline.segments[0].current_tip_hash, *headers[1].hash()); + assert_eq!(pipeline.segments[0].current_height, 1002); + assert_eq!(pipeline.segments[0].buffered_headers, headers); + assert_eq!(pipeline.segments[0].coordinator.active_count(), 0); + assert_eq!(pipeline.next_to_store, 0); + } + + #[test] + fn test_draining_active_window_exposes_next_segment_for_refill() { + let mut segments = Vec::new(); + for id in 0..ACTIVE_SEGMENT_WINDOW { + let start_hash = BlockHash::dummy(id as u32); + let mut segment = + SegmentState::new(id, id as u32, start_hash, Some(id as u32 + 1), None); + segment.complete = true; + segment.buffered_headers.push(Header::dummy(id as u32).into()); + segments.push(segment); + } + segments.push(SegmentState::new( + ACTIVE_SEGMENT_WINDOW, + ACTIVE_SEGMENT_WINDOW as u32, + BlockHash::dummy(ACTIVE_SEGMENT_WINDOW as u32), + None, + None, + )); + + let cm = create_test_checkpoint_manager(true); + let mut pipeline = HeadersPipeline::new(cm); + pipeline.set_segments_for_test(segments); + let (sender, mut rx) = create_test_request_sender(); + + assert_eq!(pipeline.send_pending(&sender).unwrap(), 0); + assert_eq!(pipeline.take_ready_to_store().len(), ACTIVE_SEGMENT_WINDOW); + assert_eq!(pipeline.next_to_store, ACTIVE_SEGMENT_WINDOW); + assert_eq!(pipeline.send_pending(&sender).unwrap(), 1); + + match rx.try_recv().expect("newly exposed segment was not requested") { + NetworkRequest::SendMessage( + dashcore::network::message::NetworkMessage::GetHeaders(request), + ) => { + assert_eq!( + request.locator_hashes[0], + BlockHash::dummy(ACTIVE_SEGMENT_WINDOW as u32) + ) + } + other => panic!("Expected GetHeaders, got {other:?}"), + } + assert!(rx.try_recv().is_err()); + } + #[test] fn test_completed_segment_does_not_steal_next_segment_headers() { // Create two segments which share the checkpoint hash boundary. @@ -502,7 +602,7 @@ mod tests { pipeline.segments[1].coordinator.mark_sent(&[shared_hash]); // Route headers should go to segment 1, not the completed segment 0 - let matched = pipeline.receive_headers(&[header]).unwrap(); + let matched = pipeline.receive_headers(&[header.into()]).unwrap(); assert_eq!(matched, Some(1), "Headers should route to segment 1, not completed segment 0"); // Segment 0 should still have no extra buffered headers @@ -530,7 +630,7 @@ mod tests { // Another peer sends the same header (prev_blockhash is old tip, first // header hash matches the segment's current tip) - let matched = pipeline.receive_headers(&[first_header]).unwrap(); + let matched = pipeline.receive_headers(&[first_header.into()]).unwrap(); assert_eq!(matched, None, "Duplicate headers should be silently ignored"); assert!(pipeline.segments[0].buffered_headers.is_empty()); } @@ -553,7 +653,7 @@ mod tests { mid.coordinator.mark_sent(&[shared_hash]); let mut mid_header = Header::dummy(2); mid_header.prev_blockhash = shared_hash; - mid.receive_headers(&[mid_header]).unwrap(); + mid.receive_headers(&[mid_header.into()]).unwrap(); let mid_preserved_tip = mid.current_tip_hash; let mid_preserved_height = mid.current_height; let mid_preserved_buffered = mid.buffered_headers.len(); diff --git a/dash-spv/src/sync/block_headers/segment_state.rs b/dash-spv/src/sync/block_headers/segment_state.rs index 9829d1250..58a332223 100644 --- a/dash-spv/src/sync/block_headers/segment_state.rs +++ b/dash-spv/src/sync/block_headers/segment_state.rs @@ -2,7 +2,7 @@ use crate::error::{SyncError, SyncResult}; use crate::network::RequestSender; use crate::sync::download_coordinator::{DownloadConfig, DownloadCoordinator}; use crate::types::HashedBlockHeader; -use dashcore::{BlockHash, Header}; +use dashcore::BlockHash; use std::time::Duration; /// Timeout for header requests. @@ -85,7 +85,7 @@ impl SegmentState { /// Process received headers for this segment. /// Returns the number of headers processed, or an error if checkpoint validation fails. - pub(super) fn receive_headers(&mut self, headers: &[Header]) -> SyncResult { + pub(super) fn receive_headers(&mut self, headers: &[HashedBlockHeader]) -> SyncResult { if headers.is_empty() { // Empty response means we've reached the peer's tip for this segment self.complete = true; @@ -110,7 +110,7 @@ impl SegmentState { } // Mark the request as received, reject if we never requested this hash - let prev_hash = headers[0].prev_blockhash; + let prev_hash = headers[0].header().prev_blockhash; if !self.coordinator.receive(&prev_hash) { return Err(SyncError::InvalidState(format!( "Segment {}: received unrequested headers (prev_hash {})", @@ -120,8 +120,7 @@ impl SegmentState { // Process headers let mut processed = 0; - for header in headers { - let hashed = HashedBlockHeader::from(*header); + for hashed in headers { let hash = *hashed.hash(); let height = self.current_height + processed as u32 + 1; @@ -135,7 +134,7 @@ impl SegmentState { self.segment_id, target_height ); - self.buffered_headers.push(hashed); + self.buffered_headers.push(hashed.clone()); processed += 1; self.complete = true; break; @@ -155,13 +154,13 @@ impl SegmentState { } } - self.buffered_headers.push(hashed); + self.buffered_headers.push(hashed.clone()); processed += 1; } // Update current tip for next request if processed > 0 { - self.current_tip_hash = headers[processed - 1].block_hash(); + self.current_tip_hash = *headers[processed - 1].hash(); self.current_height += processed as u32; } @@ -265,7 +264,7 @@ mod tests { let mut first = headers[0]; first.prev_blockhash = hash; - let processed = segment.receive_headers(&[first]).unwrap(); + let processed = segment.receive_headers(&[first.into()]).unwrap(); assert_eq!(processed, 1); assert_eq!(segment.buffered_headers.len(), 1); @@ -292,7 +291,7 @@ mod tests { assert_ne!(*actual_hash, expected_checkpoint_hash); // Receiving this header should fail with a validation error - let result = segment.receive_headers(&[header]); + let result = segment.receive_headers(&[header.into()]); assert!(result.is_err()); let err = result.unwrap_err(); @@ -323,7 +322,7 @@ mod tests { segment.coordinator.mark_sent(&[start_hash]); // Receiving this header should succeed and complete the segment - let result = segment.receive_headers(&[header]); + let result = segment.receive_headers(&[header.into()]); assert!(result.is_ok()); assert_eq!(result.unwrap(), 1); @@ -340,7 +339,7 @@ mod tests { let mut header = Header::dummy(1); header.prev_blockhash = start_hash; - let result = segment.receive_headers(&[header]); + let result = segment.receive_headers(&[header.into()]); assert!(result.is_err()); match result.unwrap_err() { SyncError::InvalidState(msg) => { @@ -365,7 +364,7 @@ mod tests { header.prev_blockhash = start_hash; // Completed segment should return an invalid state error - let result = segment.receive_headers(&[header]); + let result = segment.receive_headers(&[header.into()]); assert!(result.is_err()); match result.unwrap_err() { SyncError::InvalidState(msg) => { @@ -384,7 +383,7 @@ mod tests { let mut header = Header::dummy(1); header.prev_blockhash = start_hash; - segment.receive_headers(&[header]).unwrap(); + segment.receive_headers(&[header.into()]).unwrap(); let preserved_tip_hash = segment.current_tip_hash; let preserved_height = segment.current_height; diff --git a/dash-spv/src/sync/block_headers/sync_manager.rs b/dash-spv/src/sync/block_headers/sync_manager.rs index 600ee58e8..bb4b56d3a 100644 --- a/dash-spv/src/sync/block_headers/sync_manager.rs +++ b/dash-spv/src/sync/block_headers/sync_manager.rs @@ -104,8 +104,27 @@ impl SyncManager for BlockHeadersMana ) -> SyncResult> { match msg.inner() { NetworkMessage::Headers(headers) => { + let hashed_headers: Vec = + if let Some(hashes) = msg.header_hashes() { + if hashes.len() != headers.len() { + return Err(crate::error::SyncError::InvalidState(format!( + "headers2 hash count mismatch: {} headers, {} hashes", + headers.len(), + hashes.len() + ))); + } + headers + .iter() + .zip(hashes) + .map(|(header, hash)| { + crate::types::HashedBlockHeader::with_trusted_hash(*header, *hash) + }) + .collect() + } else { + headers.iter().map(crate::types::HashedBlockHeader::from).collect() + }; // Always route through pipeline when initialized - self.handle_headers_pipeline(headers, requests).await + self.handle_headers_pipeline(&hashed_headers, requests).await } NetworkMessage::Inv(inv) => { diff --git a/dash-spv/src/types.rs b/dash-spv/src/types.rs index ff14a9449..679af2fec 100644 --- a/dash-spv/src/types.rs +++ b/dash-spv/src/types.rs @@ -29,9 +29,9 @@ pub struct HashedBlockHeader { impl HashedBlockHeader { /// Pair a header with a known, trusted hash without recomputing it from the header. /// - /// The caller guarantees the hash is correct. This is used to anchor the chain at a - /// checkpoint whose hash is trusted but whose exact header bytes (the block version) - /// are not stored, so hashing the reconstructed header would not reproduce it. + /// The caller guarantees the hash is correct. This is used for hashes computed by the + /// network decoder and for checkpoint anchors whose reconstructed header bytes would not + /// reproduce the trusted checkpoint hash. pub(crate) fn with_trusted_hash(header: BlockHeader, hash: BlockHash) -> Self { Self { header, diff --git a/dash-spv/tests/header_dispatch_order.rs b/dash-spv/tests/header_dispatch_order.rs new file mode 100644 index 000000000..552db58b7 --- /dev/null +++ b/dash-spv/tests/header_dispatch_order.rs @@ -0,0 +1,147 @@ +use std::net::SocketAddr; +use std::sync::mpsc; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use dash_spv::client::ClientConfig; +use dash_spv::network::{MessageType, NetworkManager, PeerNetworkManager}; +use dashcore::consensus::encode::serialize; +use dashcore::network::address::Address; +use dashcore::network::constants::{ServiceFlags, PROTOCOL_VERSION}; +use dashcore::network::message::{NetworkMessage, RawNetworkMessage}; +use dashcore::network::message_headers2::{CompressionState, Headers2Message}; +use dashcore::network::message_network::VersionMessage; +use dashcore::{Header, Network}; +use tempfile::TempDir; +use tokio::io::AsyncWriteExt; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::oneshot; +use tokio::time::{sleep, timeout}; + +async fn send_message(stream: &mut TcpStream, message: NetworkMessage) { + let raw = RawNetworkMessage { + magic: Network::Regtest.magic(), + payload: message, + }; + stream.write_all(&serialize(&raw)).await.unwrap(); +} + +fn version_message(peer: SocketAddr) -> VersionMessage { + let local = "127.0.0.1:0".parse().unwrap(); + VersionMessage { + version: PROTOCOL_VERSION, + services: ServiceFlags::NODE_HEADERS_COMPRESSED, + timestamp: SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() as i64, + receiver: Address::new(&peer, ServiceFlags::NETWORK), + sender: Address::new(&local, ServiceFlags::NODE_HEADERS_COMPRESSED), + nonce: 1, + user_agent: "/header-order-test:0.1/".to_owned(), + start_height: 2, + relay: false, + mn_auth_challenge: [0; 32], + masternode_connection: false, + } +} + +#[test] +fn peer_reader_preserves_compressed_then_regular_header_order() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .max_blocking_threads(1) + .build() + .unwrap(); + + runtime.block_on(async { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let peer_addr = listener.local_addr().unwrap(); + let (send_headers_tx, send_headers_rx) = oneshot::channel(); + let (headers_sent_tx, headers_sent_rx) = oneshot::channel(); + let (close_tx, close_rx) = oneshot::channel(); + + let first = Header::dummy(1); + let second = Header::dummy(2); + let mut compression = CompressionState::default(); + let compressed = compression.compress(&first); + + let server = tokio::spawn(async move { + let (mut stream, client_addr) = listener.accept().await.unwrap(); + send_message(&mut stream, NetworkMessage::Version(version_message(client_addr))).await; + send_message(&mut stream, NetworkMessage::Verack).await; + + send_headers_rx.await.unwrap(); + let headers2 = NetworkMessage::Headers2(Headers2Message::new(vec![compressed])); + let regular = NetworkMessage::Headers(vec![second]); + let mut messages = serialize(&RawNetworkMessage { + magic: Network::Regtest.magic(), + payload: headers2, + }); + messages.extend(serialize(&RawNetworkMessage { + magic: Network::Regtest.magic(), + payload: regular, + })); + messages.extend(serialize(&RawNetworkMessage { + magic: Network::Regtest.magic(), + payload: NetworkMessage::Inv(vec![]), + })); + stream.write_all(&messages).await.unwrap(); + headers_sent_tx.send(()).unwrap(); + + let _ = close_rx.await; + }); + + let storage = TempDir::new().unwrap(); + let mut config = ClientConfig::new(Network::Regtest); + config.storage_path = storage.path().to_path_buf(); + config.max_peers = 1; + config.peers = vec![peer_addr]; + config.restrict_to_configured_peers = true; + config.enable_filters = false; + config.enable_masternodes = false; + + let mut manager = PeerNetworkManager::new(&config).await.unwrap(); + let mut headers = manager.message_receiver(&[MessageType::Headers]).await; + let mut marker = manager.message_receiver(&[MessageType::Inv]).await; + manager.start().await.unwrap(); + timeout(Duration::from_secs(5), async { + while manager.peer_count() == 0 { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("peer handshake did not complete"); + + // Occupy the only blocking worker before the peer sends both messages. Headers2 + // decompression must queue while the reader handles the following regular message. + let (blocking_started_tx, blocking_started_rx) = oneshot::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let blocker = tokio::task::spawn_blocking(move || { + blocking_started_tx.send(()).unwrap(); + release_rx.recv().unwrap(); + }); + blocking_started_rx.await.unwrap(); + + send_headers_tx.send(()).unwrap(); + headers_sent_rx.await.unwrap(); + timeout(Duration::from_secs(2), marker.recv()) + .await + .expect("peer reader did not reach the marker after both header messages") + .unwrap(); + assert!( + headers.try_recv().is_err(), + "regular headers overtook the blocked headers2 message" + ); + + release_tx.send(()).unwrap(); + blocker.await.unwrap(); + + let first_message = timeout(Duration::from_secs(2), headers.recv()).await.unwrap().unwrap(); + assert_eq!(first_message.inner(), &NetworkMessage::Headers(vec![first])); + + let second_message = + timeout(Duration::from_secs(2), headers.recv()).await.unwrap().unwrap(); + assert_eq!(second_message.inner(), &NetworkMessage::Headers(vec![second])); + + manager.shutdown().await; + close_tx.send(()).unwrap(); + server.await.unwrap(); + }); +} diff --git a/dash/src/network/message_headers2.rs b/dash/src/network/message_headers2.rs index 47e76743d..8ed995dc2 100644 --- a/dash/src/network/message_headers2.rs +++ b/dash/src/network/message_headers2.rs @@ -272,6 +272,8 @@ pub struct CompressionState { pub version_cache: Vec, /// Previous header for delta encoding pub prev_header: Option
, + /// Cached X11 hash paired with the exact `prev_header` it was computed from. + prev_header_hash: Option<(Header, BlockHash)>, } impl CompressionState { @@ -280,6 +282,7 @@ impl CompressionState { Self { version_cache: Vec::with_capacity(MAX_VERSION_CACHE_SIZE), prev_header: None, + prev_header_hash: None, } } @@ -345,6 +348,7 @@ impl CompressionState { }; self.prev_header = Some(*header); + self.prev_header_hash = None; CompressedHeader { flags, @@ -363,10 +367,10 @@ impl CompressionState { /// Version offset decoding (matching C++ DIP-0025): /// - offset = 0: version NOT in cache, read full version from message /// - offset = 1-7: version at position (offset-1) in cache - pub fn decompress( + pub fn decompress_with_hash( &mut self, compressed: &CompressedHeader, - ) -> Result { + ) -> Result<(Header, BlockHash), DecompressionError> { // Version (C++ semantics) let version = match compressed.flags.version_offset() { 0 => { @@ -389,7 +393,7 @@ impl CompressionState { let prev_blockhash = if let Some(hash) = compressed.prev_blockhash { hash } else { - self.prev_header.as_ref().ok_or(DecompressionError::MissingPreviousHeader)?.block_hash() + self.previous_header_hash().ok_or(DecompressionError::MissingPreviousHeader)? }; // Timestamp @@ -417,23 +421,48 @@ impl CompressionState { nonce: compressed.nonce, }; + let hash = header.block_hash(); self.prev_header = Some(header); + self.prev_header_hash = Some((header, hash)); + + Ok((header, hash)) + } - Ok(header) + /// Decompress a header without retaining its computed X11 hash. + pub fn decompress( + &mut self, + compressed: &CompressedHeader, + ) -> Result { + self.decompress_with_hash(compressed).map(|(header, _)| header) } pub fn process_headers( &mut self, headers: &[CompressedHeader], ) -> Result, ProcessError> { + self.process_headers_with_hashes(headers) + .map(|headers| headers.into_iter().map(|(header, _)| header).collect()) + } + + /// Decompress a self-contained headers2 batch and retain each computed X11 hash. + pub fn process_headers_with_hashes( + &mut self, + headers: &[CompressedHeader], + ) -> Result, ProcessError> { + // Dash Core resets compression history for every headers2 message. + self.version_cache.clear(); + self.prev_header = None; + self.prev_header_hash = None; + if headers.is_empty() { return Ok(Vec::new()); } let mut decompressed = Vec::with_capacity(headers.len()); for (i, compressed) in headers.iter().enumerate() { - let header = - self.decompress(compressed).map_err(|e| ProcessError::DecompressionError(i, e))?; + let header = self + .decompress_with_hash(compressed) + .map_err(|e| ProcessError::DecompressionError(i, e))?; decompressed.push(header); } @@ -471,10 +500,16 @@ impl CompressionState { /// Check if the given hash matches the hash of the previous header fn is_sequential(&self, prev_hash: &BlockHash) -> bool { - if let Some(prev) = &self.prev_header { - prev.block_hash() == *prev_hash - } else { - false + self.previous_header_hash() == Some(*prev_hash) + } + + fn previous_header_hash(&self) -> Option { + match (self.prev_header.as_ref(), self.prev_header_hash.as_ref()) { + (Some(previous), Some((cached_header, cached_hash))) if previous == cached_header => { + Some(*cached_hash) + } + (Some(previous), _) => Some(previous.block_hash()), + (None, _) => None, } } } @@ -827,6 +862,71 @@ mod tests { assert_eq!(decompressed[1], header2); } + #[test] + fn test_process_headers_with_hashes() { + let headers = create_test_chain(10); + let mut compress_state = CompressionState::new(); + let compressed: Vec<_> = + headers.iter().map(|header| compress_state.compress(header)).collect(); + + let mut decompress_state = CompressionState::new(); + let decompressed = decompress_state + .process_headers_with_hashes(&compressed) + .expect("valid headers2 batch"); + + for ((decoded, hash), expected) in decompressed.iter().zip(&headers) { + assert_eq!(decoded, expected); + assert_eq!(*hash, expected.block_hash()); + } + } + + #[test] + fn test_cached_hash_is_used_only_for_its_header() { + let original = create_test_header(1, 0); + let mut compress_state = CompressionState::new(); + let compressed = compress_state.compress(&original); + + let mut state = CompressionState::new(); + let (_, original_hash) = + state.decompress_with_hash(&compressed).expect("valid compressed header"); + + let mut following = create_test_header(2, 1); + following.prev_blockhash = original_hash; + + let mut replaced = state.clone(); + replaced.prev_header = Some(create_test_header(99, 98)); + assert_eq!( + replaced.compress(&following).prev_blockhash, + Some(original_hash), + "a hash cached for a replaced previous header must not suppress prev_blockhash" + ); + + state.prev_header = None; + assert_eq!( + state.compress(&following).prev_blockhash, + Some(original_hash), + "a hash cached for a cleared previous header must not suppress prev_blockhash" + ); + } + + #[test] + fn test_process_headers_resets_state_between_messages() { + let headers = create_test_chain(2); + let mut compress_state = CompressionState::new(); + let compressed: Vec<_> = + headers.iter().map(|header| compress_state.compress(header)).collect(); + + let mut decompress_state = CompressionState::new(); + decompress_state.process_headers(&compressed).expect("valid headers2 batch"); + + decompress_state.process_headers(&[]).expect("valid empty headers2 batch"); + assert!(decompress_state.decompress(&compressed[1]).is_err()); + + // A compressed continuation cannot use state retained from the prior message. + let result = decompress_state.process_headers(&compressed[1..]); + assert!(matches!(result, Err(ProcessError::DecompressionError(0, _)))); + } + #[test] fn test_headers2_message_capacity_overflow() { use crate::consensus::encode::deserialize;