1mod affinity;
30
31use crate::config::*;
32
33use affinity::AffinityFilter;
34use codec::{Compact, Decode, Encode, MaxEncodedLen};
35use futures::{
36 channel::oneshot,
37 future::{pending, FusedFuture},
38 prelude::*,
39 stream::FuturesUnordered,
40};
41use governor::{
42 clock::DefaultClock,
43 state::{InMemoryState, NotKeyed},
44 Quota, RateLimiter,
45};
46use prometheus_endpoint::{
47 exponential_buckets, register, Counter, Gauge, GaugeVec, Histogram, HistogramOpts, Opts,
48 PrometheusError, Registry, U64,
49};
50use rand::seq::IteratorRandom;
51use sc_network::{
52 config::{NonReservedPeerMode, SetConfig},
53 error, multiaddr,
54 peer_store::PeerStoreProvider,
55 service::{
56 traits::{NotificationEvent, NotificationService, ValidationResult},
57 NotificationMetrics,
58 },
59 types::ProtocolName,
60 utils::{interval, LruHashSet},
61 NetworkBackend, NetworkEventStream, NetworkPeers,
62};
63use sc_network_sync::{SyncEvent, SyncEventStream};
64use sc_network_types::PeerId;
65use sp_runtime::traits::Block as BlockT;
66use sp_statement_store::{
67 FilterDecision, Hash, Statement, StatementSource, StatementStore, SubmitResult,
68};
69use std::{
70 collections::{hash_map::Entry, HashMap, HashSet, VecDeque},
71 iter,
72 num::{NonZeroU32, NonZeroUsize},
73 pin::Pin,
74 sync::Arc,
75 time::Instant,
76};
77use tokio::time::timeout;
78pub mod config;
79
80pub type Statements = Vec<Statement>;
82
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85enum PeerProtocolVersion {
86 V1,
88 V2,
90}
91
92impl PeerProtocolVersion {
93 fn envelope_overhead(&self) -> usize {
95 match self {
96 PeerProtocolVersion::V1 => V1_ENVELOPE_OVERHEAD,
97 PeerProtocolVersion::V2 => V2_ENVELOPE_OVERHEAD,
98 }
99 }
100}
101
102#[derive(Debug, Encode, Decode)]
103enum StatementMessage {
104 #[codec(index = 0)]
105 Statements(Vec<Statement>),
106 #[codec(index = 1)]
108 ExplicitTopicAffinity(AffinityFilter),
109}
110
111const STATEMENTS_VARIANT_INDEX: u8 = 0;
113
114impl StatementMessage {
115 fn encode_statement_refs(statements: &[&Statement]) -> Vec<u8> {
118 let mut out = Vec::new();
119 STATEMENTS_VARIANT_INDEX.encode_to(&mut out);
120 statements.encode_to(&mut out);
121 out
122 }
123}
124
125pub type StatementImportFuture = oneshot::Receiver<SubmitResult>;
127
128mod rep {
129 use sc_network::ReputationChange as Rep;
130 pub const ANY_STATEMENT: Rep = Rep::new(-(1 << 4), "Any statement");
135 pub const ANY_STATEMENT_REFUND: Rep = Rep::new(1 << 4, "Any statement (refund)");
137 pub const GOOD_STATEMENT: Rep = Rep::new(1 << 8, "Good statement");
139 pub const INVALID_STATEMENT: Rep = Rep::new(-(1 << 12), "Invalid statement");
141 pub const DUPLICATE_STATEMENT: Rep = Rep::new(-(1 << 7), "Duplicate statement");
143 pub const STATEMENT_FLOODING: Rep = Rep::new_fatal("Statement flooding");
145 pub const BAD_MESSAGE: Rep = Rep::new(-(1 << 12), "Bad statement message");
147}
148
149const LOG_TARGET: &str = "statement-gossip";
150const STATEMENT_PROTOCOL_V2: &str = "statement/2";
153const STATEMENT_PROTOCOL_V1: &str = "statement/1";
155const SEND_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
157const INITIAL_SYNC_BURST_INTERVAL: std::time::Duration = std::time::Duration::from_millis(10);
159const PENDING_AFFINITIES_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);
161const SYNC_RECOVERY_READD_DELAY: std::time::Duration = std::time::Duration::from_secs(60);
163
164struct Metrics {
165 propagated_statements: Counter<U64>,
166 known_statements_received: Counter<U64>,
167 skipped_oversized_statements: Counter<U64>,
168 propagated_statements_chunks: Histogram,
169 pending_statements: Gauge<U64>,
170 ignored_statements: Counter<U64>,
171 peers_connected: GaugeVec<U64>,
172 statements_received: Counter<U64>,
173 bytes_sent_total: Counter<U64>,
174 bytes_received_total: Counter<U64>,
175 sent_latency_seconds: Histogram,
176 initial_sync_statements_sent: Counter<U64>,
177 initial_sync_bursts_total: Counter<U64>,
178 initial_sync_peers_active: Gauge<U64>,
179 initial_sync_duration_seconds: Histogram,
180 statement_flooding_detected: Counter<U64>,
181}
182
183impl Metrics {
184 fn register(r: &Registry) -> Result<Self, PrometheusError> {
185 let peers_connected = register(
186 GaugeVec::new(
187 Opts::new(
188 "substrate_sync_statement_peers_connected",
189 "Number of peers connected using the statement protocol by kind",
190 ),
191 &["kind"],
192 )?,
193 r,
194 )?;
195 peers_connected.with_label_values(&["full"]).set(0);
196 peers_connected.with_label_values(&["light"]).set(0);
197
198 Ok(Self {
199 propagated_statements: register(
200 Counter::new(
201 "substrate_sync_propagated_statements",
202 "Total statements propagated to peers, counted once per recipient (a statement sent to N peers increments by N)",
203 )?,
204 r,
205 )?,
206 known_statements_received: register(
207 Counter::new(
208 "substrate_sync_known_statement_received",
209 "Number of statements received via gossiping that were already in the statement store",
210 )?,
211 r,
212 )?,
213 skipped_oversized_statements: register(
214 Counter::new(
215 "substrate_sync_skipped_oversized_statements",
216 "Number of oversized statements that were skipped to be gossiped",
217 )?,
218 r,
219 )?,
220 propagated_statements_chunks: register(
221 Histogram::with_opts(
222 HistogramOpts::new(
223 "substrate_sync_propagated_statements_chunks",
224 "Distribution of chunk sizes when propagating statements",
225 )
226 .buckets(exponential_buckets(1.0, 2.0, 14)?),
227 )?,
228 r,
229 )?,
230 pending_statements: register(
231 Gauge::new(
232 "substrate_sync_pending_statement_validations",
233 "Number of pending statement validations, sampled once per propagation tick",
234 )?,
235 r,
236 )?,
237 ignored_statements: register(
238 Counter::new(
239 "substrate_sync_ignored_statements",
240 "Number of statements ignored due to exceeding MAX_PENDING_STATEMENTS limit",
241 )?,
242 r,
243 )?,
244 peers_connected,
245 statements_received: register(
246 Counter::new(
247 "substrate_sync_statements_received",
248 "Total number of statements received from peers",
249 )?,
250 r,
251 )?,
252 bytes_sent_total: register(
253 Counter::new(
254 "substrate_sync_statement_bytes_sent_total",
255 "Total bytes sent for statement protocol messages",
256 )?,
257 r,
258 )?,
259 bytes_received_total: register(
260 Counter::new(
261 "substrate_sync_statement_bytes_received_total",
262 "Total bytes received for statement protocol messages (includes bytes from notifications that are later discarded — e.g. while major-syncing)",
263 )?,
264 r,
265 )?,
266 sent_latency_seconds: register(
267 Histogram::with_opts(
268 HistogramOpts::new(
269 "substrate_sync_statement_sent_latency_seconds",
270 "Time to send statement messages to peers",
271 )
272 .buckets(vec![0.000_001, 0.000_01, 0.000_1, 0.001, 0.01, 0.1, 1.0]),
274 )?,
275 r,
276 )?,
277 initial_sync_statements_sent: register(
278 Counter::new(
279 "substrate_sync_initial_sync_statements_sent",
280 "Total statements sent during initial sync bursts to newly connected peers",
281 )?,
282 r,
283 )?,
284 initial_sync_bursts_total: register(
285 Counter::new(
286 "substrate_sync_initial_sync_bursts_total",
287 "Total initial-sync burst rounds attempted (includes rounds that return early with no hashes left)",
288 )?,
289 r,
290 )?,
291 initial_sync_peers_active: register(
292 Gauge::new(
293 "substrate_sync_initial_sync_peers_active",
294 "Number of peers currently being synced via initial sync",
295 )?,
296 r,
297 )?,
298 initial_sync_duration_seconds: register(
299 Histogram::with_opts(
300 HistogramOpts::new(
301 "substrate_sync_initial_sync_duration_seconds",
302 "Per-peer duration of initial sync from start until completion or peer disconnect (whichever comes first)",
303 )
304 .buckets(vec![0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0]),
305 )?,
306 r,
307 )?,
308 statement_flooding_detected: register(
309 Counter::new(
310 "substrate_sync_statement_flooding_detected",
311 "Number of peers disconnected for exceeding statement rate limits",
312 )?,
313 r,
314 )?,
315 })
316 }
317}
318
319pub struct StatementHandlerPrototype {
321 protocol_name: ProtocolName,
322 notification_service: Box<dyn NotificationService>,
323}
324
325impl StatementHandlerPrototype {
326 pub fn new<
328 Hash: AsRef<[u8]>,
329 Block: BlockT,
330 Net: NetworkBackend<Block, <Block as BlockT>::Hash>,
331 >(
332 genesis_hash: Hash,
333 fork_id: Option<&str>,
334 metrics: NotificationMetrics,
335 peer_store_handle: Arc<dyn PeerStoreProvider>,
336 ) -> (Self, Net::NotificationProtocolConfig) {
337 let genesis_hash = genesis_hash.as_ref();
338 let hex = array_bytes::bytes2hex("", genesis_hash);
339 let (protocol_name, fallback_name) = if let Some(fork_id) = fork_id {
340 (
341 format!("/{hex}/{fork_id}/{STATEMENT_PROTOCOL_V2}"),
342 format!("/{hex}/{fork_id}/{STATEMENT_PROTOCOL_V1}"),
343 )
344 } else {
345 (format!("/{hex}/{STATEMENT_PROTOCOL_V2}"), format!("/{hex}/{STATEMENT_PROTOCOL_V1}"))
346 };
347 let (config, notification_service) = Net::notification_config(
348 protocol_name.clone().into(),
349 vec![fallback_name.into()],
350 MAX_STATEMENT_NOTIFICATION_SIZE,
351 None,
352 SetConfig {
353 in_peers: 0,
354 out_peers: 0,
355 reserved_nodes: Vec::new(),
356 non_reserved_mode: NonReservedPeerMode::Deny,
357 },
358 metrics,
359 peer_store_handle,
360 );
361
362 (Self { protocol_name: protocol_name.into(), notification_service }, config)
363 }
364
365 pub fn build<
370 N: NetworkPeers + NetworkEventStream,
371 S: SyncEventStream + sp_consensus::SyncOracle,
372 >(
373 self,
374 network: N,
375 sync: S,
376 statement_store: Arc<dyn StatementStore>,
377 metrics_registry: Option<&Registry>,
378 executor: impl Fn(Pin<Box<dyn Future<Output = ()> + Send>>) + Send,
379 mut num_submission_workers: usize,
380 statements_per_second: u32,
381 ) -> error::Result<StatementHandler<N, S>> {
382 let sync_event_stream = sync.event_stream("statement-handler-sync");
383 let (queue_sender, queue_receiver) = async_channel::bounded(MAX_PENDING_STATEMENTS);
384
385 if num_submission_workers == 0 {
386 log::warn!(
387 target: LOG_TARGET,
388 "num_submission_workers is 0, defaulting to 1"
389 );
390 num_submission_workers = 1;
391 }
392
393 let statements_per_second = match NonZeroU32::new(statements_per_second) {
394 Some(rate) => rate,
395 None => {
396 log::warn!(
397 target: LOG_TARGET,
398 "statements_per_second is 0, defaulting to {}",
399 DEFAULT_STATEMENTS_PER_SECOND
400 );
401 NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
402 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero")
403 },
404 };
405
406 let metrics =
407 if let Some(r) = metrics_registry { Some(Metrics::register(r)?) } else { None };
408
409 for _ in 0..num_submission_workers {
410 let store = statement_store.clone();
411 let mut queue_receiver = queue_receiver.clone();
412 executor(
413 async move {
414 loop {
415 let task: Option<(Statement, oneshot::Sender<SubmitResult>)> =
416 queue_receiver.next().await;
417 match task {
418 None => return,
419 Some((statement, completion)) => {
420 let result = store.submit(statement, StatementSource::Network);
421 if completion.send(result).is_err() {
422 log::debug!(
423 target: LOG_TARGET,
424 "Error sending validation completion"
425 );
426 }
427 },
428 }
429 }
430 }
431 .boxed(),
432 );
433 }
434
435 let handler = StatementHandler {
436 protocol_name: self.protocol_name,
437 notification_service: self.notification_service,
438 propagate_timeout: (Box::pin(interval(PROPAGATE_TIMEOUT))
439 as Pin<Box<dyn Stream<Item = ()> + Send>>)
440 .fuse(),
441 pending_statements: FuturesUnordered::new(),
442 pending_statements_peers: HashMap::new(),
443 network,
444 sync,
445 sync_event_stream: sync_event_stream.fuse(),
446 peers: HashMap::new(),
447 statement_store,
448 queue_sender,
449 statements_per_second,
450 metrics,
451 initial_sync_timeout: Box::pin(tokio::time::sleep(INITIAL_SYNC_BURST_INTERVAL).fuse()),
452 pending_affinities_timeout: Box::pin(
453 tokio::time::sleep(PENDING_AFFINITIES_INTERVAL).fuse(),
454 ),
455 pending_initial_syncs: HashMap::new(),
456 initial_sync_peer_queue: VecDeque::new(),
457 deferred_peers: HashSet::new(),
458 dropped_statements_during_sync: false,
459 sync_recovery_peer: None,
460 sync_recovery_readd_timeout: Box::pin(pending().fuse()),
461 };
462
463 Ok(handler)
464 }
465}
466
467pub struct StatementHandler<
469 N: NetworkPeers + NetworkEventStream,
470 S: SyncEventStream + sp_consensus::SyncOracle,
471> {
472 protocol_name: ProtocolName,
473 propagate_timeout: stream::Fuse<Pin<Box<dyn Stream<Item = ()> + Send>>>,
475 pending_statements:
477 FuturesUnordered<Pin<Box<dyn Future<Output = (Hash, Option<SubmitResult>)> + Send>>>,
478 pending_statements_peers: HashMap<Hash, HashSet<PeerId>>,
483 network: N,
485 sync: S,
487 sync_event_stream: stream::Fuse<Pin<Box<dyn Stream<Item = SyncEvent> + Send>>>,
489 notification_service: Box<dyn NotificationService>,
491 peers: HashMap<PeerId, Peer>,
493 statement_store: Arc<dyn StatementStore>,
494 queue_sender: async_channel::Sender<(Statement, oneshot::Sender<SubmitResult>)>,
495 statements_per_second: NonZeroU32,
497 metrics: Option<Metrics>,
499 initial_sync_timeout: Pin<Box<dyn FusedFuture<Output = ()> + Send>>,
501 pending_affinities_timeout: Pin<Box<dyn FusedFuture<Output = ()> + Send>>,
503 pending_initial_syncs: HashMap<PeerId, PendingInitialSync>,
505 initial_sync_peer_queue: VecDeque<PeerId>,
507 deferred_peers: HashSet<PeerId>,
510 dropped_statements_during_sync: bool,
512 sync_recovery_peer: Option<PeerId>,
514 sync_recovery_readd_timeout: Pin<Box<dyn FusedFuture<Output = ()> + Send>>,
516}
517
518#[derive(Debug)]
523struct PeerRateLimiter {
524 limiter: RateLimiter<NotKeyed, InMemoryState, DefaultClock>,
525}
526
527impl PeerRateLimiter {
528 fn new(statements_per_second: NonZeroU32, burst: NonZeroU32) -> Self {
529 let quota = Quota::per_second(statements_per_second).allow_burst(burst);
530 Self { limiter: RateLimiter::direct(quota) }
531 }
532
533 fn is_flooding(&self, count: usize) -> bool {
535 if count > u32::MAX as usize {
536 return true;
537 }
538
539 let Some(n) = NonZeroU32::new(count as u32) else {
540 return false;
541 };
542 !matches!(self.limiter.check_n(n), Ok(Ok(())))
543 }
544}
545
546#[cfg_attr(not(any(test, feature = "test-helpers")), doc(hidden))]
548#[derive(Debug)]
549pub struct Peer {
550 known_statements: LruHashSet<Hash>,
552 rate_limiter: PeerRateLimiter,
554 protocol_version: PeerProtocolVersion,
556 topic_affinity: Option<AffinityFilter>,
559 is_light: bool,
562 pending_topic_affinity: Option<AffinityFilter>,
566}
567
568struct PendingInitialSync {
570 hashes: Vec<Hash>,
571 started_at: Instant,
572}
573
574enum ChunkResult {
576 Send(usize),
578 SkipOversized,
580}
581
582enum SendChunkResult {
584 Sent(usize),
586 Skipped,
588 Empty,
590 Failed,
592}
593
594const V1_ENVELOPE_OVERHEAD: usize = 5;
596
597const V2_ENVELOPE_OVERHEAD: usize = 1 + V1_ENVELOPE_OVERHEAD;
599
600fn max_statement_payload_size(envelope_overhead: usize) -> usize {
603 debug_assert_eq!(
604 V1_ENVELOPE_OVERHEAD,
605 Compact::<u32>::max_encoded_len(),
606 "V1_ENVELOPE_OVERHEAD must equal Compact::<u32>::max_encoded_len()"
607 );
608 MAX_STATEMENT_NOTIFICATION_SIZE as usize - envelope_overhead
609}
610
611fn find_sendable_chunk(statements: &[&Statement], envelope_overhead: usize) -> ChunkResult {
618 if statements.is_empty() {
619 return ChunkResult::Send(0);
620 }
621 let max_size = max_statement_payload_size(envelope_overhead);
622
623 let mut accumulated_size = 0;
628 let mut count = 0usize;
629
630 for stmt in &statements[0..] {
631 let stmt_size = stmt.encoded_size();
632 let new_count = count + 1;
633 let new_total = accumulated_size + stmt_size;
635 if new_total > max_size {
636 break;
637 }
638
639 accumulated_size += stmt_size;
640 count = new_count;
641 }
642
643 if count == 0 {
645 ChunkResult::SkipOversized
646 } else {
647 ChunkResult::Send(count)
648 }
649}
650
651impl Peer {
652 #[cfg(any(test, feature = "test-helpers"))]
654 pub fn new_for_testing(
655 known_statements: LruHashSet<Hash>,
656 statements_per_second: NonZeroU32,
657 burst: NonZeroU32,
658 ) -> Self {
659 Self {
660 known_statements,
661 rate_limiter: PeerRateLimiter::new(statements_per_second, burst),
662 protocol_version: PeerProtocolVersion::V1,
663 topic_affinity: None,
664 is_light: false,
665 pending_topic_affinity: None,
666 }
667 }
668
669 fn can_receive(&self) -> bool {
673 !(self.is_light &&
674 self.protocol_version == PeerProtocolVersion::V2 &&
675 self.topic_affinity.is_none())
676 }
677
678 fn kind(&self) -> &'static str {
679 if self.is_light {
680 "light"
681 } else {
682 "full"
683 }
684 }
685}
686
687impl<N, S> StatementHandler<N, S>
688where
689 N: NetworkPeers + NetworkEventStream,
690 S: SyncEventStream + sp_consensus::SyncOracle,
691{
692 #[cfg(any(test, feature = "test-helpers"))]
694 pub fn new_for_testing(
695 protocol_name: ProtocolName,
696 notification_service: Box<dyn NotificationService>,
697 propagate_timeout: stream::Fuse<Pin<Box<dyn Stream<Item = ()> + Send>>>,
698 network: N,
699 sync: S,
700 sync_event_stream: stream::Fuse<Pin<Box<dyn Stream<Item = SyncEvent> + Send>>>,
701 peers: HashMap<PeerId, Peer>,
702 statement_store: Arc<dyn StatementStore>,
703 queue_sender: async_channel::Sender<(Statement, oneshot::Sender<SubmitResult>)>,
704 statements_per_second: NonZeroU32,
705 ) -> Self {
706 Self {
707 protocol_name,
708 notification_service,
709 propagate_timeout,
710 pending_statements: FuturesUnordered::new(),
711 pending_statements_peers: HashMap::new(),
712 network,
713 sync,
714 sync_event_stream,
715 peers,
716 statement_store,
717 queue_sender,
718 statements_per_second,
719 metrics: None,
720 initial_sync_timeout: Box::pin(pending().fuse()),
721 pending_affinities_timeout: Box::pin(pending().fuse()),
722 pending_initial_syncs: HashMap::new(),
723 initial_sync_peer_queue: VecDeque::new(),
724 deferred_peers: HashSet::new(),
725 dropped_statements_during_sync: false,
726 sync_recovery_peer: None,
727 sync_recovery_readd_timeout: Box::pin(pending().fuse()),
728 }
729 }
730
731 #[cfg(any(test, feature = "test-helpers"))]
733 pub fn pending_statements_mut(
734 &mut self,
735 ) -> &mut FuturesUnordered<Pin<Box<dyn Future<Output = (Hash, Option<SubmitResult>)> + Send>>>
736 {
737 &mut self.pending_statements
738 }
739
740 pub async fn run(mut self) {
743 loop {
744 futures::select_biased! {
745 _ = self.propagate_timeout.next() => {
746 self.propagate_statements().await;
747 self.metrics.as_ref().map(|metrics| {
748 metrics.pending_statements.set(self.pending_statements.len() as u64);
749 });
750 },
751 (hash, result) = self.pending_statements.select_next_some() => {
752 if let Some(peers) = self.pending_statements_peers.remove(&hash) {
753 if let Some(result) = result {
754 peers.into_iter().for_each(|p| self.on_handle_statement_import(p, &result));
755 }
756 } else {
757 log::warn!(target: LOG_TARGET, "Inconsistent state, no peers for pending statement!");
758 }
759 },
760 sync_event = self.sync_event_stream.next() => {
761 if let Some(sync_event) = sync_event {
762 self.handle_sync_event(sync_event);
763 } else {
764 return;
766 }
767 }
768 event = self.notification_service.next_event().fuse() => {
769 if let Some(event) = event {
770 self.handle_notification_event(event).await
771 } else {
772 return
774 }
775 }
776 _ = &mut self.initial_sync_timeout => {
777 self.process_initial_sync_burst().await;
778 self.initial_sync_timeout =
779 Box::pin(tokio::time::sleep(INITIAL_SYNC_BURST_INTERVAL).fuse());
780 },
781 _ = &mut self.pending_affinities_timeout => {
782 self.process_pending_affinities();
783 self.pending_affinities_timeout =
784 Box::pin(tokio::time::sleep(PENDING_AFFINITIES_INTERVAL).fuse());
785 },
786 _ = &mut self.sync_recovery_readd_timeout => {
787 self.try_readd_sync_recovery_peer();
788 self.sync_recovery_readd_timeout = Box::pin(pending().fuse());
789 },
790 }
791
792 if !self.sync.is_major_syncing() {
793 self.drain_deferred_peers();
794 self.start_sync_recovery();
795 }
796 }
797 }
798
799 async fn send_statement_chunk(
805 &mut self,
806 peer: &PeerId,
807 statements: &[&Statement],
808 ) -> SendChunkResult {
809 let Some(peer_data) = self.peers.get(peer) else {
810 log::error!(target: LOG_TARGET, "Peer {peer} not found in peers map during send_statement_chunk");
811 return SendChunkResult::Failed;
812 };
813 let peer_version = peer_data.protocol_version;
814 let envelope_overhead = peer_version.envelope_overhead();
815 match find_sendable_chunk(statements, envelope_overhead) {
816 ChunkResult::Send(0) => SendChunkResult::Empty,
817 ChunkResult::Send(chunk_end) => {
818 let chunk = &statements[..chunk_end];
819 let encoded = match peer_version {
820 PeerProtocolVersion::V1 => chunk.encode(),
821 PeerProtocolVersion::V2 => StatementMessage::encode_statement_refs(chunk),
822 };
823 let bytes_to_send = encoded.len() as u64;
824
825 let sent_latency_timer =
826 self.metrics.as_ref().map(|m| m.sent_latency_seconds.start_timer());
827 let send_result = timeout(
828 SEND_TIMEOUT,
829 self.notification_service.send_async_notification(peer, encoded),
830 )
831 .await;
832 drop(sent_latency_timer);
833
834 if let Err(e) = send_result {
835 log::debug!(target: LOG_TARGET, "Failed to send notification to {peer}: {e:?}");
836 return SendChunkResult::Failed;
837 }
838
839 log::trace!(target: LOG_TARGET, "Sent {} statements to {}", chunk.len(), peer);
840 self.metrics.as_ref().map(|metrics| {
841 metrics.propagated_statements.inc_by(chunk.len() as u64);
842 metrics.bytes_sent_total.inc_by(bytes_to_send);
843 metrics.propagated_statements_chunks.observe(chunk.len() as f64);
844 });
845 SendChunkResult::Sent(chunk_end)
846 },
847 ChunkResult::SkipOversized => {
848 log::warn!(target: LOG_TARGET, "Statement too large, skipping");
849 self.metrics.as_ref().map(|metrics| {
850 metrics.skipped_oversized_statements.inc();
851 });
852 SendChunkResult::Skipped
853 },
854 }
855 }
856
857 fn drain_deferred_peers(&mut self) {
859 if self.deferred_peers.is_empty() {
860 return;
861 }
862
863 log::debug!(
864 target: LOG_TARGET,
865 "Major sync complete, adding {} deferred statement peers",
866 self.deferred_peers.len(),
867 );
868
869 let addrs: HashSet<multiaddr::Multiaddr> = self
870 .deferred_peers
871 .drain()
872 .map(|p| {
873 iter::once(multiaddr::Protocol::P2p(p.into())).collect::<multiaddr::Multiaddr>()
874 })
875 .collect();
876
877 if let Err(err) = self.network.add_peers_to_reserved_set(self.protocol_name.clone(), addrs)
878 {
879 log::warn!(target: LOG_TARGET, "Failed to add deferred peers: {err}");
880 }
881 }
882
883 fn start_sync_recovery(&mut self) {
888 if !self.dropped_statements_during_sync {
889 return;
890 }
891 self.dropped_statements_during_sync = false;
892
893 if self.sync_recovery_peer.is_some() {
894 return;
895 }
896
897 let Some(&peer_id) = self.peers.keys().choose(&mut rand::thread_rng()) else {
898 return;
899 };
900
901 log::trace!(
902 target: LOG_TARGET,
903 "Major sync complete, force-reconnecting {peer_id} for statement recovery",
904 );
905
906 if let Err(err) = self.network.remove_peers_from_reserved_set(
907 self.protocol_name.clone(),
908 iter::once(peer_id).collect(),
909 ) {
910 log::warn!(target: LOG_TARGET, "Failed to remove peer {peer_id} for sync recovery: {err}");
911 return;
912 }
913
914 self.sync_recovery_peer = Some(peer_id);
915 self.sync_recovery_readd_timeout =
916 Box::pin(tokio::time::sleep(SYNC_RECOVERY_READD_DELAY).fuse());
917 }
918
919 fn try_readd_sync_recovery_peer(&mut self) {
921 let Some(peer_id) = self.sync_recovery_peer.take() else { return };
922 log::trace!(
923 target: LOG_TARGET,
924 "Re-adding {peer_id} to reserved set after sync recovery window",
925 );
926 let addr =
927 iter::once(multiaddr::Protocol::P2p(peer_id.into())).collect::<multiaddr::Multiaddr>();
928 if let Err(err) = self
929 .network
930 .add_peers_to_reserved_set(self.protocol_name.clone(), iter::once(addr).collect())
931 {
932 log::warn!(target: LOG_TARGET, "Failed to re-add sync recovery peer {peer_id}: {err}");
933 }
934 }
935
936 fn handle_sync_event(&mut self, event: SyncEvent) {
937 match event {
938 SyncEvent::PeerConnected(remote) => {
939 if self.sync.is_major_syncing() {
940 log::trace!(
941 target: LOG_TARGET,
942 "Major sync in progress, deferring connection to {remote}",
943 );
944 self.deferred_peers.insert(remote);
945 return;
946 }
947 let addr = iter::once(multiaddr::Protocol::P2p(remote.into()))
948 .collect::<multiaddr::Multiaddr>();
949 let result = self.network.add_peers_to_reserved_set(
950 self.protocol_name.clone(),
951 iter::once(addr).collect(),
952 );
953 if let Err(err) = result {
954 log::error!(target: LOG_TARGET, "Add reserved peer failed: {}", err);
955 }
956 },
957 SyncEvent::PeerDisconnected(remote) => {
958 if self.deferred_peers.remove(&remote) {
959 return;
960 }
961 let result = self.network.remove_peers_from_reserved_set(
962 self.protocol_name.clone(),
963 iter::once(remote).collect(),
964 );
965 if let Err(err) = result {
966 log::error!(target: LOG_TARGET, "Failed to remove reserved peer: {err}");
967 }
968 },
969 }
970 }
971
972 async fn handle_notification_event(&mut self, event: NotificationEvent) {
973 match event {
974 NotificationEvent::ValidateInboundSubstream { peer, handshake, result_tx, .. } => {
975 let result = self
977 .network
978 .peer_role(peer, handshake)
979 .map_or(ValidationResult::Reject, |_| ValidationResult::Accept);
980 let _ = result_tx.send(result);
981 },
982 NotificationEvent::NotificationStreamOpened {
983 peer,
984 negotiated_fallback,
985 handshake,
986 ..
987 } => {
988 let protocol_version = if negotiated_fallback.is_some() {
991 PeerProtocolVersion::V1
992 } else {
993 PeerProtocolVersion::V2
994 };
995 let Some(peer_role) = self.network.peer_role(peer, handshake) else {
996 log::debug!(
997 target: LOG_TARGET,
998 "Peer {peer} connected but role could not be determined, ignoring"
999 );
1000 return;
1001 };
1002 let is_light = peer_role.is_light();
1003 log::debug!(
1004 target: LOG_TARGET,
1005 "Peer {peer} connected with statement protocol {protocol_version:?}, role={peer_role:?}"
1006 );
1007 let _was_in = self.peers.insert(
1008 peer,
1009 Peer {
1010 known_statements: LruHashSet::new(
1011 NonZeroUsize::new(MAX_KNOWN_STATEMENTS).expect("Constant is nonzero"),
1012 ),
1013 rate_limiter: PeerRateLimiter::new(
1014 self.statements_per_second,
1015 NonZeroU32::new(
1016 self.statements_per_second.get() *
1017 config::STATEMENTS_BURST_COEFFICIENT,
1018 )
1019 .expect("burst capacity is nonzero"),
1020 ),
1021 protocol_version,
1022 topic_affinity: None,
1023 is_light,
1024 pending_topic_affinity: None,
1025 },
1026 );
1027 debug_assert!(_was_in.is_none());
1028
1029 self.metrics.as_ref().map(|metrics| {
1030 if let Some(peer) = self.peers.get(&peer) {
1031 metrics.peers_connected.with_label_values(&[peer.kind()]).inc();
1032 }
1033 });
1034
1035 if self.peers.get(&peer).map_or(false, |p| p.can_receive()) {
1038 self.schedule_initial_sync_for_peer(peer);
1039 }
1040 },
1041 NotificationEvent::NotificationStreamClosed { peer } => {
1042 let removed_peer = self.peers.remove(&peer);
1043 debug_assert!(removed_peer.is_some());
1044
1045 if let Some(removed_peer) = removed_peer {
1046 self.metrics.as_ref().map(|metrics| {
1047 metrics.peers_connected.with_label_values(&[removed_peer.kind()]).dec();
1048 });
1049 }
1050
1051 if let Some(pending) = self.pending_initial_syncs.remove(&peer) {
1052 self.metrics.as_ref().map(|metrics| {
1053 metrics.initial_sync_peers_active.dec();
1054 metrics
1055 .initial_sync_duration_seconds
1056 .observe(pending.started_at.elapsed().as_secs_f64());
1057 });
1058 }
1059 self.initial_sync_peer_queue.retain(|p| *p != peer);
1060 },
1061 NotificationEvent::NotificationReceived { peer, notification } => {
1062 let bytes_received = notification.len() as u64;
1063 self.metrics.as_ref().map(|metrics| {
1064 metrics.bytes_received_total.inc_by(bytes_received);
1065 });
1066
1067 if self.sync.is_major_syncing() {
1069 log::trace!(
1070 target: LOG_TARGET,
1071 "{peer}: Ignoring statements while major syncing or offline"
1072 );
1073 self.dropped_statements_during_sync = true;
1074 return;
1075 }
1076
1077 let Some(peer_data) = self.peers.get(&peer) else {
1078 log::error!(target: LOG_TARGET, "Received notification from unknown peer {peer}");
1079 return;
1080 };
1081
1082 match peer_data.protocol_version {
1083 PeerProtocolVersion::V1 => {
1084 if let Ok(statements) =
1086 <Statements as Decode>::decode(&mut notification.as_ref())
1087 {
1088 self.on_statements(peer, statements);
1089 } else {
1090 log::debug!(
1091 target: LOG_TARGET,
1092 "Failed to decode v1 statement list from {peer}"
1093 );
1094 self.network.report_peer(peer, rep::BAD_MESSAGE);
1095 }
1096 },
1097 PeerProtocolVersion::V2 => {
1098 if let Ok(message) = StatementMessage::decode(&mut notification.as_ref()) {
1100 match message {
1101 StatementMessage::Statements(statements) => {
1102 self.on_statements(peer, statements)
1103 },
1104 StatementMessage::ExplicitTopicAffinity(filter) => {
1105 if let Some(peer_data) = self.peers.get_mut(&peer) {
1106 if peer_data.rate_limiter.is_flooding(1) {
1107 log::debug!(
1108 target: LOG_TARGET,
1109 "Rate-limiting ExplicitTopicAffinity from {peer}"
1110 );
1111 self.network.report_peer(peer, rep::BAD_MESSAGE);
1112 } else {
1113 log::debug!(
1114 target: LOG_TARGET,
1115 "Received topic affinity filter from {peer}"
1116 );
1117 peer_data.pending_topic_affinity = Some(filter);
1120 }
1121 }
1122 },
1123 }
1124 } else {
1125 log::debug!(
1126 target: LOG_TARGET,
1127 "Failed to decode v2 statement message from {peer}"
1128 );
1129 self.network.report_peer(peer, rep::BAD_MESSAGE);
1130 }
1131 },
1132 }
1133 },
1134 }
1135 }
1136
1137 #[cfg_attr(not(any(test, feature = "test-helpers")), doc(hidden))]
1139 pub fn on_statements(&mut self, who: PeerId, statements: Statements) {
1140 log::trace!(target: LOG_TARGET, "Received {} statements from {}", statements.len(), who);
1141
1142 self.metrics.as_ref().map(|metrics| {
1143 metrics.statements_received.inc_by(statements.len() as u64);
1144 });
1145
1146 if let Some(ref mut peer) = self.peers.get_mut(&who) {
1147 if peer.rate_limiter.is_flooding(statements.len()) {
1148 log::warn!(
1149 target: LOG_TARGET,
1150 "Peer {} exceeded statement rate limit ({} statements/sec). Disconnecting.",
1151 who,
1152 self.statements_per_second
1153 );
1154
1155 self.network.report_peer(who, rep::STATEMENT_FLOODING);
1156
1157 self.network.disconnect_peer(who, self.protocol_name.clone());
1159
1160 if let Some(ref metrics) = self.metrics {
1161 metrics.statement_flooding_detected.inc();
1162 }
1163
1164 return;
1165 }
1166
1167 let mut statements_left = statements.len() as u64;
1168 for s in statements {
1169 if self.pending_statements.len() > MAX_PENDING_STATEMENTS {
1170 log::debug!(
1171 target: LOG_TARGET,
1172 "Ignoring {} statements that exceed `MAX_PENDING_STATEMENTS`({}) limit",
1173 statements_left,
1174 MAX_PENDING_STATEMENTS,
1175 );
1176 self.metrics.as_ref().map(|metrics| {
1177 metrics.ignored_statements.inc_by(statements_left);
1178 });
1179 break;
1180 }
1181
1182 let hash = s.hash();
1183 peer.known_statements.insert(hash);
1184
1185 if self.statement_store.has_statement(&hash) {
1186 self.metrics.as_ref().map(|metrics| {
1187 metrics.known_statements_received.inc();
1188 });
1189
1190 if let Some(peers) = self.pending_statements_peers.get(&hash) {
1191 if peers.contains(&who) {
1192 log::trace!(
1193 target: LOG_TARGET,
1194 "Already received the statement from the same peer {who}.",
1195 );
1196 self.network.report_peer(who, rep::DUPLICATE_STATEMENT);
1197 }
1198 }
1199 continue;
1200 }
1201
1202 self.network.report_peer(who, rep::ANY_STATEMENT);
1203
1204 match self.pending_statements_peers.entry(hash) {
1205 Entry::Vacant(entry) => {
1206 let (completion_sender, completion_receiver) = oneshot::channel();
1207 match self.queue_sender.try_send((s, completion_sender)) {
1208 Ok(()) => {
1209 self.pending_statements.push(
1210 async move {
1211 let res = completion_receiver.await;
1212 (hash, res.ok())
1213 }
1214 .boxed(),
1215 );
1216 entry.insert(HashSet::from_iter([who]));
1217 },
1218 Err(async_channel::TrySendError::Full(_)) => {
1219 log::debug!(
1220 target: LOG_TARGET,
1221 "Dropped statement because validation channel is full",
1222 );
1223 },
1224 Err(async_channel::TrySendError::Closed(_)) => {
1225 log::trace!(
1226 target: LOG_TARGET,
1227 "Dropped statement because validation channel is closed",
1228 );
1229 },
1230 }
1231 },
1232 Entry::Occupied(mut entry) => {
1233 if !entry.get_mut().insert(who) {
1234 self.network.report_peer(who, rep::DUPLICATE_STATEMENT);
1236 }
1237 },
1238 }
1239
1240 statements_left -= 1;
1241 }
1242 }
1243 }
1244
1245 fn on_handle_statement_import(&mut self, who: PeerId, import: &SubmitResult) {
1246 match import {
1247 SubmitResult::New => self.network.report_peer(who, rep::GOOD_STATEMENT),
1248 SubmitResult::Known => self.network.report_peer(who, rep::ANY_STATEMENT_REFUND),
1249 SubmitResult::KnownExpired => {},
1250 SubmitResult::Rejected(_) => {},
1251 SubmitResult::Invalid(_) => self.network.report_peer(who, rep::INVALID_STATEMENT),
1252 SubmitResult::InternalError(_) => {},
1253 }
1254 }
1255
1256 pub async fn propagate_statement(&mut self, hash: &Hash) {
1258 if self.sync.is_major_syncing() {
1260 return;
1261 }
1262
1263 log::debug!(target: LOG_TARGET, "Propagating statement [{:?}]", hash);
1264 if let Ok(Some(statement)) = self.statement_store.statement(hash) {
1265 self.do_propagate_statements(&[(*hash, statement)]).await;
1266 }
1267 }
1268
1269 async fn send_statements_to_peer(&mut self, who: &PeerId, statements: &[(Hash, Statement)]) {
1274 let Some(peer) = self.peers.get_mut(who) else {
1275 return;
1276 };
1277
1278 if !peer.can_receive() {
1279 return;
1280 }
1281
1282 let to_send: Vec<_> = statements
1283 .iter()
1284 .filter_map(|(hash, stmt)| {
1285 if peer.known_statements.contains(hash) {
1286 return None;
1287 }
1288 if peer.topic_affinity.as_ref().is_some_and(|a| !a.matches_statement(stmt)) {
1292 return None;
1293 }
1294 peer.known_statements.insert(*hash);
1295 Some(stmt)
1296 })
1297 .collect();
1298
1299 log::trace!(target: LOG_TARGET, "We have {} statements that the peer doesn't know about", to_send.len());
1300
1301 if to_send.is_empty() {
1302 return;
1303 }
1304
1305 self.send_statements_in_chunks(who, &to_send).await;
1306 }
1307
1308 async fn send_statements_in_chunks(&mut self, who: &PeerId, statements: &[&Statement]) {
1310 let mut offset = 0;
1311 while offset < statements.len() {
1312 match self.send_statement_chunk(who, &statements[offset..]).await {
1313 SendChunkResult::Sent(chunk_end) => {
1314 offset += chunk_end;
1315 },
1316 SendChunkResult::Skipped => {
1317 offset += 1;
1318 },
1319 SendChunkResult::Empty | SendChunkResult::Failed => return,
1320 }
1321 }
1322 }
1323
1324 async fn do_propagate_statements(&mut self, statements: &[(Hash, Statement)]) {
1325 log::debug!(target: LOG_TARGET, "Propagating {} statements for {} peers", statements.len(), self.peers.len());
1326 let peers: Vec<_> = self.peers.keys().copied().collect();
1327 for who in peers {
1328 log::trace!(target: LOG_TARGET, "Start propagating statements for {}", who);
1329 self.send_statements_to_peer(&who, statements).await;
1330 }
1331 log::trace!(target: LOG_TARGET, "Statements propagated to all peers");
1332 }
1333
1334 async fn propagate_statements(&mut self) {
1336 if self.sync.is_major_syncing() {
1338 return;
1339 }
1340
1341 let Ok(statements) = self.statement_store.take_recent_statements() else { return };
1342 if !statements.is_empty() {
1343 self.do_propagate_statements(&statements).await;
1344 }
1345 }
1346
1347 fn schedule_initial_sync_for_peer(&mut self, peer: PeerId) {
1353 if let Some(pending) = self.pending_initial_syncs.remove(&peer) {
1355 self.record_initial_sync_completion(pending.started_at);
1356 self.initial_sync_peer_queue.retain(|p| *p != peer);
1357 }
1358 let hashes = self.statement_store.statement_hashes();
1359 if let Some(peer_data) = self.peers.get_mut(&peer) {
1364 peer_data.known_statements.clear();
1365 }
1366 if !hashes.is_empty() {
1367 self.pending_initial_syncs
1368 .insert(peer, PendingInitialSync { hashes, started_at: Instant::now() });
1369 self.initial_sync_peer_queue.push_back(peer);
1370 self.metrics.as_ref().map(|metrics| {
1371 metrics.initial_sync_peers_active.inc();
1372 });
1373 }
1374 }
1375
1376 fn process_pending_affinities(&mut self) {
1382 let ready_peers: Vec<PeerId> = self
1383 .peers
1384 .iter()
1385 .filter(|(peer_id, peer_data)| {
1386 peer_data.pending_topic_affinity.is_some() &&
1387 !self.pending_initial_syncs.contains_key(peer_id)
1388 })
1389 .map(|(peer_id, _)| *peer_id)
1390 .collect();
1391
1392 for peer_id in ready_peers {
1393 if let Some(peer_data) = self.peers.get_mut(&peer_id) {
1394 peer_data.topic_affinity = peer_data.pending_topic_affinity.take();
1395 }
1396 self.schedule_initial_sync_for_peer(peer_id);
1397 }
1398 }
1399
1400 fn record_initial_sync_completion(&self, started_at: Instant) {
1402 self.metrics.as_ref().map(|metrics| {
1403 metrics.initial_sync_peers_active.dec();
1404 metrics
1405 .initial_sync_duration_seconds
1406 .observe(started_at.elapsed().as_secs_f64());
1407 });
1408 }
1409
1410 async fn process_initial_sync_burst(&mut self) {
1412 if self.sync.is_major_syncing() {
1413 return;
1414 }
1415
1416 let Some(peer_id) = self.initial_sync_peer_queue.pop_front() else {
1417 return;
1418 };
1419
1420 let Entry::Occupied(mut entry) = self.pending_initial_syncs.entry(peer_id) else {
1421 return;
1422 };
1423
1424 self.metrics.as_ref().map(|metrics| {
1425 metrics.initial_sync_bursts_total.inc();
1426 });
1427
1428 if entry.get().hashes.is_empty() {
1429 let started_at = entry.get().started_at;
1430 entry.remove();
1431 self.record_initial_sync_completion(started_at);
1432 return;
1433 }
1434
1435 let Some(peer_data) = self.peers.get(&peer_id) else {
1440 log::error!(target: LOG_TARGET, "Peer {peer_id} has pending initial sync but is not in peers map");
1441 entry.remove();
1442 return;
1443 };
1444 let envelope_overhead = peer_data.protocol_version.envelope_overhead();
1445 let max_size = max_statement_payload_size(envelope_overhead);
1446 let mut accumulated_size = 0;
1447 let (statements, processed) = match self.statement_store.statements_by_hashes(
1448 &entry.get().hashes,
1449 &mut |hash, encoded, stmt| {
1450 if peer_data.known_statements.contains(hash) {
1454 return FilterDecision::Skip;
1455 }
1456 if peer_data.topic_affinity.as_ref().is_some_and(|a| !a.matches_statement(stmt)) {
1457 return FilterDecision::Skip;
1458 }
1459 if accumulated_size > 0 && accumulated_size + encoded.len() > max_size {
1460 return FilterDecision::Abort;
1461 }
1462 accumulated_size += encoded.len();
1463 FilterDecision::Take
1464 },
1465 ) {
1466 Ok(r) => r,
1467 Err(e) => {
1468 log::debug!(target: LOG_TARGET, "Failed to fetch statements for initial sync: {e:?}");
1469 let started_at = entry.get().started_at;
1470 entry.remove();
1471 self.record_initial_sync_completion(started_at);
1472 return;
1473 },
1474 };
1475
1476 entry.get_mut().hashes.drain(..processed);
1478 let has_more = !entry.get().hashes.is_empty();
1479 drop(entry);
1480
1481 let send_stmts: Vec<_> = statements.iter().map(|(_, stmt)| stmt).collect();
1482 match self.send_statement_chunk(&peer_id, &send_stmts).await {
1483 SendChunkResult::Failed => {
1484 if let Some(pending) = self.pending_initial_syncs.remove(&peer_id) {
1485 self.record_initial_sync_completion(pending.started_at);
1486 }
1487 return;
1488 },
1489 SendChunkResult::Sent(sent) => {
1490 debug_assert_eq!(send_stmts.len(), sent);
1491 self.metrics.as_ref().map(|metrics| {
1492 metrics.initial_sync_statements_sent.inc_by(sent as u64);
1493 });
1494 if let Some(peer) = self.peers.get_mut(&peer_id) {
1496 for (hash, _) in &statements {
1497 peer.known_statements.insert(*hash);
1498 }
1499 }
1500 },
1501 SendChunkResult::Empty | SendChunkResult::Skipped => {},
1502 }
1503
1504 if has_more {
1506 self.initial_sync_peer_queue.push_back(peer_id);
1507 } else {
1508 if let Some(pending) = self.pending_initial_syncs.remove(&peer_id) {
1509 self.record_initial_sync_completion(pending.started_at);
1510 }
1511 }
1512 }
1513}
1514
1515#[cfg(test)]
1516mod tests {
1517
1518 use super::*;
1519 use std::sync::{
1520 atomic::{AtomicBool, Ordering},
1521 Mutex,
1522 };
1523
1524 const BLOOM_SEED: u128 = 0x5EED_5EED_5EED_5EED;
1526
1527 #[derive(Clone)]
1528 struct TestNetwork {
1529 reported_peers: Arc<Mutex<Vec<(PeerId, sc_network::ReputationChange)>>>,
1530 disconnected_peers: Arc<Mutex<Vec<PeerId>>>,
1531 default_role: sc_network::ObservedRole,
1533 added_reserved: Arc<Mutex<Vec<HashSet<sc_network::Multiaddr>>>>,
1534 removed_reserved: Arc<Mutex<Vec<Vec<PeerId>>>>,
1535 }
1536
1537 impl TestNetwork {
1538 fn new() -> Self {
1539 Self {
1540 reported_peers: Arc::new(Mutex::new(Vec::new())),
1541 disconnected_peers: Arc::new(Mutex::new(Vec::new())),
1542 default_role: sc_network::ObservedRole::Full,
1543 added_reserved: Arc::new(Mutex::new(Vec::new())),
1544 removed_reserved: Arc::new(Mutex::new(Vec::new())),
1545 }
1546 }
1547
1548 fn new_light() -> Self {
1549 Self {
1550 reported_peers: Arc::new(Mutex::new(Vec::new())),
1551 disconnected_peers: Arc::new(Mutex::new(Vec::new())),
1552 default_role: sc_network::ObservedRole::Light,
1553 added_reserved: Arc::new(Mutex::new(Vec::new())),
1554 removed_reserved: Arc::new(Mutex::new(Vec::new())),
1555 }
1556 }
1557
1558 fn get_reports(&self) -> Vec<(PeerId, sc_network::ReputationChange)> {
1559 self.reported_peers.lock().unwrap().clone()
1560 }
1561
1562 fn get_disconnected_peers(&self) -> Vec<PeerId> {
1563 self.disconnected_peers.lock().unwrap().clone()
1564 }
1565
1566 fn get_added_reserved(&self) -> Vec<HashSet<sc_network::Multiaddr>> {
1567 self.added_reserved.lock().unwrap().clone()
1568 }
1569
1570 fn get_removed_reserved(&self) -> Vec<Vec<PeerId>> {
1571 self.removed_reserved.lock().unwrap().clone()
1572 }
1573 }
1574
1575 #[async_trait::async_trait]
1576 impl NetworkPeers for TestNetwork {
1577 fn set_authorized_peers(&self, _: std::collections::HashSet<PeerId>) {
1578 unimplemented!()
1579 }
1580
1581 fn set_authorized_only(&self, _: bool) {
1582 unimplemented!()
1583 }
1584
1585 fn add_known_address(&self, _: PeerId, _: sc_network::Multiaddr) {
1586 unimplemented!()
1587 }
1588
1589 fn report_peer(&self, peer_id: PeerId, cost_benefit: sc_network::ReputationChange) {
1590 self.reported_peers.lock().unwrap().push((peer_id, cost_benefit));
1591 }
1592
1593 fn peer_reputation(&self, _: &PeerId) -> i32 {
1594 unimplemented!()
1595 }
1596
1597 fn disconnect_peer(&self, peer: PeerId, _: sc_network::ProtocolName) {
1598 self.disconnected_peers.lock().unwrap().push(peer);
1599 }
1600
1601 fn accept_unreserved_peers(&self) {
1602 unimplemented!()
1603 }
1604
1605 fn deny_unreserved_peers(&self) {
1606 unimplemented!()
1607 }
1608
1609 fn add_reserved_peer(
1610 &self,
1611 _: sc_network::config::MultiaddrWithPeerId,
1612 ) -> Result<(), String> {
1613 unimplemented!()
1614 }
1615
1616 fn remove_reserved_peer(&self, _: PeerId) {
1617 unimplemented!()
1618 }
1619
1620 fn set_reserved_peers(
1621 &self,
1622 _: sc_network::ProtocolName,
1623 _: std::collections::HashSet<sc_network::Multiaddr>,
1624 ) -> Result<(), String> {
1625 unimplemented!()
1626 }
1627
1628 fn add_peers_to_reserved_set(
1629 &self,
1630 _: sc_network::ProtocolName,
1631 addrs: std::collections::HashSet<sc_network::Multiaddr>,
1632 ) -> Result<(), String> {
1633 self.added_reserved.lock().unwrap().push(addrs);
1634 Ok(())
1635 }
1636
1637 fn remove_peers_from_reserved_set(
1638 &self,
1639 _: sc_network::ProtocolName,
1640 peers: Vec<PeerId>,
1641 ) -> Result<(), String> {
1642 self.removed_reserved.lock().unwrap().push(peers);
1643 Ok(())
1644 }
1645
1646 fn sync_num_connected(&self) -> usize {
1647 unimplemented!()
1648 }
1649
1650 fn peer_role(&self, _: PeerId, _: Vec<u8>) -> Option<sc_network::ObservedRole> {
1651 Some(self.default_role)
1652 }
1653
1654 async fn reserved_peers(&self) -> Result<Vec<PeerId>, ()> {
1655 unimplemented!();
1656 }
1657 }
1658
1659 #[derive(Clone)]
1660 struct TestSync {
1661 major_syncing: Arc<AtomicBool>,
1662 }
1663
1664 impl TestSync {
1665 fn new() -> Self {
1666 Self { major_syncing: Arc::new(AtomicBool::new(false)) }
1667 }
1668
1669 fn with_syncing(initial: bool) -> (Self, Arc<AtomicBool>) {
1670 let flag = Arc::new(AtomicBool::new(initial));
1671 (Self { major_syncing: flag.clone() }, flag)
1672 }
1673 }
1674
1675 impl SyncEventStream for TestSync {
1676 fn event_stream(
1677 &self,
1678 _name: &'static str,
1679 ) -> Pin<Box<dyn Stream<Item = sc_network_sync::types::SyncEvent> + Send>> {
1680 Box::pin(futures::stream::pending())
1681 }
1682 }
1683
1684 impl sp_consensus::SyncOracle for TestSync {
1685 fn is_major_syncing(&self) -> bool {
1686 self.major_syncing.load(Ordering::Relaxed)
1687 }
1688
1689 fn is_offline(&self) -> bool {
1690 unimplemented!()
1691 }
1692 }
1693
1694 impl NetworkEventStream for TestNetwork {
1695 fn event_stream(
1696 &self,
1697 _name: &'static str,
1698 ) -> Pin<Box<dyn Stream<Item = sc_network::Event> + Send>> {
1699 unimplemented!()
1700 }
1701 }
1702
1703 #[derive(Debug, Clone)]
1704 struct TestNotificationService {
1705 sent_notifications: Arc<Mutex<Vec<(PeerId, Vec<u8>)>>>,
1706 }
1707
1708 impl TestNotificationService {
1709 fn new() -> Self {
1710 Self { sent_notifications: Arc::new(Mutex::new(Vec::new())) }
1711 }
1712
1713 fn get_sent_notifications(&self) -> Vec<(PeerId, Vec<u8>)> {
1714 self.sent_notifications.lock().unwrap().clone()
1715 }
1716
1717 fn clear_sent_notifications(&self) {
1718 self.sent_notifications.lock().unwrap().clear();
1719 }
1720 }
1721
1722 #[async_trait::async_trait]
1723 impl NotificationService for TestNotificationService {
1724 async fn open_substream(&mut self, _peer: PeerId) -> Result<(), ()> {
1725 unimplemented!()
1726 }
1727
1728 async fn close_substream(&mut self, _peer: PeerId) -> Result<(), ()> {
1729 unimplemented!()
1730 }
1731
1732 fn send_sync_notification(&mut self, peer: &PeerId, notification: Vec<u8>) {
1733 self.sent_notifications.lock().unwrap().push((*peer, notification));
1734 }
1735
1736 async fn send_async_notification(
1737 &mut self,
1738 peer: &PeerId,
1739 notification: Vec<u8>,
1740 ) -> Result<(), sc_network::error::Error> {
1741 self.sent_notifications.lock().unwrap().push((*peer, notification));
1742 Ok(())
1743 }
1744
1745 async fn set_handshake(&mut self, _handshake: Vec<u8>) -> Result<(), ()> {
1746 unimplemented!()
1747 }
1748
1749 fn try_set_handshake(&mut self, _handshake: Vec<u8>) -> Result<(), ()> {
1750 unimplemented!()
1751 }
1752
1753 async fn next_event(&mut self) -> Option<sc_network::service::traits::NotificationEvent> {
1754 None
1755 }
1756
1757 fn clone(&mut self) -> Result<Box<dyn NotificationService>, ()> {
1758 unimplemented!()
1759 }
1760
1761 fn protocol(&self) -> &sc_network::types::ProtocolName {
1762 unimplemented!()
1763 }
1764
1765 fn message_sink(
1766 &self,
1767 _peer: &PeerId,
1768 ) -> Option<Box<dyn sc_network::service::traits::MessageSink>> {
1769 unimplemented!()
1770 }
1771 }
1772
1773 #[derive(Clone)]
1774 struct TestStatementStore {
1775 statements: Arc<Mutex<HashMap<sp_statement_store::Hash, sp_statement_store::Statement>>>,
1776 recent_statements:
1777 Arc<Mutex<HashMap<sp_statement_store::Hash, sp_statement_store::Statement>>>,
1778 }
1779
1780 impl TestStatementStore {
1781 fn new() -> Self {
1782 Self { statements: Default::default(), recent_statements: Default::default() }
1783 }
1784 }
1785
1786 impl StatementStore for TestStatementStore {
1787 fn statements(
1788 &self,
1789 ) -> sp_statement_store::Result<
1790 Vec<(sp_statement_store::Hash, sp_statement_store::Statement)>,
1791 > {
1792 Ok(self.statements.lock().unwrap().iter().map(|(h, s)| (*h, s.clone())).collect())
1793 }
1794
1795 fn take_recent_statements(
1796 &self,
1797 ) -> sp_statement_store::Result<
1798 Vec<(sp_statement_store::Hash, sp_statement_store::Statement)>,
1799 > {
1800 Ok(self.recent_statements.lock().unwrap().drain().collect())
1801 }
1802
1803 fn statement(
1804 &self,
1805 _hash: &sp_statement_store::Hash,
1806 ) -> sp_statement_store::Result<Option<sp_statement_store::Statement>> {
1807 unimplemented!()
1808 }
1809
1810 fn has_statement(&self, hash: &sp_statement_store::Hash) -> bool {
1811 self.statements.lock().unwrap().contains_key(hash)
1812 }
1813
1814 fn statement_hashes(&self) -> Vec<sp_statement_store::Hash> {
1815 self.statements.lock().unwrap().keys().cloned().collect()
1816 }
1817
1818 fn statements_by_hashes(
1819 &self,
1820 hashes: &[sp_statement_store::Hash],
1821 filter: &mut dyn FnMut(
1822 &sp_statement_store::Hash,
1823 &[u8],
1824 &sp_statement_store::Statement,
1825 ) -> FilterDecision,
1826 ) -> sp_statement_store::Result<(
1827 Vec<(sp_statement_store::Hash, sp_statement_store::Statement)>,
1828 usize,
1829 )> {
1830 let statements = self.statements.lock().unwrap();
1831 let mut result = Vec::new();
1832 let mut processed = 0;
1833 for hash in hashes {
1834 let Some(stmt) = statements.get(hash) else {
1835 processed += 1;
1836 continue;
1837 };
1838 let encoded = stmt.encode();
1839 match filter(hash, &encoded, stmt) {
1840 FilterDecision::Skip => {
1841 processed += 1;
1842 },
1843 FilterDecision::Take => {
1844 processed += 1;
1845 result.push((*hash, stmt.clone()));
1846 },
1847 FilterDecision::Abort => break,
1848 }
1849 }
1850 Ok((result, processed))
1851 }
1852
1853 fn broadcasts(
1854 &self,
1855 _match_all_topics: &[sp_statement_store::Topic],
1856 ) -> sp_statement_store::Result<Vec<Vec<u8>>> {
1857 unimplemented!()
1858 }
1859
1860 fn posted(
1861 &self,
1862 _match_all_topics: &[sp_statement_store::Topic],
1863 _dest: [u8; 32],
1864 ) -> sp_statement_store::Result<Vec<Vec<u8>>> {
1865 unimplemented!()
1866 }
1867
1868 fn posted_clear(
1869 &self,
1870 _match_all_topics: &[sp_statement_store::Topic],
1871 _dest: [u8; 32],
1872 ) -> sp_statement_store::Result<Vec<Vec<u8>>> {
1873 unimplemented!()
1874 }
1875
1876 fn broadcasts_stmt(
1877 &self,
1878 _match_all_topics: &[sp_statement_store::Topic],
1879 ) -> sp_statement_store::Result<Vec<Vec<u8>>> {
1880 unimplemented!()
1881 }
1882
1883 fn posted_stmt(
1884 &self,
1885 _match_all_topics: &[sp_statement_store::Topic],
1886 _dest: [u8; 32],
1887 ) -> sp_statement_store::Result<Vec<Vec<u8>>> {
1888 unimplemented!()
1889 }
1890
1891 fn posted_clear_stmt(
1892 &self,
1893 _match_all_topics: &[sp_statement_store::Topic],
1894 _dest: [u8; 32],
1895 ) -> sp_statement_store::Result<Vec<Vec<u8>>> {
1896 unimplemented!()
1897 }
1898
1899 fn submit(
1900 &self,
1901 _statement: sp_statement_store::Statement,
1902 _source: sp_statement_store::StatementSource,
1903 ) -> sp_statement_store::SubmitResult {
1904 unimplemented!()
1905 }
1906
1907 fn remove(&self, _hash: &sp_statement_store::Hash) -> sp_statement_store::Result<()> {
1908 unimplemented!()
1909 }
1910
1911 fn remove_by(&self, _who: [u8; 32]) -> sp_statement_store::Result<()> {
1912 unimplemented!()
1913 }
1914 }
1915
1916 fn build_handler(
1917 num_peers: usize,
1918 ) -> (
1919 StatementHandler<TestNetwork, TestSync>,
1920 TestStatementStore,
1921 TestNetwork,
1922 TestNotificationService,
1923 async_channel::Receiver<(Statement, oneshot::Sender<SubmitResult>)>,
1924 Vec<PeerId>,
1925 ) {
1926 let statement_store = TestStatementStore::new();
1927 let (queue_sender, queue_receiver) = async_channel::bounded(100);
1928 let network = TestNetwork::new();
1929 let notification_service = TestNotificationService::new();
1930 let mut peers = HashMap::new();
1931 let mut peer_ids = Vec::with_capacity(num_peers);
1932
1933 for _ in 0..num_peers {
1934 let peer_id = PeerId::random();
1935 peer_ids.push(peer_id);
1936 peers.insert(
1937 peer_id,
1938 Peer {
1939 known_statements: LruHashSet::new(NonZeroUsize::new(1000).unwrap()),
1940 rate_limiter: PeerRateLimiter::new(
1941 NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
1942 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero"),
1943 NonZeroU32::new(
1944 DEFAULT_STATEMENTS_PER_SECOND * config::STATEMENTS_BURST_COEFFICIENT,
1945 )
1946 .expect("burst capacity is nonzero"),
1947 ),
1948 protocol_version: PeerProtocolVersion::V1,
1949 topic_affinity: None,
1950 is_light: false,
1951 pending_topic_affinity: None,
1952 },
1953 );
1954 }
1955
1956 let handler = StatementHandler {
1957 protocol_name: format!("/{STATEMENT_PROTOCOL_V1}").into(),
1958 notification_service: Box::new(notification_service.clone()),
1959 propagate_timeout: (Box::pin(futures::stream::pending())
1960 as Pin<Box<dyn Stream<Item = ()> + Send>>)
1961 .fuse(),
1962 pending_statements: FuturesUnordered::new(),
1963 pending_statements_peers: HashMap::new(),
1964 network: network.clone(),
1965 sync: TestSync::new(),
1966 sync_event_stream: (Box::pin(futures::stream::pending())
1967 as Pin<Box<dyn Stream<Item = sc_network_sync::types::SyncEvent> + Send>>)
1968 .fuse(),
1969 peers,
1970 statement_store: Arc::new(statement_store.clone()),
1971 queue_sender,
1972 statements_per_second: NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
1973 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero"),
1974 metrics: None,
1975 initial_sync_timeout: Box::pin(futures::future::pending()),
1976 pending_affinities_timeout: Box::pin(futures::future::pending()),
1977 pending_initial_syncs: HashMap::new(),
1978 initial_sync_peer_queue: VecDeque::new(),
1979 deferred_peers: HashSet::new(),
1980 dropped_statements_during_sync: false,
1981 sync_recovery_peer: None,
1982 sync_recovery_readd_timeout: Box::pin(futures::future::pending()),
1983 };
1984 (handler, statement_store, network, notification_service, queue_receiver, peer_ids)
1985 }
1986
1987 fn get_peer_hashes(sent: &[(PeerId, Vec<u8>)], peer: PeerId) -> Vec<sp_statement_store::Hash> {
1988 sent.iter()
1989 .filter(|(p, _)| *p == peer)
1990 .flat_map(|(_, notification)| {
1991 <Statements as Decode>::decode(&mut notification.as_slice()).unwrap()
1992 })
1993 .map(|s| s.hash())
1994 .collect()
1995 }
1996
1997 async fn dispatch_disconnects(
2000 handler: &mut StatementHandler<TestNetwork, TestSync>,
2001 network: &TestNetwork,
2002 ) {
2003 for peer in network.get_disconnected_peers() {
2004 handler
2005 .handle_notification_event(NotificationEvent::NotificationStreamClosed { peer })
2006 .await;
2007 }
2008 }
2009
2010 #[tokio::test]
2011 async fn test_skips_processing_statements_that_already_in_store() {
2012 let (mut handler, statement_store, _network, _notification_service, queue_receiver, _) =
2013 build_handler(1);
2014
2015 let mut statement1 = Statement::new();
2016 statement1.set_plain_data(b"statement1".to_vec());
2017 let hash1 = statement1.hash();
2018
2019 statement_store.statements.lock().unwrap().insert(hash1, statement1.clone());
2020
2021 let mut statement2 = Statement::new();
2022 statement2.set_plain_data(b"statement2".to_vec());
2023 let hash2 = statement2.hash();
2024
2025 let peer_id = *handler.peers.keys().next().unwrap();
2026
2027 handler.on_statements(peer_id, vec![statement1, statement2]);
2028
2029 let to_submit = queue_receiver.try_recv();
2030 assert_eq!(to_submit.unwrap().0.hash(), hash2, "Expected only statement2 to be queued");
2031
2032 let no_more = queue_receiver.try_recv();
2033 assert!(no_more.is_err(), "Expected only one statement to be queued");
2034 }
2035
2036 #[tokio::test]
2037 async fn test_reports_for_duplicate_statements() {
2038 let (mut handler, statement_store, network, _notification_service, queue_receiver, _) =
2039 build_handler(1);
2040
2041 let peer_id = *handler.peers.keys().next().unwrap();
2042
2043 let mut statement1 = Statement::new();
2044 statement1.set_plain_data(b"statement1".to_vec());
2045
2046 handler.on_statements(peer_id, vec![statement1.clone()]);
2047 {
2048 let (s, _) = queue_receiver.try_recv().unwrap();
2050 let _ = statement_store.statements.lock().unwrap().insert(s.hash(), s);
2051 handler.network.report_peer(peer_id, rep::ANY_STATEMENT_REFUND);
2052 }
2053
2054 handler.on_statements(peer_id, vec![statement1]);
2055
2056 let reports = network.get_reports();
2057 assert_eq!(
2058 reports,
2059 vec![
2060 (peer_id, rep::ANY_STATEMENT), (peer_id, rep::ANY_STATEMENT_REFUND), (peer_id, rep::DUPLICATE_STATEMENT) ],
2064 "Expected ANY_STATEMENT, ANY_STATEMENT_REFUND, DUPLICATE_STATEMENT reputation change, but got: {:?}",
2065 reports
2066 );
2067 }
2068
2069 #[tokio::test]
2070 async fn test_splits_large_batches_into_smaller_chunks() {
2071 let (mut handler, statement_store, _network, notification_service, _queue_receiver, _) =
2072 build_handler(1);
2073
2074 let num_statements = 30;
2075 let statement_size = 100 * 1024; for i in 0..num_statements {
2077 let mut statement = Statement::new();
2078 let mut data = vec![0u8; statement_size];
2079 data[0] = i as u8;
2080 statement.set_plain_data(data);
2081 let hash = statement.hash();
2082 statement_store.recent_statements.lock().unwrap().insert(hash, statement);
2083 }
2084
2085 handler.propagate_statements().await;
2086
2087 let sent = notification_service.get_sent_notifications();
2088 let mut total_statements_sent = 0;
2089 assert!(
2090 sent.len() == 3,
2091 "Expected batch to be split into 3 chunks, but got {} chunks",
2092 sent.len()
2093 );
2094 for (_peer, notification) in sent.iter() {
2095 assert!(
2096 notification.len() <= MAX_STATEMENT_NOTIFICATION_SIZE as usize,
2097 "Notification size {} exceeds limit {}",
2098 notification.len(),
2099 MAX_STATEMENT_NOTIFICATION_SIZE
2100 );
2101 if let Ok(stmts) = <Statements as Decode>::decode(&mut notification.as_slice()) {
2102 total_statements_sent += stmts.len();
2103 }
2104 }
2105
2106 assert_eq!(
2107 total_statements_sent, num_statements,
2108 "Expected all {} statements to be sent, but only {} were sent",
2109 num_statements, total_statements_sent
2110 );
2111 }
2112
2113 #[tokio::test]
2114 async fn test_skips_only_oversized_statements() {
2115 let (mut handler, statement_store, _network, notification_service, _queue_receiver, _) =
2116 build_handler(1);
2117
2118 let mut statement1 = Statement::new();
2119 statement1.set_plain_data(vec![1u8; 100]);
2120 let hash1 = statement1.hash();
2121 statement_store
2122 .recent_statements
2123 .lock()
2124 .unwrap()
2125 .insert(hash1, statement1.clone());
2126
2127 let mut oversized1 = Statement::new();
2128 oversized1.set_plain_data(vec![2u8; MAX_STATEMENT_NOTIFICATION_SIZE as usize * 100]);
2129 let hash_oversized1 = oversized1.hash();
2130 statement_store
2131 .recent_statements
2132 .lock()
2133 .unwrap()
2134 .insert(hash_oversized1, oversized1);
2135
2136 let mut statement2 = Statement::new();
2137 statement2.set_plain_data(vec![3u8; 100]);
2138 let hash2 = statement2.hash();
2139 statement_store
2140 .recent_statements
2141 .lock()
2142 .unwrap()
2143 .insert(hash2, statement2.clone());
2144
2145 let mut oversized2 = Statement::new();
2146 oversized2.set_plain_data(vec![4u8; MAX_STATEMENT_NOTIFICATION_SIZE as usize]);
2147 let hash_oversized2 = oversized2.hash();
2148 statement_store
2149 .recent_statements
2150 .lock()
2151 .unwrap()
2152 .insert(hash_oversized2, oversized2);
2153
2154 let mut statement3 = Statement::new();
2155 statement3.set_plain_data(vec![5u8; 100]);
2156 let hash3 = statement3.hash();
2157 statement_store
2158 .recent_statements
2159 .lock()
2160 .unwrap()
2161 .insert(hash3, statement3.clone());
2162
2163 handler.propagate_statements().await;
2164
2165 let sent = notification_service.get_sent_notifications();
2166
2167 let mut sent_hashes = sent
2168 .iter()
2169 .flat_map(|(_peer, notification)| {
2170 <Statements as Decode>::decode(&mut notification.as_slice()).unwrap()
2171 })
2172 .map(|s| s.hash())
2173 .collect::<Vec<_>>();
2174 sent_hashes.sort();
2175 let mut expected_hashes = vec![hash1, hash2, hash3];
2176 expected_hashes.sort();
2177 assert_eq!(sent_hashes, expected_hashes, "Only small statements should be sent");
2178 }
2179
2180 fn build_handler_no_peers() -> (
2181 StatementHandler<TestNetwork, TestSync>,
2182 TestStatementStore,
2183 TestNetwork,
2184 TestNotificationService,
2185 ) {
2186 let statement_store = TestStatementStore::new();
2187 let (queue_sender, _queue_receiver) = async_channel::bounded(2);
2188 let network = TestNetwork::new();
2189 let notification_service = TestNotificationService::new();
2190
2191 let handler = StatementHandler {
2192 protocol_name: format!("/{STATEMENT_PROTOCOL_V1}").into(),
2193 notification_service: Box::new(notification_service.clone()),
2194 propagate_timeout: (Box::pin(futures::stream::pending())
2195 as Pin<Box<dyn Stream<Item = ()> + Send>>)
2196 .fuse(),
2197 pending_statements: FuturesUnordered::new(),
2198 pending_statements_peers: HashMap::new(),
2199 network: network.clone(),
2200 sync: TestSync::new(),
2201 sync_event_stream: (Box::pin(futures::stream::pending())
2202 as Pin<Box<dyn Stream<Item = sc_network_sync::types::SyncEvent> + Send>>)
2203 .fuse(),
2204 peers: HashMap::new(),
2205 statement_store: Arc::new(statement_store.clone()),
2206 queue_sender,
2207 statements_per_second: NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
2208 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero"),
2209 metrics: None,
2210 initial_sync_timeout: Box::pin(futures::future::pending()),
2211 pending_affinities_timeout: Box::pin(futures::future::pending()),
2212 pending_initial_syncs: HashMap::new(),
2213 initial_sync_peer_queue: VecDeque::new(),
2214 deferred_peers: HashSet::new(),
2215 dropped_statements_during_sync: false,
2216 sync_recovery_peer: None,
2217 sync_recovery_readd_timeout: Box::pin(futures::future::pending()),
2218 };
2219 (handler, statement_store, network, notification_service)
2220 }
2221
2222 fn build_handler_no_peers_light() -> (
2224 StatementHandler<TestNetwork, TestSync>,
2225 TestStatementStore,
2226 TestNetwork,
2227 TestNotificationService,
2228 ) {
2229 let statement_store = TestStatementStore::new();
2230 let (queue_sender, _queue_receiver) = async_channel::bounded(2);
2231 let network = TestNetwork::new_light();
2232 let notification_service = TestNotificationService::new();
2233
2234 let handler = StatementHandler {
2235 protocol_name: format!("/{STATEMENT_PROTOCOL_V1}").into(),
2236 notification_service: Box::new(notification_service.clone()),
2237 propagate_timeout: (Box::pin(futures::stream::pending())
2238 as Pin<Box<dyn Stream<Item = ()> + Send>>)
2239 .fuse(),
2240 pending_statements: FuturesUnordered::new(),
2241 pending_statements_peers: HashMap::new(),
2242 network: network.clone(),
2243 sync: TestSync::new(),
2244 sync_event_stream: (Box::pin(futures::stream::pending())
2245 as Pin<Box<dyn Stream<Item = sc_network_sync::types::SyncEvent> + Send>>)
2246 .fuse(),
2247 peers: HashMap::new(),
2248 statement_store: Arc::new(statement_store.clone()),
2249 queue_sender,
2250 statements_per_second: NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
2251 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero"),
2252 metrics: None,
2253 initial_sync_timeout: Box::pin(futures::future::pending()),
2254 pending_affinities_timeout: Box::pin(futures::future::pending()),
2255 pending_initial_syncs: HashMap::new(),
2256 initial_sync_peer_queue: VecDeque::new(),
2257 deferred_peers: HashSet::new(),
2258 dropped_statements_during_sync: false,
2259 sync_recovery_peer: None,
2260 sync_recovery_readd_timeout: Box::pin(futures::future::pending()),
2261 };
2262 (handler, statement_store, network, notification_service)
2263 }
2264
2265 #[tokio::test]
2266 async fn test_initial_sync_burst_single_peer() {
2267 let (mut handler, statement_store, _network, notification_service, _, _) = build_handler(0);
2268
2269 let num_statements = 200;
2272 let statement_size = 100 * 1024; let mut expected_hashes = Vec::new();
2274 for i in 0..num_statements {
2275 let mut statement = Statement::new();
2276 let mut data = vec![0u8; statement_size];
2277 data[0] = (i % 256) as u8;
2279 data[1] = (i / 256) as u8;
2280 statement.set_plain_data(data);
2281 let hash = statement.hash();
2282 expected_hashes.push(hash);
2283 statement_store.statements.lock().unwrap().insert(hash, statement);
2284 }
2285
2286 let peer_id = PeerId::random();
2288
2289 handler
2290 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
2291 peer: peer_id,
2292 direction: sc_network::service::traits::Direction::Inbound,
2293 handshake: vec![],
2294 negotiated_fallback: Some(format!("/{STATEMENT_PROTOCOL_V1}").into()),
2295 })
2296 .await;
2297
2298 assert!(handler.peers.contains_key(&peer_id));
2300 assert!(handler.pending_initial_syncs.contains_key(&peer_id));
2301 assert_eq!(handler.initial_sync_peer_queue.len(), 1);
2302
2303 let mut burst_count = 0;
2305 while handler.pending_initial_syncs.contains_key(&peer_id) {
2306 handler.process_initial_sync_burst().await;
2307 burst_count += 1;
2308 assert!(burst_count <= 300, "Too many bursts, possible infinite loop");
2310 }
2311
2312 assert!(
2315 burst_count >= 10,
2316 "Expected multiple bursts for 200 statements of 100KB each, got {}",
2317 burst_count
2318 );
2319
2320 let sent = notification_service.get_sent_notifications();
2322 let mut sent_hashes: Vec<_> = sent
2323 .iter()
2324 .flat_map(|(peer, notification)| {
2325 assert_eq!(*peer, peer_id);
2326 <Statements as Decode>::decode(&mut notification.as_slice()).unwrap()
2327 })
2328 .map(|s| s.hash())
2329 .collect();
2330 sent_hashes.sort();
2331 expected_hashes.sort();
2332
2333 assert_eq!(
2334 sent_hashes.len(),
2335 expected_hashes.len(),
2336 "Expected {} statements to be sent, got {}",
2337 expected_hashes.len(),
2338 sent_hashes.len()
2339 );
2340 assert_eq!(sent_hashes, expected_hashes, "All statements should be sent");
2341
2342 assert!(!handler.pending_initial_syncs.contains_key(&peer_id));
2344 assert!(handler.initial_sync_peer_queue.is_empty());
2345 }
2346
2347 #[tokio::test]
2348 async fn test_initial_sync_burst_multiple_peers_round_robin() {
2349 let (mut handler, statement_store, _network, notification_service, _, _) = build_handler(0);
2350
2351 let num_statements = 200;
2353 let statement_size = 100 * 1024; let mut expected_hashes = Vec::new();
2355 for i in 0..num_statements {
2356 let mut statement = Statement::new();
2357 let mut data = vec![0u8; statement_size];
2358 data[0] = (i % 256) as u8;
2359 data[1] = (i / 256) as u8;
2360 statement.set_plain_data(data);
2361 let hash = statement.hash();
2362 expected_hashes.push(hash);
2363 statement_store.statements.lock().unwrap().insert(hash, statement);
2364 }
2365
2366 let peer1 = PeerId::random();
2368 let peer2 = PeerId::random();
2369 let peer3 = PeerId::random();
2370
2371 for peer in [peer1, peer2, peer3] {
2373 handler
2374 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
2375 peer,
2376 direction: sc_network::service::traits::Direction::Inbound,
2377 handshake: vec![],
2378 negotiated_fallback: Some(format!("/{STATEMENT_PROTOCOL_V1}").into()),
2379 })
2380 .await;
2381 }
2382
2383 assert_eq!(handler.peers.len(), 3);
2385 assert_eq!(handler.pending_initial_syncs.len(), 3);
2386 assert_eq!(handler.initial_sync_peer_queue.len(), 3);
2387
2388 let mut peer_burst_order = Vec::new();
2390 let mut burst_count = 0;
2391
2392 while !handler.pending_initial_syncs.is_empty() {
2393 if let Some(&next_peer) = handler.initial_sync_peer_queue.front() {
2395 peer_burst_order.push(next_peer);
2396 }
2397 handler.process_initial_sync_burst().await;
2398 burst_count += 1;
2399 assert!(burst_count <= 500, "Too many bursts, possible infinite loop");
2401 }
2402
2403 assert!(
2406 burst_count >= 30,
2407 "Expected many bursts for 3 peers with 200 statements each, got {}",
2408 burst_count
2409 );
2410
2411 assert!(peer_burst_order.len() >= 9, "Expected at least 9 bursts");
2413 assert_eq!(peer_burst_order[0], peer1, "First burst should be peer1");
2415 assert_eq!(peer_burst_order[1], peer2, "Second burst should be peer2");
2416 assert_eq!(peer_burst_order[2], peer3, "Third burst should be peer3");
2417 assert_eq!(peer_burst_order[3], peer1, "Fourth burst should be peer1");
2419 assert_eq!(peer_burst_order[4], peer2, "Fifth burst should be peer2");
2420 assert_eq!(peer_burst_order[5], peer3, "Sixth burst should be peer3");
2421
2422 let sent = notification_service.get_sent_notifications();
2424 let mut peer1_hashes = get_peer_hashes(&sent, peer1);
2425 let mut peer2_hashes = get_peer_hashes(&sent, peer2);
2426 let mut peer3_hashes = get_peer_hashes(&sent, peer3);
2427
2428 peer1_hashes.sort();
2429 peer2_hashes.sort();
2430 peer3_hashes.sort();
2431 expected_hashes.sort();
2432
2433 assert_eq!(peer1_hashes, expected_hashes, "Peer1 should receive all statements");
2434 assert_eq!(peer2_hashes, expected_hashes, "Peer2 should receive all statements");
2435 assert_eq!(peer3_hashes, expected_hashes, "Peer3 should receive all statements");
2436
2437 assert!(handler.pending_initial_syncs.is_empty());
2439 assert!(handler.initial_sync_peer_queue.is_empty());
2440 }
2441
2442 #[tokio::test]
2443 async fn test_send_statements_in_chunks_exact_max_size() {
2444 let (mut handler, statement_store, _network, notification_service, _queue_receiver, _) =
2445 build_handler(1);
2446
2447 let max_size = MAX_STATEMENT_NOTIFICATION_SIZE as usize - Compact::<u32>::max_encoded_len();
2461 let num_statements: usize = 100;
2462 let per_statement_overhead = 1 + 1 + 8 + 1 + 2; let total_overhead = per_statement_overhead * num_statements;
2464 let total_data_size = max_size - total_overhead;
2465 let per_statement_data_size = total_data_size / num_statements;
2466 let remainder = total_data_size % num_statements;
2467
2468 let mut expected_hashes = Vec::with_capacity(num_statements);
2469 let mut total_encoded_size = 0;
2470
2471 for i in 0..num_statements {
2472 let mut statement = Statement::new();
2473 let extra = if i < remainder { 1 } else { 0 };
2475 let mut data = vec![42u8; per_statement_data_size + extra];
2476 data[0] = i as u8;
2478 data[1] = (i >> 8) as u8;
2479 statement.set_plain_data(data);
2480
2481 total_encoded_size += statement.encoded_size();
2482
2483 let hash = statement.hash();
2484 expected_hashes.push(hash);
2485 statement_store.recent_statements.lock().unwrap().insert(hash, statement);
2486 }
2487
2488 assert!(
2490 total_encoded_size == max_size,
2491 "Total encoded size {} should be <= max_size {}",
2492 total_encoded_size,
2493 max_size
2494 );
2495
2496 handler.propagate_statements().await;
2497
2498 let sent = notification_service.get_sent_notifications();
2499
2500 assert_eq!(
2502 sent.len(),
2503 1,
2504 "Expected 1 notification for all {} statements, but got {}",
2505 num_statements,
2506 sent.len()
2507 );
2508
2509 let (_peer, notification) = &sent[0];
2510 assert!(
2511 notification.len() <= MAX_STATEMENT_NOTIFICATION_SIZE as usize,
2512 "Notification size {} exceeds limit {}",
2513 notification.len(),
2514 MAX_STATEMENT_NOTIFICATION_SIZE
2515 );
2516
2517 let decoded = <Statements as Decode>::decode(&mut notification.as_slice()).unwrap();
2518 assert_eq!(
2519 decoded.len(),
2520 num_statements,
2521 "Expected {} statements in the notification",
2522 num_statements
2523 );
2524
2525 let mut received_hashes: Vec<_> = decoded.iter().map(|s| s.hash()).collect();
2527 expected_hashes.sort();
2528 received_hashes.sort();
2529 assert_eq!(expected_hashes, received_hashes, "All statement hashes should match");
2530 }
2531
2532 #[tokio::test]
2533 async fn test_initial_sync_burst_size_limit_consistency() {
2534 let (mut handler, statement_store, _network, notification_service, _, _) = build_handler(0);
2545
2546 let payload_limit = max_statement_payload_size(V1_ENVELOPE_OVERHEAD);
2548
2549 let first_stmt_data_size = payload_limit / 2 + 10;
2551 let mut stmt1 = Statement::new();
2552 stmt1.set_plain_data(vec![1u8; first_stmt_data_size]);
2553 let stmt1_encoded_size = stmt1.encoded_size();
2554
2555 let remaining = payload_limit.saturating_sub(stmt1_encoded_size);
2558 let target_stmt2_encoded = remaining + 3; let stmt2_data_size = target_stmt2_encoded.saturating_sub(4); let mut stmt2 = Statement::new();
2561 stmt2.set_plain_data(vec![2u8; stmt2_data_size]);
2562 let stmt2_encoded_size = stmt2.encoded_size();
2563
2564 let total_encoded = stmt1_encoded_size + stmt2_encoded_size;
2565
2566 assert!(
2568 total_encoded > payload_limit,
2569 "Total {} should exceed payload_limit {} so filter rejects second statement",
2570 total_encoded,
2571 payload_limit
2572 );
2573
2574 let hash1 = stmt1.hash();
2575 let hash2 = stmt2.hash();
2576 statement_store.statements.lock().unwrap().insert(hash1, stmt1);
2577 statement_store.statements.lock().unwrap().insert(hash2, stmt2);
2578
2579 let peer_id = PeerId::random();
2581
2582 handler
2583 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
2584 peer: peer_id,
2585 direction: sc_network::service::traits::Direction::Inbound,
2586 handshake: vec![],
2587 negotiated_fallback: Some(format!("/{STATEMENT_PROTOCOL_V1}").into()),
2588 })
2589 .await;
2590
2591 assert!(handler.pending_initial_syncs.contains_key(&peer_id));
2593 assert_eq!(handler.pending_initial_syncs.get(&peer_id).unwrap().hashes.len(), 2);
2594
2595 handler.process_initial_sync_burst().await;
2597
2598 let sent = notification_service.get_sent_notifications();
2601 assert_eq!(sent.len(), 1, "First burst should send one notification");
2602
2603 let decoded = <Statements as Decode>::decode(&mut sent[0].1.as_slice()).unwrap();
2604 assert_eq!(decoded.len(), 1, "First notification should contain one statement");
2605
2606 let sent_hash = decoded[0].hash();
2608 assert!(
2609 sent_hash == hash1 || sent_hash == hash2,
2610 "Sent statement should be one of the two created"
2611 );
2612
2613 assert!(handler.pending_initial_syncs.contains_key(&peer_id));
2615 assert_eq!(handler.pending_initial_syncs.get(&peer_id).unwrap().hashes.len(), 1);
2616
2617 handler.process_initial_sync_burst().await;
2619
2620 let sent = notification_service.get_sent_notifications();
2621 assert_eq!(sent.len(), 2, "Second burst should send another notification");
2622
2623 let mut sent_hashes: Vec<_> = sent
2625 .iter()
2626 .flat_map(|(_, notification)| {
2627 <Statements as Decode>::decode(&mut notification.as_slice()).unwrap()
2628 })
2629 .map(|s| s.hash())
2630 .collect();
2631 sent_hashes.sort();
2632 let mut expected_hashes = vec![hash1, hash2];
2633 expected_hashes.sort();
2634 assert_eq!(sent_hashes, expected_hashes, "Both statements should be sent");
2635
2636 assert!(!handler.pending_initial_syncs.contains_key(&peer_id));
2638 }
2639
2640 #[tokio::test]
2641 async fn test_peer_disconnected_on_flooding() {
2642 let (mut handler, _statement_store, network, _notification_service, _queue_receiver, _) =
2643 build_handler(1);
2644
2645 let peer_id = *handler.peers.keys().next().unwrap();
2646
2647 let mut flood_statements = Vec::new();
2648 for i in 0..600_000 {
2649 let mut statement = Statement::new();
2650 statement.set_plain_data(vec![i as u8, (i >> 8) as u8, (i >> 16) as u8]);
2651 flood_statements.push(statement);
2652 }
2653
2654 handler.on_statements(peer_id, flood_statements);
2655
2656 let reports = network.get_reports();
2657 assert!(
2658 reports
2659 .iter()
2660 .any(|(id, rep)| *id == peer_id && *rep == rep::STATEMENT_FLOODING),
2661 "Expected STATEMENT_FLOODING reputation change, but got: {:?}",
2662 reports
2663 );
2664
2665 let disconnected = network.get_disconnected_peers();
2666 assert!(
2667 disconnected.contains(&peer_id),
2668 "Expected peer {} to be disconnected, but it wasn't. Disconnected peers: {:?}",
2669 peer_id,
2670 disconnected
2671 );
2672
2673 dispatch_disconnects(&mut handler, &network).await;
2674
2675 assert!(!handler.peers.contains_key(&peer_id), "Peer should be removed from peers map");
2677 assert!(
2678 !handler.pending_initial_syncs.contains_key(&peer_id),
2679 "Peer should be removed from pending_initial_syncs"
2680 );
2681 assert!(
2682 !handler.initial_sync_peer_queue.contains(&peer_id),
2683 "Peer should be removed from initial_sync_peer_queue"
2684 );
2685 }
2686
2687 #[tokio::test]
2688 async fn test_legitimate_traffic_not_flagged() {
2689 let (mut handler, _statement_store, network, _notification_service, _queue_receiver, _) =
2690 build_handler(1);
2691
2692 let peer_id = *handler.peers.keys().next().unwrap();
2693
2694 let start = std::time::Instant::now();
2695 let duration = std::time::Duration::from_secs(5);
2696 let mut counter = 0u32;
2697
2698 while start.elapsed() < duration {
2699 let mut statements = Vec::new();
2700 for i in 0..5_000 {
2701 let mut statement = Statement::new();
2702 statement.set_plain_data(vec![
2703 counter as u8,
2704 (counter >> 8) as u8,
2705 (counter >> 16) as u8,
2706 i as u8,
2707 ]);
2708 statements.push(statement);
2709 counter = counter.wrapping_add(1);
2710 }
2711
2712 handler.on_statements(peer_id, statements);
2713
2714 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
2715 }
2716
2717 let reports = network.get_reports();
2718 assert!(
2719 !reports
2720 .iter()
2721 .any(|(id, rep)| *id == peer_id && *rep == rep::STATEMENT_FLOODING),
2722 "Legitimate traffic should not trigger flooding detection. Reports: {:?}",
2723 reports
2724 );
2725
2726 let disconnected = network.get_disconnected_peers();
2727 assert!(
2728 !disconnected.contains(&peer_id),
2729 "Legitimate traffic should not cause disconnection. Disconnected peers: {:?}",
2730 disconnected
2731 );
2732
2733 assert!(handler.peers.contains_key(&peer_id), "Peer should still be connected");
2734 }
2735
2736 #[tokio::test]
2737 async fn test_just_over_rate_limit_triggers_flooding() {
2738 let (mut handler, _statement_store, network, _notification_service, _queue_receiver, _) =
2739 build_handler(1);
2740
2741 let peer_id = *handler.peers.keys().next().unwrap();
2742
2743 let mut statements = Vec::new();
2744 for i in 0..260_000 {
2745 let mut statement = Statement::new();
2746 statement.set_plain_data(vec![
2747 i as u8,
2748 (i >> 8) as u8,
2749 (i >> 16) as u8,
2750 (i >> 24) as u8,
2751 ]);
2752 statements.push(statement);
2753 }
2754
2755 handler.on_statements(peer_id, statements);
2756
2757 let reports = network.get_reports();
2758 let expected_burst = DEFAULT_STATEMENTS_PER_SECOND * config::STATEMENTS_BURST_COEFFICIENT;
2759 assert!(
2760 reports
2761 .iter()
2762 .any(|(id, rep)| *id == peer_id && *rep == rep::STATEMENT_FLOODING),
2763 "Sending 260,000 statements should trigger flooding (burst limit: {}). Reports: {:?}",
2764 expected_burst,
2765 reports
2766 );
2767
2768 let disconnected = network.get_disconnected_peers();
2769 assert!(
2770 disconnected.contains(&peer_id),
2771 "Peer should be disconnected after exceeding rate limit. Disconnected: {:?}",
2772 disconnected
2773 );
2774
2775 dispatch_disconnects(&mut handler, &network).await;
2776
2777 assert!(!handler.peers.contains_key(&peer_id), "Peer should be removed from peers map");
2778 }
2779
2780 #[tokio::test]
2781 async fn test_burst_of_250k_statements_allowed() {
2782 let (mut handler, _statement_store, network, _notification_service, _queue_receiver, _) =
2783 build_handler(1);
2784
2785 let peer_id = *handler.peers.keys().next().unwrap();
2786
2787 let mut statements = Vec::new();
2788 for i in 0..250_000 {
2789 let mut statement = Statement::new();
2790 statement.set_plain_data(vec![
2791 i as u8,
2792 (i >> 8) as u8,
2793 (i >> 16) as u8,
2794 (i >> 24) as u8,
2795 ]);
2796 statements.push(statement);
2797 }
2798
2799 handler.on_statements(peer_id, statements);
2800
2801 let reports = network.get_reports();
2802 assert!(
2803 !reports
2804 .iter()
2805 .any(|(id, rep)| *id == peer_id && *rep == rep::STATEMENT_FLOODING),
2806 "250k burst should be allowed (burst = rate × 5). Reports: {:?}",
2807 reports
2808 );
2809
2810 assert!(
2811 handler.peers.contains_key(&peer_id),
2812 "Peer should still be connected after 250k burst"
2813 );
2814 }
2815
2816 #[tokio::test]
2817 async fn test_sustained_rate_above_limit_triggers_flooding() {
2818 let (mut handler, _statement_store, network, _notification_service, _queue_receiver, _) =
2819 build_handler(1);
2820
2821 let peer_id = *handler.peers.keys().next().unwrap();
2822
2823 let mut counter = 0u32;
2824
2825 let start = std::time::Instant::now();
2826 let duration = std::time::Duration::from_secs(5);
2827
2828 let mut flooding_detected = false;
2829 while start.elapsed() < duration {
2830 let mut statements = Vec::new();
2831 for i in 0..30_000 {
2832 let mut statement = Statement::new();
2833 statement.set_plain_data(vec![
2834 counter as u8,
2835 (counter >> 8) as u8,
2836 (counter >> 16) as u8,
2837 i as u8,
2838 ]);
2839 statements.push(statement);
2840 counter = counter.wrapping_add(1);
2841 }
2842
2843 handler.on_statements(peer_id, statements);
2844
2845 let reports = network.get_reports();
2847 if reports
2848 .iter()
2849 .any(|(id, rep)| *id == peer_id && *rep == rep::STATEMENT_FLOODING)
2850 {
2851 flooding_detected = true;
2852 break;
2853 }
2854
2855 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
2856 }
2857
2858 assert!(flooding_detected, "Sustained rate of 300k/sec should trigger flooding");
2859
2860 let disconnected = network.get_disconnected_peers();
2861 assert!(
2862 disconnected.contains(&peer_id),
2863 "Peer should be disconnected after sustained high rate. Disconnected: {:?}",
2864 disconnected
2865 );
2866
2867 dispatch_disconnects(&mut handler, &network).await;
2868
2869 assert!(!handler.peers.contains_key(&peer_id), "Peer should be removed from peers map");
2870 }
2871
2872 #[tokio::test]
2873 async fn test_v2_peer_detected_when_no_fallback() {
2874 let (mut handler, _statement_store, _network, _notification_service) =
2875 build_handler_no_peers();
2876
2877 let peer_id = PeerId::random();
2878
2879 handler
2881 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
2882 peer: peer_id,
2883 direction: sc_network::service::traits::Direction::Inbound,
2884 handshake: vec![],
2885 negotiated_fallback: None,
2886 })
2887 .await;
2888
2889 assert_eq!(
2890 handler.peers.get(&peer_id).unwrap().protocol_version,
2891 PeerProtocolVersion::V2,
2892 "Peer should be detected as v2 when no fallback is negotiated"
2893 );
2894 }
2895
2896 #[tokio::test]
2897 async fn test_v1_peer_detected_when_fallback_negotiated() {
2898 let (mut handler, _statement_store, _network, _notification_service) =
2899 build_handler_no_peers();
2900
2901 let peer_id = PeerId::random();
2902
2903 handler
2905 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
2906 peer: peer_id,
2907 direction: sc_network::service::traits::Direction::Inbound,
2908 handshake: vec![],
2909 negotiated_fallback: Some(format!("/{STATEMENT_PROTOCOL_V1}").into()),
2910 })
2911 .await;
2912
2913 assert_eq!(
2914 handler.peers.get(&peer_id).unwrap().protocol_version,
2915 PeerProtocolVersion::V1,
2916 "Peer should be detected as v1 when fallback is negotiated"
2917 );
2918 }
2919
2920 #[tokio::test]
2921 async fn test_v1_peer_decodes_raw_statements() {
2922 let (mut handler, _statement_store, _network, _notification_service) =
2923 build_handler_no_peers();
2924
2925 let peer_id = PeerId::random();
2926 let (queue_sender, queue_receiver) = async_channel::bounded(10);
2927 handler.queue_sender = queue_sender;
2928
2929 handler
2931 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
2932 peer: peer_id,
2933 direction: sc_network::service::traits::Direction::Inbound,
2934 handshake: vec![],
2935 negotiated_fallback: Some(format!("/{STATEMENT_PROTOCOL_V1}").into()),
2936 })
2937 .await;
2938
2939 let mut statement = Statement::new();
2941 statement.set_plain_data(b"v1 statement".to_vec());
2942 let hash = statement.hash();
2943 let raw_encoded = vec![statement].encode();
2944
2945 handler
2946 .handle_notification_event(NotificationEvent::NotificationReceived {
2947 peer: peer_id,
2948 notification: raw_encoded.into(),
2949 })
2950 .await;
2951
2952 let (received, _) = queue_receiver.try_recv().unwrap();
2953 assert_eq!(received.hash(), hash, "V1 peer's raw statement should be decoded correctly");
2954 }
2955
2956 #[tokio::test]
2957 async fn test_v2_peer_decodes_statement_message() {
2958 let (mut handler, _statement_store, _network, _notification_service) =
2959 build_handler_no_peers();
2960
2961 let peer_id = PeerId::random();
2962 let (queue_sender, queue_receiver) = async_channel::bounded(10);
2963 handler.queue_sender = queue_sender;
2964
2965 handler
2967 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
2968 peer: peer_id,
2969 direction: sc_network::service::traits::Direction::Inbound,
2970 handshake: vec![],
2971 negotiated_fallback: None,
2972 })
2973 .await;
2974
2975 let mut statement = Statement::new();
2977 statement.set_plain_data(b"v2 statement".to_vec());
2978 let hash = statement.hash();
2979 let msg = StatementMessage::Statements(vec![statement]);
2980 let encoded = msg.encode();
2981
2982 handler
2983 .handle_notification_event(NotificationEvent::NotificationReceived {
2984 peer: peer_id,
2985 notification: encoded.into(),
2986 })
2987 .await;
2988
2989 let (received, _) = queue_receiver.try_recv().unwrap();
2990 assert_eq!(received.hash(), hash, "V2 peer's StatementMessage should be decoded correctly");
2991 }
2992
2993 #[tokio::test]
2994 async fn test_v2_peer_topic_affinity_stored() {
2995 let (mut handler, _statement_store, _network, _notification_service) =
2996 build_handler_no_peers();
2997
2998 let peer_id = PeerId::random();
2999
3000 handler
3002 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
3003 peer: peer_id,
3004 direction: sc_network::service::traits::Direction::Inbound,
3005 handshake: vec![],
3006 negotiated_fallback: None,
3007 })
3008 .await;
3009
3010 assert!(
3011 handler.peers.get(&peer_id).unwrap().topic_affinity.is_none(),
3012 "Topic affinity should be None initially"
3013 );
3014
3015 let topic: [u8; 32] = [0xAA; 32];
3017 let mut filter = AffinityFilter::new(BLOOM_SEED, 0.01, 100);
3018 filter.insert(&topic);
3019 let msg = StatementMessage::ExplicitTopicAffinity(filter);
3020 let encoded = msg.encode();
3021
3022 handler
3023 .handle_notification_event(NotificationEvent::NotificationReceived {
3024 peer: peer_id,
3025 notification: encoded.into(),
3026 })
3027 .await;
3028
3029 handler.process_pending_affinities();
3031
3032 let peer_data = handler.peers.get(&peer_id).unwrap();
3033 assert!(
3034 peer_data.topic_affinity.is_some(),
3035 "Topic affinity should be set after receiving ExplicitTopicAffinity"
3036 );
3037 assert!(
3039 peer_data.topic_affinity.as_ref().unwrap().contains(&topic),
3040 "Stored affinity filter should match the topic"
3041 );
3042 }
3043
3044 #[tokio::test]
3045 async fn test_topic_affinity_filters_propagation() {
3046 let (mut handler, statement_store, _network, notification_service) =
3047 build_handler_no_peers();
3048
3049 let peer_id = PeerId::random();
3050
3051 handler
3053 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
3054 peer: peer_id,
3055 direction: sc_network::service::traits::Direction::Inbound,
3056 handshake: vec![],
3057 negotiated_fallback: None,
3058 })
3059 .await;
3060
3061 let topic_aa: [u8; 32] = [0xAA; 32];
3063 let topic_bb: [u8; 32] = [0xBB; 32];
3064 let mut filter = AffinityFilter::new(BLOOM_SEED, 0.01, 100);
3065 filter.insert(&topic_aa);
3066 let msg = StatementMessage::ExplicitTopicAffinity(filter);
3067 let encoded = msg.encode();
3068 handler
3069 .handle_notification_event(NotificationEvent::NotificationReceived {
3070 peer: peer_id,
3071 notification: encoded.into(),
3072 })
3073 .await;
3074
3075 handler.process_pending_affinities();
3077
3078 let mut stmt_matching = Statement::new();
3080 stmt_matching.set_plain_data(b"matching".to_vec());
3081 stmt_matching.set_topic(0, topic_aa.into());
3082 let hash_matching = stmt_matching.hash();
3083
3084 let mut stmt_not_matching = Statement::new();
3085 stmt_not_matching.set_plain_data(b"not matching".to_vec());
3086 stmt_not_matching.set_topic(0, topic_bb.into());
3087 let hash_not_matching = stmt_not_matching.hash();
3088
3089 let mut stmt_no_topic = Statement::new();
3090 stmt_no_topic.set_plain_data(b"no topic".to_vec());
3091 let hash_no_topic = stmt_no_topic.hash();
3092
3093 statement_store
3094 .recent_statements
3095 .lock()
3096 .unwrap()
3097 .insert(hash_matching, stmt_matching);
3098 statement_store
3099 .recent_statements
3100 .lock()
3101 .unwrap()
3102 .insert(hash_not_matching, stmt_not_matching);
3103 statement_store
3104 .recent_statements
3105 .lock()
3106 .unwrap()
3107 .insert(hash_no_topic, stmt_no_topic);
3108
3109 handler.propagate_statements().await;
3110
3111 let sent = notification_service.get_sent_notifications();
3112 let mut sent_hashes: Vec<_> = sent
3113 .iter()
3114 .flat_map(|(_, notification)| {
3115 match StatementMessage::decode(&mut notification.as_slice()).unwrap() {
3117 StatementMessage::Statements(stmts) => stmts,
3118 _ => panic!("Expected StatementMessage::Statements"),
3119 }
3120 })
3121 .map(|s| s.hash())
3122 .collect();
3123 sent_hashes.sort();
3124
3125 assert!(
3127 sent_hashes.contains(&hash_matching),
3128 "Statement matching topic affinity should be propagated"
3129 );
3130 assert!(
3131 sent_hashes.contains(&hash_no_topic),
3132 "Statement with no topics should be propagated (broadcast)"
3133 );
3134 assert!(
3135 !sent_hashes.contains(&hash_not_matching),
3136 "Statement NOT matching topic affinity should be filtered out"
3137 );
3138 }
3139
3140 #[tokio::test]
3141 async fn test_v1_peer_no_topic_filtering() {
3142 let (mut handler, statement_store, _network, notification_service) =
3143 build_handler_no_peers();
3144
3145 let peer_id = PeerId::random();
3146
3147 handler
3149 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
3150 peer: peer_id,
3151 direction: sc_network::service::traits::Direction::Inbound,
3152 handshake: vec![],
3153 negotiated_fallback: Some(format!("/{STATEMENT_PROTOCOL_V1}").into()),
3154 })
3155 .await;
3156
3157 let topic_aa: [u8; 32] = [0xAA; 32];
3159 let mut stmt_with_topic = Statement::new();
3160 stmt_with_topic.set_plain_data(b"with topic".to_vec());
3161 stmt_with_topic.set_topic(0, topic_aa.into());
3162 let hash_with_topic = stmt_with_topic.hash();
3163
3164 let mut stmt_no_topic = Statement::new();
3165 stmt_no_topic.set_plain_data(b"no topic".to_vec());
3166 let hash_no_topic = stmt_no_topic.hash();
3167
3168 statement_store
3169 .recent_statements
3170 .lock()
3171 .unwrap()
3172 .insert(hash_with_topic, stmt_with_topic);
3173 statement_store
3174 .recent_statements
3175 .lock()
3176 .unwrap()
3177 .insert(hash_no_topic, stmt_no_topic);
3178
3179 handler.propagate_statements().await;
3180
3181 let sent = notification_service.get_sent_notifications();
3182 let sent_hashes: Vec<_> = sent
3183 .iter()
3184 .flat_map(|(_, notification)| {
3185 <Statements as Decode>::decode(&mut notification.as_slice()).unwrap()
3186 })
3187 .map(|s| s.hash())
3188 .collect();
3189
3190 assert_eq!(
3191 sent_hashes.len(),
3192 2,
3193 "V1 peer should receive all statements regardless of topics"
3194 );
3195 assert!(sent_hashes.contains(&hash_with_topic));
3196 assert!(sent_hashes.contains(&hash_no_topic));
3197 }
3198
3199 #[tokio::test]
3200 async fn test_affinity_change_triggers_resync() {
3201 let (mut handler, statement_store, _network, notification_service) =
3202 build_handler_no_peers_light();
3203
3204 let peer_id = PeerId::random();
3205
3206 let topic_aa: [u8; 32] = [0xAA; 32];
3208 let topic_bb: [u8; 32] = [0xBB; 32];
3209
3210 let mut stmt_aa = Statement::new();
3211 stmt_aa.set_plain_data(b"stmt_aa".to_vec());
3212 stmt_aa.set_topic(0, topic_aa.into());
3213 let hash_aa = stmt_aa.hash();
3214
3215 let mut stmt_bb = Statement::new();
3216 stmt_bb.set_plain_data(b"stmt_bb".to_vec());
3217 stmt_bb.set_topic(0, topic_bb.into());
3218 let hash_bb = stmt_bb.hash();
3219
3220 let mut stmt_no_topic = Statement::new();
3221 stmt_no_topic.set_plain_data(b"no topic".to_vec());
3222 let hash_no_topic = stmt_no_topic.hash();
3223
3224 statement_store.statements.lock().unwrap().insert(hash_aa, stmt_aa);
3225 statement_store.statements.lock().unwrap().insert(hash_bb, stmt_bb);
3226 statement_store.statements.lock().unwrap().insert(hash_no_topic, stmt_no_topic);
3227
3228 handler
3230 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
3231 peer: peer_id,
3232 direction: sc_network::service::traits::Direction::Inbound,
3233 handshake: vec![],
3234 negotiated_fallback: None,
3235 })
3236 .await;
3237
3238 assert!(
3240 !handler.pending_initial_syncs.contains_key(&peer_id),
3241 "Light V2 peer should NOT have initial sync scheduled on connect"
3242 );
3243
3244 let mut filter = AffinityFilter::new(BLOOM_SEED, 0.01, 100);
3246 filter.insert(&topic_aa);
3247 let msg = StatementMessage::ExplicitTopicAffinity(filter);
3248 let encoded = msg.encode();
3249 handler
3250 .handle_notification_event(NotificationEvent::NotificationReceived {
3251 peer: peer_id,
3252 notification: encoded.into(),
3253 })
3254 .await;
3255
3256 handler.process_pending_affinities();
3258
3259 assert!(
3260 handler.pending_initial_syncs.contains_key(&peer_id),
3261 "Initial sync should be scheduled after setting affinity"
3262 );
3263
3264 while handler.pending_initial_syncs.contains_key(&peer_id) {
3266 handler.process_initial_sync_burst().await;
3267 }
3268
3269 let sent = notification_service.get_sent_notifications();
3270 let sent_hashes: HashSet<_> = sent
3271 .iter()
3272 .flat_map(|(_, notification)| {
3273 match StatementMessage::decode(&mut notification.as_slice()).unwrap() {
3274 StatementMessage::Statements(stmts) => stmts,
3275 _ => panic!("Expected StatementMessage::Statements"),
3276 }
3277 })
3278 .map(|s| s.hash())
3279 .collect();
3280 assert!(sent_hashes.contains(&hash_aa), "stmt_aa should be sent (matches affinity)");
3281 assert!(
3282 sent_hashes.contains(&hash_no_topic),
3283 "stmt_no_topic should be sent (broadcast, no topic)"
3284 );
3285 assert!(!sent_hashes.contains(&hash_bb), "stmt_bb should NOT be sent (filtered)");
3286
3287 let mut filter = AffinityFilter::new(BLOOM_SEED, 0.01, 100);
3289 filter.insert(&topic_bb);
3290 let msg = StatementMessage::ExplicitTopicAffinity(filter);
3291 let encoded = msg.encode();
3292 handler
3293 .handle_notification_event(NotificationEvent::NotificationReceived {
3294 peer: peer_id,
3295 notification: encoded.into(),
3296 })
3297 .await;
3298
3299 handler.process_pending_affinities();
3301
3302 assert!(
3303 handler.pending_initial_syncs.contains_key(&peer_id),
3304 "Initial sync should be re-scheduled after affinity change"
3305 );
3306
3307 notification_service.clear_sent_notifications();
3308 while handler.pending_initial_syncs.contains_key(&peer_id) {
3309 handler.process_initial_sync_burst().await;
3310 }
3311
3312 let sent_after_bb = notification_service.get_sent_notifications();
3313 let sent_hashes_bb: HashSet<_> = sent_after_bb
3314 .iter()
3315 .flat_map(|(_, notification)| {
3316 match StatementMessage::decode(&mut notification.as_slice()).unwrap() {
3317 StatementMessage::Statements(stmts) => stmts,
3318 _ => panic!("Expected StatementMessage::Statements"),
3319 }
3320 })
3321 .map(|s| s.hash())
3322 .collect();
3323 assert!(
3325 sent_hashes_bb.contains(&hash_bb),
3326 "stmt_bb should now be sent after affinity changed to topic_bb"
3327 );
3328 assert!(
3330 sent_hashes_bb.contains(&hash_no_topic),
3331 "stmt_no_topic should be re-sent (known_statements cleared on affinity change)"
3332 );
3333 }
3334
3335 #[tokio::test]
3336 async fn test_affinity_change_sends_previously_filtered_statements() {
3337 let (mut handler, statement_store, _network, notification_service) =
3342 build_handler_no_peers_light();
3343
3344 let peer_id = PeerId::random();
3345
3346 let topic_aa: [u8; 32] = [0xAA; 32];
3347 let topic_bb: [u8; 32] = [0xBB; 32];
3348
3349 let mut stmt_aa = Statement::new();
3350 stmt_aa.set_plain_data(b"stmt_aa".to_vec());
3351 stmt_aa.set_topic(0, topic_aa.into());
3352 let hash_aa = stmt_aa.hash();
3353
3354 let mut stmt_bb = Statement::new();
3355 stmt_bb.set_plain_data(b"stmt_bb".to_vec());
3356 stmt_bb.set_topic(0, topic_bb.into());
3357 let hash_bb = stmt_bb.hash();
3358
3359 statement_store.statements.lock().unwrap().insert(hash_aa, stmt_aa.clone());
3360 statement_store.statements.lock().unwrap().insert(hash_bb, stmt_bb.clone());
3361
3362 statement_store.recent_statements.lock().unwrap().insert(hash_aa, stmt_aa);
3364 statement_store.recent_statements.lock().unwrap().insert(hash_bb, stmt_bb);
3365
3366 handler
3368 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
3369 peer: peer_id,
3370 direction: sc_network::service::traits::Direction::Inbound,
3371 handshake: vec![],
3372 negotiated_fallback: None,
3373 })
3374 .await;
3375
3376 let mut filter = AffinityFilter::new(BLOOM_SEED, 0.01, 100);
3378 filter.insert(&topic_aa);
3379 let msg = StatementMessage::ExplicitTopicAffinity(filter);
3380 let encoded = msg.encode();
3381 handler
3382 .handle_notification_event(NotificationEvent::NotificationReceived {
3383 peer: peer_id,
3384 notification: encoded.into(),
3385 })
3386 .await;
3387
3388 handler.process_pending_affinities();
3390
3391 while handler.pending_initial_syncs.contains_key(&peer_id) {
3393 handler.process_initial_sync_burst().await;
3394 }
3395
3396 let sent = notification_service.get_sent_notifications();
3397 let sent_hashes: HashSet<_> = sent
3398 .iter()
3399 .flat_map(|(_, notification)| {
3400 match StatementMessage::decode(&mut notification.as_slice()).unwrap() {
3401 StatementMessage::Statements(stmts) => stmts,
3402 _ => panic!("Expected StatementMessage::Statements"),
3403 }
3404 })
3405 .map(|s| s.hash())
3406 .collect();
3407 assert!(sent_hashes.contains(&hash_aa), "stmt_aa should be sent (matches affinity)");
3408 assert!(
3409 !sent_hashes.contains(&hash_bb),
3410 "stmt_bb should NOT be sent (filtered by affinity)"
3411 );
3412
3413 handler.propagate_statements().await;
3416
3417 let peer = handler.peers.get(&peer_id).unwrap();
3419 assert!(
3420 !peer.known_statements.contains(&hash_bb),
3421 "stmt_bb should NOT be in known_statements (filtered by affinity)"
3422 );
3423 assert!(peer.known_statements.contains(&hash_aa), "stmt_aa should be in known_statements");
3424
3425 let mut filter = AffinityFilter::new(BLOOM_SEED, 0.01, 100);
3427 filter.insert(&topic_aa);
3428 filter.insert(&topic_bb);
3429 let msg = StatementMessage::ExplicitTopicAffinity(filter);
3430 let encoded = msg.encode();
3431
3432 notification_service.clear_sent_notifications();
3433 handler
3434 .handle_notification_event(NotificationEvent::NotificationReceived {
3435 peer: peer_id,
3436 notification: encoded.into(),
3437 })
3438 .await;
3439
3440 handler.process_pending_affinities();
3442
3443 while handler.pending_initial_syncs.contains_key(&peer_id) {
3445 handler.process_initial_sync_burst().await;
3446 }
3447
3448 let sent = notification_service.get_sent_notifications();
3449 let sent_hashes: HashSet<_> = sent
3450 .iter()
3451 .flat_map(|(_, notification)| {
3452 match StatementMessage::decode(&mut notification.as_slice()).unwrap() {
3453 StatementMessage::Statements(stmts) => stmts,
3454 _ => panic!("Expected StatementMessage::Statements"),
3455 }
3456 })
3457 .map(|s| s.hash())
3458 .collect();
3459 assert!(
3460 sent_hashes.contains(&hash_bb),
3461 "stmt_bb should now be sent after affinity expanded to include topic_bb"
3462 );
3463 assert!(
3465 sent_hashes.contains(&hash_aa),
3466 "stmt_aa should be re-sent (known_statements cleared on affinity change)"
3467 );
3468 }
3469
3470 #[test]
3471 fn test_encode_statement_refs_matches_derive_encoding() {
3472 let mut stmt1 = Statement::new();
3473 stmt1.set_plain_data(b"first".to_vec());
3474 let mut stmt2 = Statement::new();
3475 stmt2.set_plain_data(b"second".to_vec());
3476
3477 let refs: Vec<&Statement> = vec![&stmt1, &stmt2];
3478 let hand_rolled = StatementMessage::encode_statement_refs(&refs);
3479 let derive_encoded = StatementMessage::Statements(vec![stmt1, stmt2]).encode();
3480
3481 assert_eq!(
3482 hand_rolled, derive_encoded,
3483 "encode_statement_refs must produce identical bytes to derive Encode"
3484 );
3485 }
3486
3487 #[test]
3488 fn test_encode_statement_refs_empty() {
3489 let refs: Vec<&Statement> = vec![];
3490 let hand_rolled = StatementMessage::encode_statement_refs(&refs);
3491 let derive_encoded = StatementMessage::Statements(vec![]).encode();
3492
3493 assert_eq!(hand_rolled, derive_encoded);
3494 }
3495
3496 #[test]
3497 fn test_can_receive_all_combinations() {
3498 let make_peer = |is_light: bool, version: PeerProtocolVersion, has_affinity: bool| {
3499 let topic_affinity = has_affinity.then(|| AffinityFilter::new(BLOOM_SEED, 0.01, 10));
3500 Peer {
3501 known_statements: LruHashSet::new(NonZeroUsize::new(10).unwrap()),
3502 rate_limiter: PeerRateLimiter::new(
3503 NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND).expect("nonzero"),
3504 NonZeroU32::new(
3505 DEFAULT_STATEMENTS_PER_SECOND * config::STATEMENTS_BURST_COEFFICIENT,
3506 )
3507 .expect("nonzero"),
3508 ),
3509 protocol_version: version,
3510 topic_affinity,
3511 is_light,
3512 pending_topic_affinity: None,
3513 }
3514 };
3515
3516 assert!(make_peer(false, PeerProtocolVersion::V1, false).can_receive());
3518 assert!(make_peer(false, PeerProtocolVersion::V2, false).can_receive());
3520 assert!(make_peer(true, PeerProtocolVersion::V1, false).can_receive());
3522 assert!(!make_peer(true, PeerProtocolVersion::V2, false).can_receive());
3524 assert!(make_peer(true, PeerProtocolVersion::V2, true).can_receive());
3526 assert!(make_peer(false, PeerProtocolVersion::V2, true).can_receive());
3528 }
3529
3530 #[tokio::test]
3531 async fn test_send_chunk_v1_vs_v2_encoding() {
3532 let (mut handler, _statement_store, _network, notification_service) =
3533 build_handler_no_peers();
3534
3535 let v1_peer = PeerId::random();
3536 let v2_peer = PeerId::random();
3537
3538 handler
3540 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
3541 peer: v1_peer,
3542 direction: sc_network::service::traits::Direction::Inbound,
3543 handshake: vec![],
3544 negotiated_fallback: Some(format!("/{STATEMENT_PROTOCOL_V1}").into()),
3545 })
3546 .await;
3547
3548 handler
3550 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
3551 peer: v2_peer,
3552 direction: sc_network::service::traits::Direction::Inbound,
3553 handshake: vec![],
3554 negotiated_fallback: None,
3555 })
3556 .await;
3557
3558 let mut stmt = Statement::new();
3559 stmt.set_plain_data(b"encoding test".to_vec());
3560
3561 notification_service.clear_sent_notifications();
3563 handler.send_statement_chunk(&v1_peer, &[&stmt]).await;
3564 let v1_sent = notification_service.get_sent_notifications();
3565 assert_eq!(v1_sent.len(), 1);
3566 let v1_bytes = &v1_sent[0].1;
3567 let decoded_v1 = <Statements as Decode>::decode(&mut v1_bytes.as_slice())
3569 .expect("V1 peer should receive raw Vec<Statement> encoding");
3570 assert_eq!(decoded_v1.len(), 1);
3571
3572 notification_service.clear_sent_notifications();
3574 handler.send_statement_chunk(&v2_peer, &[&stmt]).await;
3575 let v2_sent = notification_service.get_sent_notifications();
3576 assert_eq!(v2_sent.len(), 1);
3577 let v2_bytes = &v2_sent[0].1;
3578 let decoded_v2 = StatementMessage::decode(&mut v2_bytes.as_slice())
3580 .expect("V2 peer should receive StatementMessage encoding");
3581 match decoded_v2 {
3582 StatementMessage::Statements(stmts) => assert_eq!(stmts.len(), 1),
3583 _ => panic!("Expected StatementMessage::Statements for V2 peer"),
3584 }
3585
3586 assert_ne!(v1_bytes, v2_bytes, "V1 and V2 encodings should differ");
3588 }
3589
3590 #[tokio::test]
3591 async fn test_schedule_initial_sync_replaces_existing() {
3592 let (mut handler, statement_store, _network, _notification_service) =
3593 build_handler_no_peers();
3594
3595 let peer_id = PeerId::random();
3596
3597 let mut stmt1 = Statement::new();
3599 stmt1.set_plain_data(b"stmt1".to_vec());
3600 let hash1 = stmt1.hash();
3601 statement_store.statements.lock().unwrap().insert(hash1, stmt1);
3602
3603 handler
3605 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
3606 peer: peer_id,
3607 direction: sc_network::service::traits::Direction::Inbound,
3608 handshake: vec![],
3609 negotiated_fallback: Some(format!("/{STATEMENT_PROTOCOL_V1}").into()),
3610 })
3611 .await;
3612
3613 assert!(handler.pending_initial_syncs.contains_key(&peer_id));
3615 assert_eq!(
3616 handler.initial_sync_peer_queue.iter().filter(|p| **p == peer_id).count(),
3617 1,
3618 "Peer should appear exactly once in the queue"
3619 );
3620
3621 let mut stmt2 = Statement::new();
3623 stmt2.set_plain_data(b"stmt2".to_vec());
3624 let hash2 = stmt2.hash();
3625 statement_store.statements.lock().unwrap().insert(hash2, stmt2);
3626
3627 handler.schedule_initial_sync_for_peer(peer_id);
3628
3629 assert_eq!(
3631 handler.initial_sync_peer_queue.iter().filter(|p| **p == peer_id).count(),
3632 1,
3633 "Peer should NOT be duplicated in the queue after re-schedule"
3634 );
3635 let pending = handler.pending_initial_syncs.get(&peer_id).unwrap();
3637 assert!(pending.hashes.contains(&hash1));
3638 assert!(pending.hashes.contains(&hash2));
3639 }
3640
3641 #[tokio::test]
3642 async fn test_initial_sync_queued_during_major_sync_processed_after() {
3643 let statement_store = TestStatementStore::new();
3644 let (queue_sender, _queue_receiver) = async_channel::bounded(2);
3645 let network = TestNetwork::new();
3646 let notification_service = TestNotificationService::new();
3647 let sync = TestSync::new();
3648 sync.major_syncing.store(true, Ordering::Relaxed);
3650
3651 let mut handler = StatementHandler {
3652 protocol_name: format!("/{STATEMENT_PROTOCOL_V1}").into(),
3653 notification_service: Box::new(notification_service.clone()),
3654 propagate_timeout: (Box::pin(futures::stream::pending())
3655 as Pin<Box<dyn Stream<Item = ()> + Send>>)
3656 .fuse(),
3657 pending_statements: FuturesUnordered::new(),
3658 pending_statements_peers: HashMap::new(),
3659 network: network.clone(),
3660 sync: sync.clone(),
3661 sync_event_stream: (Box::pin(futures::stream::pending())
3662 as Pin<Box<dyn Stream<Item = sc_network_sync::types::SyncEvent> + Send>>)
3663 .fuse(),
3664 peers: HashMap::new(),
3665 statement_store: Arc::new(statement_store.clone()),
3666 queue_sender,
3667 statements_per_second: NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
3668 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero"),
3669 metrics: None,
3670 initial_sync_timeout: Box::pin(futures::future::pending()),
3671 pending_affinities_timeout: Box::pin(futures::future::pending()),
3672 pending_initial_syncs: HashMap::new(),
3673 initial_sync_peer_queue: VecDeque::new(),
3674 deferred_peers: HashSet::new(),
3675 dropped_statements_during_sync: false,
3676 sync_recovery_peer: None,
3677 sync_recovery_readd_timeout: Box::pin(futures::future::pending()),
3678 };
3679
3680 let mut stmt = Statement::new();
3682 stmt.set_plain_data(b"during major sync".to_vec());
3683 let hash = stmt.hash();
3684 statement_store.statements.lock().unwrap().insert(hash, stmt);
3685
3686 let peer_id = PeerId::random();
3688 handler.peers.insert(
3689 peer_id,
3690 Peer::new_for_testing(
3691 LruHashSet::new(NonZeroUsize::new(100).unwrap()),
3692 NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND).unwrap(),
3693 NonZeroU32::new(
3694 DEFAULT_STATEMENTS_PER_SECOND * config::STATEMENTS_BURST_COEFFICIENT,
3695 )
3696 .unwrap(),
3697 ),
3698 );
3699
3700 handler.schedule_initial_sync_for_peer(peer_id);
3702
3703 assert!(
3704 handler.pending_initial_syncs.contains_key(&peer_id),
3705 "Initial sync should be queued even during major sync"
3706 );
3707 assert_eq!(handler.initial_sync_peer_queue.len(), 1);
3708
3709 handler.process_initial_sync_burst().await;
3711 assert!(
3712 handler.pending_initial_syncs.contains_key(&peer_id),
3713 "Pending sync should remain untouched during major sync"
3714 );
3715
3716 sync.major_syncing.store(false, Ordering::Relaxed);
3718 handler.process_initial_sync_burst().await;
3719 assert!(
3720 handler.initial_sync_peer_queue.is_empty(),
3721 "Peer should have been processed after major sync ended"
3722 );
3723 }
3724
3725 #[tokio::test]
3726 async fn test_schedule_initial_sync_resends_all_matching() {
3727 let (mut handler, statement_store, _network, _notification_service) =
3728 build_handler_no_peers();
3729
3730 let peer_id = PeerId::random();
3731
3732 let mut stmt1 = Statement::new();
3734 stmt1.set_plain_data(b"known".to_vec());
3735 let hash1 = stmt1.hash();
3736 let mut stmt2 = Statement::new();
3737 stmt2.set_plain_data(b"unknown".to_vec());
3738 let hash2 = stmt2.hash();
3739
3740 statement_store.statements.lock().unwrap().insert(hash1, stmt1);
3741 statement_store.statements.lock().unwrap().insert(hash2, stmt2);
3742
3743 let mut known = LruHashSet::new(NonZeroUsize::new(100).unwrap());
3745 known.insert(hash1);
3746 handler.peers.insert(
3747 peer_id,
3748 Peer {
3749 known_statements: known,
3750 rate_limiter: PeerRateLimiter::new(
3751 NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND).unwrap(),
3752 NonZeroU32::new(
3753 DEFAULT_STATEMENTS_PER_SECOND * config::STATEMENTS_BURST_COEFFICIENT,
3754 )
3755 .unwrap(),
3756 ),
3757 protocol_version: PeerProtocolVersion::V1,
3758 topic_affinity: None,
3759 is_light: false,
3760 pending_topic_affinity: None,
3761 },
3762 );
3763
3764 handler.schedule_initial_sync_for_peer(peer_id);
3765
3766 let pending = handler.pending_initial_syncs.get(&peer_id).unwrap();
3767 assert!(
3769 pending.hashes.contains(&hash1),
3770 "Previously known hash should be included after affinity change"
3771 );
3772 assert!(pending.hashes.contains(&hash2), "Unknown hash should be included in initial sync");
3773 let peer_data = handler.peers.get(&peer_id).unwrap();
3775 assert!(
3776 !peer_data.known_statements.contains(&hash1),
3777 "known_statements should be cleared after schedule_initial_sync_for_peer"
3778 );
3779 }
3780
3781 #[tokio::test]
3782 async fn test_malformed_v2_message_does_not_panic() {
3783 let (mut handler, _statement_store, _network, _notification_service) =
3784 build_handler_no_peers();
3785
3786 let peer_id = PeerId::random();
3787
3788 handler
3790 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
3791 peer: peer_id,
3792 direction: sc_network::service::traits::Direction::Inbound,
3793 handshake: vec![],
3794 negotiated_fallback: None,
3795 })
3796 .await;
3797
3798 handler
3800 .handle_notification_event(NotificationEvent::NotificationReceived {
3801 peer: peer_id,
3802 notification: vec![0xFF, 0xFE, 0xFD].into(),
3803 })
3804 .await;
3805
3806 let mut stmt = Statement::new();
3808 stmt.set_plain_data(b"v1 encoded".to_vec());
3809 let v1_encoded = vec![stmt].encode();
3810 handler
3811 .handle_notification_event(NotificationEvent::NotificationReceived {
3812 peer: peer_id,
3813 notification: v1_encoded.into(),
3814 })
3815 .await;
3816
3817 assert!(handler.peers.contains_key(&peer_id), "Peer should still be connected");
3819 }
3820
3821 #[test]
3822 fn test_find_sendable_chunk_v2_overhead() {
3823 let v1_max = max_statement_payload_size(V1_ENVELOPE_OVERHEAD);
3824 let v2_max = max_statement_payload_size(V2_ENVELOPE_OVERHEAD);
3825
3826 assert!(
3828 v2_max < v1_max,
3829 "V2 payload capacity ({v2_max}) should be less than V1 ({v1_max})"
3830 );
3831 assert_eq!(v1_max - v2_max, 1, "V2 overhead is exactly 1 byte more than V1");
3832
3833 let stmts: Vec<Statement> = (0..1000)
3835 .map(|i| {
3836 let mut s = Statement::new();
3837 s.set_plain_data(format!("stmt-{i}").into_bytes());
3838 s
3839 })
3840 .collect();
3841 let refs: Vec<&Statement> = stmts.iter().collect();
3842
3843 let v1_chunk = find_sendable_chunk(&refs, V1_ENVELOPE_OVERHEAD);
3844 let v2_chunk = find_sendable_chunk(&refs, V2_ENVELOPE_OVERHEAD);
3845
3846 let v1_count = match v1_chunk {
3848 ChunkResult::Send(n) => n,
3849 _ => panic!("Expected Send for V1"),
3850 };
3851 let v2_count = match v2_chunk {
3852 ChunkResult::Send(n) => n,
3853 _ => panic!("Expected Send for V2"),
3854 };
3855 assert!(
3856 v2_count <= v1_count,
3857 "V2 ({v2_count}) should fit at most as many statements as V1 ({v1_count})"
3858 );
3859 }
3860
3861 #[tokio::test]
3862 async fn test_full_node_v2_gets_initial_sync_immediately() {
3863 let (mut handler, statement_store, _network, _notification_service) =
3864 build_handler_no_peers();
3865
3866 let mut stmt = Statement::new();
3868 stmt.set_plain_data(b"full node v2".to_vec());
3869 let hash = stmt.hash();
3870 statement_store.statements.lock().unwrap().insert(hash, stmt);
3871
3872 let peer_id = PeerId::random();
3873
3874 handler
3876 .handle_notification_event(NotificationEvent::NotificationStreamOpened {
3877 peer: peer_id,
3878 direction: sc_network::service::traits::Direction::Inbound,
3879 handshake: vec![],
3880 negotiated_fallback: None,
3881 })
3882 .await;
3883
3884 assert!(
3886 handler.pending_initial_syncs.contains_key(&peer_id),
3887 "Full-node V2 peer should have initial sync scheduled immediately"
3888 );
3889 assert_eq!(handler.peers.get(&peer_id).unwrap().protocol_version, PeerProtocolVersion::V2);
3890 assert!(!handler.peers.get(&peer_id).unwrap().is_light);
3891 }
3892
3893 #[tokio::test]
3894 async fn test_propagation_reaches_all_connected_peers() {
3895 let (
3896 mut handler,
3897 statement_store,
3898 _network,
3899 notification_service,
3900 _queue_receiver,
3901 peer_ids,
3902 ) = build_handler(5);
3903
3904 let mut expected_hashes = Vec::new();
3906 for i in 0..3u8 {
3907 let mut statement = Statement::new();
3908 statement.set_plain_data(vec![i; 100]);
3909 let hash = statement.hash();
3910 expected_hashes.push(hash);
3911 statement_store.recent_statements.lock().unwrap().insert(hash, statement);
3912 }
3913 expected_hashes.sort();
3914
3915 handler.propagate_statements().await;
3916
3917 let sent = notification_service.get_sent_notifications();
3918
3919 for peer_id in &peer_ids {
3921 let mut received_hashes = get_peer_hashes(&sent, *peer_id);
3922 received_hashes.sort();
3923
3924 assert_eq!(
3925 received_hashes, expected_hashes,
3926 "Peer {peer_id} should have received all 3 statements"
3927 );
3928 }
3929
3930 assert!(statement_store.recent_statements.lock().unwrap().is_empty());
3932 }
3933
3934 #[tokio::test]
3935 async fn test_known_statement_filtering_per_peer() {
3936 let (
3937 mut handler,
3938 statement_store,
3939 _network,
3940 notification_service,
3941 _queue_receiver,
3942 peer_ids,
3943 ) = build_handler(3);
3944
3945 let peer_a = peer_ids[0];
3946 let peer_b = peer_ids[1];
3947 let peer_c = peer_ids[2];
3948
3949 let mut hashes = Vec::new();
3951 for i in 0..5u8 {
3952 let mut statement = Statement::new();
3953 statement.set_plain_data(vec![i; 100]);
3954 let hash = statement.hash();
3955 hashes.push(hash);
3956 statement_store.recent_statements.lock().unwrap().insert(hash, statement);
3957 }
3958
3959 handler.peers.get_mut(&peer_a).unwrap().known_statements.insert(hashes[0]);
3961 handler.peers.get_mut(&peer_a).unwrap().known_statements.insert(hashes[1]);
3962 handler.peers.get_mut(&peer_b).unwrap().known_statements.insert(hashes[2]);
3963
3964 handler.propagate_statements().await;
3965
3966 let sent = notification_service.get_sent_notifications();
3967
3968 let peer_a_hashes = get_peer_hashes(&sent, peer_a);
3969 let peer_b_hashes = get_peer_hashes(&sent, peer_b);
3970 let peer_c_hashes = get_peer_hashes(&sent, peer_c);
3971
3972 assert_eq!(peer_a_hashes.len(), 3, "peer_a should get 3 statements");
3974 assert!(!peer_a_hashes.contains(&hashes[0]), "peer_a already knows s1");
3975 assert!(!peer_a_hashes.contains(&hashes[1]), "peer_a already knows s2");
3976 assert!(peer_a_hashes.contains(&hashes[2]));
3977 assert!(peer_a_hashes.contains(&hashes[3]));
3978 assert!(peer_a_hashes.contains(&hashes[4]));
3979
3980 assert_eq!(peer_b_hashes.len(), 4, "peer_b should get 4 statements");
3982 assert!(!peer_b_hashes.contains(&hashes[2]), "peer_b already knows s3");
3983 assert!(peer_b_hashes.contains(&hashes[0]));
3984 assert!(peer_b_hashes.contains(&hashes[1]));
3985 assert!(peer_b_hashes.contains(&hashes[3]));
3986 assert!(peer_b_hashes.contains(&hashes[4]));
3987
3988 let mut sorted_peer_c: Vec<_> = peer_c_hashes.into_iter().collect();
3990 sorted_peer_c.sort();
3991 let mut all_hashes = hashes.clone();
3992 all_hashes.sort();
3993 assert_eq!(sorted_peer_c, all_hashes, "peer_c should get all 5 statements");
3994 }
3995
3996 #[test]
3999 fn major_sync_defers_peers_and_handles_disconnect() {
4000 let (sync, _flag) = TestSync::with_syncing(true);
4001 let network = TestNetwork::new();
4002 let notification_service = TestNotificationService::new();
4003 let statement_store = TestStatementStore::new();
4004 let (queue_sender, _queue_receiver) = async_channel::bounded(100);
4005
4006 let mut handler = StatementHandler {
4007 protocol_name: "/statement/1".into(),
4008 notification_service: Box::new(notification_service),
4009 propagate_timeout: (Box::pin(futures::stream::pending())
4010 as Pin<Box<dyn Stream<Item = ()> + Send>>)
4011 .fuse(),
4012 pending_statements: FuturesUnordered::new(),
4013 pending_statements_peers: HashMap::new(),
4014 network: network.clone(),
4015 sync,
4016 sync_event_stream: (Box::pin(futures::stream::pending())
4017 as Pin<Box<dyn Stream<Item = sc_network_sync::types::SyncEvent> + Send>>)
4018 .fuse(),
4019 peers: HashMap::new(),
4020 statement_store: Arc::new(statement_store),
4021 queue_sender,
4022 statements_per_second: NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
4023 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero"),
4024 metrics: None,
4025 initial_sync_timeout: Box::pin(futures::future::pending()),
4026 pending_affinities_timeout: Box::pin(futures::future::pending()),
4027 pending_initial_syncs: HashMap::new(),
4028 initial_sync_peer_queue: VecDeque::new(),
4029 deferred_peers: HashSet::new(),
4030 dropped_statements_during_sync: false,
4031 sync_recovery_peer: None,
4032 sync_recovery_readd_timeout: Box::pin(pending().fuse()),
4033 };
4034
4035 let peer1 = PeerId::random();
4036 let peer2 = PeerId::random();
4037 let peer3 = PeerId::random();
4038
4039 handler.handle_sync_event(SyncEvent::PeerConnected(peer1));
4040 handler.handle_sync_event(SyncEvent::PeerConnected(peer2));
4041 handler.handle_sync_event(SyncEvent::PeerConnected(peer3));
4042
4043 assert!(network.get_added_reserved().is_empty());
4045 assert!(network.get_removed_reserved().is_empty());
4046 assert_eq!(handler.deferred_peers.len(), 3);
4047
4048 handler.handle_sync_event(SyncEvent::PeerDisconnected(peer1));
4050 assert_eq!(handler.deferred_peers.len(), 2);
4051 assert!(!handler.deferred_peers.contains(&peer1), "disconnected peer must leave buffer");
4052 assert!(handler.deferred_peers.contains(&peer2));
4053 assert!(handler.deferred_peers.contains(&peer3));
4054 assert!(network.get_removed_reserved().is_empty(), "no remove call for buffered peer");
4055 }
4056
4057 #[test]
4058 fn deferred_peers_flushed_on_sync_end_without_remove() {
4059 let (sync, flag) = TestSync::with_syncing(true);
4060 let network = TestNetwork::new();
4061 let notification_service = TestNotificationService::new();
4062 let statement_store = TestStatementStore::new();
4063 let (queue_sender, _queue_receiver) = async_channel::bounded(100);
4064
4065 let peer1 = PeerId::random();
4066 let peer2 = PeerId::random();
4067 let mut deferred = HashSet::new();
4068 deferred.insert(peer1);
4069 deferred.insert(peer2);
4070
4071 let mut handler = StatementHandler {
4072 protocol_name: "/statement/1".into(),
4073 notification_service: Box::new(notification_service),
4074 propagate_timeout: (Box::pin(futures::stream::pending())
4075 as Pin<Box<dyn Stream<Item = ()> + Send>>)
4076 .fuse(),
4077 pending_statements: FuturesUnordered::new(),
4078 pending_statements_peers: HashMap::new(),
4079 network: network.clone(),
4080 sync,
4081 sync_event_stream: (Box::pin(futures::stream::pending())
4082 as Pin<Box<dyn Stream<Item = sc_network_sync::types::SyncEvent> + Send>>)
4083 .fuse(),
4084 peers: HashMap::new(),
4085 statement_store: Arc::new(statement_store),
4086 queue_sender,
4087 statements_per_second: NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
4088 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero"),
4089 metrics: None,
4090 initial_sync_timeout: Box::pin(futures::future::pending()),
4091 pending_affinities_timeout: Box::pin(futures::future::pending()),
4092 pending_initial_syncs: HashMap::new(),
4093 initial_sync_peer_queue: VecDeque::new(),
4094 deferred_peers: deferred,
4095 dropped_statements_during_sync: false,
4096 sync_recovery_peer: None,
4097 sync_recovery_readd_timeout: Box::pin(pending().fuse()),
4098 };
4099
4100 flag.store(false, std::sync::atomic::Ordering::Relaxed);
4101 handler.drain_deferred_peers();
4102
4103 assert!(handler.deferred_peers.is_empty());
4104
4105 let added = network.get_added_reserved();
4106 assert_eq!(added.len(), 1);
4107 let added_addrs = &added[0];
4108 let expected_addr1: sc_network::Multiaddr =
4109 iter::once(multiaddr::Protocol::P2p(peer1.into())).collect();
4110 let expected_addr2: sc_network::Multiaddr =
4111 iter::once(multiaddr::Protocol::P2p(peer2.into())).collect();
4112 assert!(added_addrs.contains(&expected_addr1), "peer1 must be in added set");
4113 assert!(added_addrs.contains(&expected_addr2), "peer2 must be in added set");
4114
4115 assert!(network.get_removed_reserved().is_empty());
4116 }
4117
4118 #[tokio::test]
4119 async fn sync_recovery_schedules_remove_for_one_connected_peer() {
4120 let network = TestNetwork::new();
4121 let notification_service = TestNotificationService::new();
4122 let (sync, _flag) = TestSync::with_syncing(false);
4123 let (queue_sender, _) = async_channel::bounded(2);
4124 let statement_store = TestStatementStore::new();
4125
4126 let connected_peer = PeerId::random();
4127
4128 let mut peers = HashMap::new();
4129 peers.insert(
4130 connected_peer,
4131 Peer {
4132 known_statements: LruHashSet::new(NonZeroUsize::new(1024).unwrap()),
4133 rate_limiter: PeerRateLimiter::new(
4134 NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
4135 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero"),
4136 NonZeroU32::new(
4137 DEFAULT_STATEMENTS_PER_SECOND * config::STATEMENTS_BURST_COEFFICIENT,
4138 )
4139 .expect("burst capacity is nonzero"),
4140 ),
4141 protocol_version: PeerProtocolVersion::V1,
4142 topic_affinity: None,
4143 is_light: false,
4144 pending_topic_affinity: None,
4145 },
4146 );
4147
4148 let mut handler = StatementHandler {
4149 protocol_name: format!("/{STATEMENT_PROTOCOL_V1}").into(),
4150 notification_service: Box::new(notification_service),
4151 propagate_timeout: (Box::pin(futures::stream::pending())
4152 as Pin<Box<dyn Stream<Item = ()> + Send>>)
4153 .fuse(),
4154 pending_statements: FuturesUnordered::new(),
4155 pending_statements_peers: HashMap::new(),
4156 network: network.clone(),
4157 sync,
4158 sync_event_stream: (Box::pin(futures::stream::pending())
4159 as Pin<Box<dyn Stream<Item = sc_network_sync::types::SyncEvent> + Send>>)
4160 .fuse(),
4161 peers,
4162 statement_store: Arc::new(statement_store),
4163 queue_sender,
4164 statements_per_second: NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
4165 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero"),
4166 metrics: None,
4167 initial_sync_timeout: Box::pin(futures::future::pending()),
4168 pending_affinities_timeout: Box::pin(futures::future::pending()),
4169 pending_initial_syncs: HashMap::new(),
4170 initial_sync_peer_queue: VecDeque::new(),
4171 deferred_peers: HashSet::new(),
4172 dropped_statements_during_sync: true,
4173 sync_recovery_peer: None,
4174 sync_recovery_readd_timeout: Box::pin(futures::future::pending()),
4175 };
4176
4177 handler.start_sync_recovery();
4178
4179 {
4181 let removed = network.removed_reserved.lock().unwrap();
4182 assert_eq!(
4183 removed.len(),
4184 1,
4185 "Expected exactly one remove_peers_from_reserved_set call"
4186 );
4187 assert!(removed[0].contains(&connected_peer));
4188 }
4189
4190 assert_eq!(handler.sync_recovery_peer, Some(connected_peer));
4192
4193 handler.try_readd_sync_recovery_peer();
4196 assert!(handler.sync_recovery_peer.is_none());
4197 {
4198 let added = network.added_reserved.lock().unwrap();
4199 assert_eq!(added.len(), 1);
4200 let expected_addr: multiaddr::Multiaddr =
4201 iter::once(multiaddr::Protocol::P2p(connected_peer.into())).collect();
4202 assert!(added[0].contains(&expected_addr));
4203 }
4204
4205 {
4208 let peer2 = PeerId::random();
4209 handler.sync_recovery_peer = Some(peer2);
4210 handler.start_sync_recovery();
4211 assert_eq!(
4212 handler.sync_recovery_peer,
4213 Some(peer2),
4214 "Re-entry guard: recovery peer must not change on second call"
4215 );
4216 assert_eq!(
4217 network.removed_reserved.lock().unwrap().len(),
4218 1,
4219 "Re-entry guard: no extra remove call while recovery is in flight"
4220 );
4221 }
4222 }
4223
4224 #[tokio::test]
4225 async fn sync_recovery_gated_by_dropped_statements_flag() {
4226 let make_peer = || Peer {
4227 known_statements: LruHashSet::new(NonZeroUsize::new(1024).unwrap()),
4228 rate_limiter: PeerRateLimiter::new(
4229 NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
4230 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero"),
4231 NonZeroU32::new(
4232 DEFAULT_STATEMENTS_PER_SECOND * config::STATEMENTS_BURST_COEFFICIENT,
4233 )
4234 .expect("burst capacity is nonzero"),
4235 ),
4236 protocol_version: PeerProtocolVersion::V1,
4237 topic_affinity: None,
4238 is_light: false,
4239 pending_topic_affinity: None,
4240 };
4241
4242 let make_handler =
4243 |network: TestNetwork, dropped: bool| -> StatementHandler<TestNetwork, TestSync> {
4244 let (sync, _) = TestSync::with_syncing(false);
4245 let (queue_sender, _) = async_channel::bounded(2);
4246 let mut peers = HashMap::new();
4247 peers.insert(PeerId::random(), make_peer());
4248 StatementHandler {
4249 protocol_name: format!("/{STATEMENT_PROTOCOL_V1}").into(),
4250 notification_service: Box::new(TestNotificationService::new()),
4251 propagate_timeout: (Box::pin(futures::stream::pending())
4252 as Pin<Box<dyn Stream<Item = ()> + Send>>)
4253 .fuse(),
4254 pending_statements: FuturesUnordered::new(),
4255 pending_statements_peers: HashMap::new(),
4256 network,
4257 sync,
4258 sync_event_stream: (Box::pin(futures::stream::pending())
4259 as Pin<Box<dyn Stream<Item = sc_network_sync::types::SyncEvent> + Send>>)
4260 .fuse(),
4261 peers,
4262 statement_store: Arc::new(TestStatementStore::new()),
4263 queue_sender,
4264 statements_per_second: NonZeroU32::new(DEFAULT_STATEMENTS_PER_SECOND)
4265 .expect("DEFAULT_STATEMENTS_PER_SECOND is nonzero"),
4266 metrics: None,
4267 initial_sync_timeout: Box::pin(futures::future::pending()),
4268 pending_affinities_timeout: Box::pin(futures::future::pending()),
4269 pending_initial_syncs: HashMap::new(),
4270 initial_sync_peer_queue: VecDeque::new(),
4271 deferred_peers: HashSet::new(),
4272 dropped_statements_during_sync: dropped,
4273 sync_recovery_peer: None,
4274 sync_recovery_readd_timeout: Box::pin(pending().fuse()),
4275 }
4276 };
4277
4278 let net = TestNetwork::new();
4280 let mut handler = make_handler(net.clone(), false);
4281 handler.start_sync_recovery();
4282 assert!(handler.sync_recovery_peer.is_none());
4283 assert!(net.get_removed_reserved().is_empty());
4284
4285 let net2 = TestNetwork::new();
4287 let mut handler2 = make_handler(net2.clone(), true);
4288 handler2.start_sync_recovery();
4289 assert!(handler2.sync_recovery_peer.is_some());
4290 assert_eq!(net2.get_removed_reserved().len(), 1);
4291 }
4292}