1mod trace;
4
5use std::{
6 collections::{HashMap, HashSet},
7 future,
8 io::{Cursor, Read},
9 net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
10 path::PathBuf,
11 sync::{
12 atomic::{AtomicU64, Ordering},
13 Arc, Mutex as StdMutex,
14 },
15 time::Duration,
16};
17
18use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
19use futures::StreamExt as _;
20use iroh::Watcher as _;
21use iroh::{
22 endpoint::{
23 Connection, ConnectionType, Endpoint, RecvStream, SendStream, TransportConfig, VarInt,
24 },
25 protocol::{AcceptError, ProtocolHandler, Router},
26 NodeAddr, NodeId, SecretKey,
27};
28use rand::{rngs::OsRng, RngCore};
29use thiserror::Error;
30use tokio::{
31 sync::{broadcast, mpsc, oneshot, watch, Mutex, OwnedSemaphorePermit, Semaphore},
32 task::{AbortHandle, JoinHandle, JoinSet},
33 time::{timeout, Instant},
34};
35use tokio_stream::StreamMap;
36use tokio_util::sync::CancellationToken;
37use zakura_chain::{
38 block::{self, Block, CountedHeader},
39 parameters::{Network, NetworkKind},
40 serialization::{CompactSizeMessage, ZcashDeserialize, MAX_HEADERS_PER_MESSAGE},
41 transaction::Transaction,
42};
43
44use self::trace::ZakuraConnTrace;
45use super::discovery::{self, native_dial_supervised, spawn_native_bootstrap_dialer, RedialPolicy};
46use super::trace::{reject_reason_label, ZakuraTrace};
47#[cfg(any(test, feature = "zakura-testkit"))]
48use crate::zakura::drive_header_sync_actions;
49#[cfg(any(test, feature = "zakura-testkit"))]
50use crate::zakura::HeaderSyncAction;
51use crate::{peer_registry::PeerRegistry, BoxError, Config, MAX_TX_INV_IN_SENT_MESSAGE};
52use crate::{
53 protocol::external::InventoryHash,
54 zakura::{
55 canonical_ip, direct_endpoint_builder, spawn_block_sync_reactor, spawn_header_sync_reactor,
56 AuthenticatedPeerRegistration, BlockSyncAction, BlockSyncFrontiers, BlockSyncHandle,
57 BlockSyncService, BlockSyncStartup, BoxRunFuture, Clock, CloseCause, Frame, FramedRecv,
58 FramedSend, FullStateFrontiers, HeaderSyncPassthroughService, HeaderSyncService,
59 HeaderSyncStartup, OrderedSessionDemand, OrderedStreamOpening, OrderedStreamPolicy, Peer,
60 RealClock, Service, ServicePeerDirection, ServiceRegistry, ServiceStream, SinkReject,
61 Stream, StreamMode, StreamPrelude, ZakuraAcceptedLimits, ZakuraBlockSyncConfig,
62 ZakuraConnId, ZakuraControlAck, ZakuraControlHello, ZakuraControlRole,
63 ZakuraControlValidation, ZakuraHandshakeConfig, ZakuraHandshakePath,
64 ZakuraHeaderSyncConfig, ZakuraInitialLimits, ZakuraLimits, ZakuraPeerId,
65 ZakuraPeerSupervisor, ZakuraProtocolError, ZakuraRejectReason, ZakuraServiceId,
66 ZakuraUpgradeDialStart, CONTROL_ACK_MAGIC, CONTROL_HELLO_MAGIC, CONTROL_VERSION,
67 FRAME_HEADER_BYTES, MAX_CONTROL_PAYLOAD_BYTES, P2P_V2_ALPN, STREAM_PRELUDE_MAGIC,
68 TRANSCRIPT_HASH_BYTES, ZAKURA_CAP_HEADER_SYNC, ZAKURA_HEADER_SYNC_STREAM_VERSION,
69 ZAKURA_PROTOCOL_VERSION_1, ZAKURA_STREAM_BLOCK_SYNC, ZAKURA_STREAM_HEADER_SYNC,
70 },
71};
72
73pub const DEFAULT_ZAKURA_MAX_CONNECTIONS: usize = 256;
75pub const DEFAULT_ZAKURA_MAX_CONNS_PER_IP: usize = 16;
77pub const DEFAULT_ZAKURA_MAX_PENDING_HANDSHAKES: usize = 32;
79pub const DEFAULT_ZAKURA_STREAM_OPEN_RATE_PER_SECOND: u32 = 32;
81pub const DEFAULT_ZAKURA_MESSAGE_RATE_PER_SECOND: u32 = 2048;
91pub const DEFAULT_ZAKURA_LISTEN_ADDR: SocketAddr =
93 SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 8234));
94pub const DEFAULT_MAINNET_ZAKURA_BOOTSTRAP_PEERS: &[&str] = &[
96 "1398f62c6d1a457c51ba6a4b5f3dbd2f69fca93216218dc8997e416bd17d93ca@165.22.54.66:8234",
97 "fd1724385aa0c75b64fb78cd602fa1d991fdebf76b13c58ed702eac835e9f618@104.131.184.123:8234",
98 "9ec67ad6834bc2ca0d659c240e042d3446c37cabcc092b527d459c87d938b4a4@159.65.183.89:8234",
99 "bd3dc5d2a3d44c6bf90e364bf446231dbf9737e38a562ccf9e91ea631ea59b22@143.244.184.176:8234",
100 "14ab98fa0c4b07d40119e1dbc9f3c36d20c8f226ae5ba4216218a2034f148e57@159.203.38.10:8234",
101 "681d21b18644cd82ec13256a97f92bec1fff815683ef6f65dc7c993f098a4fe5@64.227.44.93:8234",
102 "058b3f20dc9bef7bb447f94d7663d793cfbc036720f97e52d7f13661b21818e1@161.35.156.226:8234",
103 "291323d78eb7186c3fa225ef5e305e95363e0ef06d42dca91bd4ef0254aed1ae@139.59.64.115:8234",
104 "85e425233a68697d4be91dd5d542305a8a327cd06d992d53c0913cef2fa75084@168.144.173.250:8234",
105];
106pub const DEFAULT_TESTNET_ZAKURA_BOOTSTRAP_PEERS: &[&str] = &[
108 "57ad39fad4f0bca46cf1ea831772a99d5027b372fef2be5a0ea68e1b5bb4da49@167.99.103.111:8234",
109 "2bbb907b5d90598ef49f2e637066586b311a64587479be6ed43e8388587fcd2a@164.92.209.78:8234",
110 "50999835f48f4a048c0e9042e5332844c9673943d7fab1f7e993bae698c27ea3@206.189.148.0:8234",
111];
112pub const DEFAULT_ZAKURA_BOOTSTRAP_PEERS: &[&str] = DEFAULT_MAINNET_ZAKURA_BOOTSTRAP_PEERS;
114pub const DEFAULT_ZAKURA_PRELUDE_TIMEOUT: Duration = Duration::from_secs(3);
116pub const DEFAULT_ZAKURA_CONTROL_TIMEOUT: Duration = Duration::from_secs(10);
118pub const DEFAULT_ZAKURA_QUIC_IDLE_TIMEOUT: Duration = Duration::from_secs(150);
126pub const DEFAULT_ZAKURA_KEEP_ALIVE_INTERVAL: Duration = Duration::from_secs(10);
128pub const ZAKURA_DUPLICATE_EVICT_MIN_AGE: Duration = Duration::from_secs(300);
138pub const ZAKURA_SAME_IP_DUPLICATE_EVICT_MIN_AGE: Duration = Duration::from_secs(5);
150pub const DEFAULT_ZAKURA_STREAM_RECEIVE_WINDOW: u32 = 32 * 1024 * 1024;
152pub const DEFAULT_ZAKURA_RECEIVE_WINDOW: u32 = 32 * 1024 * 1024;
154pub const DEFAULT_ZAKURA_SEND_WINDOW: u64 = 32 * 1024 * 1024;
156pub const DEFAULT_ZAKURA_REDIAL_INITIAL_BACKOFF: Duration = Duration::from_secs(1);
158pub const DEFAULT_ZAKURA_REDIAL_MAX_BACKOFF: Duration = Duration::from_secs(30);
160const CONTROL_LENGTH_BYTES: usize = 4;
161const STREAM_PRELUDE_FIXED_BYTES: usize = 4 + 2 + 2 + 1;
162const STREAM_PRELUDE_REQUEST_ID_FLAG_OFFSET: usize = STREAM_PRELUDE_FIXED_BYTES - 1;
163const STREAM_PRELUDE_REQUEST_ID_BYTES: usize = 8;
164const STREAM_PRELUDE_CAP_BYTES: usize = 4;
165const STREAM_WORKER_DRAIN_TIMEOUT: Duration = Duration::from_secs(1);
166const ORDERED_STREAM_REOPEN_BACKOFF: Duration = Duration::from_millis(250);
167const ORDERED_STREAM_REOPEN_BACKOFF_CAP: Duration = Duration::from_secs(8);
168const OUTBOUND_STREAM_WRITE_TIMEOUT: Duration = Duration::from_secs(10);
169const OUTBOUND_REQUEST_RESPONSE_TIMEOUT: Duration = Duration::from_secs(30);
170const LEGACY_GOSSIP_STREAM_KIND: u16 = 2;
173const LEGACY_REQUEST_STREAM_KIND: u16 = 3;
174const DISCOVERY_STREAM_KIND: u16 = 4;
175const HEADER_SYNC_STREAM_KIND: u16 = ZAKURA_STREAM_HEADER_SYNC;
176const LEGACY_REQUEST_BLOCKS_BY_HASH: u16 = 3;
177const LEGACY_REQUEST_TRANSACTIONS_BY_ID: u16 = 4;
178const LEGACY_RESPONSE_BLOCK: u16 = 5;
179const LEGACY_RESPONSE_TRANSACTION: u16 = 6;
180const LEGACY_RESPONSE_MISSING_BLOCKS: u16 = 7;
181const LEGACY_RESPONSE_MISSING_TRANSACTIONS: u16 = 8;
182const LEGACY_REQUEST_FIND_BLOCKS: u16 = 9;
183const LEGACY_REQUEST_FIND_HEADERS: u16 = 10;
184const LEGACY_REQUEST_MEMPOOL_TRANSACTION_IDS: u16 = 11;
185const LEGACY_REQUEST_PING: u16 = 12;
186const LEGACY_REQUEST_PUSH_TRANSACTION: u16 = 13;
187const LEGACY_RESPONSE_BLOCK_HASHES: u16 = 14;
188const LEGACY_RESPONSE_BLOCK_HEADERS: u16 = 15;
189const LEGACY_RESPONSE_TRANSACTION_IDS: u16 = 16;
190const LEGACY_RESPONSE_PONG: u16 = 17;
191const LEGACY_RESPONSE_NIL: u16 = 18;
192const LEGACY_RESPONSE_REQUEST_ID_BYTES: usize = 8;
193const LEGACY_RESPONSE_CHUNK_HEADER_BYTES: usize = LEGACY_RESPONSE_REQUEST_ID_BYTES + 1;
194const LEGACY_COMPACT_SIZE_PREFIX_BYTES: usize = 9;
195const LEGACY_BLOCK_HASH_BYTES: usize = 32;
196const LEGACY_INVENTORY_HASH_BYTES: usize = 36;
197const LEGACY_RESPONSE_MAX_FRAMES_PER_ITEM: usize = 8;
198const LEGACY_RESPONSE_MAX_AGGREGATE_BYTES: usize =
211 8 * zakura_chain::serialization::MAX_PROTOCOL_MESSAGE_LEN;
212const _: () = assert!(LEGACY_GOSSIP_STREAM_KIND == super::legacy_gossip::ZAKURA_STREAM_GOSSIP);
213const _: () =
214 assert!(LEGACY_REQUEST_STREAM_KIND == super::legacy_gossip::ZAKURA_STREAM_LEGACY_REQUESTS);
215const _: () = assert!(DISCOVERY_STREAM_KIND == super::discovery::ZAKURA_STREAM_DISCOVERY);
216const _: () = assert!(HEADER_SYNC_STREAM_KIND == super::header_sync::ZAKURA_STREAM_HEADER_SYNC);
217const _: () = assert!(ZAKURA_HEADER_SYNC_STREAM_VERSION == 8);
218const _: () =
219 assert!(LEGACY_REQUEST_BLOCKS_BY_HASH == super::legacy_gossip::MSG_REQUEST_BLOCKS_BY_HASH);
220const _: () = assert!(
221 LEGACY_REQUEST_TRANSACTIONS_BY_ID == super::legacy_gossip::MSG_REQUEST_TRANSACTIONS_BY_ID
222);
223const _: () = assert!(LEGACY_RESPONSE_BLOCK == super::legacy_gossip::MSG_RESPONSE_BLOCK);
224const _: () =
225 assert!(LEGACY_RESPONSE_TRANSACTION == super::legacy_gossip::MSG_RESPONSE_TRANSACTION);
226const _: () =
227 assert!(LEGACY_RESPONSE_MISSING_BLOCKS == super::legacy_gossip::MSG_RESPONSE_MISSING_BLOCKS);
228const _: () = assert!(
229 LEGACY_RESPONSE_MISSING_TRANSACTIONS == super::legacy_gossip::MSG_RESPONSE_MISSING_TRANSACTIONS
230);
231const _: () = assert!(LEGACY_REQUEST_FIND_BLOCKS == super::legacy_gossip::MSG_REQUEST_FIND_BLOCKS);
232const _: () =
233 assert!(LEGACY_REQUEST_FIND_HEADERS == super::legacy_gossip::MSG_REQUEST_FIND_HEADERS);
234const _: () = assert!(
235 LEGACY_REQUEST_MEMPOOL_TRANSACTION_IDS
236 == super::legacy_gossip::MSG_REQUEST_MEMPOOL_TRANSACTION_IDS
237);
238const _: () = assert!(LEGACY_REQUEST_PING == super::legacy_gossip::MSG_REQUEST_PING);
239const _: () =
240 assert!(LEGACY_REQUEST_PUSH_TRANSACTION == super::legacy_gossip::MSG_REQUEST_PUSH_TRANSACTION);
241const _: () =
242 assert!(LEGACY_RESPONSE_BLOCK_HASHES == super::legacy_gossip::MSG_RESPONSE_BLOCK_HASHES);
243const _: () =
244 assert!(LEGACY_RESPONSE_BLOCK_HEADERS == super::legacy_gossip::MSG_RESPONSE_BLOCK_HEADERS);
245const _: () =
246 assert!(LEGACY_RESPONSE_TRANSACTION_IDS == super::legacy_gossip::MSG_RESPONSE_TRANSACTION_IDS);
247const _: () = assert!(LEGACY_RESPONSE_PONG == super::legacy_gossip::MSG_RESPONSE_PONG);
248const _: () = assert!(LEGACY_RESPONSE_NIL == super::legacy_gossip::MSG_RESPONSE_NIL);
249const ZAKURA_CLOSE_NEUTRAL: u32 = 0;
250const ZAKURA_CLOSE_RESOURCE: u32 = 1;
251const ZAKURA_CLOSE_BAD_PRELUDE: u32 = 2;
252const ZAKURA_CLOSE_RATE_LIMIT: u32 = 3;
253const ZAKURA_CLOSE_OVERSIZE: u32 = 4;
254const ZAKURA_CLOSE_UNKNOWN_STREAM: u32 = 5;
259
260#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
262#[serde(deny_unknown_fields, default)]
263pub struct ZakuraConfig {
264 pub bootstrap_peers: Vec<String>,
270 pub listen_addr: Option<SocketAddr>,
280 pub max_connections: usize,
282 pub max_connections_per_ip: usize,
289 pub max_pending_handshakes: usize,
291 pub stream_open_rate_per_second: u32,
293 pub message_rate_per_second: u32,
295 pub trace_dir: Option<PathBuf>,
302 pub header_sync: ZakuraHeaderSyncConfig,
304 pub block_sync: ZakuraBlockSyncConfig,
306 pub dev_network: Option<String>,
319}
320
321impl Default for ZakuraConfig {
322 fn default() -> Self {
323 Self {
324 bootstrap_peers: DEFAULT_ZAKURA_BOOTSTRAP_PEERS
325 .iter()
326 .map(ToString::to_string)
327 .collect(),
328 listen_addr: Some(DEFAULT_ZAKURA_LISTEN_ADDR),
329 max_connections: DEFAULT_ZAKURA_MAX_CONNECTIONS,
330 max_connections_per_ip: DEFAULT_ZAKURA_MAX_CONNS_PER_IP,
331 max_pending_handshakes: DEFAULT_ZAKURA_MAX_PENDING_HANDSHAKES,
332 stream_open_rate_per_second: DEFAULT_ZAKURA_STREAM_OPEN_RATE_PER_SECOND,
333 message_rate_per_second: DEFAULT_ZAKURA_MESSAGE_RATE_PER_SECOND,
334 trace_dir: None,
335 header_sync: ZakuraHeaderSyncConfig::default(),
336 block_sync: ZakuraBlockSyncConfig::default(),
337 dev_network: None,
338 }
339 }
340}
341
342impl ZakuraConfig {
343 pub fn max_connections_per_ip(&self) -> usize {
345 if self.max_connections_per_ip == 0 {
346 DEFAULT_ZAKURA_MAX_CONNS_PER_IP
347 } else {
348 self.max_connections_per_ip
349 }
350 }
351
352 pub fn default_bootstrap_peers_for_network(network: &Network) -> Vec<String> {
354 match network.kind() {
355 NetworkKind::Mainnet => {
356 bootstrap_peers_to_strings(DEFAULT_MAINNET_ZAKURA_BOOTSTRAP_PEERS)
357 }
358 NetworkKind::Testnet => {
359 bootstrap_peers_to_strings(DEFAULT_TESTNET_ZAKURA_BOOTSTRAP_PEERS)
360 }
361 NetworkKind::Regtest => Vec::new(),
362 }
363 }
364
365 pub fn apply_network_defaults(&mut self, network: &Network) {
370 if self.bootstrap_peers
371 != bootstrap_peers_to_strings(DEFAULT_MAINNET_ZAKURA_BOOTSTRAP_PEERS)
372 {
373 return;
374 }
375
376 self.bootstrap_peers = Self::default_bootstrap_peers_for_network(network);
377 }
378}
379
380fn bootstrap_peers_to_strings(peers: &[&str]) -> Vec<String> {
381 peers.iter().map(ToString::to_string).collect()
382}
383
384#[derive(Clone, Debug)]
386pub struct ZakuraLocalLimits {
387 pub max_connections: usize,
389 pub max_pending_handshakes: usize,
391 pub quic_idle_timeout: Duration,
393 pub keep_alive_interval: Duration,
395 pub prelude_timeout: Duration,
397 pub control_timeout: Duration,
399 pub stream_open_rate_per_second: u32,
401 pub message_rate_per_second: u32,
403 pub max_frame_bytes: u32,
405 pub max_message_bytes: u32,
407 pub max_open_streams: u16,
409 pub max_inbound_queue_depth: u16,
411}
412
413impl ZakuraLocalLimits {
414 pub fn from_config(config: &Config) -> Self {
416 let handshake = ZakuraHandshakeConfig::for_network(&config.network);
417 Self {
418 max_connections: config.zakura.max_connections.max(1),
419 max_pending_handshakes: config.zakura.max_pending_handshakes.max(1),
420 quic_idle_timeout: DEFAULT_ZAKURA_QUIC_IDLE_TIMEOUT,
421 keep_alive_interval: DEFAULT_ZAKURA_KEEP_ALIVE_INTERVAL,
422 prelude_timeout: DEFAULT_ZAKURA_PRELUDE_TIMEOUT,
423 control_timeout: DEFAULT_ZAKURA_CONTROL_TIMEOUT,
424 stream_open_rate_per_second: config.zakura.stream_open_rate_per_second.max(1),
425 message_rate_per_second: config.zakura.message_rate_per_second.max(1),
426 max_frame_bytes: handshake.max_message_bytes,
427 max_message_bytes: handshake.max_message_bytes,
428 max_open_streams: handshake.max_open_streams,
429 max_inbound_queue_depth: handshake.max_inbound_queue_depth,
430 }
431 }
432
433 pub fn clamp(&self, negotiated: &ZakuraAcceptedLimits) -> ZakuraConnectionLimits {
435 let max_open_streams = negotiated
436 .max_open_streams
437 .min(self.max_open_streams)
438 .max(1);
439 let idle_timeout = Duration::from_millis(
440 u64::from(negotiated.idle_timeout_millis)
441 .min(self.quic_idle_timeout.as_millis().saturating_sub(1) as u64)
442 .max(1),
443 );
444
445 ZakuraConnectionLimits {
446 max_frame_bytes: negotiated.max_frame_bytes.min(self.max_frame_bytes).max(1),
447 max_message_bytes: negotiated
448 .max_message_bytes
449 .min(self.max_message_bytes)
450 .max(1),
451 max_open_streams,
452 max_inbound_queue_depth: negotiated
453 .max_inbound_queue_depth
454 .min(self.max_inbound_queue_depth)
455 .max(1),
456 idle_timeout,
457 prelude_timeout: self.prelude_timeout,
458 control_timeout: self.control_timeout,
459 stream_open_rate_per_second: self.stream_open_rate_per_second,
460 message_rate_per_second: self.message_rate_per_second,
461 }
462 }
463
464 pub fn initial_limits(&self) -> ZakuraInitialLimits {
466 ZakuraLimits {
467 max_frame_bytes: self.max_frame_bytes,
468 max_message_bytes: self.max_message_bytes,
469 max_open_streams: self.max_open_streams,
470 max_inbound_queue_depth: self.max_inbound_queue_depth,
471 idle_timeout_millis: self.quic_idle_timeout.as_millis().saturating_sub(1) as u32,
472 }
473 }
474
475 pub fn transport_config(&self) -> TransportConfig {
477 let mut transport = TransportConfig::default();
478 transport
479 .max_concurrent_bidi_streams(VarInt::from_u32(u32::from(self.max_open_streams)))
480 .max_concurrent_uni_streams(VarInt::from_u32(0))
481 .stream_receive_window(VarInt::from_u32(DEFAULT_ZAKURA_STREAM_RECEIVE_WINDOW))
482 .receive_window(VarInt::from_u32(DEFAULT_ZAKURA_RECEIVE_WINDOW))
483 .send_window(DEFAULT_ZAKURA_SEND_WINDOW)
484 .max_idle_timeout(Some(
485 self.quic_idle_timeout
486 .try_into()
487 .expect("default Zakura idle timeout is a valid QUIC idle timeout"),
488 ))
489 .keep_alive_interval(Some(self.keep_alive_interval))
490 .datagram_receive_buffer_size(None)
491 .datagram_send_buffer_size(0);
492 transport
493 }
494}
495
496#[derive(Copy, Clone, Debug, Eq, PartialEq)]
498pub struct ZakuraConnectionLimits {
499 pub max_frame_bytes: u32,
501 pub max_message_bytes: u32,
503 pub max_open_streams: u16,
505 pub max_inbound_queue_depth: u16,
507 pub idle_timeout: Duration,
509 pub prelude_timeout: Duration,
511 pub control_timeout: Duration,
513 pub stream_open_rate_per_second: u32,
515 pub message_rate_per_second: u32,
517}
518
519#[derive(Clone, Debug)]
521pub struct CustomService {
522 pub service: Arc<dyn Service>,
524
525 pub provides: Vec<ZakuraServiceId>,
527
528 pub seeks: Vec<ZakuraServiceId>,
530}
531
532impl From<CustomService> for Arc<dyn Service> {
533 fn from(service: CustomService) -> Self {
534 service.service
535 }
536}
537
538#[derive(Debug, Clone)]
540pub struct ZakuraEndpoint {
541 router: Router,
542 supervisor: ZakuraSupervisorHandle,
543 handler: ZakuraProtocolHandler,
544 header_sync: Option<super::HeaderSyncHandle>,
545 block_sync: Option<BlockSyncHandle>,
546 header_sync_tasks: Option<Arc<HeaderSyncBackgroundTasks>>,
547 header_sync_fatal_events:
548 Option<Arc<Mutex<Option<mpsc::UnboundedReceiver<super::HeaderSyncFatalEvent>>>>>,
549 #[cfg(any(test, feature = "zakura-testkit"))]
550 header_sync_actions: Option<Arc<Mutex<Option<mpsc::Receiver<HeaderSyncAction>>>>>,
551 block_sync_actions: Option<Arc<Mutex<Option<mpsc::Receiver<BlockSyncAction>>>>>,
552 upgrade_dials: Arc<StdMutex<HashMap<ZakuraPeerId, UpgradeDialOwnership>>>,
556}
557
558#[derive(Debug)]
559struct UpgradeDialOwnership {
560 dial_abort: AbortHandle,
561 lifetime: CancellationToken,
562}
563
564#[derive(Debug)]
565struct HeaderSyncBackgroundTasks {
566 shutdown: CancellationToken,
567 tasks: Mutex<Vec<JoinHandle<()>>>,
568}
569
570#[derive(Clone, Debug)]
572pub struct ZakuraHeaderSyncDriverStartup {
573 pub frontiers: FullStateFrontiers,
575 pub best_header_tip: Option<(block::Height, block::Hash)>,
577 pub verified_block_tip_hash: block::Hash,
579 pub committed_snapshots: watch::Receiver<Option<zakura_header_chain::EngineSnapshot>>,
582 pub committed_views: watch::Receiver<Option<zakura_header_chain::CommittedHeaderChainView>>,
584 pub service_demand: watch::Receiver<zakura_node_services::sync_lifecycle::SyncServiceDemand>,
586 pub vct_root_repairs: Option<watch::Receiver<zakura_header_chain::VctRootRepairStatus>>,
588 pub header_chain_port: Arc<dyn zakura_node_services::header_chain::Port>,
590}
591
592impl ZakuraEndpoint {
593 pub(crate) fn local_node_id(&self) -> NodeId {
595 self.router.endpoint().node_id()
596 }
597
598 pub fn connector(&self) -> super::ZakuraHandshakeConnector {
600 super::ZakuraHandshakeConnector::new_with_endpoint(self.clone())
601 }
602
603 pub(crate) async fn local_upgrade_hints(&self) -> (Vec<u8>, Vec<Vec<u8>>) {
609 let endpoint = self.router.endpoint();
610 let node_id = endpoint.node_id().as_bytes().to_vec();
611 let node_addr = endpoint.node_addr().initialized().await;
612 let direct_addresses = node_addr
613 .direct_addresses()
614 .take(super::MAX_IROH_DIRECT_ADDRESSES)
615 .map(|addr| addr.to_string().into_bytes())
616 .collect();
617 (node_id, direct_addresses)
618 }
619
620 pub fn supervisor(&self) -> ZakuraSupervisorHandle {
622 self.supervisor.clone()
623 }
624
625 pub fn trace(&self) -> ZakuraTrace {
627 self.handler.trace.clone()
628 }
629
630 pub fn header_sync(&self) -> Option<super::HeaderSyncHandle> {
632 self.header_sync.clone()
633 }
634
635 pub fn block_sync(&self) -> Option<BlockSyncHandle> {
637 self.block_sync.clone()
638 }
639
640 #[cfg(any(test, feature = "zakura-testkit"))]
642 pub async fn take_header_sync_actions(&self) -> Option<mpsc::Receiver<HeaderSyncAction>> {
643 let actions = self.header_sync_actions.as_ref()?;
644 actions.lock().await.take()
645 }
646
647 pub async fn take_block_sync_actions(&self) -> Option<mpsc::Receiver<BlockSyncAction>> {
649 let actions = self.block_sync_actions.as_ref()?;
650 actions.lock().await.take()
651 }
652
653 pub fn header_sync_shutdown(&self) -> Option<CancellationToken> {
655 self.header_sync_tasks
656 .as_ref()
657 .map(|tasks| tasks.shutdown.clone())
658 }
659
660 pub async fn take_header_sync_fatal_events(
662 &self,
663 ) -> Option<mpsc::UnboundedReceiver<super::HeaderSyncFatalEvent>> {
664 let events = self.header_sync_fatal_events.as_ref()?;
665 events.lock().await.take()
666 }
667
668 pub async fn shutdown_sync_tasks(&self) {
674 if let Some(tasks) = &self.header_sync_tasks {
675 tasks.shutdown.cancel();
676 let mut tasks = tasks.tasks.lock().await;
677 for task in tasks.drain(..) {
678 let _ = task.await;
679 }
680 }
681 }
682
683 pub(crate) fn background_shutdown_token(&self) -> CancellationToken {
693 self.header_sync_tasks
694 .as_ref()
695 .map(|tasks| tasks.shutdown.clone())
696 .unwrap_or_default()
697 }
698
699 pub async fn push_header_sync_task(&self, task: JoinHandle<()>) {
701 if let Some(tasks) = self.header_sync_tasks.as_ref() {
702 tasks.tasks.lock().await.push(task);
703 }
704 }
705
706 pub async fn push_block_sync_task(&self, task: JoinHandle<()>) {
708 self.push_header_sync_task(task).await;
709 }
710
711 pub async fn node_addr(&self) -> NodeAddr {
713 self.router.endpoint().node_addr().initialized().await
714 }
715
716 pub fn add_node_addr(
718 &self,
719 node_addr: NodeAddr,
720 ) -> Result<(), iroh::endpoint::AddNodeAddrError> {
721 self.router.endpoint().add_node_addr(node_addr)
722 }
723
724 pub fn spawn_native_dial(&self, node_addr: NodeAddr) -> tokio::task::JoinHandle<()> {
736 let endpoint = self.clone();
737 let limits = self.handler.limits.clone();
738 let policy = RedialPolicy::maintain(
739 DEFAULT_ZAKURA_REDIAL_INITIAL_BACKOFF,
740 DEFAULT_ZAKURA_REDIAL_MAX_BACKOFF,
741 );
742 tokio::spawn(native_dial_supervised(endpoint, node_addr, limits, policy))
743 }
744
745 pub(crate) fn start_upgrade_native_dial(&self, node_addr: NodeAddr) -> ZakuraUpgradeDialStart {
752 let Ok(peer_id) = ZakuraPeerId::new(node_addr.node_id.as_bytes().to_vec()) else {
753 return ZakuraUpgradeDialStart::InvalidPeerId;
754 };
755
756 let mut upgrade_dials = self
762 .upgrade_dials
763 .lock()
764 .expect("Zakura upgrade dial registry mutex is never poisoned");
765 if upgrade_dials.contains_key(&peer_id) {
766 return ZakuraUpgradeDialStart::AlreadyRunning;
767 }
768
769 let endpoint = self.clone();
770 let limits = self.handler.limits.clone();
771 let policy = RedialPolicy::maintain(
772 DEFAULT_ZAKURA_REDIAL_INITIAL_BACKOFF,
773 DEFAULT_ZAKURA_REDIAL_MAX_BACKOFF,
774 );
775 let lifetime = CancellationToken::new();
776 let task_lifetime = lifetime.clone();
777 let task_peer_id = peer_id.clone();
778 let dial = tokio::spawn(async move {
779 let _lifetime_guard = task_lifetime.drop_guard();
780 native_dial_supervised(endpoint.clone(), node_addr, limits, policy).await;
781 endpoint
782 .supervisor
783 .forget_retained_native_metadata(&task_peer_id);
784 endpoint
785 .upgrade_dials
786 .lock()
787 .expect("Zakura upgrade dial registry mutex is never poisoned")
788 .remove(&task_peer_id);
789 });
790 upgrade_dials.insert(
791 peer_id,
792 UpgradeDialOwnership {
793 dial_abort: dial.abort_handle(),
794 lifetime,
795 },
796 );
797 ZakuraUpgradeDialStart::Started
798 }
799
800 pub(crate) fn upgrade_dial_lifetime(
803 &self,
804 peer_id: &ZakuraPeerId,
805 ) -> Option<CancellationToken> {
806 self.upgrade_dials
807 .lock()
808 .expect("Zakura upgrade dial registry mutex is never poisoned")
809 .get(peer_id)
810 .map(|ownership| ownership.lifetime.clone())
811 }
812
813 pub(crate) fn cancel_upgrade_native_dial(&self, peer_id: &ZakuraPeerId) {
825 let ownership = self
826 .upgrade_dials
827 .lock()
828 .expect("Zakura upgrade dial registry mutex is never poisoned")
829 .remove(peer_id);
830 if let Some(ownership) = ownership {
831 ownership.lifetime.cancel();
832 ownership.dial_abort.abort();
833 self.supervisor.forget_retained_native_metadata(peer_id);
834 }
835 }
836
837 pub(crate) fn has_native_admission_capacity(&self) -> bool {
842 self.handler.admission.available_permits() > 0
843 }
844
845 pub async fn shutdown(&self) {
847 if let Some(tasks) = &self.header_sync_tasks {
848 tasks.shutdown.cancel();
849 let mut tasks = tasks.tasks.lock().await;
850 for mut task in tasks.drain(..) {
851 if timeout(Duration::from_secs(1), &mut task).await.is_err() {
852 task.abort();
853 let _ = task.await;
854 }
855 }
856 }
857 self.supervisor.shutdown();
858 let _ = self.router.shutdown().await;
859 }
860
861 #[cfg(any(test, feature = "zakura-testkit"))]
862 pub(crate) fn from_parts(
863 router: Router,
864 supervisor: ZakuraSupervisorHandle,
865 handler: ZakuraProtocolHandler,
866 ) -> Self {
867 Self {
868 router,
869 supervisor,
870 handler,
871 header_sync: None,
872 block_sync: None,
873 header_sync_tasks: None,
874 header_sync_fatal_events: None,
875 header_sync_actions: None,
876 block_sync_actions: None,
877 upgrade_dials: Arc::new(StdMutex::new(HashMap::new())),
878 }
879 }
880
881 #[cfg(any(test, feature = "zakura-testkit"))]
882 #[allow(dead_code)]
883 pub(crate) fn from_parts_with_header_sync(
884 router: Router,
885 supervisor: ZakuraSupervisorHandle,
886 handler: ZakuraProtocolHandler,
887 header_sync: super::HeaderSyncHandle,
888 shutdown: CancellationToken,
889 tasks: Vec<JoinHandle<()>>,
890 actions: Option<mpsc::Receiver<HeaderSyncAction>>,
891 ) -> Self {
892 Self {
893 router,
894 supervisor,
895 handler,
896 header_sync: Some(header_sync),
897 block_sync: None,
898 header_sync_tasks: Some(Arc::new(HeaderSyncBackgroundTasks {
899 shutdown,
900 tasks: Mutex::new(tasks),
901 })),
902 header_sync_fatal_events: None,
903 header_sync_actions: actions.map(|actions| Arc::new(Mutex::new(Some(actions)))),
904 block_sync_actions: None,
905 upgrade_dials: Arc::new(StdMutex::new(HashMap::new())),
906 }
907 }
908
909 #[cfg(any(test, feature = "zakura-testkit"))]
910 #[allow(clippy::too_many_arguments)]
911 pub(crate) fn from_parts_with_sync_services(
912 router: Router,
913 supervisor: ZakuraSupervisorHandle,
914 handler: ZakuraProtocolHandler,
915 header_sync: super::HeaderSyncHandle,
916 block_sync: BlockSyncHandle,
917 shutdown: CancellationToken,
918 tasks: Vec<JoinHandle<()>>,
919 header_sync_actions: Option<mpsc::Receiver<HeaderSyncAction>>,
920 block_sync_actions: Option<mpsc::Receiver<BlockSyncAction>>,
921 ) -> Self {
922 Self {
923 router,
924 supervisor,
925 handler,
926 header_sync: Some(header_sync),
927 block_sync: Some(block_sync),
928 header_sync_tasks: Some(Arc::new(HeaderSyncBackgroundTasks {
929 shutdown,
930 tasks: Mutex::new(tasks),
931 })),
932 header_sync_fatal_events: None,
933 header_sync_actions: header_sync_actions
934 .map(|actions| Arc::new(Mutex::new(Some(actions)))),
935 block_sync_actions: block_sync_actions
936 .map(|actions| Arc::new(Mutex::new(Some(actions)))),
937 upgrade_dials: Arc::new(StdMutex::new(HashMap::new())),
938 }
939 }
940}
941
942#[derive(Clone, Debug)]
944pub struct ZakuraSupervisorHandle {
945 id: u64,
946 inner: Arc<Mutex<ZakuraSupervisorState>>,
947 shutdown: CancellationToken,
948 peer_set_tx: watch::Sender<Vec<ZakuraPeerId>>,
949 registration_tx: broadcast::Sender<ZakuraConnectionRegistration>,
950 pending_handoffs: Arc<StdMutex<HashMap<ZakuraPeerId, ZakuraHandoffId>>>,
951 peer_registry: Option<PeerRegistry>,
952}
953
954#[derive(Clone, Debug, Eq, PartialEq)]
956struct ZakuraConnectionRegistration {
957 peer_id: ZakuraPeerId,
958 conn_id: ZakuraConnId,
959}
960
961type ZakuraHandoffId = u64;
963
964const ZAKURA_REGISTRATION_EVENT_CAPACITY: usize = 1024;
967
968pub(crate) struct ZakuraConnectionRegistrationWait {
970 peer_id: ZakuraPeerId,
971 handoff_id: ZakuraHandoffId,
972 registrations: broadcast::Receiver<ZakuraConnectionRegistration>,
973 pending_handoffs: Arc<StdMutex<HashMap<ZakuraPeerId, ZakuraHandoffId>>>,
974}
975
976pub(crate) enum ZakuraConnectionRegistrationWaitStart {
978 Waiting(ZakuraConnectionRegistrationWait),
980 Duplicate,
982 Unavailable,
984}
985
986impl ZakuraConnectionRegistrationWait {
987 pub(crate) async fn wait(mut self, registration_timeout: Duration) -> Option<ZakuraConnId> {
990 timeout(registration_timeout, async {
991 loop {
992 match self.registrations.recv().await {
993 Ok(registration) if registration.peer_id == self.peer_id => {
994 return Some(registration.conn_id);
995 }
996 Ok(_) => {}
997 Err(broadcast::error::RecvError::Lagged(_))
998 | Err(broadcast::error::RecvError::Closed) => return None,
999 }
1000 }
1001 })
1002 .await
1003 .unwrap_or(None)
1004 }
1005}
1006
1007impl Drop for ZakuraConnectionRegistrationWait {
1008 fn drop(&mut self) {
1009 let peer_id = self.peer_id.clone();
1010 let handoff_id = self.handoff_id;
1011 let mut pending_handoffs = self
1012 .pending_handoffs
1013 .lock()
1014 .expect("Zakura pending handoff mutex is never poisoned");
1015 if pending_handoffs.get(&peer_id) == Some(&handoff_id) {
1017 pending_handoffs.remove(&peer_id);
1018 }
1019 }
1020}
1021
1022static NEXT_SUPERVISOR_ID: AtomicU64 = AtomicU64::new(1);
1023
1024#[derive(Debug)]
1025struct ZakuraSupervisorState {
1026 supervisor: ZakuraPeerSupervisor,
1027 active_by_peer: HashMap<ZakuraPeerId, ZakuraPeerConnectionEntry>,
1028 active_by_ip: HashMap<IpAddr, usize>,
1029 next_handoff_id: ZakuraHandoffId,
1030 max_connections_per_ip: usize,
1031 next_registration_id: ZakuraConnId,
1032}
1033
1034#[derive(Debug)]
1035struct ZakuraPeerConnectionEntry {
1036 conn_id: ZakuraConnId,
1039 outbound_handle: ZakuraPeerHandle,
1040 disconnect_token: CancellationToken,
1041 registered_at: Instant,
1042 remote_ip: Option<IpAddr>,
1043}
1044
1045impl ZakuraSupervisorState {
1046 fn increment_ip(&mut self, remote_ip: Option<IpAddr>) {
1047 if let Some(remote_ip) = remote_ip {
1048 *self.active_by_ip.entry(remote_ip).or_default() += 1;
1049 }
1050 }
1051
1052 fn decrement_ip(&mut self, remote_ip: Option<IpAddr>) {
1053 if let Some(remote_ip) = remote_ip {
1054 if let Some(count) = self.active_by_ip.get_mut(&remote_ip) {
1055 *count = count.saturating_sub(1);
1056 if *count == 0 {
1057 self.active_by_ip.remove(&remote_ip);
1058 }
1059 }
1060 }
1061 }
1062
1063 #[cfg(debug_assertions)]
1064 fn debug_assert_accounting(&self) {
1065 let active_by_ip_total: usize = self.active_by_ip.values().sum();
1066 debug_assert_eq!(
1067 active_by_ip_total,
1068 self.active_by_peer
1069 .values()
1070 .filter(|entry| entry.remote_ip.is_some())
1071 .count(),
1072 "Zakura active_by_ip totals must match active peer registrations with known IPs",
1073 );
1074 }
1075
1076 #[cfg(not(debug_assertions))]
1077 fn debug_assert_accounting(&self) {}
1078}
1079
1080#[derive(Clone, Debug)]
1082pub struct ZakuraPeerHandle {
1083 peer_id: ZakuraPeerId,
1084 sender: mpsc::Sender<ZakuraOutboundFrame>,
1085}
1086
1087impl ZakuraPeerHandle {
1088 #[cfg(test)]
1089 pub(crate) fn new_for_tests(
1090 peer_id: ZakuraPeerId,
1091 sender: mpsc::Sender<ZakuraOutboundFrame>,
1092 ) -> Self {
1093 Self { peer_id, sender }
1094 }
1095
1096 pub fn peer_id(&self) -> &ZakuraPeerId {
1098 &self.peer_id
1099 }
1100
1101 fn has_outbound_capacity(&self) -> bool {
1102 self.sender.capacity() > 0
1103 }
1104
1105 pub async fn request(
1107 &self,
1108 stream_kind: u16,
1109 request_id: u64,
1110 message_type: u16,
1111 flags: u16,
1112 payload: Vec<u8>,
1113 ) -> Result<Vec<Frame>, BoxError> {
1114 let (completion, completed) = oneshot::channel();
1115 let frame = ZakuraOutboundFrame::Request {
1116 stream_kind,
1117 request_id,
1118 message_type,
1119 flags,
1120 payload,
1121 completion,
1122 };
1123 self.sender
1124 .send(frame)
1125 .await
1126 .map_err(|_| -> BoxError { "Zakura outbound peer queue closed".into() })?;
1127 completed
1128 .await
1129 .map_err(|_| -> BoxError { "Zakura outbound completion dropped".into() })?
1130 }
1131}
1132
1133#[derive(Debug)]
1135pub enum ZakuraOutboundFrame {
1136 Request {
1138 stream_kind: u16,
1140 request_id: u64,
1142 message_type: u16,
1144 flags: u16,
1146 payload: Vec<u8>,
1148 completion: oneshot::Sender<Result<Vec<Frame>, BoxError>>,
1150 },
1151}
1152
1153impl ZakuraSupervisorHandle {
1154 pub fn new(max_connections_per_ip: usize) -> Self {
1156 Self::new_inner(max_connections_per_ip, None)
1157 }
1158
1159 pub(crate) fn new_with_peer_registry(
1160 max_connections_per_ip: usize,
1161 peer_registry: PeerRegistry,
1162 ) -> Self {
1163 Self::new_inner(max_connections_per_ip, Some(peer_registry))
1164 }
1165
1166 fn new_inner(max_connections_per_ip: usize, peer_registry: Option<PeerRegistry>) -> Self {
1167 let (registration_tx, _) = broadcast::channel(ZAKURA_REGISTRATION_EVENT_CAPACITY);
1168 Self {
1169 id: NEXT_SUPERVISOR_ID.fetch_add(1, Ordering::Relaxed),
1170 inner: Arc::new(Mutex::new(ZakuraSupervisorState {
1171 supervisor: ZakuraPeerSupervisor::default(),
1172 active_by_peer: HashMap::new(),
1173 active_by_ip: HashMap::new(),
1174 next_handoff_id: 1,
1175 max_connections_per_ip: max_connections_per_ip.max(1),
1176 next_registration_id: 1,
1177 })),
1178 shutdown: CancellationToken::new(),
1179 peer_set_tx: watch::channel(Vec::new()).0,
1180 registration_tx,
1181 pending_handoffs: Arc::new(StdMutex::new(HashMap::new())),
1182 peer_registry,
1183 }
1184 }
1185
1186 pub async fn registered_ids(&self) -> Vec<ZakuraPeerId> {
1188 let state = self.inner.lock().await;
1189 state.active_by_peer.keys().cloned().collect()
1190 }
1191
1192 pub(crate) fn id(&self) -> u64 {
1193 self.id
1194 }
1195
1196 pub async fn outbound_peer_handles(&self) -> Vec<ZakuraPeerHandle> {
1198 let state = self.inner.lock().await;
1199 state
1200 .active_by_peer
1201 .values()
1202 .map(|entry| &entry.outbound_handle)
1203 .filter(|handle| handle.has_outbound_capacity())
1204 .cloned()
1205 .collect()
1206 }
1207
1208 pub fn subscribe(&self) -> watch::Receiver<Vec<ZakuraPeerId>> {
1210 self.peer_set_tx.subscribe()
1211 }
1212
1213 pub(crate) async fn begin_connection_registration_wait(
1219 &self,
1220 peer_id: &ZakuraPeerId,
1221 ) -> ZakuraConnectionRegistrationWaitStart {
1222 let mut state = self.inner.lock().await;
1223 let mut pending_handoffs = self
1224 .pending_handoffs
1225 .lock()
1226 .expect("Zakura pending handoff mutex is never poisoned");
1227 if state.active_by_peer.contains_key(peer_id) || pending_handoffs.contains_key(peer_id) {
1228 return ZakuraConnectionRegistrationWaitStart::Duplicate;
1229 }
1230 if state.next_handoff_id == u64::MAX {
1231 return ZakuraConnectionRegistrationWaitStart::Unavailable;
1232 }
1233 let handoff_id = state.next_handoff_id;
1234 state.next_handoff_id += 1;
1235 pending_handoffs.insert(peer_id.clone(), handoff_id);
1236 let registrations = self.registration_tx.subscribe();
1237 ZakuraConnectionRegistrationWaitStart::Waiting(ZakuraConnectionRegistrationWait {
1238 peer_id: peer_id.clone(),
1239 handoff_id,
1240 registrations,
1241 pending_handoffs: self.pending_handoffs.clone(),
1242 })
1243 }
1244
1245 pub async fn disconnect_peer(&self, peer_id: &ZakuraPeerId) -> bool {
1247 let token = {
1248 let state = self.inner.lock().await;
1249 state
1250 .active_by_peer
1251 .get(peer_id)
1252 .map(|entry| entry.disconnect_token.clone())
1253 };
1254
1255 if let Some(token) = token {
1256 token.cancel();
1257 true
1258 } else {
1259 false
1260 }
1261 }
1262
1263 #[allow(clippy::too_many_arguments)]
1264 async fn register(
1265 &self,
1266 _trace_conn_id: ZakuraConnId,
1267 peer_id: ZakuraPeerId,
1268 remote_ip: Option<IpAddr>,
1269 transcript_hash: [u8; TRANSCRIPT_HASH_BYTES],
1270 outbound_handle: ZakuraPeerHandle,
1271 disconnect_token: CancellationToken,
1272 _accepted_capabilities: u64,
1273 ) -> ZakuraRegistration {
1274 let remote_ip = remote_ip.map(canonical_ip);
1275 let mut state = self.inner.lock().await;
1276 let same_ip_duplicate_redial = state
1285 .active_by_peer
1286 .get(&peer_id)
1287 .is_some_and(|entry| entry.remote_ip == remote_ip);
1288 if let Some(remote_ip) = remote_ip {
1289 if !same_ip_duplicate_redial {
1290 let ip_count = state
1291 .active_by_ip
1292 .get(&remote_ip)
1293 .copied()
1294 .unwrap_or_default();
1295 if ip_count >= state.max_connections_per_ip {
1296 metrics::counter!("zakura.p2p.conn.rejected.admission").increment(1);
1297 return ZakuraRegistration::Rejected(ZakuraRejectReason::ResourceLimit);
1298 }
1299 }
1300 }
1301 if state.next_registration_id == u64::MAX {
1302 metrics::counter!("zakura.p2p.conn.rejected.admission").increment(1);
1303 return ZakuraRegistration::Rejected(ZakuraRejectReason::TemporaryUnavailable);
1304 }
1305
1306 match state
1307 .supervisor
1308 .register_authenticated(peer_id.clone(), transcript_hash)
1309 {
1310 AuthenticatedPeerRegistration::Registered => {
1311 let conn_id = state.next_registration_id;
1312 state.next_registration_id += 1;
1313 self.pending_handoffs
1314 .lock()
1315 .expect("Zakura pending handoff mutex is never poisoned")
1316 .remove(&peer_id);
1317 let entry = ZakuraPeerConnectionEntry {
1318 conn_id,
1319 outbound_handle,
1320 disconnect_token,
1321 registered_at: Instant::now(),
1322 remote_ip,
1323 };
1324 if let Some(old_entry) = state.active_by_peer.insert(peer_id.clone(), entry) {
1325 state.decrement_ip(old_entry.remote_ip);
1326 old_entry.disconnect_token.cancel();
1327 metrics::counter!("zakura.p2p.conn.duplicate.evicted_upgraded").increment(1);
1328 }
1329 state.increment_ip(remote_ip);
1330 state.debug_assert_accounting();
1331 let registered_ids: Vec<_> = state.active_by_peer.keys().cloned().collect();
1332 set_active_connection_gauge(registered_ids.len());
1333 self.peer_set_tx.send_replace(registered_ids);
1334 if let Some(peer_registry) = &self.peer_registry {
1335 peer_registry.register_native_connection(peer_id.clone(), conn_id);
1336 }
1337 let _ = self.registration_tx.send(ZakuraConnectionRegistration {
1338 peer_id: peer_id.clone(),
1339 conn_id,
1340 });
1341 let disconnect_token = state
1342 .active_by_peer
1343 .get(&peer_id)
1344 .map(|entry| entry.disconnect_token.clone())
1345 .expect("disconnect token exists because this peer was just registered");
1346 ZakuraRegistration::Registered {
1347 conn_id,
1348 peer_id,
1349 remote_ip,
1350 disconnect_token,
1351 }
1352 }
1353 AuthenticatedPeerRegistration::Duplicate => {
1354 let evict_min_age = if same_ip_duplicate_redial {
1372 ZAKURA_SAME_IP_DUPLICATE_EVICT_MIN_AGE
1373 } else {
1374 ZAKURA_DUPLICATE_EVICT_MIN_AGE
1375 };
1376 if let Some(entry) = state.active_by_peer.get(&peer_id) {
1377 if entry.registered_at.elapsed() >= evict_min_age
1378 && !entry.disconnect_token.is_cancelled()
1379 {
1380 entry.disconnect_token.cancel();
1381 metrics::counter!("zakura.p2p.conn.duplicate.evicted_stale").increment(1);
1382 }
1383 }
1384 ZakuraRegistration::Duplicate { peer_id }
1385 }
1386 }
1387 }
1388
1389 async fn deregister(&self, peer_id: &ZakuraPeerId, conn_id: ZakuraConnId) {
1390 let mut state = self.inner.lock().await;
1391 let Some(entry) = state.active_by_peer.get(peer_id) else {
1392 state.debug_assert_accounting();
1393 return;
1394 };
1395 if entry.conn_id != conn_id {
1396 state.debug_assert_accounting();
1397 return;
1398 }
1399 let entry = state
1400 .active_by_peer
1401 .remove(peer_id)
1402 .expect("peer entry exists because it was just checked");
1403 state.decrement_ip(entry.remote_ip);
1404 state.supervisor.deregister_authenticated(peer_id);
1405 state.debug_assert_accounting();
1406 let registered_ids: Vec<_> = state.active_by_peer.keys().cloned().collect();
1407 set_active_connection_gauge(registered_ids.len());
1408 self.peer_set_tx.send_replace(registered_ids);
1409 if let Some(peer_registry) = &self.peer_registry {
1410 peer_registry.deregister_native_connection(peer_id, conn_id);
1411 }
1412 }
1413
1414 fn forget_retained_native_metadata(&self, peer_id: &ZakuraPeerId) {
1415 if let Some(peer_registry) = &self.peer_registry {
1416 peer_registry.forget_retained_native_metadata(peer_id);
1417 }
1418 }
1419
1420 fn shutdown(&self) {
1421 self.shutdown.cancel();
1422 }
1423
1424 pub(crate) async fn can_accept_remote_ip_with_in_flight(
1429 &self,
1430 remote_ip: IpAddr,
1431 in_flight_count: usize,
1432 ) -> bool {
1433 let remote_ip = canonical_ip(remote_ip);
1434 let state = self.inner.lock().await;
1435 let active_count = state
1436 .active_by_ip
1437 .get(&remote_ip)
1438 .copied()
1439 .unwrap_or_default();
1440 active_count.saturating_add(in_flight_count) < state.max_connections_per_ip
1441 }
1442}
1443
1444fn set_active_connection_gauge(active_connections: usize) {
1445 metrics::gauge!("zakura.p2p.conn.active").set(active_connections as f64);
1448}
1449
1450#[derive(Debug)]
1451enum ZakuraRegistration {
1452 Registered {
1453 conn_id: ZakuraConnId,
1454 peer_id: ZakuraPeerId,
1455 remote_ip: Option<IpAddr>,
1456 disconnect_token: CancellationToken,
1457 },
1458 Duplicate {
1459 peer_id: ZakuraPeerId,
1460 },
1461 Rejected(ZakuraRejectReason),
1462}
1463
1464#[derive(Debug)]
1465struct RegisteredPeerCleanupGuard {
1466 supervisor: ZakuraSupervisorHandle,
1467 registry: Arc<ServiceRegistry>,
1468 peer_id: ZakuraPeerId,
1469 conn_id: ZakuraConnId,
1470 disconnect_token: CancellationToken,
1471 admitted_capabilities: u64,
1472 armed: bool,
1473}
1474
1475impl RegisteredPeerCleanupGuard {
1476 fn new(
1477 supervisor: ZakuraSupervisorHandle,
1478 registry: Arc<ServiceRegistry>,
1479 peer_id: ZakuraPeerId,
1480 conn_id: ZakuraConnId,
1481 disconnect_token: CancellationToken,
1482 ) -> Self {
1483 Self {
1484 supervisor,
1485 registry,
1486 peer_id,
1487 conn_id,
1488 disconnect_token,
1489 admitted_capabilities: 0,
1490 armed: true,
1491 }
1492 }
1493
1494 fn add_admitted_capabilities(&mut self, capabilities: u64) {
1495 self.admitted_capabilities |= capabilities;
1496 }
1497
1498 fn remove_admitted_services(&mut self) {
1499 if self.admitted_capabilities != 0 {
1500 self.registry
1501 .remove_peer(&self.peer_id, self.conn_id, self.admitted_capabilities);
1502 self.admitted_capabilities = 0;
1503 }
1504 }
1505
1506 async fn cleanup_registered_peer(mut self) {
1507 if self.armed {
1508 self.disconnect_token.cancel();
1509 self.remove_admitted_services();
1510 self.supervisor
1511 .deregister(&self.peer_id, self.conn_id)
1512 .await;
1513 self.armed = false;
1514 }
1515 }
1516}
1517
1518impl Drop for RegisteredPeerCleanupGuard {
1519 fn drop(&mut self) {
1520 if !self.armed {
1521 return;
1522 }
1523
1524 self.disconnect_token.cancel();
1525 self.remove_admitted_services();
1526
1527 let supervisor = self.supervisor.clone();
1528 let peer_id = self.peer_id.clone();
1529 let conn_id = self.conn_id;
1530 if let Ok(handle) = tokio::runtime::Handle::try_current() {
1531 let _task = handle.spawn(async move {
1532 supervisor.deregister(&peer_id, conn_id).await;
1533 });
1534 }
1535 }
1536}
1537
1538struct StreamAdmission<'a> {
1539 conn: ZakuraConnTrace,
1540 peer_id: &'a ZakuraPeerId,
1541 stream_sem: &'a Arc<Semaphore>,
1542 open_limiter: &'a mut TokenBucket,
1543 message_buckets: &'a mut MessageRateBuckets,
1544 workers: &'a mut JoinSet<()>,
1545 limits: ZakuraConnectionLimits,
1546 accepted_capabilities: u64,
1547 connection_token: CancellationToken,
1548 close_cause: CloseCause,
1549 freshness_tx: watch::Sender<Instant>,
1550}
1551
1552struct ConnectionServeContext {
1553 limits: ZakuraConnectionLimits,
1554 accepted_capabilities: u64,
1555 role: &'static str,
1556 direction: ServicePeerDirection,
1557 transcript_hash: [u8; TRANSCRIPT_HASH_BYTES],
1558 i_open_collision_winner: bool,
1563 conn: ZakuraConnTrace,
1564}
1565
1566struct RegisteredConnectionServeContext {
1567 limits: ZakuraConnectionLimits,
1568 conn: ZakuraConnTrace,
1569 conn_id: ZakuraConnId,
1570 connection_token: CancellationToken,
1571 close_cause: CloseCause,
1572 accepted_capabilities: u64,
1573 is_initiator: bool,
1577 i_open_collision_winner: bool,
1579 direction: ServicePeerDirection,
1580}
1581
1582struct StreamWorkerContext {
1583 conn: ZakuraConnTrace,
1584 peer_id: ZakuraPeerId,
1585 stream_id: u64,
1586 _permit: OwnedSemaphorePermit,
1587 limits: ZakuraConnectionLimits,
1588 inbound_frame_cap: u32,
1589 outbound_frame_cap: u32,
1590 message_bucket: SharedMessageBucket,
1591 connection_token: CancellationToken,
1592 stream_token: CancellationToken,
1593 close_cause: CloseCause,
1594 freshness_tx: watch::Sender<Instant>,
1595}
1596
1597struct AdmittedOrderedSession {
1598 kind: u16,
1599 version: u16,
1600 session_id: u64,
1601 recv: FramedRecv,
1602 send: FramedSend,
1603 cancel_token: CancellationToken,
1604}
1605
1606#[derive(Copy, Clone, Debug)]
1607struct OrderedSessionExit {
1608 stream: Stream,
1609 session_id: u64,
1610 opened_locally: bool,
1611}
1612
1613#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1614enum OrderedSessionWaitReason {
1615 Demand,
1616 Transport,
1617}
1618
1619#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
1620enum OrderedSessionReopenState {
1621 #[default]
1622 Idle,
1623 Waiting(OrderedSessionWaitReason),
1624 Retired,
1625}
1626
1627type OrderedSessionWait = futures::stream::Once<BoxRunFuture<'static, ()>>;
1628type OrderedSessionWaits = StreamMap<u16, OrderedSessionWait>;
1629
1630#[derive(Debug)]
1632struct OrderedSessionState {
1633 stream: Stream,
1634 local_session_id: Option<u64>,
1635 remote_session_id: Option<u64>,
1636 reopen_attempts: u32,
1640 reopen_state: OrderedSessionReopenState,
1641}
1642
1643impl OrderedSessionState {
1644 fn new(stream: Stream) -> Self {
1645 Self {
1646 stream,
1647 local_session_id: None,
1648 remote_session_id: None,
1649 reopen_attempts: 0,
1650 reopen_state: OrderedSessionReopenState::Idle,
1651 }
1652 }
1653
1654 fn has_active_session(&self) -> bool {
1655 self.local_session_id.is_some() || self.remote_session_id.is_some()
1656 }
1657
1658 fn remove_active_session(&mut self, opened_locally: bool, session_id: u64) -> bool {
1659 let active_session = if opened_locally {
1660 &mut self.local_session_id
1661 } else {
1662 &mut self.remote_session_id
1663 };
1664 if *active_session != Some(session_id) {
1665 return false;
1666 }
1667
1668 *active_session = None;
1669 true
1670 }
1671
1672 fn schedule_demand(&mut self, waits: &mut OrderedSessionWaits, demand: OrderedSessionDemand) {
1673 match demand {
1674 OrderedSessionDemand::OpenNow => {
1675 self.cancel_wait(waits);
1676 self.install_wait(
1677 waits,
1678 OrderedSessionWaitReason::Demand,
1679 Box::pin(future::ready(())),
1680 );
1681 }
1682 OrderedSessionDemand::RetryAt(at) => {
1683 self.install_demand_wait(
1684 waits,
1685 Box::pin(async move {
1686 tokio::time::sleep_until(at.into()).await;
1687 }),
1688 );
1689 }
1690 OrderedSessionDemand::WaitForChange(changed) => {
1691 self.install_demand_wait(waits, changed);
1692 }
1693 OrderedSessionDemand::Retire => {
1694 waits.remove(&self.stream.kind);
1695 self.reopen_state = OrderedSessionReopenState::Retired;
1696 }
1697 }
1698 }
1699
1700 fn schedule_transport_backoff(&mut self, waits: &mut OrderedSessionWaits) -> Option<Duration> {
1701 if self.reopen_state != OrderedSessionReopenState::Idle {
1702 return None;
1703 }
1704 debug_assert!(
1705 !waits.contains_key(&self.stream.kind),
1706 "idle ordered session cannot already have a reopen wait"
1707 );
1708
1709 let delay = ordered_session_reopen_backoff(self.reopen_attempts);
1710 self.reopen_attempts = self.reopen_attempts.saturating_add(1);
1711 self.install_wait(
1712 waits,
1713 OrderedSessionWaitReason::Transport,
1714 Box::pin(tokio::time::sleep(delay)),
1715 );
1716 Some(delay)
1717 }
1718
1719 fn finish_wait(&mut self, waits: &mut OrderedSessionWaits) {
1720 let removed = waits.remove(&self.stream.kind);
1721 debug_assert!(
1722 removed.is_some(),
1723 "yielded ordered session wait remains keyed until it is handled"
1724 );
1725 if matches!(self.reopen_state, OrderedSessionReopenState::Waiting(_)) {
1726 self.reopen_state = OrderedSessionReopenState::Idle;
1727 }
1728 }
1729
1730 fn cancel_wait(&mut self, waits: &mut OrderedSessionWaits) {
1731 waits.remove(&self.stream.kind);
1732 if matches!(self.reopen_state, OrderedSessionReopenState::Waiting(_)) {
1733 self.reopen_state = OrderedSessionReopenState::Idle;
1734 }
1735 }
1736
1737 fn install_demand_wait(
1738 &mut self,
1739 waits: &mut OrderedSessionWaits,
1740 wait: BoxRunFuture<'static, ()>,
1741 ) {
1742 if self.reopen_state == OrderedSessionReopenState::Waiting(OrderedSessionWaitReason::Demand)
1743 {
1744 return;
1745 }
1746
1747 self.cancel_wait(waits);
1748 self.install_wait(waits, OrderedSessionWaitReason::Demand, wait);
1749 }
1750
1751 fn install_wait(
1752 &mut self,
1753 waits: &mut OrderedSessionWaits,
1754 reason: OrderedSessionWaitReason,
1755 wait: BoxRunFuture<'static, ()>,
1756 ) {
1757 let replaced = waits.insert(self.stream.kind, futures::stream::once(wait));
1758 debug_assert!(
1759 replaced.is_none(),
1760 "ordered session wait is cancelled before replacement"
1761 );
1762 self.reopen_state = OrderedSessionReopenState::Waiting(reason);
1763 }
1764}
1765
1766fn may_open_ordered_stream(policy: OrderedStreamPolicy, is_initiator: bool) -> bool {
1772 is_initiator || policy.opening == OrderedStreamOpening::EitherSide
1773}
1774
1775fn opens_ordered_stream_locally(
1782 policy: OrderedStreamPolicy,
1783 is_initiator: bool,
1784 i_open_collision_winner: bool,
1785) -> bool {
1786 match policy.opening {
1787 OrderedStreamOpening::InitiatorOnly => is_initiator,
1788 OrderedStreamOpening::EitherSide => i_open_collision_winner,
1789 }
1790}
1791
1792fn should_reopen_ordered_session(
1796 exited: OrderedSessionExit,
1797 policy: OrderedStreamPolicy,
1798 is_initiator: bool,
1799 i_open_collision_winner: bool,
1800 connection_cancelled: bool,
1801) -> bool {
1802 exited.stream.mode == StreamMode::Ordered
1803 && policy.reopen
1804 && opens_ordered_stream_locally(policy, is_initiator, i_open_collision_winner)
1805 && !connection_cancelled
1806}
1807
1808fn ordered_session_reopen_backoff(attempts: u32) -> Duration {
1811 ORDERED_STREAM_REOPEN_BACKOFF
1813 .saturating_mul(1u32 << attempts.min(8))
1814 .min(ORDERED_STREAM_REOPEN_BACKOFF_CAP)
1815}
1816
1817#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1818enum InboundMessageAdmission {
1819 Admit,
1820 Oversize,
1821 Throttled,
1822}
1823
1824fn admit_inbound_message(
1825 payload_len: usize,
1826 context: &StreamWorkerContext,
1827 stream_kind: u16,
1828) -> InboundMessageAdmission {
1829 let stream_kind = stream_kind_label(stream_kind);
1830 let max_message_bytes = usize::try_from(context.limits.max_message_bytes)
1831 .expect("u32 message byte limit fits in usize");
1832 if payload_len > max_message_bytes {
1833 metrics::counter!(
1834 "zakura.p2p.ratelimit.message.oversize",
1835 "stream_kind" => stream_kind,
1836 )
1837 .increment(1);
1838 context.conn.trace_rate_limit(
1839 "message.oversize",
1840 context.stream_id,
1841 stream_kind,
1842 None,
1843 None,
1844 None,
1845 );
1846 return InboundMessageAdmission::Oversize;
1847 }
1848
1849 let admitted = {
1850 let mut bucket = context
1851 .message_bucket
1852 .lock()
1853 .expect("Zakura message-rate bucket mutex is never poisoned");
1854 bucket.try_take()
1855 };
1856 if !admitted {
1857 metrics::counter!(
1858 "zakura.p2p.ratelimit.message.throttled",
1859 "stream_kind" => stream_kind,
1860 )
1861 .increment(1);
1862 context.conn.trace_rate_limit(
1863 "message.throttled",
1864 context.stream_id,
1865 stream_kind,
1866 None,
1867 None,
1868 None,
1869 );
1870 return InboundMessageAdmission::Throttled;
1871 }
1872
1873 InboundMessageAdmission::Admit
1874}
1875
1876#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1877pub(crate) struct NativeHandshakeNegotiated {
1878 pub(crate) limits: ZakuraAcceptedLimits,
1879 pub(crate) accepted_capabilities: u64,
1880}
1881
1882#[allow(clippy::too_many_arguments)]
1883pub(crate) fn service_registry(
1884 _supervisor: &ZakuraSupervisorHandle,
1885 header_sync: Option<super::HeaderSyncHandle>,
1886 block_sync: Option<BlockSyncHandle>,
1887 block_sync_config: ZakuraBlockSyncConfig,
1888 legacy_service: Arc<dyn Service>,
1889 discovery_service: Arc<super::DiscoveryService>,
1890 service_demand: Option<
1891 watch::Receiver<zakura_node_services::sync_lifecycle::SyncServiceDemand>,
1892 >,
1893 custom_services: Vec<CustomService>,
1894) -> Result<Arc<ServiceRegistry>, BoxError> {
1895 let mut services = vec![legacy_service.clone()];
1896 let header_sync_service = if let Some(header_sync) = &header_sync {
1897 Arc::new(
1898 HeaderSyncService::new(header_sync.clone()).with_service_demand(service_demand.clone()),
1899 ) as Arc<dyn Service>
1900 } else {
1901 Arc::new(HeaderSyncPassthroughService::new(legacy_service.clone())) as Arc<dyn Service>
1902 };
1903 services.push(header_sync_service.clone());
1904 let block_sync = match block_sync {
1905 Some(block_sync) => BlockSyncService::new_with_handle(block_sync_config, block_sync),
1906 None => match header_sync.as_ref() {
1907 Some(header_sync) => BlockSyncService::new_with_header_tip(
1908 block_sync_config,
1909 header_sync.subscribe_tip(),
1910 ),
1911 None => BlockSyncService::new(block_sync_config),
1912 },
1913 };
1914 let block_sync = Arc::new(block_sync.with_service_demand(service_demand)) as Arc<dyn Service>;
1915 let custom_services = custom_services
1916 .into_iter()
1917 .map(Into::into)
1918 .collect::<Vec<_>>();
1919 let mut connection_owners = vec![legacy_service, header_sync_service, block_sync.clone()];
1920 connection_owners.extend(custom_services.iter().cloned());
1921 discovery_service.set_connection_owners(connection_owners);
1922 services.push(discovery_service as Arc<dyn Service>);
1923 services.push(block_sync);
1924 services.extend(custom_services);
1925
1926 Ok(Arc::new(
1927 ServiceRegistry::new(services).map_err(|error| -> BoxError { Box::new(error) })?,
1928 ))
1929}
1930
1931#[derive(Debug, Clone)]
1933pub struct ZakuraProtocolHandler {
1934 supervisor: ZakuraSupervisorHandle,
1935 handshake_config: ZakuraHandshakeConfig,
1936 supported_capabilities: Arc<AtomicU64>,
1937 limits: ZakuraLocalLimits,
1938 registry: Arc<ServiceRegistry>,
1939 trace: ZakuraTrace,
1940 next_conn_id: Arc<AtomicU64>,
1941 next_stream_id: Arc<AtomicU64>,
1942 admission: Arc<Semaphore>,
1943 pending_handshakes: Arc<Semaphore>,
1944 shutdown: CancellationToken,
1945 endpoint: Option<Endpoint>,
1953}
1954
1955fn random_stream_session_seed() -> u64 {
1956 let mut rng = OsRng;
1957 loop {
1958 let seed = rng.next_u64();
1959 if seed != 0 {
1960 return seed;
1961 }
1962 }
1963}
1964
1965fn confirmed_remote_ip(endpoint: &Endpoint, node_id: NodeId) -> Option<IpAddr> {
1981 if let Some(mut conn_type) = endpoint.conn_type(node_id) {
1982 match conn_type.get() {
1983 ConnectionType::Direct(addr) | ConnectionType::Mixed(addr, _) => {
1984 return Some(addr.ip())
1985 }
1986 ConnectionType::Relay(_) | ConnectionType::None => {}
1987 }
1988 }
1989 endpoint.remote_info(node_id).and_then(|info| {
1990 info.addrs
1991 .into_iter()
1992 .filter(|addr| addr.last_payload.is_some())
1993 .min_by_key(|addr| addr.last_payload)
1996 .map(|addr| addr.addr.ip())
1997 })
1998}
1999
2000fn native_connection_transcript_hash(
2001 direction: ServicePeerDirection,
2002 local_node_id: &NodeId,
2003 remote_node_id: &NodeId,
2004) -> [u8; TRANSCRIPT_HASH_BYTES] {
2005 let initiator = match direction {
2006 ServicePeerDirection::Inbound => remote_node_id,
2007 ServicePeerDirection::Outbound => local_node_id,
2008 };
2009 *initiator.as_bytes()
2010}
2011
2012fn i_open_collision_winner(local_node_id: &NodeId, remote_node_id: &NodeId) -> bool {
2017 local_node_id.as_bytes() < remote_node_id.as_bytes()
2018}
2019
2020impl ZakuraProtocolHandler {
2021 pub fn new(
2023 supervisor: ZakuraSupervisorHandle,
2024 network: Network,
2025 handshake_config: ZakuraHandshakeConfig,
2026 limits: ZakuraLocalLimits,
2027 ) -> Self {
2028 Self::new_with_registry(
2029 supervisor,
2030 network,
2031 handshake_config,
2032 limits,
2033 Arc::new(ServiceRegistry::default()),
2034 )
2035 }
2036
2037 pub fn new_with_registry(
2039 supervisor: ZakuraSupervisorHandle,
2040 network: Network,
2041 handshake_config: ZakuraHandshakeConfig,
2042 limits: ZakuraLocalLimits,
2043 registry: Arc<ServiceRegistry>,
2044 ) -> Self {
2045 Self::new_with_registry_and_trace(
2046 supervisor,
2047 network,
2048 handshake_config,
2049 limits,
2050 registry,
2051 ZakuraTrace::noop(),
2052 )
2053 }
2054
2055 pub fn new_with_registry_and_trace(
2057 supervisor: ZakuraSupervisorHandle,
2058 _network: Network,
2059 handshake_config: ZakuraHandshakeConfig,
2060 limits: ZakuraLocalLimits,
2061 registry: Arc<ServiceRegistry>,
2062 trace: ZakuraTrace,
2063 ) -> Self {
2064 let mut handshake_config = handshake_config;
2065 handshake_config.supported_capabilities = registry.supported_capabilities();
2066 let supported_capabilities =
2067 Arc::new(AtomicU64::new(handshake_config.supported_capabilities));
2068 Self {
2069 supervisor,
2070 handshake_config,
2071 supported_capabilities,
2072 registry,
2073 trace,
2074 next_conn_id: Arc::new(AtomicU64::new(1)),
2075 next_stream_id: Arc::new(AtomicU64::new(random_stream_session_seed())),
2079 admission: Arc::new(Semaphore::new(limits.max_connections)),
2080 pending_handshakes: Arc::new(Semaphore::new(limits.max_pending_handshakes)),
2081 shutdown: CancellationToken::new(),
2082 limits,
2083 endpoint: None,
2084 }
2085 }
2086
2087 #[cfg(any(test, feature = "zakura-testkit"))]
2092 pub(crate) fn with_supported_capabilities(mut self, supported_capabilities: u64) -> Self {
2093 self.handshake_config.supported_capabilities &= supported_capabilities;
2094 self.supported_capabilities.store(
2095 self.handshake_config.supported_capabilities,
2096 Ordering::Relaxed,
2097 );
2098 self
2099 }
2100
2101 fn current_handshake_config(&self) -> ZakuraHandshakeConfig {
2102 let mut config = self.handshake_config;
2103 config.supported_capabilities = self.supported_capabilities.load(Ordering::Relaxed);
2104 config
2105 }
2106
2107 fn set_header_sync_enabled(&self, enabled: bool) {
2108 if enabled {
2109 self.supported_capabilities
2110 .fetch_or(ZAKURA_CAP_HEADER_SYNC, Ordering::Relaxed);
2111 } else {
2112 self.supported_capabilities
2113 .fetch_and(!ZAKURA_CAP_HEADER_SYNC, Ordering::Relaxed);
2114 }
2115 }
2116
2117 pub fn with_endpoint(mut self, endpoint: Endpoint) -> Self {
2120 self.endpoint = Some(endpoint);
2121 self
2122 }
2123
2124 async fn accept_connection(&self, connection: Connection) -> Result<(), AcceptError> {
2125 let conn_id = self.next_conn_id.fetch_add(1, Ordering::Relaxed);
2126 let Ok(_admission) = self.admission.clone().try_acquire_owned() else {
2127 metrics::counter!("zakura.p2p.conn.rejected.admission").increment(1);
2128 let conn = ZakuraConnTrace::without_peer_on(&self.trace, conn_id);
2129 conn.trace_connection(
2130 "rejected.admission",
2131 None,
2132 Some("inbound"),
2133 Some("admission"),
2134 );
2135 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"admission");
2136 return Ok(());
2137 };
2138
2139 let remote_node_id = connection.remote_node_id()?;
2140 let remote_peer_id =
2141 ZakuraPeerId::new(remote_node_id.as_bytes().to_vec()).map_err(AcceptError::from_err)?;
2142 let conn = ZakuraConnTrace::new(&self.trace, conn_id, &remote_peer_id);
2143
2144 let negotiated = match self
2145 .run_native_responder_handshake_with_permit(&connection, &remote_peer_id, &conn)
2146 .await
2147 {
2148 Ok(negotiated) => negotiated,
2149 Err(ZakuraHandlerError::ResourceLimit("pending handshake")) => return Ok(()),
2150 Err(error) => {
2151 debug!(?error, "Zakura control handshake failed");
2152 connection.close(VarInt::from_u32(ZAKURA_CLOSE_NEUTRAL), b"control handshake");
2153 return Ok(());
2154 }
2155 };
2156
2157 let conn_limits = self.limits.clamp(&negotiated.limits);
2158 let remote_ip = self
2167 .endpoint
2168 .as_ref()
2169 .and_then(|endpoint| confirmed_remote_ip(endpoint, remote_node_id));
2170 let local_node_id = self
2171 .endpoint
2172 .as_ref()
2173 .map(|endpoint| endpoint.node_id())
2174 .unwrap_or(remote_node_id);
2175 let direction = ServicePeerDirection::Inbound;
2176 self.register_and_serve(
2177 connection,
2178 remote_peer_id,
2179 remote_ip,
2180 ConnectionServeContext {
2181 limits: conn_limits,
2182 accepted_capabilities: negotiated.accepted_capabilities,
2183 role: "responder",
2184 direction,
2185 transcript_hash: native_connection_transcript_hash(
2186 direction,
2187 &local_node_id,
2188 &remote_node_id,
2189 ),
2190 i_open_collision_winner: i_open_collision_winner(&local_node_id, &remote_node_id),
2191 conn,
2192 },
2193 )
2194 .await
2195 .map_err(AcceptError::from_err)
2196 }
2197
2198 async fn run_native_responder_handshake_with_permit(
2199 &self,
2200 connection: &Connection,
2201 remote_peer_id: &ZakuraPeerId,
2202 conn: &ZakuraConnTrace,
2203 ) -> Result<NativeHandshakeNegotiated, ZakuraHandlerError> {
2204 let Ok(_handshake) = self.pending_handshakes.clone().try_acquire_owned() else {
2205 metrics::counter!("zakura.p2p.conn.rejected.pending_handshake").increment(1);
2206 conn.trace_connection(
2207 "rejected.admission",
2208 None,
2209 Some("inbound"),
2210 Some("pending_handshake"),
2211 );
2212 connection.close(
2213 VarInt::from_u32(ZAKURA_CLOSE_RESOURCE),
2214 b"pending handshake",
2215 );
2216 return Err(ZakuraHandlerError::ResourceLimit("pending handshake"));
2217 };
2218
2219 self.run_native_responder_handshake(connection, remote_peer_id, conn)
2220 .await
2221 }
2222
2223 async fn run_native_responder_handshake(
2224 &self,
2225 connection: &Connection,
2226 remote_peer_id: &ZakuraPeerId,
2227 conn: &ZakuraConnTrace,
2228 ) -> Result<NativeHandshakeNegotiated, ZakuraHandlerError> {
2229 let handshake_config = self.current_handshake_config();
2230 conn.trace_handshake(
2231 "control.started",
2232 "responder",
2233 None,
2234 handshake_config.network_label(),
2235 );
2236 let (mut send, mut recv) = timeout(self.limits.control_timeout, connection.accept_bi())
2237 .await
2238 .map_err(|_| ZakuraHandlerError::Timeout("accept control stream"))??;
2239
2240 let hello_bytes = read_control_payload(
2241 &mut recv,
2242 handshake_config.max_control_frame_bytes,
2243 self.limits.control_timeout,
2244 )
2245 .await?;
2246 let hello = ZakuraControlHello::decode(&hello_bytes)?;
2247 let expected = ZakuraControlValidation {
2248 local: &handshake_config,
2249 authenticated_remote_id: remote_peer_id.as_bytes(),
2250 selected_zakura_protocol: ZAKURA_PROTOCOL_VERSION_1,
2251 handshake_path: ZakuraHandshakePath::Native,
2252 remote_role: ZakuraControlRole::Initiator,
2253 initiator_upgrade_nonce: [0; 32],
2254 responder_upgrade_nonce: [0; 32],
2255 legacy_upgrade_transcript: [0; 32],
2256 };
2257 hello.validate(&expected)?;
2258
2259 let mut local_nonce = [0; 32];
2260 OsRng.fill_bytes(&mut local_nonce);
2261
2262 let accepted_limits = self.accepted_limits_for(&hello.initial_limits);
2263 let ack = ZakuraControlAck {
2264 magic: CONTROL_ACK_MAGIC,
2265 control_version: CONTROL_VERSION,
2266 selected_zakura_protocol: hello.selected_zakura_protocol,
2267 peer_nonce: local_nonce,
2268 remote_peer_nonce: hello.peer_nonce,
2269 accepted_capabilities: hello.capabilities & handshake_config.supported_capabilities,
2270 accepted_channels: hello.required_channels & handshake_config.supported_channels,
2271 accepted_limits,
2272 };
2273 write_control_payload(&mut send, &ack.encode()?, self.limits.control_timeout).await?;
2274 conn.trace_handshake(
2275 "control.succeeded",
2276 "responder",
2277 Some(ack.selected_zakura_protocol),
2278 handshake_config.network_label(),
2279 );
2280 Ok(NativeHandshakeNegotiated {
2281 limits: accepted_limits,
2282 accepted_capabilities: ack.accepted_capabilities,
2283 })
2284 }
2285
2286 fn accepted_limits_for(&self, remote_limits: &ZakuraInitialLimits) -> ZakuraAcceptedLimits {
2287 ZakuraLimits {
2288 max_frame_bytes: remote_limits
2289 .max_frame_bytes
2290 .min(self.limits.max_frame_bytes),
2291 max_message_bytes: remote_limits
2292 .max_message_bytes
2293 .min(self.limits.max_message_bytes),
2294 max_open_streams: remote_limits
2295 .max_open_streams
2296 .min(self.limits.max_open_streams),
2297 max_inbound_queue_depth: remote_limits
2298 .max_inbound_queue_depth
2299 .min(self.limits.max_inbound_queue_depth),
2300 idle_timeout_millis: remote_limits
2301 .idle_timeout_millis
2302 .min(self.limits.initial_limits().idle_timeout_millis),
2303 }
2304 }
2305
2306 async fn serve_connection(
2307 &self,
2308 connection: Connection,
2309 peer_id: ZakuraPeerId,
2310 remote_ip: Option<IpAddr>,
2311 mut outbound_rx: mpsc::Receiver<ZakuraOutboundFrame>,
2312 context: RegisteredConnectionServeContext,
2313 mut cleanup_guard: RegisteredPeerCleanupGuard,
2314 ) -> Result<(), ZakuraHandlerError> {
2315 let limits = context.limits;
2316 let conn = context.conn;
2317 let conn_id = context.conn_id;
2318 let connection_token = context.connection_token;
2319 let close_cause = context.close_cause;
2320 let accepted_capabilities = context.accepted_capabilities;
2321 let stream_sem = Arc::new(Semaphore::new(usize::from(limits.max_open_streams)));
2322 let mut workers = JoinSet::new();
2323 let (ordered_session_exit_tx, mut ordered_session_exit_rx) = mpsc::unbounded_channel();
2324 let mut ordered_session_waits = OrderedSessionWaits::new();
2325 let mut open_limiter = TokenBucket::new(limits.stream_open_rate_per_second);
2326 let mut message_buckets = MessageRateBuckets::new();
2327 let (freshness_tx, freshness_rx) = watch::channel(Instant::now());
2328 let negotiated_ordered_streams = self
2329 .registry
2330 .ordered_streams_for_negotiated(accepted_capabilities);
2331 let mut ordered_sessions: HashMap<u16, OrderedSessionState> = negotiated_ordered_streams
2332 .iter()
2333 .copied()
2334 .map(|stream| (stream.kind, OrderedSessionState::new(stream)))
2335 .collect();
2336 let mut ordered_streams = Vec::new();
2342 let mut deferred_ordered_streams = Vec::new();
2343 for stream in negotiated_ordered_streams.iter().copied() {
2344 let policy = self.registry.ordered_stream_policy(stream.kind);
2345 if !opens_ordered_stream_locally(
2346 policy,
2347 context.is_initiator,
2348 context.i_open_collision_winner,
2349 ) {
2350 continue;
2351 }
2352
2353 match self.registry.ordered_session_demand(
2354 stream.kind,
2355 conn_id,
2356 accepted_capabilities,
2357 &peer_id,
2358 context.direction,
2359 ) {
2360 OrderedSessionDemand::OpenNow => ordered_streams.push(stream),
2361 demand => deferred_ordered_streams.push((stream, demand)),
2362 }
2363 }
2364 let request_response_stream_count = self
2365 .registry
2366 .request_response_streams_for_negotiated(accepted_capabilities)
2367 .len();
2368 if ordered_streams.len() > usize::from(limits.max_open_streams) {
2369 debug!(
2370 max_open_streams = limits.max_open_streams,
2371 ordered_stream_count = ordered_streams.len(),
2372 "closing Zakura peer because negotiated ordered streams exceed max-open-streams"
2373 );
2374 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"ordered streams");
2375 close_cause.record("resource_ordered_streams");
2376 connection_token.cancel();
2377 } else if !negotiated_ordered_streams.is_empty()
2378 && usize::from(limits.max_inbound_queue_depth) < negotiated_ordered_streams.len()
2379 {
2380 debug!(
2381 max_inbound_queue_depth = limits.max_inbound_queue_depth,
2382 ordered_stream_count = negotiated_ordered_streams.len(),
2383 "closing Zakura peer because inbound queue depth cannot be split across ordered streams"
2384 );
2385 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"queue split");
2386 close_cause.record("resource_queue_split");
2387 connection_token.cancel();
2388 }
2389 let queue_split_stream_count = negotiated_ordered_streams.len().max(ordered_streams.len());
2390 let per_stream_queue_depth = per_stream_inbound_queue_depth(
2391 limits.max_inbound_queue_depth,
2392 queue_split_stream_count,
2393 );
2394 let mut service_streams = HashMap::new();
2395 let run_freshness_reaper =
2396 should_run_freshness_reaper(queue_split_stream_count, request_response_stream_count);
2397
2398 if negotiated_ordered_streams.is_empty() {
2399 self.registry
2400 .add_peer(Peer::new_with_conn_id_and_direction_and_close_cause(
2401 conn_id,
2402 peer_id.clone(),
2403 remote_ip,
2404 accepted_capabilities,
2405 context.direction,
2406 HashMap::new(),
2407 connection_token.clone(),
2408 close_cause.clone(),
2409 ));
2410 cleanup_guard.add_admitted_capabilities(accepted_capabilities);
2411 } else if !connection_token.is_cancelled() {
2412 let mut opened_capabilities = 0;
2413 for stream in ordered_streams {
2414 opened_capabilities |= stream.capability;
2415 let admitted = match self
2416 .open_ordered_service_stream(
2417 &connection,
2418 stream,
2419 &mut workers,
2420 &stream_sem,
2421 &mut message_buckets,
2422 limits,
2423 per_stream_queue_depth,
2424 connection_token.clone(),
2425 close_cause.clone(),
2426 freshness_tx.clone(),
2427 conn.clone(),
2428 peer_id.clone(),
2429 ordered_session_exit_tx.clone(),
2430 )
2431 .await
2432 {
2433 Ok(admitted) => admitted,
2434 Err(error) => {
2435 debug!(
2436 ?error,
2437 stream_kind = stream.kind,
2438 "closing Zakura peer after ordered stream setup failed"
2439 );
2440 connection.close(
2441 VarInt::from_u32(ZAKURA_CLOSE_RESOURCE),
2442 b"ordered stream setup",
2443 );
2444 close_cause.record("stream_setup_failed");
2445 connection_token.cancel();
2446 break;
2447 }
2448 };
2449 ordered_sessions
2450 .get_mut(&admitted.kind)
2451 .expect("opened ordered stream was selected from negotiated session state")
2452 .local_session_id = Some(admitted.session_id);
2453 service_streams.insert(
2454 admitted.kind,
2455 ServiceStream::new(
2456 admitted.session_id,
2457 admitted.version,
2458 admitted.recv,
2459 admitted.send,
2460 admitted.cancel_token,
2461 ),
2462 );
2463 }
2464 if !connection_token.is_cancelled() {
2465 let admitted_capabilities =
2469 self.registry
2470 .add_escalated_peer(Peer::new_with_service_streams(
2471 conn_id,
2472 peer_id.clone(),
2473 remote_ip,
2474 opened_capabilities,
2475 context.direction,
2476 std::mem::take(&mut service_streams),
2477 connection_token.clone(),
2478 close_cause.clone(),
2479 ));
2480 cleanup_guard.add_admitted_capabilities(admitted_capabilities);
2481 }
2482 }
2483
2484 if !connection_token.is_cancelled() {
2485 for (stream, demand) in deferred_ordered_streams {
2486 debug!(
2487 ?peer_id,
2488 stream_kind = stream.kind,
2489 ?demand,
2490 "deferring ordered service session according to reactor demand"
2491 );
2492 ordered_sessions
2493 .get_mut(&stream.kind)
2494 .expect("deferred ordered stream has negotiated session state")
2495 .schedule_demand(&mut ordered_session_waits, demand);
2496 }
2497 }
2498
2499 loop {
2500 tokio::select! {
2501 biased;
2502 _ = connection_token.cancelled() => break,
2503 _ = freshness_reaper(freshness_rx.clone(), limits.idle_timeout), if run_freshness_reaper => {
2504 connection.close(VarInt::from_u32(ZAKURA_CLOSE_NEUTRAL), b"idle");
2505 close_cause.record("idle_timeout");
2506 break;
2507 }
2508 Some(exited) = ordered_session_exit_rx.recv() => {
2509 let Some(session) = ordered_sessions.get_mut(&exited.stream.kind) else {
2513 continue;
2514 };
2515 let removed_active_session =
2516 session.remove_active_session(exited.opened_locally, exited.session_id);
2517 let policy = self.registry.ordered_stream_policy(exited.stream.kind);
2518 if removed_active_session
2519 && !session.has_active_session()
2520 && should_reopen_ordered_session(
2521 exited,
2522 policy,
2523 context.is_initiator,
2524 context.i_open_collision_winner,
2525 connection_token.is_cancelled(),
2526 )
2527 {
2528 if let Some(delay) =
2529 session.schedule_transport_backoff(&mut ordered_session_waits)
2530 {
2531 debug!(
2532 stream_kind = exited.stream.kind,
2533 stream_version = exited.stream.version,
2534 session_id = exited.session_id,
2535 ?delay,
2536 "scheduling Zakura ordered stream reopen"
2537 );
2538 }
2539 }
2540 }
2541 Some((kind, ())) = ordered_session_waits.next(), if !ordered_session_waits.is_empty() => {
2542 let Some(session) = ordered_sessions.get_mut(&kind) else {
2543 continue;
2544 };
2545 session.finish_wait(&mut ordered_session_waits);
2546 let stream = session.stream;
2547 let policy = self.registry.ordered_stream_policy(stream.kind);
2548 if connection_token.is_cancelled()
2549 || !opens_ordered_stream_locally(
2550 policy,
2551 context.is_initiator,
2552 context.i_open_collision_winner,
2553 )
2554 || session.has_active_session()
2555 {
2556 continue;
2557 }
2558
2559 let demand = self.registry.ordered_session_demand(
2560 stream.kind,
2561 conn_id,
2562 accepted_capabilities,
2563 &peer_id,
2564 context.direction,
2565 );
2566 match demand {
2567 OrderedSessionDemand::OpenNow => {}
2568 demand => {
2569 session.schedule_demand(&mut ordered_session_waits, demand);
2570 continue;
2571 }
2572 }
2573
2574 match self
2575 .open_ordered_service_stream(
2576 &connection,
2577 stream,
2578 &mut workers,
2579 &stream_sem,
2580 &mut message_buckets,
2581 limits,
2582 per_stream_queue_depth,
2583 connection_token.clone(),
2584 close_cause.clone(),
2585 freshness_tx.clone(),
2586 conn.clone(),
2587 peer_id.clone(),
2588 ordered_session_exit_tx.clone(),
2589 )
2590 .await
2591 {
2592 Ok(admitted) => {
2593 let session = ordered_sessions
2594 .get_mut(&admitted.kind)
2595 .expect("opened ordered stream has negotiated session state");
2596 session.cancel_wait(&mut ordered_session_waits);
2597 session.local_session_id = Some(admitted.session_id);
2598 let service_streams = HashMap::from([(
2599 admitted.kind,
2600 ServiceStream::new(
2601 admitted.session_id,
2602 admitted.version,
2603 admitted.recv,
2604 admitted.send,
2605 admitted.cancel_token,
2606 ),
2607 )]);
2608 let admitted_capabilities = self.registry.add_escalated_peer(
2609 Peer::new_with_service_streams(
2610 conn_id,
2611 peer_id.clone(),
2612 remote_ip,
2613 accepted_capabilities,
2614 context.direction,
2615 service_streams,
2616 connection_token.clone(),
2617 close_cause.clone(),
2618 ),
2619 );
2620 cleanup_guard.add_admitted_capabilities(admitted_capabilities);
2621 }
2622 Err(error) => {
2623 let delay = ordered_sessions
2624 .get_mut(&stream.kind)
2625 .expect("failed ordered stream open has negotiated session state")
2626 .schedule_transport_backoff(&mut ordered_session_waits);
2627 debug!(
2628 ?error,
2629 stream_kind = stream.kind,
2630 ?delay,
2631 "retrying Zakura ordered stream reopen after local failure"
2632 );
2633 }
2634 }
2635 }
2636 Some(joined) = workers.join_next() => {
2637 if let Err(error) = joined {
2638 debug!(?error, "Zakura stream worker exited unexpectedly");
2639 }
2640 }
2641 accepted = connection.accept_bi() => {
2642 match accepted {
2643 Ok((send, recv)) => {
2644 let mut admission = StreamAdmission {
2645 conn: conn.clone(),
2646 peer_id: &peer_id,
2647 stream_sem: &stream_sem,
2648 open_limiter: &mut open_limiter,
2649 message_buckets: &mut message_buckets,
2650 workers: &mut workers,
2651 limits,
2652 accepted_capabilities,
2653 connection_token: connection_token.clone(),
2654 close_cause: close_cause.clone(),
2655 freshness_tx: freshness_tx.clone(),
2656 };
2657 if let Some(admitted) = self
2658 .admit_bi_stream(
2659 send,
2660 recv,
2661 &mut admission,
2662 per_stream_queue_depth,
2663 ordered_session_exit_tx.clone(),
2664 )
2665 .await
2666 {
2667 let kind = admitted.kind;
2668
2669 if !ordered_sessions.contains_key(&kind) {
2673 debug!(
2674 stream_kind = kind,
2675 "closing peer after unexpected ordered stream"
2676 );
2677 close_cause.record("unexpected_stream");
2678 connection_token.cancel();
2679 continue;
2680 }
2681 let policy = self.registry.ordered_stream_policy(kind);
2685 if !may_open_ordered_stream(policy, !context.is_initiator) {
2686 debug!(
2687 stream_kind = kind,
2688 "closing peer after unexpected ordered stream"
2689 );
2690 close_cause.record("unexpected_stream");
2691 connection_token.cancel();
2692 continue;
2693 }
2694
2695 let is_collision =
2696 ordered_sessions[&kind].local_session_id.is_some();
2697 if is_collision && context.i_open_collision_winner {
2701 debug!(
2702 stream_kind = kind,
2703 "winning ordered-stream collision; parking peer's duplicate"
2704 );
2705 admitted.cancel_token.cancel();
2706 continue;
2707 }
2708
2709 if ordered_sessions[&kind].remote_session_id.is_some() {
2713 debug!(
2714 stream_kind = kind,
2715 "closing peer after duplicate ordered stream"
2716 );
2717 close_cause.record("duplicate_stream");
2718 connection_token.cancel();
2719 continue;
2720 }
2721 ordered_sessions
2722 .get_mut(&kind)
2723 .expect("accepted ordered stream has negotiated session state")
2724 .remote_session_id = Some(admitted.session_id);
2725
2726 let demand = (!is_collision).then(|| {
2730 self.registry.ordered_session_demand(
2731 kind,
2732 conn_id,
2733 accepted_capabilities,
2734 &peer_id,
2735 context.direction,
2736 )
2737 });
2738 if demand
2739 .as_ref()
2740 .is_some_and(|demand| {
2741 !matches!(demand, OrderedSessionDemand::OpenNow)
2742 })
2743 {
2744 metrics::counter!(
2745 "zakura.p2p.stream.parked.no_demand",
2746 "stream_kind" => stream_kind_label(kind),
2747 )
2748 .increment(1);
2749 info!(
2750 ?peer_id,
2751 stream_kind = kind,
2752 "locally parking ordered service stream because the service has no demand"
2753 );
2754 let session = ordered_sessions
2758 .get_mut(&kind)
2759 .expect("parked ordered stream has negotiated session state");
2760 session.remove_active_session(false, admitted.session_id);
2761 if policy.reopen
2765 && opens_ordered_stream_locally(
2766 policy,
2767 context.is_initiator,
2768 context.i_open_collision_winner,
2769 )
2770 {
2771 session.schedule_demand(
2772 &mut ordered_session_waits,
2773 demand.expect("non-open demand exists because this branch checked it"),
2774 );
2775 }
2776 admitted.cancel_token.cancel();
2777 continue;
2778 }
2779
2780 ordered_sessions
2781 .get_mut(&kind)
2782 .expect("adopted ordered stream has negotiated session state")
2783 .cancel_wait(&mut ordered_session_waits);
2784 let service_streams = HashMap::from([(
2785 kind,
2786 ServiceStream::new(
2787 admitted.session_id,
2788 admitted.version,
2789 admitted.recv,
2790 admitted.send,
2791 admitted.cancel_token.clone(),
2792 ),
2793 )]);
2794 let admitted_capabilities = self.registry.add_escalated_peer(
2807 Peer::new_with_service_streams(
2808 conn_id,
2809 peer_id.clone(),
2810 remote_ip,
2811 accepted_capabilities,
2812 context.direction,
2813 service_streams,
2814 connection_token.clone(),
2815 close_cause.clone(),
2816 ),
2817 );
2818 cleanup_guard.add_admitted_capabilities(admitted_capabilities);
2819 }
2820 }
2821 Err(error) => {
2822 debug!(?error, "Zakura connection stopped accepting streams");
2823 close_cause.record("accept_failed");
2824 break;
2825 }
2826 }
2827 }
2828 outbound = outbound_rx.recv() => {
2829 let Some(outbound) = outbound else {
2830 close_cause.record("outbound_closed");
2831 break;
2832 };
2833 match outbound {
2834 ZakuraOutboundFrame::Request {
2835 stream_kind,
2836 request_id,
2837 message_type,
2838 flags,
2839 payload,
2840 mut completion,
2841 } => {
2842 let Some(stream) = self.registry.stream_for_kind(stream_kind) else {
2843 let _ = completion.send(Err(format!(
2844 "Zakura outbound request stream kind {stream_kind} is not registered"
2845 )
2846 .into()));
2847 continue;
2848 };
2849 let result = tokio::select! {
2850 biased;
2851 _ = completion.closed() => {
2852 Err(OutboundRequestError::Local(
2853 "Zakura outbound request receiver dropped".into(),
2854 ))
2855 }
2856 result = write_outbound_request_frame(
2857 &connection,
2858 limits,
2859 stream,
2860 request_id,
2861 message_type,
2862 flags,
2863 payload,
2864 ) => result,
2865 };
2866 match result {
2867 Ok(frames) => {
2868 let _ = completion.send(Ok(frames));
2869 }
2870 Err(OutboundRequestError::Local(error)) => {
2871 let _ = completion.send(Err(error));
2872 }
2873 Err(OutboundRequestError::Fatal(error)) => {
2874 connection.close(
2875 VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE),
2876 b"malformed response",
2877 );
2878 close_cause.record("bad_response");
2879 connection_token.cancel();
2880 let _ = completion.send(Err(error));
2881 }
2882 }
2883 }
2884 }
2885 }
2886 }
2887 }
2888
2889 connection_token.cancel();
2890 while let Some(joined) = timeout(STREAM_WORKER_DRAIN_TIMEOUT, workers.join_next())
2891 .await
2892 .ok()
2893 .flatten()
2894 {
2895 if let Err(error) = joined {
2896 debug!(?error, "Zakura stream worker failed during shutdown");
2897 }
2898 }
2899 workers.abort_all();
2900 cleanup_guard.cleanup_registered_peer().await;
2901 metrics::counter!("zakura.p2p.conn.closed.neutral").increment(1);
2902 conn.trace_connection(
2903 "closed.neutral",
2904 None,
2905 None,
2906 Some(close_cause.get_or("cancelled")),
2907 );
2908 Ok(())
2909 }
2910
2911 #[allow(clippy::too_many_arguments)]
2912 async fn open_ordered_service_stream(
2913 &self,
2914 connection: &Connection,
2915 stream: Stream,
2916 workers: &mut JoinSet<()>,
2917 stream_sem: &Arc<Semaphore>,
2918 message_buckets: &mut MessageRateBuckets,
2919 limits: ZakuraConnectionLimits,
2920 per_stream_queue_depth: usize,
2921 connection_token: CancellationToken,
2922 close_cause: CloseCause,
2923 freshness_tx: watch::Sender<Instant>,
2924 conn: ZakuraConnTrace,
2925 peer_id: ZakuraPeerId,
2926 ordered_session_exit_tx: mpsc::UnboundedSender<OrderedSessionExit>,
2927 ) -> Result<AdmittedOrderedSession, ZakuraHandlerError> {
2928 let stream_id = self.next_stream_id.fetch_add(1, Ordering::Relaxed);
2929 let permit = stream_sem
2930 .clone()
2931 .try_acquire_owned()
2932 .map_err(|_| ZakuraHandlerError::ResourceLimit("ordered stream permit"))?;
2933 let (mut send, recv) = timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, connection.open_bi())
2934 .await
2935 .map_err(|_| ZakuraHandlerError::Timeout("open ordered service stream"))??;
2936 let prelude = StreamPrelude {
2937 magic: STREAM_PRELUDE_MAGIC,
2938 stream_kind: stream.kind,
2939 stream_version: stream.version,
2940 request_id: None,
2941 max_frame_bytes: inbound_frame_cap_for_stream(&limits, stream),
2942 };
2943 let prelude_bytes = prelude.encode()?;
2944 timeout(
2945 OUTBOUND_STREAM_WRITE_TIMEOUT,
2946 send.write_all(&prelude_bytes),
2947 )
2948 .await
2949 .map_err(|_| ZakuraHandlerError::Timeout("ordered stream prelude write"))??;
2950
2951 let message_bucket = message_bucket_for(
2952 message_buckets,
2953 stream.kind,
2954 limits.message_rate_per_second,
2955 RealClock,
2956 );
2957 let stream_token = connection_token.child_token();
2958 let context = StreamWorkerContext {
2959 conn: conn.clone(),
2960 peer_id,
2961 stream_id,
2962 _permit: permit,
2963 limits,
2964 inbound_frame_cap: prelude.max_frame_bytes,
2965 outbound_frame_cap: application_frame_cap(&limits, stream),
2966 message_bucket,
2967 connection_token,
2968 stream_token,
2969 close_cause,
2970 freshness_tx,
2971 };
2972
2973 metrics::counter!(
2974 "zakura.p2p.stream.accepted",
2975 "stream_kind" => stream_kind_label(stream.kind),
2976 )
2977 .increment(1);
2978 conn.trace_stream("accepted", stream_id, Some(stream_kind_label(stream.kind)));
2979
2980 Ok(spawn_persistent_stream_worker(
2981 workers,
2982 send,
2983 recv,
2984 stream,
2985 prelude,
2986 context,
2987 per_stream_queue_depth,
2988 true,
2989 ordered_session_exit_tx,
2990 ))
2991 }
2992
2993 async fn admit_bi_stream(
2994 &self,
2995 mut send: SendStream,
2996 mut recv: RecvStream,
2997 admission: &mut StreamAdmission<'_>,
2998 per_stream_queue_depth: usize,
2999 ordered_session_exit_tx: mpsc::UnboundedSender<OrderedSessionExit>,
3000 ) -> Option<AdmittedOrderedSession> {
3001 let stream_id = self.next_stream_id.fetch_add(1, Ordering::Relaxed);
3002 let Ok(permit) = admission.stream_sem.clone().try_acquire_owned() else {
3003 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE));
3004 metrics::counter!("zakura.p2p.stream.rejected.semaphore").increment(1);
3005 admission
3006 .conn
3007 .trace_stream("rejected.semaphore", stream_id, None);
3008 return None;
3009 };
3010
3011 if !admission.open_limiter.try_take() {
3019 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RATE_LIMIT));
3020 metrics::counter!("zakura.p2p.stream.rejected.open_rate").increment(1);
3021 admission
3022 .conn
3023 .trace_stream("rejected.open_rate", stream_id, None);
3024 return None;
3025 }
3026
3027 let prelude = match read_stream_prelude(&mut recv, admission.limits.prelude_timeout).await {
3028 Ok(prelude) => prelude,
3029 Err(error) => {
3030 debug!(?error, "rejecting Zakura stream with bad prelude");
3031 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
3032 metrics::counter!("zakura.p2p.stream.rejected.prelude").increment(1);
3033 admission
3034 .conn
3035 .trace_stream("rejected.prelude", stream_id, None);
3036 return None;
3037 }
3038 };
3039 let stream_kind = stream_kind_label(prelude.stream_kind);
3040
3041 let Some(stream) = self
3042 .registry
3043 .stream(prelude.stream_kind, prelude.stream_version)
3044 else {
3045 debug!(
3046 stream_kind = prelude.stream_kind,
3047 stream_version = prelude.stream_version,
3048 "rejecting Zakura stream with unknown kind or unsupported version"
3049 );
3050 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_UNKNOWN_STREAM));
3051 metrics::counter!(
3052 "zakura.p2p.stream.rejected.unknown_kind",
3053 "stream_kind" => stream_kind,
3054 )
3055 .increment(1);
3056 admission
3057 .conn
3058 .trace_stream("rejected.unknown_kind", stream_id, Some(stream_kind));
3059 return None;
3060 };
3061
3062 if admission.accepted_capabilities & stream.capability != stream.capability {
3063 debug!(
3064 stream_kind = prelude.stream_kind,
3065 stream_version = prelude.stream_version,
3066 accepted_capabilities = admission.accepted_capabilities,
3067 "rejecting Zakura stream that was not negotiated for this peer"
3068 );
3069 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_UNKNOWN_STREAM));
3070 metrics::counter!(
3071 "zakura.p2p.stream.rejected.unnegotiated_capability",
3072 "stream_kind" => stream_kind,
3073 )
3074 .increment(1);
3075 admission.conn.trace_stream(
3076 "rejected.unnegotiated_capability",
3077 stream_id,
3078 Some(stream_kind),
3079 );
3080 return None;
3081 }
3082
3083 if stream.mode == StreamMode::Ordered
3084 && self
3085 .registry
3086 .ordered_streams_for_negotiated(admission.accepted_capabilities)
3087 .into_iter()
3088 .find(|selected| selected.kind == stream.kind)
3089 .is_some_and(|selected| selected.version != stream.version)
3090 {
3091 debug!(
3092 stream_kind = prelude.stream_kind,
3093 stream_version = prelude.stream_version,
3094 "rejecting a lower ordered-stream version than the mutually selected version"
3095 );
3096 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_UNKNOWN_STREAM));
3097 metrics::counter!(
3098 "zakura.p2p.stream.rejected.unselected_version",
3099 "stream_kind" => stream_kind,
3100 )
3101 .increment(1);
3102 admission.conn.trace_stream(
3103 "rejected.unselected_version",
3104 stream_id,
3105 Some(stream_kind),
3106 );
3107 return None;
3108 }
3109
3110 if stream.mode != StreamMode::RequestResponse && prelude.request_id.is_some() {
3111 debug!("rejecting non-request Zakura stream with request id");
3112 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
3113 admission.close_cause.record("unexpected_request_id");
3114 admission.connection_token.cancel();
3115 metrics::counter!("zakura.p2p.stream.rejected.unexpected_request_id").increment(1);
3116 admission.conn.trace_stream(
3117 "rejected.unexpected_request_id",
3118 stream_id,
3119 Some(stream_kind),
3120 );
3121 return None;
3122 }
3123
3124 if stream.mode == StreamMode::RequestResponse && prelude.request_id.is_none() {
3125 debug!("rejecting Zakura request stream without request id");
3126 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
3127 admission.close_cause.record("request_without_id");
3128 admission.connection_token.cancel();
3129 metrics::counter!("zakura.p2p.stream.rejected.request_without_id").increment(1);
3130 admission.conn.trace_stream(
3131 "rejected.request_without_id",
3132 stream_id,
3133 Some(stream_kind),
3134 );
3135 return None;
3136 }
3137
3138 metrics::counter!(
3139 "zakura.p2p.stream.accepted",
3140 "stream_kind" => stream_kind,
3141 )
3142 .increment(1);
3143 admission
3144 .conn
3145 .trace_stream("accepted", stream_id, Some(stream_kind));
3146
3147 let message_bucket = message_bucket_for(
3148 admission.message_buckets,
3149 prelude.stream_kind,
3150 admission.limits.message_rate_per_second,
3151 RealClock,
3152 );
3153 let stream_token = admission.connection_token.child_token();
3154
3155 let context = StreamWorkerContext {
3156 conn: admission.conn.clone(),
3157 peer_id: admission.peer_id.clone(),
3158 stream_id,
3159 _permit: permit,
3160 limits: admission.limits,
3161 inbound_frame_cap: inbound_frame_cap_for_stream(&admission.limits, stream),
3162 outbound_frame_cap: peer_accepted_frame_cap(
3163 &admission.limits,
3164 stream,
3165 prelude.max_frame_bytes,
3166 ),
3167 message_bucket,
3168 connection_token: admission.connection_token.clone(),
3169 stream_token,
3170 close_cause: admission.close_cause.clone(),
3171 freshness_tx: admission.freshness_tx.clone(),
3172 };
3173
3174 if stream.mode == StreamMode::RequestResponse {
3175 admission.workers.spawn(request_stream_worker(
3176 send,
3177 recv,
3178 prelude,
3179 context,
3180 self.registry.clone(),
3181 ));
3182 None
3183 } else {
3184 Some(spawn_persistent_stream_worker(
3185 admission.workers,
3186 send,
3187 recv,
3188 stream,
3189 prelude,
3190 context,
3191 per_stream_queue_depth,
3192 false,
3193 ordered_session_exit_tx,
3194 ))
3195 }
3196 }
3197
3198 async fn register_and_serve(
3199 &self,
3200 connection: Connection,
3201 peer_id: ZakuraPeerId,
3202 remote_ip: Option<IpAddr>,
3203 context: ConnectionServeContext,
3204 ) -> Result<(), ZakuraHandlerError> {
3205 let ordered_stream_count = self
3208 .registry
3209 .ordered_streams_for_escalation(
3210 context.accepted_capabilities,
3211 &peer_id,
3212 context.direction,
3213 )
3214 .len();
3215 if ordered_stream_count > usize::from(context.limits.max_open_streams) {
3216 debug!(
3217 max_open_streams = context.limits.max_open_streams,
3218 ordered_stream_count,
3219 "rejecting Zakura peer before registration because negotiated ordered streams exceed max-open-streams"
3220 );
3221 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"ordered streams");
3222 return Ok(());
3223 }
3224 if ordered_stream_count > 0
3225 && usize::from(context.limits.max_inbound_queue_depth) < ordered_stream_count
3226 {
3227 debug!(
3228 max_inbound_queue_depth = context.limits.max_inbound_queue_depth,
3229 ordered_stream_count,
3230 "rejecting Zakura peer before registration because inbound queue depth cannot be split"
3231 );
3232 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"queue split");
3233 return Ok(());
3234 }
3235
3236 let (outbound_tx, outbound_rx) =
3237 mpsc::channel(usize::from(context.limits.max_inbound_queue_depth));
3238 let close_cause = CloseCause::new();
3239 let outbound_handle = ZakuraPeerHandle {
3240 peer_id: peer_id.clone(),
3241 sender: outbound_tx,
3242 };
3243 let connection_token = self.shutdown.child_token();
3244 let registration = self
3245 .supervisor
3246 .register(
3247 context.conn.id,
3248 peer_id,
3249 remote_ip,
3250 context.transcript_hash,
3251 outbound_handle,
3252 connection_token.clone(),
3253 context.accepted_capabilities,
3254 )
3255 .await;
3256
3257 match registration {
3258 ZakuraRegistration::Registered {
3259 conn_id,
3260 peer_id,
3261 remote_ip,
3262 disconnect_token,
3263 } => {
3264 metrics::counter!("zakura.p2p.conn.accepted", "role" => context.role).increment(1);
3265 context.conn.trace_connection(
3266 "accepted",
3267 Some(context.role),
3268 Some(context.direction.trace_label()),
3269 None,
3270 );
3271 let cleanup_guard = RegisteredPeerCleanupGuard::new(
3272 self.supervisor.clone(),
3273 self.registry.clone(),
3274 peer_id.clone(),
3275 conn_id,
3276 disconnect_token.clone(),
3277 );
3278 self.serve_connection(
3279 connection,
3280 peer_id,
3281 remote_ip,
3282 outbound_rx,
3283 RegisteredConnectionServeContext {
3284 limits: context.limits,
3285 conn: context.conn,
3286 conn_id,
3287 connection_token: disconnect_token,
3288 close_cause,
3289 accepted_capabilities: context.accepted_capabilities,
3290 is_initiator: context.role == "initiator",
3291 i_open_collision_winner: context.i_open_collision_winner,
3292 direction: context.direction,
3293 },
3294 cleanup_guard,
3295 )
3296 .await
3297 }
3298 ZakuraRegistration::Duplicate { peer_id } => {
3299 debug!(?peer_id, "closing duplicate Zakura peer neutrally");
3300 metrics::counter!("zakura.p2p.conn.duplicate").increment(1);
3301 context.conn.trace_connection(
3302 "duplicate",
3303 Some(context.role),
3304 Some(context.direction.trace_label()),
3305 None,
3306 );
3307 connection.close(VarInt::from_u32(ZAKURA_CLOSE_NEUTRAL), b"duplicate");
3308 Ok(())
3309 }
3310 ZakuraRegistration::Rejected(reason) => {
3311 debug!(
3312 ?reason,
3313 "closing Zakura peer neutrally after registration rejection"
3314 );
3315 context.conn.trace_connection(
3316 "rejected.admission",
3317 Some(context.role),
3318 Some(context.direction.trace_label()),
3319 Some(reject_reason_label(reason)),
3320 );
3321 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"registration");
3322 Ok(())
3323 }
3324 }
3325 }
3326}
3327
3328impl ProtocolHandler for ZakuraProtocolHandler {
3329 async fn accept(&self, connection: Connection) -> Result<(), AcceptError> {
3330 self.accept_connection(connection).await
3331 }
3332
3333 async fn shutdown(&self) {
3334 self.shutdown.cancel();
3335 }
3336}
3337
3338const ZAKURA_LOOPBACK_BIND_V4: SocketAddrV4 = SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0);
3342const ZAKURA_LOOPBACK_BIND_V6: SocketAddrV6 = SocketAddrV6::new(Ipv6Addr::LOCALHOST, 0, 0, 0);
3344
3345fn bind_native_endpoint(
3355 builder: iroh::endpoint::Builder,
3356 listen_addr: Option<SocketAddr>,
3357) -> iroh::endpoint::Builder {
3358 match listen_addr {
3359 Some(SocketAddr::V4(addr)) => builder.bind_addr_v4(addr),
3360 Some(SocketAddr::V6(addr)) => builder.bind_addr_v6(addr),
3361 None => builder
3362 .bind_addr_v4(ZAKURA_LOOPBACK_BIND_V4)
3363 .bind_addr_v6(ZAKURA_LOOPBACK_BIND_V6),
3364 }
3365}
3366
3367fn discovery_direct_addrs(config: &Config, local_node_id: NodeId) -> Vec<SocketAddr> {
3368 let Some(listen_addr) = config.zakura.listen_addr else {
3369 return Vec::new();
3370 };
3371 let mut direct_addrs = Vec::new();
3372
3373 if !listen_addr.ip().is_unspecified() {
3374 direct_addrs.push(listen_addr);
3375 }
3376 if let Some(external_addr) = config.external_addr {
3377 direct_addrs.push(SocketAddr::new(external_addr.ip(), listen_addr.port()));
3378 }
3379 for entry in &config.zakura.bootstrap_peers {
3380 let Ok(node_addr) = super::discovery::parse_bootstrap_peer(entry) else {
3381 continue;
3382 };
3383 if node_addr.node_id == local_node_id {
3384 direct_addrs.extend(node_addr.direct_addresses().copied());
3385 }
3386 }
3387
3388 direct_addrs.sort_unstable();
3389 direct_addrs.dedup();
3390 direct_addrs
3391}
3392
3393fn remote_bootstrap_peer_count(bootstrap_peers: &[String], local_node_id: NodeId) -> usize {
3394 bootstrap_peers
3395 .iter()
3396 .filter_map(|entry| super::discovery::parse_bootstrap_peer(entry).ok())
3397 .filter(|node_addr| node_addr.node_id != local_node_id)
3398 .map(|node_addr| node_addr.node_id)
3399 .collect::<HashSet<_>>()
3400 .len()
3401}
3402
3403#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
3404struct HeaderCapabilityEpochs {
3405 applied_ready: Option<zakura_node_services::sync_lifecycle::LifecycleEpoch>,
3406}
3407
3408impl HeaderCapabilityEpochs {
3409 fn from_initial(demand: &zakura_node_services::sync_lifecycle::SyncServiceDemand) -> Self {
3410 let applied_ready = demand.header.capability_epoch();
3411 Self { applied_ready }
3412 }
3413
3414 fn observe(
3415 &mut self,
3416 demand: &zakura_node_services::sync_lifecycle::SyncServiceDemand,
3417 ) -> Option<zakura_node_services::sync_lifecycle::LifecycleEpoch> {
3418 let epoch = demand.header.capability_epoch()?;
3419 if self.applied_ready.is_some_and(|applied| epoch <= applied) {
3420 return None;
3421 }
3422 self.applied_ready = Some(epoch);
3423 Some(epoch)
3424 }
3425}
3426
3427async fn enable_header_sync_and_renegotiate(
3433 handler: &ZakuraProtocolHandler,
3434 supervisor: &ZakuraSupervisorHandle,
3435 capability_epoch: zakura_node_services::sync_lifecycle::LifecycleEpoch,
3436) -> usize {
3437 handler.set_header_sync_enabled(true);
3438
3439 let peers = supervisor.registered_ids().await;
3440 let mut disconnected = 0;
3441 for peer in &peers {
3442 disconnected += usize::from(supervisor.disconnect_peer(peer).await);
3443 }
3444
3445 if disconnected > 0 {
3446 metrics::counter!("zakura.p2p.header_sync.capability_reconnects")
3447 .increment(u64::try_from(disconnected).expect("peer count always fits in u64"));
3448 info!(
3449 disconnected,
3450 capability_epoch = capability_epoch.get(),
3451 "reconnecting peers to negotiate header sync after checkpoint bootstrap"
3452 );
3453 }
3454
3455 disconnected
3456}
3457
3458fn provided_services_with_custom(custom_services: &[CustomService]) -> Vec<ZakuraServiceId> {
3459 let mut provided_services = discovery::default_advertised_services();
3460 provided_services.extend(
3461 custom_services
3462 .iter()
3463 .flat_map(|custom| custom.provides.iter().cloned()),
3464 );
3465 provided_services
3466}
3467
3468fn sought_services_with_custom(custom_services: &[CustomService]) -> Vec<ZakuraServiceId> {
3469 let mut sought_services: Vec<_> = custom_services
3470 .iter()
3471 .flat_map(|custom| custom.seeks.iter().cloned())
3472 .collect();
3473 sought_services.sort_unstable();
3474 sought_services.dedup();
3475 sought_services
3476}
3477
3478pub async fn spawn_zakura_endpoint(
3480 config: &Config,
3481 sink_factory: impl FnOnce(ZakuraSupervisorHandle, ZakuraTrace) -> Arc<dyn Service>,
3482) -> Result<Option<ZakuraEndpoint>, BoxError> {
3483 spawn_zakura_endpoint_with_header_sync_driver(config, sink_factory, None).await
3484}
3485
3486pub async fn spawn_zakura_endpoint_with_header_sync_driver(
3488 config: &Config,
3489 sink_factory: impl FnOnce(ZakuraSupervisorHandle, ZakuraTrace) -> Arc<dyn Service>,
3490 header_sync_driver_startup: Option<ZakuraHeaderSyncDriverStartup>,
3491) -> Result<Option<ZakuraEndpoint>, BoxError> {
3492 spawn_zakura_endpoint_with_services(
3493 config,
3494 sink_factory,
3495 header_sync_driver_startup,
3496 Vec::new(),
3497 )
3498 .await
3499}
3500
3501pub async fn spawn_zakura_endpoint_with_services(
3503 config: &Config,
3504 sink_factory: impl FnOnce(ZakuraSupervisorHandle, ZakuraTrace) -> Arc<dyn Service>,
3505 header_sync_driver_startup: Option<ZakuraHeaderSyncDriverStartup>,
3506 custom_services: Vec<CustomService>,
3507) -> Result<Option<ZakuraEndpoint>, BoxError> {
3508 spawn_zakura_endpoint_inner(
3509 config,
3510 sink_factory,
3511 header_sync_driver_startup,
3512 custom_services,
3513 None,
3514 )
3515 .await
3516}
3517
3518pub(crate) async fn spawn_zakura_endpoint_with_peer_registry(
3519 config: &Config,
3520 sink_factory: impl FnOnce(ZakuraSupervisorHandle, ZakuraTrace) -> Arc<dyn Service>,
3521 header_sync_driver_startup: Option<ZakuraHeaderSyncDriverStartup>,
3522 custom_services: Vec<CustomService>,
3523 peer_registry: PeerRegistry,
3524) -> Result<Option<ZakuraEndpoint>, BoxError> {
3525 spawn_zakura_endpoint_inner(
3526 config,
3527 sink_factory,
3528 header_sync_driver_startup,
3529 custom_services,
3530 Some(peer_registry),
3531 )
3532 .await
3533}
3534
3535async fn spawn_zakura_endpoint_inner(
3536 config: &Config,
3537 sink_factory: impl FnOnce(ZakuraSupervisorHandle, ZakuraTrace) -> Arc<dyn Service>,
3538 header_sync_driver_startup: Option<ZakuraHeaderSyncDriverStartup>,
3539 custom_services: Vec<CustomService>,
3540 peer_registry: Option<PeerRegistry>,
3541) -> Result<Option<ZakuraEndpoint>, BoxError> {
3542 if !config.v2_p2p() {
3543 return Ok(None);
3544 }
3545
3546 let limits = ZakuraLocalLimits::from_config(config);
3547 validate_idle_invariant(&limits)?;
3548 let secret_key = zakura_secret_key(config)?;
3549 let local_node_id = secret_key.public();
3550 let discovery_secret_key = secret_key.clone();
3551 let builder = direct_endpoint_builder(secret_key).transport_config(limits.transport_config());
3552 let builder = bind_native_endpoint(builder, config.zakura.listen_addr);
3556 let endpoint = builder.bind().await?;
3557 let supervisor = match peer_registry {
3558 Some(peer_registry) => ZakuraSupervisorHandle::new_with_peer_registry(
3559 config.zakura.max_connections_per_ip(),
3560 peer_registry,
3561 ),
3562 None => ZakuraSupervisorHandle::new(config.zakura.max_connections_per_ip()),
3563 };
3564 let tracer = config
3565 .zakura
3566 .trace_dir
3567 .clone()
3568 .map(zakura_jsonl_trace::JsonlTracer::spawn)
3569 .unwrap_or_else(zakura_jsonl_trace::JsonlTracer::noop);
3570 let trace = ZakuraTrace::new(tracer, zakura_jsonl_trace::node_id());
3571 let handshake_config = ZakuraHandshakeConfig::for_network_with_dev_cohort(
3572 &config.network,
3573 config.zakura.dev_network.as_deref(),
3574 );
3575 let sought_services = sought_services_with_custom(&custom_services);
3576 let discovery = discovery::build_discovery_handle(
3577 discovery_secret_key,
3578 discovery_direct_addrs(config, local_node_id),
3579 provided_services_with_custom(&custom_services),
3580 &handshake_config,
3581 config.zakura.max_connections,
3582 remote_bootstrap_peer_count(&config.zakura.bootstrap_peers, local_node_id),
3583 supervisor.subscribe(),
3584 )?;
3585 let anchor = config.zakura.header_sync.anchor(&config.network)?;
3586 let frontiers = header_sync_driver_startup.as_ref().map_or(
3587 FullStateFrontiers {
3588 finalized_height: anchor.0,
3589 verified_block_tip: anchor.0,
3590 verified_block_hash: anchor.1,
3591 },
3592 |startup| startup.frontiers,
3593 );
3594 let best_header_tip = header_sync_driver_startup
3595 .as_ref()
3596 .map_or(Some(anchor), |startup| startup.best_header_tip);
3597 let mut startup = HeaderSyncStartup::new(
3598 config.network.clone(),
3599 anchor,
3600 frontiers,
3601 best_header_tip,
3602 config.zakura.header_sync.clone(),
3603 limits.max_frame_bytes,
3604 );
3605 let (header_sync_fatal_tx, header_sync_fatal_rx) = mpsc::unbounded_channel();
3606 startup.fatal_events = Some(header_sync_fatal_tx);
3607 startup.use_direct_port();
3608 startup.trace = trace.clone();
3609 startup.committed_snapshots = header_sync_driver_startup
3610 .as_ref()
3611 .map(|driver| driver.committed_snapshots.clone());
3612 startup.vct_root_repairs = header_sync_driver_startup
3613 .as_ref()
3614 .and_then(|driver| driver.vct_root_repairs.clone());
3615 if let Some(driver) = header_sync_driver_startup.as_ref() {
3616 startup.header_chain_port = driver.header_chain_port.clone();
3617 }
3618 let header_sync_shutdown = CancellationToken::new();
3619 startup.shutdown = header_sync_shutdown.clone();
3620 let (header_sync, header_sync_actions, header_sync_task) = spawn_header_sync_reactor(startup)?;
3621 let block_sync_driver_enabled = header_sync_driver_startup.is_some();
3622 let (block_sync, block_sync_actions, block_sync_task) =
3623 if let Some(driver_startup) = header_sync_driver_startup.as_ref() {
3624 let best_header_tip = driver_startup.best_header_tip.unwrap_or(anchor);
3625 let mut startup = BlockSyncStartup::new_with_committed_views(
3626 BlockSyncFrontiers {
3627 finalized_height: driver_startup.frontiers.finalized_height,
3628 verified_block_tip: driver_startup.frontiers.verified_block_tip,
3629 verified_block_hash: driver_startup.verified_block_tip_hash,
3630 },
3631 best_header_tip,
3632 driver_startup.committed_views.clone(),
3633 config.zakura.block_sync.clone(),
3634 );
3635 startup.shutdown = header_sync_shutdown.clone();
3636 startup.trace = trace.clone();
3637 let (handle, actions, task) = spawn_block_sync_reactor(startup);
3638 (Some(handle), Some(actions), Some(task))
3639 } else {
3640 (None, None, None)
3641 };
3642 let discovery_service = Arc::new(super::DiscoveryService::with_sync_services(
3643 discovery.clone(),
3644 header_sync.clone(),
3645 block_sync.clone(),
3646 ));
3647 let legacy_service = sink_factory(supervisor.clone(), trace.clone());
3648 let service_demand = header_sync_driver_startup
3649 .as_ref()
3650 .map(|startup| startup.service_demand.clone());
3651 let registry = service_registry(
3652 &supervisor,
3653 Some(header_sync.clone()),
3654 block_sync.clone(),
3655 config.zakura.block_sync.clone(),
3656 legacy_service,
3657 discovery_service,
3658 service_demand.clone(),
3659 custom_services,
3660 );
3661 let registry = match registry {
3662 Ok(registry) => registry,
3663 Err(error) => {
3664 header_sync_shutdown.cancel();
3665 return Err(error);
3666 }
3667 };
3668 let mut tasks = vec![header_sync_task];
3669 if let Some(task) = block_sync_task {
3670 tasks.push(task);
3671 }
3672 #[cfg(any(test, feature = "zakura-testkit"))]
3673 let header_sync_actions = if block_sync_driver_enabled {
3674 Some(Arc::new(Mutex::new(Some(header_sync_actions))))
3675 } else {
3676 let action_driver_task = tokio::spawn(drive_header_sync_actions(
3677 header_sync_actions,
3678 header_sync.clone(),
3679 supervisor.clone(),
3680 header_sync_shutdown.clone(),
3681 ));
3682 tasks.push(action_driver_task);
3683 None
3684 };
3685 #[cfg(not(any(test, feature = "zakura-testkit")))]
3686 drop(header_sync_actions);
3687 let block_sync_actions = block_sync_actions
3688 .filter(|_| block_sync_driver_enabled)
3689 .map(|actions| Arc::new(Mutex::new(Some(actions))));
3690 let header_sync_tasks = Arc::new(HeaderSyncBackgroundTasks {
3691 shutdown: header_sync_shutdown,
3692 tasks: Mutex::new(tasks),
3693 });
3694 let header_sync_ready = service_demand
3695 .as_ref()
3696 .is_some_and(|demand| demand.borrow().header.is_enabled());
3697 let handler = ZakuraProtocolHandler::new_with_registry_and_trace(
3698 supervisor.clone(),
3699 config.network.clone(),
3700 handshake_config,
3701 limits.clone(),
3702 registry,
3703 trace,
3704 )
3705 .with_endpoint(endpoint.clone());
3708 handler.set_header_sync_enabled(header_sync_ready);
3709 let router = Router::builder(endpoint)
3710 .accept(P2P_V2_ALPN, handler.clone())
3711 .spawn();
3712 let endpoint = ZakuraEndpoint {
3713 router,
3714 supervisor,
3715 handler,
3716 header_sync: Some(header_sync),
3717 block_sync,
3718 header_sync_tasks: Some(header_sync_tasks),
3719 header_sync_fatal_events: Some(Arc::new(Mutex::new(Some(header_sync_fatal_rx)))),
3720 #[cfg(any(test, feature = "zakura-testkit"))]
3721 header_sync_actions,
3722 block_sync_actions,
3723 upgrade_dials: Arc::new(StdMutex::new(HashMap::new())),
3724 };
3725 let startup_shutdown = endpoint.background_shutdown_token().drop_guard();
3726
3727 if let Some(mut service_demand) = service_demand {
3728 let capability_handler = endpoint.handler.clone();
3729 let capability_supervisor = endpoint.supervisor.clone();
3730 let shutdown = endpoint.background_shutdown_token();
3731 let task = tokio::spawn(async move {
3732 let mut epochs = HeaderCapabilityEpochs::from_initial(&service_demand.borrow());
3733 loop {
3734 let ready_epoch = epochs.observe(&service_demand.borrow());
3735 if let Some(epoch) = ready_epoch {
3736 enable_header_sync_and_renegotiate(
3737 &capability_handler,
3738 &capability_supervisor,
3739 epoch,
3740 )
3741 .await;
3742 }
3743 tokio::select! {
3744 biased;
3745 _ = shutdown.cancelled() => break,
3746 changed = service_demand.changed() => {
3747 if changed.is_err() {
3748 break;
3749 }
3750 }
3751 }
3752 }
3753 });
3754 endpoint.push_header_sync_task(task).await;
3755 }
3756
3757 {
3762 let shutdown = endpoint.background_shutdown_token();
3763 let log_endpoint = endpoint.clone();
3764 let task = tokio::spawn(async move {
3765 tokio::select! {
3766 biased;
3767 _ = shutdown.cancelled() => {}
3768 node_addr = log_endpoint.node_addr() => {
3769 let direct_addresses: Vec<String> = node_addr
3770 .direct_addresses()
3771 .map(|addr| addr.to_string())
3772 .collect();
3773 info!(
3774 node_id = %node_addr.node_id,
3775 ?direct_addresses,
3776 "Zakura P2P endpoint ready; advertise <node_id>@<direct_addr> as a bootstrap peer",
3777 );
3778 }
3779 }
3780 });
3781 endpoint.push_header_sync_task(task).await;
3782 }
3783
3784 super::discovery::insert_static_bootstrap_candidates(
3785 &discovery,
3786 &config.zakura.bootstrap_peers,
3787 )
3788 .await;
3789 for task in spawn_native_bootstrap_dialer(
3795 endpoint.clone(),
3796 config.zakura.bootstrap_peers.clone(),
3797 limits.clone(),
3798 ) {
3799 endpoint.push_header_sync_task(task).await;
3800 }
3801 let discovery_dialer = discovery::spawn_native_discovery_dialer(
3802 endpoint.clone(),
3803 discovery,
3804 limits,
3805 sought_services,
3806 );
3807 endpoint.push_header_sync_task(discovery_dialer).await;
3808 startup_shutdown.disarm();
3809 Ok(Some(endpoint))
3810}
3811
3812pub(crate) async fn serve_native_dial_connection(
3813 endpoint: &ZakuraEndpoint,
3814 node_addr: NodeAddr,
3815 limits: &ZakuraLocalLimits,
3816) -> Result<(), ZakuraHandlerError> {
3817 let conn_id = endpoint
3818 .handler
3819 .next_conn_id
3820 .fetch_add(1, Ordering::Relaxed);
3821 let _admission = endpoint
3822 .handler
3823 .admission
3824 .clone()
3825 .try_acquire_owned()
3826 .map_err(|_| ZakuraHandlerError::ResourceLimit("admission"))?;
3827 let connection = timeout(
3828 limits.control_timeout,
3829 endpoint.router.endpoint().connect(node_addr, P2P_V2_ALPN),
3830 )
3831 .await
3832 .map_err(|_| ZakuraHandlerError::Timeout("native dial"))??;
3833 let remote_node_id = connection.remote_node_id()?;
3834 let peer_id = ZakuraPeerId::new(remote_node_id.as_bytes().to_vec())?;
3835 let conn = ZakuraConnTrace::new(&endpoint.handler.trace, conn_id, &peer_id);
3836 let local_node_id = endpoint.router.endpoint().node_id();
3837 let local_peer_id = ZakuraPeerId::new(local_node_id.as_bytes().to_vec())?;
3838 let negotiated = {
3839 let _handshake = endpoint
3840 .handler
3841 .pending_handshakes
3842 .clone()
3843 .try_acquire_owned()
3844 .map_err(|_| ZakuraHandlerError::ResourceLimit("pending handshake"))?;
3845 let handshake_config = endpoint.handler.current_handshake_config();
3846 run_native_initiator_handshake(
3847 &connection,
3848 limits,
3849 &handshake_config,
3850 &local_peer_id,
3851 &endpoint.handler.trace,
3852 &conn,
3853 )
3854 .await?
3855 };
3856 let conn_limits = limits.clamp(&negotiated.limits);
3857 let remote_ip = confirmed_remote_ip(endpoint.router.endpoint(), remote_node_id);
3863 let direction = ServicePeerDirection::Outbound;
3864 endpoint
3865 .handler
3866 .register_and_serve(
3867 connection,
3868 peer_id,
3869 remote_ip,
3870 ConnectionServeContext {
3871 limits: conn_limits,
3872 accepted_capabilities: negotiated.accepted_capabilities,
3873 role: "initiator",
3874 direction,
3875 transcript_hash: native_connection_transcript_hash(
3876 direction,
3877 &local_node_id,
3878 &remote_node_id,
3879 ),
3880 i_open_collision_winner: i_open_collision_winner(&local_node_id, &remote_node_id),
3881 conn,
3882 },
3883 )
3884 .await
3885}
3886
3887#[cfg(any(test, feature = "zakura-testkit"))]
3888pub(crate) async fn run_native_initiator_handshake_without_trace(
3889 connection: &Connection,
3890 limits: &ZakuraLocalLimits,
3891 handshake_config: &ZakuraHandshakeConfig,
3892 local_peer_id: &ZakuraPeerId,
3893) -> Result<NativeHandshakeNegotiated, ZakuraHandlerError> {
3894 run_native_initiator_handshake(
3895 connection,
3896 limits,
3897 handshake_config,
3898 local_peer_id,
3899 &ZakuraTrace::noop(),
3900 &ZakuraConnTrace::placeholder(),
3901 )
3902 .await
3903}
3904
3905async fn run_native_initiator_handshake(
3906 connection: &Connection,
3907 limits: &ZakuraLocalLimits,
3908 handshake_config: &ZakuraHandshakeConfig,
3909 local_peer_id: &ZakuraPeerId,
3910 _trace: &ZakuraTrace,
3911 conn: &ZakuraConnTrace,
3912) -> Result<NativeHandshakeNegotiated, ZakuraHandlerError> {
3913 conn.trace_handshake(
3914 "control.started",
3915 "initiator",
3916 None,
3917 handshake_config.network_label(),
3918 );
3919 let (mut send, mut recv) = timeout(limits.control_timeout, connection.open_bi())
3920 .await
3921 .map_err(|_| ZakuraHandlerError::Timeout("open control stream"))??;
3922 let mut local_nonce = [0; 32];
3923 OsRng.fill_bytes(&mut local_nonce);
3924
3925 let hello = ZakuraControlHello {
3926 magic: CONTROL_HELLO_MAGIC,
3927 control_version: CONTROL_VERSION,
3928 selected_zakura_protocol: ZAKURA_PROTOCOL_VERSION_1,
3929 handshake_path: ZakuraHandshakePath::Native,
3930 role: ZakuraControlRole::Initiator,
3931 network_id: handshake_config.network_id,
3932 chain_id: handshake_config.chain_id,
3933 iroh_node_id: local_peer_id.as_bytes().to_vec(),
3934 peer_nonce: local_nonce,
3935 initiator_upgrade_nonce: [0; 32],
3936 responder_upgrade_nonce: [0; 32],
3937 legacy_upgrade_transcript: [0; 32],
3938 capabilities: handshake_config.supported_capabilities,
3939 required_channels: 0,
3940 initial_limits: limits.initial_limits(),
3941 };
3942
3943 write_control_payload(&mut send, &hello.encode()?, limits.control_timeout).await?;
3944 let ack_bytes = read_control_payload(
3945 &mut recv,
3946 handshake_config.max_control_frame_bytes,
3947 limits.control_timeout,
3948 )
3949 .await?;
3950 let ack = ZakuraControlAck::decode(&ack_bytes)?;
3951 ack.validate(
3952 ZAKURA_PROTOCOL_VERSION_1,
3953 local_nonce,
3954 ack.peer_nonce,
3955 &limits.initial_limits(),
3956 handshake_config,
3957 )?;
3958 conn.trace_handshake(
3959 "control.succeeded",
3960 "initiator",
3961 Some(ack.selected_zakura_protocol),
3962 handshake_config.network_label(),
3963 );
3964 Ok(NativeHandshakeNegotiated {
3965 limits: ack.accepted_limits,
3966 accepted_capabilities: ack.accepted_capabilities,
3967 })
3968}
3969
3970#[allow(clippy::too_many_arguments)]
3971fn spawn_persistent_stream_worker(
3972 workers: &mut JoinSet<()>,
3973 send: SendStream,
3974 recv: RecvStream,
3975 stream: Stream,
3976 prelude: StreamPrelude,
3977 context: StreamWorkerContext,
3978 queue_depth: usize,
3979 opened_locally: bool,
3980 ordered_session_exit_tx: mpsc::UnboundedSender<OrderedSessionExit>,
3981) -> AdmittedOrderedSession {
3982 let (to_service_tx, to_service_rx) = mpsc::channel(queue_depth);
3983 let (from_service_tx, from_service_rx) = mpsc::channel(queue_depth);
3984 let admitted = AdmittedOrderedSession {
3985 kind: prelude.stream_kind,
3986 version: prelude.stream_version,
3987 session_id: context.stream_id,
3988 recv: FramedRecv::new(to_service_rx),
3989 send: FramedSend::new(from_service_tx),
3990 cancel_token: context.stream_token.clone(),
3991 };
3992
3993 let exit = OrderedSessionExit {
3994 stream,
3995 session_id: admitted.session_id,
3996 opened_locally,
3997 };
3998 workers.spawn(async move {
3999 persistent_stream_worker(
4000 send,
4001 recv,
4002 prelude,
4003 context,
4004 to_service_tx,
4005 from_service_rx,
4006 queue_depth,
4007 )
4008 .await;
4009 let _ = ordered_session_exit_tx.send(exit);
4010 });
4011
4012 admitted
4013}
4014
4015async fn persistent_stream_worker(
4016 mut send: SendStream,
4017 recv: RecvStream,
4018 prelude: StreamPrelude,
4019 context: StreamWorkerContext,
4020 inbound_tx: mpsc::Sender<Frame>,
4021 outbound_rx: mpsc::Receiver<Frame>,
4022 queue_depth_limit: usize,
4023) {
4024 let context = Arc::new(context);
4025 let stream_kind = prelude.stream_kind;
4026
4027 let (frame_tx, mut frame_rx) = mpsc::channel::<Result<Frame, ZakuraHandlerError>>(1);
4039 let reader_context = Arc::clone(&context);
4040 let reader = tokio::spawn(async move {
4041 let mut recv = recv;
4042 loop {
4043 let frame = tokio::select! {
4044 biased;
4045 _ = reader_context.connection_token.cancelled() => break,
4046 _ = reader_context.stream_token.cancelled() => break,
4047 frame = read_frame(
4048 &mut recv,
4049 reader_context.inbound_frame_cap,
4050 reader_context.limits.idle_timeout,
4051 None,
4057 ) => frame,
4058 };
4059 let message = match frame {
4067 Ok(frame) => {
4068 let _ = reader_context.freshness_tx.send(Instant::now());
4069 match admit_inbound_message(frame.payload.len(), &reader_context, stream_kind) {
4070 InboundMessageAdmission::Admit => Ok(frame),
4071 InboundMessageAdmission::Oversize => Err(ZakuraHandlerError::Oversize),
4072 InboundMessageAdmission::Throttled => Err(ZakuraHandlerError::RateLimited),
4073 }
4074 }
4075 Err(error) => {
4076 if let Some((payload_len, frame_len, max_frame_bytes)) =
4080 error.oversize_frame_details()
4081 {
4082 reader_context.conn.trace_rate_limit(
4083 "frame.oversize",
4084 reader_context.stream_id,
4085 stream_kind_label(stream_kind),
4086 Some(payload_len),
4087 Some(frame_len),
4088 Some(max_frame_bytes),
4089 );
4090 }
4091 Err(error)
4092 }
4093 };
4094 let is_terminal = message.is_err();
4101 let must_disconnect =
4102 matches!(&message, Err(error) if !matches!(error, ZakuraHandlerError::Closed));
4103 let forward_failed = frame_tx.send(message).await.is_err();
4104 if must_disconnect {
4105 reader_context.close_cause.record("ordered_read_error");
4106 reader_context.connection_token.cancel();
4107 }
4108 if forward_failed || is_terminal {
4109 break;
4110 }
4111 }
4112 });
4113
4114 let mut outbound_rx = Some(outbound_rx);
4115 loop {
4116 tokio::select! {
4117 biased;
4118 _ = context.connection_token.cancelled() => break,
4119 _ = context.stream_token.cancelled() => break,
4120 outbound = async {
4121 match outbound_rx.as_mut() {
4122 Some(outbound_rx) => outbound_rx.recv().await,
4123 None => future::pending().await,
4124 }
4125 } => {
4126 match outbound {
4127 Some(frame) => {
4128 if let Err(error) = write_ordered_frame(
4129 &mut send,
4130 frame,
4131 context.limits,
4132 context.outbound_frame_cap,
4133 ).await {
4134 if ordered_stream_write_was_stopped(&error) {
4135 debug!(?error, "closing Zakura ordered stream after peer stopped receiving");
4136 break;
4137 }
4138 debug!(?error, "closing Zakura ordered stream writer");
4139 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
4140 context.close_cause.record("ordered_write_error");
4141 context.connection_token.cancel();
4142 break;
4143 }
4144 }
4145 None => {
4146 outbound_rx = None;
4147 }
4148 }
4149 }
4150 inbound = frame_rx.recv() => {
4151 match inbound {
4152 Some(Ok(frame)) => {
4154 if inbound_tx.send(frame).await.is_err() {
4155 debug!(
4156 stream_kind,
4157 "closing Zakura ordered stream after local service receiver dropped"
4158 );
4159 break;
4160 }
4161 metrics::gauge!(
4162 "zakura.p2p.queue.depth",
4163 "stream_kind" => stream_kind_label(stream_kind),
4164 )
4165 .set(queue_depth_limit.saturating_sub(inbound_tx.capacity()) as f64);
4166 }
4167 Some(Err(ZakuraHandlerError::Oversize)) => {
4169 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_OVERSIZE));
4170 context.close_cause.record("ordered_oversize");
4171 context.connection_token.cancel();
4172 break;
4173 }
4174 Some(Err(ZakuraHandlerError::RateLimited)) => {
4175 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RATE_LIMIT));
4176 context.close_cause.record("ordered_rate_limited");
4177 context.connection_token.cancel();
4178 break;
4179 }
4180 Some(Err(ZakuraHandlerError::Closed)) | None => {
4181 break;
4182 }
4183 Some(Err(error)) => {
4185 debug!(?error, "closing Zakura stream worker");
4186 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
4187 context.close_cause.record("ordered_read_error");
4188 context.connection_token.cancel();
4189 break;
4190 }
4191 }
4192 }
4193 }
4194 }
4195
4196 reader.abort();
4200}
4201
4202fn ordered_stream_write_was_stopped(error: &BoxError) -> bool {
4203 matches!(
4204 error.downcast_ref::<iroh::endpoint::WriteError>(),
4205 Some(iroh::endpoint::WriteError::Stopped(_))
4206 )
4207}
4208
4209async fn request_stream_worker(
4210 mut send: SendStream,
4211 mut recv: RecvStream,
4212 prelude: StreamPrelude,
4213 context: StreamWorkerContext,
4214 registry: Arc<ServiceRegistry>,
4215) {
4216 let Some(request_id) = prelude.request_id else {
4217 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
4218 context.close_cause.record("request_without_id");
4219 context.connection_token.cancel();
4220 return;
4221 };
4222
4223 let frame = tokio::select! {
4224 biased;
4225 _ = context.connection_token.cancelled() => return,
4226 frame = read_frame(
4227 &mut recv,
4228 context.inbound_frame_cap,
4229 context.limits.idle_timeout,
4230 Some(context.limits.idle_timeout),
4234 ) => frame,
4235 };
4236
4237 let frame = match frame {
4238 Ok(frame) => frame,
4239 Err(error) => {
4240 debug!(
4241 ?error,
4242 "closing Zakura request stream with invalid request frame"
4243 );
4244 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
4245 context.close_cause.record("request_read_error");
4246 context.connection_token.cancel();
4247 return;
4248 }
4249 };
4250
4251 let _ = context.freshness_tx.send(Instant::now());
4252 match admit_inbound_message(frame.payload.len(), &context, prelude.stream_kind) {
4253 InboundMessageAdmission::Admit => {}
4254 InboundMessageAdmission::Oversize => {
4255 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_OVERSIZE));
4256 context.close_cause.record("request_oversize");
4257 context.connection_token.cancel();
4258 return;
4259 }
4260 InboundMessageAdmission::Throttled => {
4261 }
4266 }
4267
4268 let response_frame_cap = context.outbound_frame_cap;
4269 let response_frames = match registry
4270 .request(
4271 context.peer_id.clone(),
4272 prelude.stream_kind,
4273 request_id,
4274 response_frame_cap,
4275 context.limits.max_message_bytes,
4276 frame,
4277 )
4278 .await
4279 {
4280 Ok(frames) => frames,
4281 Err(SinkReject::Protocol(error)) => {
4282 debug!(
4283 ?error,
4284 "Zakura inbound sink rejected protocol-invalid request"
4285 );
4286 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
4287 context.close_cause.record("request_protocol_reject");
4288 context.connection_token.cancel();
4289 return;
4290 }
4291 Err(SinkReject::Local(error)) => {
4292 debug!(
4293 ?error,
4294 "Zakura inbound sink could not answer request locally"
4295 );
4296 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE));
4297 return;
4298 }
4299 };
4300
4301 for frame in response_frames {
4302 if let Err(error) =
4303 write_response_frame(&mut send, frame, context.limits, response_frame_cap).await
4304 {
4305 debug!(?error, "failed to write Zakura request response frame");
4306 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
4307 return;
4308 }
4309 }
4310
4311 let _ = send.finish();
4312}
4313
4314async fn freshness_reaper(mut freshness_rx: watch::Receiver<Instant>, idle_timeout: Duration) {
4315 loop {
4316 let last = *freshness_rx.borrow_and_update();
4317 let elapsed = last.elapsed();
4318 if elapsed >= idle_timeout {
4319 return;
4320 }
4321 tokio::select! {
4322 _ = tokio::time::sleep(idle_timeout - elapsed) => return,
4323 changed = freshness_rx.changed() => {
4324 if changed.is_err() {
4325 return;
4326 }
4327 }
4328 }
4329 }
4330}
4331
4332async fn read_stream_prelude(
4333 recv: &mut RecvStream,
4334 prelude_timeout: Duration,
4335) -> Result<StreamPrelude, ZakuraHandlerError> {
4336 let mut fixed = [0; STREAM_PRELUDE_FIXED_BYTES];
4337 timeout(prelude_timeout, recv.read_exact(&mut fixed))
4338 .await
4339 .map_err(|_| ZakuraHandlerError::Timeout("stream prelude"))??;
4340
4341 let mut fixed_reader = &fixed[..];
4342 let mut magic = [0; 4];
4343 fixed_reader.read_exact(&mut magic)?;
4344 if magic != STREAM_PRELUDE_MAGIC {
4345 return Err(ZakuraProtocolError::InvalidMagic.into());
4346 }
4347 let stream_kind = fixed_reader.read_u16::<LittleEndian>()?;
4348 let stream_version = fixed_reader.read_u16::<LittleEndian>()?;
4349 let request_id = match fixed[STREAM_PRELUDE_REQUEST_ID_FLAG_OFFSET] {
4350 0 => None,
4351 1 => {
4352 let mut bytes = [0; STREAM_PRELUDE_REQUEST_ID_BYTES];
4353 timeout(prelude_timeout, recv.read_exact(&mut bytes))
4354 .await
4355 .map_err(|_| ZakuraHandlerError::Timeout("stream request id"))??;
4356 let mut request_id_reader = &bytes[..];
4357 Some(request_id_reader.read_u64::<LittleEndian>()?)
4358 }
4359 flag => return Err(ZakuraProtocolError::InvalidFlag(flag).into()),
4360 };
4361
4362 let mut cap = [0; STREAM_PRELUDE_CAP_BYTES];
4363 timeout(prelude_timeout, recv.read_exact(&mut cap))
4364 .await
4365 .map_err(|_| ZakuraHandlerError::Timeout("stream prelude cap"))??;
4366 let mut cap_reader = &cap[..];
4367 let max_frame_bytes = cap_reader.read_u32::<LittleEndian>()?;
4368
4369 Ok(StreamPrelude {
4370 magic,
4371 stream_kind,
4372 stream_version,
4373 request_id,
4374 max_frame_bytes,
4375 })
4376}
4377
4378async fn read_frame(
4398 recv: &mut RecvStream,
4399 max_frame_bytes: u32,
4400 read_timeout: Duration,
4401 first_byte_timeout: Option<Duration>,
4402) -> Result<Frame, ZakuraHandlerError> {
4403 let mut header = [0; FRAME_HEADER_BYTES];
4404 let first_byte = recv.read_exact(&mut header[..1]);
4409 match first_byte_timeout {
4410 Some(first_byte_timeout) => match timeout(first_byte_timeout, first_byte).await {
4411 Ok(Ok(())) => {}
4412 Ok(Err(_)) => return Err(ZakuraHandlerError::Closed),
4413 Err(_) => return Err(ZakuraHandlerError::Timeout("frame header")),
4414 },
4415 None => match first_byte.await {
4416 Ok(()) => {}
4417 Err(_) => return Err(ZakuraHandlerError::Closed),
4418 },
4419 }
4420 match timeout(read_timeout, recv.read_exact(&mut header[1..])).await {
4423 Ok(Ok(())) => {}
4424 Ok(Err(_)) => return Err(ZakuraHandlerError::Closed),
4425 Err(_) => return Err(ZakuraHandlerError::Timeout("frame header")),
4426 }
4427 let mut reader = &header[..];
4428 let message_type = reader.read_u16::<LittleEndian>()?;
4429 let flags = reader.read_u16::<LittleEndian>()?;
4430 let payload_len = usize::try_from(reader.read_u32::<LittleEndian>()?)
4431 .expect("u32 payload lengths fit usize on supported targets");
4432 let max_frame_bytes =
4433 usize::try_from(max_frame_bytes).expect("u32 frame cap fits usize on supported targets");
4434 let frame_len = FRAME_HEADER_BYTES.saturating_add(payload_len);
4435 if frame_len > max_frame_bytes {
4436 metrics::counter!("zakura.p2p.ratelimit.frame.oversize").increment(1);
4437 return Err(ZakuraHandlerError::OversizeFrame {
4438 payload_len,
4439 frame_len,
4440 max_frame_bytes,
4441 });
4442 }
4443 let mut payload = vec![0; payload_len];
4444 timeout(read_timeout, recv.read_exact(&mut payload))
4445 .await
4446 .map_err(|_| ZakuraHandlerError::Timeout("frame payload"))??;
4447 Ok(Frame {
4448 message_type,
4449 flags,
4450 payload,
4451 })
4452}
4453
4454async fn read_control_payload(
4455 recv: &mut RecvStream,
4456 max_bytes: u32,
4457 read_timeout: Duration,
4458) -> Result<Vec<u8>, ZakuraHandlerError> {
4459 let max_bytes = max_bytes.min(MAX_CONTROL_PAYLOAD_BYTES as u32);
4468 let mut len_bytes = [0; CONTROL_LENGTH_BYTES];
4469 timeout(read_timeout, recv.read_exact(&mut len_bytes))
4470 .await
4471 .map_err(|_| ZakuraHandlerError::Timeout("control length"))??;
4472 let len = (&len_bytes[..]).read_u32::<LittleEndian>()?;
4473 if len == 0 || len > max_bytes {
4474 return Err(ZakuraHandlerError::Oversize);
4475 }
4476 let mut bytes = vec![0; len as usize];
4477 timeout(read_timeout, recv.read_exact(&mut bytes))
4478 .await
4479 .map_err(|_| ZakuraHandlerError::Timeout("control payload"))??;
4480 Ok(bytes)
4481}
4482
4483async fn write_control_payload(
4484 send: &mut SendStream,
4485 bytes: &[u8],
4486 write_timeout: Duration,
4487) -> Result<(), ZakuraHandlerError> {
4488 let mut len = Vec::with_capacity(CONTROL_LENGTH_BYTES);
4489 len.write_u32::<LittleEndian>(bytes.len() as u32)?;
4490 timeout(write_timeout, send.write_all(&len))
4491 .await
4492 .map_err(|_| ZakuraHandlerError::Timeout("control length write"))??;
4493 timeout(write_timeout, send.write_all(bytes))
4494 .await
4495 .map_err(|_| ZakuraHandlerError::Timeout("control payload write"))??;
4496 let _ = send.finish();
4497 Ok(())
4498}
4499
4500async fn write_ordered_frame(
4501 send: &mut SendStream,
4502 frame: Frame,
4503 limits: ZakuraConnectionLimits,
4504 max_frame_bytes: u32,
4505) -> Result<(), BoxError> {
4506 if frame.payload.len() > limits.max_message_bytes as usize {
4513 return Err(format!(
4514 "Zakura outbound ordered frame payload {} exceeds negotiated max_message_bytes {}",
4515 frame.payload.len(),
4516 limits.max_message_bytes,
4517 )
4518 .into());
4519 }
4520 let frame = frame.encode(max_frame_bytes)?;
4521 timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, send.write_all(&frame))
4522 .await
4523 .map_err(|_| -> BoxError { "Zakura outbound frame write timed out".into() })??;
4524 Ok(())
4525}
4526
4527async fn write_outbound_request_frame(
4528 connection: &Connection,
4529 limits: ZakuraConnectionLimits,
4530 stream: Stream,
4531 request_id: u64,
4532 message_type: u16,
4533 flags: u16,
4534 payload: Vec<u8>,
4535) -> Result<Vec<Frame>, OutboundRequestError> {
4536 timeout(
4537 OUTBOUND_REQUEST_RESPONSE_TIMEOUT,
4538 write_outbound_request_frame_inner(
4539 connection,
4540 limits,
4541 stream,
4542 request_id,
4543 message_type,
4544 flags,
4545 payload,
4546 ),
4547 )
4548 .await
4549 .map_err(|_| OutboundRequestError::Local("Zakura outbound request/response timed out".into()))?
4550}
4551
4552async fn write_outbound_request_frame_inner(
4553 connection: &Connection,
4554 limits: ZakuraConnectionLimits,
4555 stream: Stream,
4556 request_id: u64,
4557 message_type: u16,
4558 flags: u16,
4559 payload: Vec<u8>,
4560) -> Result<Vec<Frame>, OutboundRequestError> {
4561 let mut legacy_state = LegacyResponseReadState::new(LegacyResponseBudget::from_request(
4563 message_type,
4564 &payload,
4565 limits,
4566 )?);
4567 let (mut send, mut recv) = timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, connection.open_bi())
4568 .await
4569 .map_err(|_| -> BoxError { "Zakura outbound request stream open timed out".into() })
4570 .map_err(OutboundRequestError::Local)?
4571 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
4572 let outbound_frame_cap = application_frame_cap(&limits, stream);
4573 let inbound_frame_cap = inbound_frame_cap_for_stream(&limits, stream);
4574 let prelude = StreamPrelude {
4575 magic: STREAM_PRELUDE_MAGIC,
4576 stream_kind: stream.kind,
4577 stream_version: stream.version,
4578 request_id: Some(request_id),
4579 max_frame_bytes: inbound_frame_cap,
4580 };
4581 let frame = Frame {
4582 message_type,
4583 flags,
4584 payload,
4585 };
4586 let prelude = prelude.encode().map_err(|error| {
4587 OutboundRequestError::Local(BoxError::from(format!("failed to encode prelude: {error}")))
4588 })?;
4589 let frame = frame
4590 .encode(outbound_frame_cap)
4591 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
4592 timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, send.write_all(&prelude))
4593 .await
4594 .map_err(|_| -> BoxError { "Zakura outbound request prelude write timed out".into() })
4595 .map_err(OutboundRequestError::Local)?
4596 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
4597 timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, send.write_all(&frame))
4598 .await
4599 .map_err(|_| -> BoxError { "Zakura outbound request frame write timed out".into() })
4600 .map_err(OutboundRequestError::Local)?
4601 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
4602 let _ = send.finish();
4603
4604 let mut frames = Vec::new();
4605 loop {
4606 match read_frame(
4607 &mut recv,
4608 inbound_frame_cap,
4609 limits.idle_timeout,
4610 Some(limits.idle_timeout),
4614 )
4615 .await
4616 {
4617 Ok(frame) => {
4618 legacy_state.validate_frame(request_id, &frame)?;
4619 frames.push(frame);
4620 }
4621 Err(ZakuraHandlerError::Closed) => {
4622 legacy_state.finish()?;
4623 return Ok(frames);
4624 }
4625 Err(ZakuraHandlerError::Timeout(_)) => {
4626 return Err(OutboundRequestError::Local(Box::new(
4627 ZakuraHandlerError::Timeout("outbound response"),
4628 )));
4629 }
4630 Err(error @ ZakuraHandlerError::OversizeFrame { .. })
4631 | Err(error @ ZakuraHandlerError::Oversize) => {
4632 return Err(OutboundRequestError::Fatal(Box::new(error)));
4633 }
4634 Err(error) => return Err(OutboundRequestError::Fatal(Box::new(error))),
4635 }
4636 }
4637}
4638
4639#[derive(Debug)]
4640enum OutboundRequestError {
4641 Local(BoxError),
4642 Fatal(BoxError),
4643}
4644
4645#[derive(Copy, Clone, Debug, Eq, PartialEq)]
4646enum LegacyResponseKind {
4647 Blocks,
4648 Transactions,
4649 BlockHashes,
4650 BlockHeaders,
4651 TransactionIds,
4652 Pong,
4653 Nil,
4654}
4655
4656#[derive(Copy, Clone, Debug)]
4657struct LegacyResponseBudget {
4658 kind: LegacyResponseKind,
4659 max_items: usize,
4660 max_frames: usize,
4661 max_bytes: usize,
4662 max_message_bytes: usize,
4663}
4664
4665impl LegacyResponseBudget {
4666 fn from_request(
4667 message_type: u16,
4668 payload: &[u8],
4669 limits: ZakuraConnectionLimits,
4670 ) -> Result<Self, OutboundRequestError> {
4671 let kind = match message_type {
4672 LEGACY_REQUEST_BLOCKS_BY_HASH => LegacyResponseKind::Blocks,
4673 LEGACY_REQUEST_TRANSACTIONS_BY_ID => LegacyResponseKind::Transactions,
4674 LEGACY_REQUEST_FIND_BLOCKS => LegacyResponseKind::BlockHashes,
4675 LEGACY_REQUEST_FIND_HEADERS => LegacyResponseKind::BlockHeaders,
4676 LEGACY_REQUEST_MEMPOOL_TRANSACTION_IDS => LegacyResponseKind::TransactionIds,
4677 LEGACY_REQUEST_PING => LegacyResponseKind::Pong,
4678 LEGACY_REQUEST_PUSH_TRANSACTION => LegacyResponseKind::Nil,
4679 _ => {
4680 return Err(OutboundRequestError::Local(
4681 format!("unsupported legacy request message type: {message_type}").into(),
4682 ));
4683 }
4684 };
4685 let max_message_bytes = usize::try_from(limits.max_message_bytes)
4686 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
4687 let max_inventory_items = usize::try_from(MAX_TX_INV_IN_SENT_MESSAGE)
4688 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
4689 let (max_items, max_frames, max_bytes) = match kind {
4690 LegacyResponseKind::Blocks | LegacyResponseKind::Transactions => {
4691 let item_count = legacy_inventory_count(payload, kind)?;
4692 let max_frames = item_count
4693 .saturating_mul(LEGACY_RESPONSE_MAX_FRAMES_PER_ITEM)
4694 .saturating_add(1);
4695 let item_bytes = match kind {
4696 LegacyResponseKind::Blocks => LEGACY_BLOCK_HASH_BYTES,
4697 LegacyResponseKind::Transactions => LEGACY_INVENTORY_HASH_BYTES,
4698 _ => unreachable!("matched inventory response kind"),
4699 };
4700 let missing_bytes = LEGACY_RESPONSE_REQUEST_ID_BYTES
4701 .saturating_add(LEGACY_COMPACT_SIZE_PREFIX_BYTES)
4702 .saturating_add(item_count.saturating_mul(item_bytes));
4703 let max_bytes = item_count
4704 .saturating_mul(max_message_bytes)
4705 .saturating_add(missing_bytes)
4706 .max(LEGACY_RESPONSE_CHUNK_HEADER_BYTES);
4707 (item_count, max_frames, max_bytes)
4708 }
4709 LegacyResponseKind::BlockHashes => {
4710 let max_bytes = LEGACY_RESPONSE_REQUEST_ID_BYTES
4711 .saturating_add(LEGACY_COMPACT_SIZE_PREFIX_BYTES)
4712 .saturating_add(max_inventory_items.saturating_mul(LEGACY_BLOCK_HASH_BYTES));
4713 (max_inventory_items, 1, max_bytes)
4714 }
4715 LegacyResponseKind::BlockHeaders => (
4716 MAX_HEADERS_PER_MESSAGE,
4717 1,
4718 LEGACY_RESPONSE_REQUEST_ID_BYTES.saturating_add(max_message_bytes),
4719 ),
4720 LegacyResponseKind::TransactionIds => {
4721 let max_bytes = LEGACY_RESPONSE_REQUEST_ID_BYTES
4722 .saturating_add(LEGACY_COMPACT_SIZE_PREFIX_BYTES)
4723 .saturating_add(
4724 max_inventory_items.saturating_mul(LEGACY_INVENTORY_HASH_BYTES),
4725 );
4726 (max_inventory_items, 1, max_bytes)
4727 }
4728 LegacyResponseKind::Pong | LegacyResponseKind::Nil => {
4729 (1, 1, LEGACY_RESPONSE_REQUEST_ID_BYTES)
4730 }
4731 };
4732
4733 let max_bytes = max_bytes.min(LEGACY_RESPONSE_MAX_AGGREGATE_BYTES);
4742
4743 Ok(Self {
4744 kind,
4745 max_items,
4746 max_frames,
4747 max_bytes,
4748 max_message_bytes,
4749 })
4750 }
4751}
4752
4753#[derive(Debug)]
4754struct LegacyResponseReadState {
4755 budget: LegacyResponseBudget,
4756 frames: usize,
4757 bytes: usize,
4758 items: usize,
4759 active_chunk_type: Option<u16>,
4760 active_chunk: Vec<u8>,
4761}
4762
4763impl LegacyResponseReadState {
4764 fn new(budget: LegacyResponseBudget) -> Self {
4765 Self {
4766 budget,
4767 frames: 0,
4768 bytes: 0,
4769 items: 0,
4770 active_chunk_type: None,
4771 active_chunk: Vec::new(),
4772 }
4773 }
4774
4775 fn validate_frame(
4776 &mut self,
4777 request_id: u64,
4778 frame: &Frame,
4779 ) -> Result<(), OutboundRequestError> {
4780 if frame.flags != 0 {
4781 return Err(OutboundRequestError::Fatal(
4782 format!("unsupported legacy response flags: {}", frame.flags).into(),
4783 ));
4784 }
4785
4786 self.frames = self.frames.saturating_add(1);
4787 if self.frames > self.budget.max_frames {
4788 return Err(OutboundRequestError::Fatal(
4789 "too many legacy response frames".into(),
4790 ));
4791 }
4792
4793 self.bytes = self.bytes.saturating_add(frame.payload.len());
4794 if self.bytes > self.budget.max_bytes {
4795 return Err(OutboundRequestError::Fatal(
4796 "legacy response exceeded cumulative byte budget".into(),
4797 ));
4798 }
4799
4800 match frame.message_type {
4801 LEGACY_RESPONSE_BLOCK => {
4802 self.validate_available_chunk(request_id, frame, LegacyResponseKind::Blocks)
4803 }
4804 LEGACY_RESPONSE_TRANSACTION => {
4805 self.validate_available_chunk(request_id, frame, LegacyResponseKind::Transactions)
4806 }
4807 LEGACY_RESPONSE_MISSING_BLOCKS => {
4808 self.validate_missing(request_id, &frame.payload, LegacyResponseKind::Blocks)
4809 }
4810 LEGACY_RESPONSE_MISSING_TRANSACTIONS => {
4811 self.validate_missing(request_id, &frame.payload, LegacyResponseKind::Transactions)
4812 }
4813 LEGACY_RESPONSE_BLOCK_HASHES => self.validate_id_prefixed_hashes(
4814 request_id,
4815 &frame.payload,
4816 LegacyResponseKind::BlockHashes,
4817 ),
4818 LEGACY_RESPONSE_BLOCK_HEADERS => self.validate_headers(request_id, &frame.payload),
4819 LEGACY_RESPONSE_TRANSACTION_IDS => self.validate_id_prefixed_hashes(
4820 request_id,
4821 &frame.payload,
4822 LegacyResponseKind::TransactionIds,
4823 ),
4824 LEGACY_RESPONSE_PONG => {
4825 self.validate_id_only(request_id, &frame.payload, LegacyResponseKind::Pong)
4826 }
4827 LEGACY_RESPONSE_NIL => self.validate_nil(request_id, &frame.payload),
4828 message_type => Err(OutboundRequestError::Fatal(
4829 format!("unknown legacy response message type: {message_type}").into(),
4830 )),
4831 }
4832 }
4833
4834 fn finish(self) -> Result<(), OutboundRequestError> {
4835 if self.active_chunk_type.is_some() {
4836 return Err(OutboundRequestError::Fatal(
4837 "incomplete legacy response chunk".into(),
4838 ));
4839 }
4840 Ok(())
4841 }
4842
4843 fn validate_available_chunk(
4844 &mut self,
4845 request_id: u64,
4846 frame: &Frame,
4847 kind: LegacyResponseKind,
4848 ) -> Result<(), OutboundRequestError> {
4849 if self.budget.kind != kind {
4850 return Err(OutboundRequestError::Fatal(
4851 "legacy response kind does not match request".into(),
4852 ));
4853 }
4854
4855 let (response_id, is_last, bytes) = legacy_response_chunk_header(&frame.payload)?;
4856 if response_id != request_id {
4857 return Err(OutboundRequestError::Fatal(
4858 format!(
4859 "wrong legacy response request id: expected {request_id}, got {response_id}"
4860 )
4861 .into(),
4862 ));
4863 }
4864
4865 match self.active_chunk_type {
4866 Some(active) if active != frame.message_type => {
4867 return Err(OutboundRequestError::Fatal(
4868 "interleaved legacy response chunks".into(),
4869 ));
4870 }
4871 None => self.active_chunk_type = Some(frame.message_type),
4872 _ => {}
4873 }
4874
4875 let new_len = self.active_chunk.len().saturating_add(bytes.len());
4876 if new_len > self.budget.max_message_bytes {
4877 return Err(OutboundRequestError::Fatal(
4878 "legacy response item exceeded message byte cap".into(),
4879 ));
4880 }
4881 self.active_chunk.extend_from_slice(bytes);
4882
4883 if is_last {
4884 self.validate_completed_item(kind)?;
4885 self.active_chunk_type = None;
4886 self.active_chunk.clear();
4887 self.add_items(1)?;
4888 }
4889
4890 Ok(())
4891 }
4892
4893 fn validate_completed_item(
4894 &self,
4895 kind: LegacyResponseKind,
4896 ) -> Result<(), OutboundRequestError> {
4897 match kind {
4898 LegacyResponseKind::Blocks => {
4899 Block::zcash_deserialize(&mut Cursor::new(self.active_chunk.as_slice()))
4900 .map(|_| ())
4901 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))
4902 }
4903 LegacyResponseKind::Transactions => {
4904 Transaction::zcash_deserialize(&mut Cursor::new(self.active_chunk.as_slice()))
4905 .map(|_| ())
4906 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))
4907 }
4908 LegacyResponseKind::BlockHashes
4909 | LegacyResponseKind::BlockHeaders
4910 | LegacyResponseKind::TransactionIds
4911 | LegacyResponseKind::Pong
4912 | LegacyResponseKind::Nil => unreachable!("non-chunk legacy response kind"),
4913 }
4914 }
4915
4916 fn validate_missing(
4917 &mut self,
4918 request_id: u64,
4919 payload: &[u8],
4920 kind: LegacyResponseKind,
4921 ) -> Result<(), OutboundRequestError> {
4922 let payload = self.begin_id_prefixed_response(kind, request_id, payload, "missing")?;
4923 let count = legacy_inventory_count(payload, kind)?;
4924 self.add_items(count)
4925 }
4926
4927 fn validate_id_prefixed_hashes(
4928 &mut self,
4929 request_id: u64,
4930 payload: &[u8],
4931 kind: LegacyResponseKind,
4932 ) -> Result<(), OutboundRequestError> {
4933 let payload = self.begin_id_prefixed_response(kind, request_id, payload, "list")?;
4934 let count = legacy_inventory_count(payload, kind)?;
4935 self.add_items(count)
4936 }
4937
4938 fn validate_headers(
4939 &mut self,
4940 request_id: u64,
4941 payload: &[u8],
4942 ) -> Result<(), OutboundRequestError> {
4943 let payload = self.begin_id_prefixed_response(
4944 LegacyResponseKind::BlockHeaders,
4945 request_id,
4946 payload,
4947 "headers",
4948 )?;
4949 let count = legacy_header_count(payload)?;
4950 self.add_items(count)
4951 }
4952
4953 fn validate_id_only(
4954 &mut self,
4955 request_id: u64,
4956 payload: &[u8],
4957 kind: LegacyResponseKind,
4958 ) -> Result<(), OutboundRequestError> {
4959 let payload =
4960 self.begin_id_prefixed_response(kind, request_id, payload, "acknowledgement")?;
4961 if !payload.is_empty() {
4962 return Err(OutboundRequestError::Fatal(
4963 "legacy acknowledgement response has trailing bytes".into(),
4964 ));
4965 }
4966 self.add_items(1)
4967 }
4968
4969 fn validate_nil(
4981 &mut self,
4982 request_id: u64,
4983 payload: &[u8],
4984 ) -> Result<(), OutboundRequestError> {
4985 match self.budget.kind {
4986 LegacyResponseKind::BlockHashes
4987 | LegacyResponseKind::BlockHeaders
4988 | LegacyResponseKind::TransactionIds
4989 | LegacyResponseKind::Nil => {}
4990 LegacyResponseKind::Blocks
4991 | LegacyResponseKind::Transactions
4992 | LegacyResponseKind::Pong => {
4993 return Err(OutboundRequestError::Fatal(
4994 "unexpected legacy nil response for inventory or ping request".into(),
4995 ));
4996 }
4997 }
4998 if self.active_chunk_type.is_some() {
4999 return Err(OutboundRequestError::Fatal(
5000 "legacy nil response interleaved with response chunk".into(),
5001 ));
5002 }
5003 let (response_id, payload) = legacy_response_id(payload)?;
5004 if response_id != request_id {
5005 return Err(OutboundRequestError::Fatal(
5006 format!(
5007 "wrong legacy nil response request id: expected {request_id}, got {response_id}"
5008 )
5009 .into(),
5010 ));
5011 }
5012 if !payload.is_empty() {
5013 return Err(OutboundRequestError::Fatal(
5014 "legacy nil response has trailing bytes".into(),
5015 ));
5016 }
5017 self.add_items(1)
5018 }
5019
5020 fn begin_id_prefixed_response<'a>(
5021 &self,
5022 kind: LegacyResponseKind,
5023 request_id: u64,
5024 payload: &'a [u8],
5025 label: &'static str,
5026 ) -> Result<&'a [u8], OutboundRequestError> {
5027 if self.budget.kind != kind {
5028 return Err(OutboundRequestError::Fatal(
5029 format!("legacy {label} response kind does not match request").into(),
5030 ));
5031 }
5032 if self.active_chunk_type.is_some() {
5033 return Err(OutboundRequestError::Fatal(
5034 format!("legacy {label} response interleaved with response chunk").into(),
5035 ));
5036 }
5037 let (response_id, payload) = legacy_response_id(payload)?;
5038 if response_id != request_id {
5039 return Err(OutboundRequestError::Fatal(
5040 format!(
5041 "wrong legacy {label} response request id: expected {request_id}, got {response_id}"
5042 )
5043 .into(),
5044 ));
5045 }
5046 Ok(payload)
5047 }
5048
5049 fn add_items(&mut self, count: usize) -> Result<(), OutboundRequestError> {
5050 self.items = self.items.saturating_add(count);
5051 if self.items > self.budget.max_items {
5052 return Err(OutboundRequestError::Fatal(
5053 "legacy response contained more items than requested".into(),
5054 ));
5055 }
5056 Ok(())
5057 }
5058}
5059
5060fn legacy_inventory_count(
5061 payload: &[u8],
5062 kind: LegacyResponseKind,
5063) -> Result<usize, OutboundRequestError> {
5064 let mut reader = Cursor::new(payload);
5065 let count = usize::from(
5066 CompactSizeMessage::zcash_deserialize(&mut reader)
5067 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?,
5068 );
5069
5070 match kind {
5071 LegacyResponseKind::Blocks | LegacyResponseKind::BlockHashes => {
5072 let consumed = usize::try_from(reader.position())
5073 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?;
5074 let expected_len =
5075 consumed.saturating_add(count.saturating_mul(LEGACY_BLOCK_HASH_BYTES));
5076 if expected_len != payload.len() {
5077 return Err(OutboundRequestError::Fatal(
5078 "legacy block inventory has trailing or truncated bytes".into(),
5079 ));
5080 }
5081 }
5082 LegacyResponseKind::Transactions | LegacyResponseKind::TransactionIds => {
5083 for _ in 0..count {
5084 let inventory = InventoryHash::zcash_deserialize(&mut reader)
5085 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?;
5086 if inventory.unmined_tx_id().is_none() {
5087 return Err(OutboundRequestError::Fatal(
5088 "legacy transaction inventory contained a non-transaction item".into(),
5089 ));
5090 }
5091 }
5092 reject_legacy_inventory_trailing(payload, &reader)?;
5093 }
5094 LegacyResponseKind::BlockHeaders | LegacyResponseKind::Pong | LegacyResponseKind::Nil => {
5095 return Err(OutboundRequestError::Local(
5096 "legacy response kind does not use inventory counts".into(),
5097 ));
5098 }
5099 }
5100
5101 Ok(count)
5102}
5103
5104fn legacy_header_count(payload: &[u8]) -> Result<usize, OutboundRequestError> {
5105 let mut reader = Cursor::new(payload);
5106 let count = usize::from(
5107 CompactSizeMessage::zcash_deserialize(&mut reader)
5108 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?,
5109 );
5110 if count > MAX_HEADERS_PER_MESSAGE {
5111 return Err(OutboundRequestError::Fatal(
5112 "legacy headers response exceeded header count cap".into(),
5113 ));
5114 }
5115 for _ in 0..count {
5116 CountedHeader::zcash_deserialize(&mut reader)
5117 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?;
5118 }
5119 reject_legacy_inventory_trailing(payload, &reader)?;
5120 Ok(count)
5121}
5122
5123fn reject_legacy_inventory_trailing(
5124 payload: &[u8],
5125 reader: &Cursor<&[u8]>,
5126) -> Result<(), OutboundRequestError> {
5127 if usize::try_from(reader.position())
5128 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?
5129 != payload.len()
5130 {
5131 return Err(OutboundRequestError::Fatal(
5132 "legacy transaction inventory has trailing bytes".into(),
5133 ));
5134 }
5135 Ok(())
5136}
5137
5138fn legacy_response_id(payload: &[u8]) -> Result<(u64, &[u8]), OutboundRequestError> {
5139 if payload.len() < LEGACY_RESPONSE_REQUEST_ID_BYTES {
5140 return Err(OutboundRequestError::Fatal(
5141 "truncated legacy response id".into(),
5142 ));
5143 }
5144 let mut id = [0; LEGACY_RESPONSE_REQUEST_ID_BYTES];
5145 id.copy_from_slice(&payload[..LEGACY_RESPONSE_REQUEST_ID_BYTES]);
5146 Ok((
5147 u64::from_le_bytes(id),
5148 &payload[LEGACY_RESPONSE_REQUEST_ID_BYTES..],
5149 ))
5150}
5151
5152fn legacy_response_chunk_header(
5153 payload: &[u8],
5154) -> Result<(u64, bool, &[u8]), OutboundRequestError> {
5155 if payload.len() < LEGACY_RESPONSE_CHUNK_HEADER_BYTES {
5156 return Err(OutboundRequestError::Fatal(
5157 "truncated legacy response chunk".into(),
5158 ));
5159 }
5160 let (request_id, payload) = legacy_response_id(payload)?;
5161 Ok((request_id, payload[0] != 0, &payload[1..]))
5162}
5163
5164async fn write_response_frame(
5165 send: &mut SendStream,
5166 frame: Frame,
5167 limits: ZakuraConnectionLimits,
5168 max_frame_bytes: u32,
5169) -> Result<(), ZakuraHandlerError> {
5170 if frame.payload.len() > limits.max_message_bytes as usize {
5171 return Err(ZakuraHandlerError::Oversize);
5172 }
5173 let frame = frame.encode(max_frame_bytes)?;
5174 timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, send.write_all(&frame))
5175 .await
5176 .map_err(|_| ZakuraHandlerError::Timeout("frame write"))??;
5177 Ok(())
5178}
5179
5180fn validate_idle_invariant(limits: &ZakuraLocalLimits) -> Result<(), ZakuraHandlerError> {
5181 if limits.keep_alive_interval >= limits.quic_idle_timeout {
5182 return Err(ZakuraHandlerError::InvalidLocalLimits);
5183 }
5184 if limits.initial_limits().idle_timeout_millis as u128 >= limits.quic_idle_timeout.as_millis() {
5185 return Err(ZakuraHandlerError::InvalidLocalLimits);
5186 }
5187 Ok(())
5188}
5189
5190fn zakura_secret_key(config: &Config) -> Result<SecretKey, ZakuraHandlerError> {
5191 config
5194 .zakura_secret_key()
5195 .map_err(|_| ZakuraHandlerError::InvalidSecretKey)
5196}
5197
5198fn stream_kind_label(stream_kind: u16) -> &'static str {
5199 match stream_kind {
5200 0 => "control",
5201 1 => "request",
5202 LEGACY_GOSSIP_STREAM_KIND => "gossip",
5203 LEGACY_REQUEST_STREAM_KIND => "legacy_request",
5204 DISCOVERY_STREAM_KIND => "discovery",
5205 HEADER_SYNC_STREAM_KIND => "header_sync",
5206 ZAKURA_STREAM_BLOCK_SYNC => "block_sync",
5207 _ => "unknown",
5208 }
5209}
5210
5211fn application_frame_cap(limits: &ZakuraConnectionLimits, stream: Stream) -> u32 {
5212 limits.max_frame_bytes.min(stream.frame_cap)
5213}
5214
5215fn peer_accepted_frame_cap(
5216 limits: &ZakuraConnectionLimits,
5217 stream: Stream,
5218 peer_max_frame_bytes: u32,
5219) -> u32 {
5220 application_frame_cap(limits, stream).min(peer_max_frame_bytes)
5221}
5222
5223fn inbound_frame_cap_for_stream(limits: &ZakuraConnectionLimits, stream: Stream) -> u32 {
5233 let frame_header_bytes =
5234 u32::try_from(FRAME_HEADER_BYTES).expect("frame header byte count fits in u32");
5235 application_frame_cap(limits, stream)
5236 .min(limits.max_message_bytes.saturating_add(frame_header_bytes))
5237}
5238
5239fn per_stream_inbound_queue_depth(
5240 max_inbound_queue_depth: u16,
5241 ordered_stream_count: usize,
5242) -> usize {
5243 let total = usize::from(max_inbound_queue_depth).max(1);
5244 if ordered_stream_count == 0 {
5245 return total;
5246 }
5247
5248 total.saturating_div(ordered_stream_count).max(1)
5249}
5250
5251fn should_run_freshness_reaper(
5252 ordered_stream_count: usize,
5253 request_response_stream_count: usize,
5254) -> bool {
5255 ordered_stream_count > 0 || request_response_stream_count == 0
5256}
5257
5258#[cfg(test)]
5260const ZAKURA_STREAM_VERSION_1: u16 = 1;
5261
5262#[cfg(test)]
5270fn is_supported_stream(registry: &ServiceRegistry, stream_kind: u16, stream_version: u16) -> bool {
5271 registry
5272 .capability_for_stream(stream_kind, stream_version)
5273 .is_some()
5274}
5275
5276type SharedMessageBucket<C = RealClock> = Arc<std::sync::Mutex<TokenBucket<C>>>;
5283
5284type MessageRateBuckets<C = RealClock> = HashMap<u16, SharedMessageBucket<C>>;
5288
5289fn message_bucket_for<C: Clock>(
5296 buckets: &mut MessageRateBuckets<C>,
5297 stream_kind: u16,
5298 message_rate_per_second: u32,
5299 clock: C,
5300) -> SharedMessageBucket<C> {
5301 buckets
5302 .entry(stream_kind)
5303 .or_insert_with(|| {
5304 Arc::new(std::sync::Mutex::new(TokenBucket::with_clock(
5305 message_rate_per_second,
5306 clock,
5307 )))
5308 })
5309 .clone()
5310}
5311
5312#[derive(Clone, Debug)]
5313struct TokenBucket<C = RealClock> {
5314 capacity: u32,
5315 tokens: u32,
5316 refill_per_second: u32,
5317 last_refill: Instant,
5318 clock: C,
5319}
5320
5321impl TokenBucket<RealClock> {
5322 fn new(refill_per_second: u32) -> Self {
5323 Self::with_clock(refill_per_second, RealClock)
5324 }
5325}
5326
5327impl<C: Clock> TokenBucket<C> {
5328 fn with_clock(refill_per_second: u32, clock: C) -> Self {
5329 let capacity = refill_per_second.max(1);
5330 Self {
5331 capacity,
5332 tokens: capacity,
5333 refill_per_second: capacity,
5334 last_refill: clock.now(),
5335 clock,
5336 }
5337 }
5338
5339 fn try_take(&mut self) -> bool {
5340 self.refill(self.clock.now());
5341 if self.tokens == 0 {
5342 return false;
5343 }
5344 self.tokens -= 1;
5345 true
5346 }
5347
5348 fn refill(&mut self, now: Instant) {
5349 let elapsed = now.saturating_duration_since(self.last_refill);
5350 let elapsed_nanos = elapsed.as_nanos();
5351 let new_tokens =
5352 elapsed_nanos.saturating_mul(u128::from(self.refill_per_second)) / 1_000_000_000;
5353 if new_tokens == 0 {
5354 return;
5355 }
5356 let new_tokens = new_tokens.min(u128::from(u32::MAX)) as u32;
5357 self.tokens = self.capacity.min(self.tokens.saturating_add(new_tokens));
5358 self.last_refill = now;
5359 }
5360}
5361
5362#[derive(Debug, Error)]
5364pub enum ZakuraHandlerError {
5365 #[error("Zakura {0} timed out")]
5367 Timeout(&'static str),
5368 #[error("Zakura payload exceeded its cap")]
5370 Oversize,
5371 #[error(
5373 "Zakura frame length {frame_len} exceeded cap {max_frame_bytes} \
5374 (payload length {payload_len})"
5375 )]
5376 OversizeFrame {
5377 payload_len: usize,
5379 frame_len: usize,
5381 max_frame_bytes: usize,
5383 },
5384 #[error("Zakura stream closed")]
5386 Closed,
5387 #[error("invalid Zakura bootstrap peer")]
5389 InvalidBootstrapPeer,
5390 #[error("invalid Zakura iroh secret key")]
5392 InvalidSecretKey,
5393 #[error("invalid Zakura local limits")]
5395 InvalidLocalLimits,
5396 #[error("Zakura resource limit exceeded: {0}")]
5398 ResourceLimit(&'static str),
5399 #[error("Zakura message rate exceeded")]
5401 RateLimited,
5402 #[error(transparent)]
5404 IrohConnection(#[from] iroh::endpoint::ConnectionError),
5405 #[error(transparent)]
5407 IrohConnect(#[from] iroh::endpoint::ConnectError),
5408 #[error(transparent)]
5410 IrohRemoteId(#[from] iroh::endpoint::RemoteNodeIdError),
5411 #[error(transparent)]
5413 IrohWrite(#[from] iroh::endpoint::WriteError),
5414 #[error(transparent)]
5416 IrohRead(#[from] iroh::endpoint::ReadExactError),
5417 #[error(transparent)]
5419 IrohClosedStream(#[from] iroh::endpoint::ClosedStream),
5420 #[error(transparent)]
5422 Protocol(#[from] ZakuraProtocolError),
5423 #[error(transparent)]
5425 Validation(#[from] super::ZakuraValidationError),
5426 #[error(transparent)]
5428 Io(#[from] std::io::Error),
5429}
5430
5431impl ZakuraHandlerError {
5432 fn oversize_frame_details(&self) -> Option<(u64, u64, u64)> {
5433 let Self::OversizeFrame {
5434 payload_len,
5435 frame_len,
5436 max_frame_bytes,
5437 } = self
5438 else {
5439 return None;
5440 };
5441
5442 Some((
5443 u64::try_from(*payload_len).unwrap_or(u64::MAX),
5444 u64::try_from(*frame_len).unwrap_or(u64::MAX),
5445 u64::try_from(*max_frame_bytes).unwrap_or(u64::MAX),
5446 ))
5447 }
5448}
5449
5450#[cfg(test)]
5451mod tests {
5452 use super::*;
5453 use crate::{
5454 protocol::internal::{InventoryResponse, Response},
5455 zakura::{
5456 legacy_gossip::{LegacyRequestFrame, LegacyRequestKind, LegacyResponseCodec},
5457 testkit::{await_until, LocalEndpointFactory, ZakuraTestNode},
5458 Event, HeaderSyncMisbehavior, PeerSession, ServicePeerLimits,
5459 LOCAL_MAX_CONTROL_FRAME_BYTES, LOCAL_MAX_MESSAGE_BYTES, MAX_BS_FRAME_BYTES,
5460 MAX_HS_MESSAGE_BYTES, ZAKURA_BLOCK_SYNC_STREAM_VERSION, ZAKURA_CAP_BLOCK_SYNC,
5461 ZAKURA_CAP_DISCOVERY, ZAKURA_CAP_HEADER_SYNC, ZAKURA_CAP_LEGACY_GOSSIP,
5462 ZAKURA_HEADER_SYNC_STREAM_VERSION,
5463 },
5464 P2pStack,
5465 };
5466 use iroh::{
5467 endpoint::Connection,
5468 protocol::{AcceptError, ProtocolHandler},
5469 };
5470 use zakura_chain::{
5471 block::{self, Block},
5472 serialization::{ZcashDeserialize, MAX_PROTOCOL_MESSAGE_LEN},
5473 transaction::{self, UnminedTxId},
5474 };
5475 use zakura_test::vectors::{BLOCK_MAINNET_GENESIS_BYTES, BLOCK_TESTNET_141042_BYTES};
5476
5477 #[tokio::test]
5496 async fn parked_block_sync_peer_gets_a_stream_when_its_cooldown_lapses() -> Result<(), BoxError>
5497 {
5498 const COOLDOWN: Duration = Duration::from_secs(3);
5499
5500 let _guard = zakura_test::init();
5501
5502 let genesis = Block::zcash_deserialize(&BLOCK_MAINNET_GENESIS_BYTES[..])?;
5503 let anchor = (block::Height(0), genesis.hash());
5504 let frontiers = FullStateFrontiers {
5505 finalized_height: block::Height(0),
5506 verified_block_tip: block::Height(0),
5507 verified_block_hash: genesis.hash(),
5508 };
5509 let node = |seed| {
5510 ZakuraTestNode::builder(seed)
5511 .header_sync_driver(Config::default().network, anchor, frontiers, Some(anchor))
5512 .spawn()
5513 };
5514 let dialer = node(140).await?;
5515 let listener = node(141).await?;
5516
5517 let listener_peer =
5518 ZakuraPeerId::new(listener.node_addr().await.node_id.as_bytes().to_vec())?;
5519 let block_sync = dialer
5520 .block_sync()
5521 .expect("the header-sync driver spawns the block-sync reactor");
5522
5523 block_sync.park_peer_for_test(&listener_peer, COOLDOWN);
5526 dialer
5527 .connect_native(&listener, Duration::from_secs(10))
5528 .await?;
5529
5530 assert_eq!(
5533 block_sync.peer_snapshot().outbound_peers,
5534 0,
5535 "a parked peer must not be given a block-sync stream while its cooldown runs",
5536 );
5537
5538 await_until(
5541 "block sync opens a stream to the peer whose park expired",
5542 Duration::from_secs(30),
5543 || block_sync.peer_snapshot().outbound_peers == 1,
5544 )
5545 .await?;
5546
5547 dialer.shutdown().await;
5548 listener.shutdown().await;
5549 Ok(())
5550 }
5551
5552 #[test]
5553 fn header_capability_mask_tracks_explicit_runtime_enablement() {
5554 let handler = ZakuraProtocolHandler::new(
5555 ZakuraSupervisorHandle::new(1),
5556 Network::Mainnet,
5557 ZakuraHandshakeConfig::for_network(&Network::Mainnet),
5558 ZakuraLocalLimits::from_config(&Config::default()),
5559 );
5560 handler.supported_capabilities.store(
5561 ZAKURA_CAP_DISCOVERY | ZAKURA_CAP_HEADER_SYNC,
5562 Ordering::Relaxed,
5563 );
5564
5565 handler.set_header_sync_enabled(false);
5566 assert_eq!(
5567 handler.current_handshake_config().supported_capabilities,
5568 ZAKURA_CAP_DISCOVERY
5569 );
5570
5571 handler.set_header_sync_enabled(true);
5572 assert_eq!(
5573 handler.current_handshake_config().supported_capabilities,
5574 ZAKURA_CAP_DISCOVERY | ZAKURA_CAP_HEADER_SYNC
5575 );
5576 }
5577
5578 #[test]
5579 fn header_capability_epoch_is_applied_once_and_rejects_stale_demand() {
5580 use zakura_node_services::sync_lifecycle::{
5581 BlockServiceDemand, HeaderServiceDemand, LifecycleEpoch, SyncServiceDemand,
5582 };
5583
5584 let disabled = SyncServiceDemand {
5585 header: HeaderServiceDemand::Disabled {
5586 runtime_epoch: LifecycleEpoch::INITIAL,
5587 },
5588 block: BlockServiceDemand::ServingAndApplying {
5589 apply_epoch: LifecycleEpoch::INITIAL,
5590 },
5591 };
5592 let mut epochs = HeaderCapabilityEpochs::from_initial(&disabled);
5593 let ready = SyncServiceDemand {
5594 header: HeaderServiceDemand::Enabled {
5595 capability_epoch: LifecycleEpoch::new(2),
5596 },
5597 ..disabled
5598 };
5599 assert_eq!(epochs.observe(&ready), Some(LifecycleEpoch::new(2)));
5600 assert_eq!(epochs.observe(&ready), None);
5601 let stale = SyncServiceDemand {
5602 header: HeaderServiceDemand::Enabled {
5603 capability_epoch: LifecycleEpoch::new(1),
5604 },
5605 ..disabled
5606 };
5607 assert_eq!(epochs.observe(&stale), None);
5608 }
5609
5610 #[tokio::test]
5611 async fn enabling_header_sync_reconnects_peers_negotiated_without_it() {
5612 let supervisor = ZakuraSupervisorHandle::new(4);
5613 let handler = ZakuraProtocolHandler::new(
5614 supervisor.clone(),
5615 Network::Mainnet,
5616 ZakuraHandshakeConfig::for_network(&Network::Mainnet),
5617 ZakuraLocalLimits::from_config(&Config::default()),
5618 );
5619 handler.set_header_sync_enabled(false);
5620
5621 let first_token = CancellationToken::new();
5622 let second_token = CancellationToken::new();
5623 register_test_peer(&supervisor, test_peer(31), first_token.clone()).await;
5624 register_test_peer(&supervisor, test_peer(32), second_token.clone()).await;
5625
5626 assert_eq!(
5627 enable_header_sync_and_renegotiate(
5628 &handler,
5629 &supervisor,
5630 zakura_node_services::sync_lifecycle::LifecycleEpoch::new(1),
5631 )
5632 .await,
5633 2
5634 );
5635 assert_ne!(
5636 handler.current_handshake_config().supported_capabilities & ZAKURA_CAP_HEADER_SYNC,
5637 0
5638 );
5639 assert!(first_token.is_cancelled());
5640 assert!(second_token.is_cancelled());
5641 }
5642
5643 #[tokio::test]
5649 async fn unset_listen_addr_binds_loopback_not_unspecified() {
5650 let builder = direct_endpoint_builder(SecretKey::generate(OsRng));
5651 let builder = bind_native_endpoint(builder, None);
5652 let endpoint = builder.bind().await.expect("loopback bind should succeed");
5653
5654 let sockets = endpoint.bound_sockets();
5655 assert!(
5656 !sockets.is_empty(),
5657 "endpoint should bind at least one socket"
5658 );
5659 for socket in &sockets {
5660 assert!(
5661 socket.ip().is_loopback(),
5662 "unset listen_addr must bind loopback only, but bound {socket} \
5663 (exposes P2P_V2_ALPN on all interfaces)"
5664 );
5665 }
5666
5667 endpoint.close().await;
5668 }
5669
5670 #[test]
5671 fn discovery_uses_external_ip_with_the_native_listen_port() {
5672 let secret_key = SecretKey::generate(OsRng);
5673 let mut config = Config::default();
5674 config.zakura.listen_addr = Some("0.0.0.0:8234".parse().expect("test address parses"));
5675 config.external_addr = Some("203.0.113.42:8233".parse().expect("test address parses"));
5676 config.zakura.bootstrap_peers.clear();
5677
5678 assert_eq!(
5679 discovery_direct_addrs(&config, secret_key.public()),
5680 vec!["203.0.113.42:8234".parse().expect("test address parses")]
5681 );
5682 }
5683
5684 #[test]
5685 fn discovery_uses_matching_local_bootstrap_address_and_counts_only_remote_peers() {
5686 let local_secret_key = SecretKey::generate(OsRng);
5687 let remote_secret_key = SecretKey::generate(OsRng);
5688 let local_node_id = local_secret_key.public();
5689 let remote_node_id = remote_secret_key.public();
5690 let local_entry = format!("{local_node_id}@198.51.100.7:8234");
5691 let remote_entry = format!("{remote_node_id}@198.51.100.8:8234");
5692 let mut config = Config::default();
5693 config.zakura.listen_addr = Some("0.0.0.0:8234".parse().expect("test address parses"));
5694 config.external_addr = None;
5695 config.zakura.bootstrap_peers = vec![
5696 local_entry,
5697 remote_entry.clone(),
5698 remote_entry,
5699 "malformed".to_string(),
5700 ];
5701
5702 assert_eq!(
5703 discovery_direct_addrs(&config, local_node_id),
5704 vec!["198.51.100.7:8234".parse().expect("test address parses")]
5705 );
5706 assert_eq!(
5707 remote_bootstrap_peer_count(&config.zakura.bootstrap_peers, local_node_id),
5708 1
5709 );
5710 }
5711
5712 #[derive(Debug, Clone)]
5713 struct CaptureConnection {
5714 connection_tx: mpsc::Sender<Connection>,
5715 stream_tx: mpsc::Sender<(SendStream, RecvStream)>,
5716 }
5717
5718 impl ProtocolHandler for CaptureConnection {
5719 async fn accept(&self, connection: Connection) -> Result<(), AcceptError> {
5720 let _ = self.connection_tx.send(connection.clone()).await;
5721 for _ in 0..2 {
5722 let Ok(streams) = connection.accept_bi().await else {
5723 break;
5724 };
5725 let _ = self.stream_tx.send(streams).await;
5726 }
5727 Ok(())
5728 }
5729 }
5730
5731 #[derive(Debug)]
5732 struct NoopService;
5733
5734 impl Service for NoopService {
5735 fn name(&self) -> &'static str {
5736 "noop"
5737 }
5738
5739 fn streams(&self) -> &[Stream] {
5740 &[]
5741 }
5742
5743 fn add_peer(&self, _peer: Peer) {}
5744
5745 fn remove_peer(&self, _peer: &ZakuraPeerId, _conn_id: ZakuraConnId) {}
5746 }
5747
5748 #[derive(Debug)]
5749 struct DeclaredStreamService {
5750 streams: Vec<Stream>,
5751 }
5752
5753 impl Service for DeclaredStreamService {
5754 fn name(&self) -> &'static str {
5755 "declared-stream"
5756 }
5757
5758 fn streams(&self) -> &[Stream] {
5759 &self.streams
5760 }
5761
5762 fn add_peer(&self, _peer: Peer) {}
5763
5764 fn remove_peer(&self, _peer: &ZakuraPeerId, _conn_id: ZakuraConnId) {}
5765 }
5766
5767 #[derive(Debug)]
5768 struct OrderedStreamService {
5769 stream: Stream,
5770 sessions: mpsc::UnboundedSender<(FramedRecv, FramedSend)>,
5771 }
5772
5773 impl Service for OrderedStreamService {
5774 fn name(&self) -> &'static str {
5775 "ordered-stream"
5776 }
5777
5778 fn streams(&self) -> &[Stream] {
5779 std::slice::from_ref(&self.stream)
5780 }
5781
5782 fn add_peer(&self, mut peer: Peer) {
5783 if let Some(session) = peer.take_stream(self.stream.kind) {
5784 let _ = self.sessions.send(session);
5785 }
5786 }
5787
5788 fn remove_peer(&self, _peer: &ZakuraPeerId, _conn_id: ZakuraConnId) {}
5789 }
5790
5791 #[derive(Debug)]
5792 struct GenerationGuardedRecordingService {
5793 streams: Vec<Stream>,
5794 active: std::sync::Mutex<HashMap<ZakuraPeerId, (ZakuraConnId, CancellationToken)>>,
5795 disconnected: std::sync::Mutex<Vec<(ZakuraPeerId, ZakuraConnId)>>,
5796 }
5797
5798 impl GenerationGuardedRecordingService {
5799 fn new(streams: Vec<Stream>) -> Arc<Self> {
5800 Arc::new(Self {
5801 streams,
5802 active: std::sync::Mutex::new(HashMap::new()),
5803 disconnected: std::sync::Mutex::new(Vec::new()),
5804 })
5805 }
5806
5807 fn active_conn(&self, peer: &ZakuraPeerId) -> Option<ZakuraConnId> {
5808 self.active
5809 .lock()
5810 .expect("recording service active map is never poisoned")
5811 .get(peer)
5812 .map(|(conn_id, _token)| *conn_id)
5813 }
5814
5815 fn active_token_cancelled(&self, peer: &ZakuraPeerId) -> Option<bool> {
5816 self.active
5817 .lock()
5818 .expect("recording service active map is never poisoned")
5819 .get(peer)
5820 .map(|(_conn_id, token)| token.is_cancelled())
5821 }
5822
5823 fn disconnected(&self) -> Vec<(ZakuraPeerId, ZakuraConnId)> {
5824 self.disconnected
5825 .lock()
5826 .expect("recording service disconnected list is never poisoned")
5827 .clone()
5828 }
5829 }
5830
5831 impl Service for GenerationGuardedRecordingService {
5832 fn name(&self) -> &'static str {
5833 "generation-recording"
5834 }
5835
5836 fn streams(&self) -> &[Stream] {
5837 &self.streams
5838 }
5839
5840 fn add_peer(&self, peer: Peer) {
5841 let mut active = self
5842 .active
5843 .lock()
5844 .expect("recording service active map is never poisoned");
5845 if active
5846 .get(&peer.id)
5847 .is_some_and(|(active_conn_id, _token)| *active_conn_id > peer.conn_id)
5848 {
5849 peer.service_cancel_token().cancel();
5850 return;
5851 }
5852 if let Some((_old_conn_id, old_token)) =
5853 active.insert(peer.id.clone(), (peer.conn_id, peer.service_cancel_token()))
5854 {
5855 old_token.cancel();
5856 }
5857 }
5858
5859 fn remove_peer(&self, peer: &ZakuraPeerId, conn_id: ZakuraConnId) {
5860 let removed = {
5861 let mut active = self
5862 .active
5863 .lock()
5864 .expect("recording service active map is never poisoned");
5865 if active
5866 .get(peer)
5867 .is_some_and(|(active_conn_id, _token)| *active_conn_id == conn_id)
5868 {
5869 active.remove(peer)
5870 } else {
5871 None
5872 }
5873 };
5874
5875 if let Some((_active_conn_id, token)) = removed {
5876 token.cancel();
5877 self.disconnected
5878 .lock()
5879 .expect("recording service disconnected list is never poisoned")
5880 .push((peer.clone(), conn_id));
5881 }
5882 }
5883 }
5884
5885 fn test_peer(byte: u8) -> ZakuraPeerId {
5886 ZakuraPeerId::new(vec![byte; 32]).expect("32-byte node id is valid")
5887 }
5888
5889 fn test_conn_id() -> ZakuraConnId {
5890 static NEXT_TEST_CONN_ID: AtomicU64 = AtomicU64::new(1);
5891 NEXT_TEST_CONN_ID.fetch_add(1, Ordering::Relaxed)
5892 }
5893
5894 fn registered_conn_id(registration: ZakuraRegistration) -> ZakuraConnId {
5895 match registration {
5896 ZakuraRegistration::Registered { conn_id, .. } => conn_id,
5897 other => panic!("expected peer registration to succeed, got {other:?}"),
5898 }
5899 }
5900
5901 #[tokio::test]
5902 async fn handoff_wait_returns_exact_generation_and_rejects_duplicates() {
5903 let supervisor = ZakuraSupervisorHandle::new(1);
5904 let peer_id = test_peer(6);
5905 let ZakuraConnectionRegistrationWaitStart::Waiting(registration_wait) = supervisor
5906 .begin_connection_registration_wait(&peer_id)
5907 .await
5908 else {
5909 panic!("the first handoff must reserve the peer identity");
5910 };
5911 assert!(
5912 matches!(
5913 supervisor
5914 .begin_connection_registration_wait(&peer_id)
5915 .await,
5916 ZakuraConnectionRegistrationWaitStart::Duplicate,
5917 ),
5918 "a concurrent handoff cannot share one connection registration",
5919 );
5920
5921 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
5922 let registration = supervisor
5923 .register(
5924 test_conn_id(),
5925 peer_id.clone(),
5926 Some("192.0.2.6".parse().expect("test IP parses")),
5927 [6; TRANSCRIPT_HASH_BYTES],
5928 ZakuraPeerHandle::new_for_tests(peer_id.clone(), outbound_tx),
5929 CancellationToken::new(),
5930 ZAKURA_CAP_LEGACY_GOSSIP,
5931 )
5932 .await;
5933 let conn_id = registered_conn_id(registration);
5934
5935 assert_eq!(
5936 registration_wait.wait(Duration::from_secs(1)).await,
5937 Some(conn_id),
5938 );
5939 assert!(
5940 matches!(
5941 supervisor
5942 .begin_connection_registration_wait(&peer_id)
5943 .await,
5944 ZakuraConnectionRegistrationWaitStart::Duplicate,
5945 ),
5946 "an incumbent connection makes a later legacy handoff a duplicate",
5947 );
5948 }
5949
5950 #[test]
5951 fn native_duplicate_tie_breaker_converges_for_simultaneous_open() {
5952 let node_a = LocalEndpointFactory::secret_key(1).public();
5953 let node_b = LocalEndpointFactory::secret_key(2).public();
5954 let a_outbound =
5955 native_connection_transcript_hash(ServicePeerDirection::Outbound, &node_a, &node_b);
5956 let a_inbound =
5957 native_connection_transcript_hash(ServicePeerDirection::Inbound, &node_a, &node_b);
5958 let b_outbound =
5959 native_connection_transcript_hash(ServicePeerDirection::Outbound, &node_b, &node_a);
5960 let b_inbound =
5961 native_connection_transcript_hash(ServicePeerDirection::Inbound, &node_b, &node_a);
5962
5963 assert_eq!(a_outbound, b_inbound);
5964 assert_eq!(a_inbound, b_outbound);
5965 assert_ne!(a_outbound, a_inbound);
5966
5967 let winning_key = a_outbound.min(a_inbound);
5968 let losing_key = a_outbound.max(a_inbound);
5969 let peer = test_peer(7);
5970 let mut supervisor_a = ZakuraPeerSupervisor::default();
5971 let mut supervisor_b = ZakuraPeerSupervisor::default();
5972 assert!(matches!(
5973 supervisor_a.register_authenticated(peer.clone(), a_outbound),
5974 AuthenticatedPeerRegistration::Registered
5975 ));
5976 let _ = supervisor_a.register_authenticated(peer.clone(), a_inbound);
5977 assert!(matches!(
5978 supervisor_b.register_authenticated(peer.clone(), b_inbound),
5979 AuthenticatedPeerRegistration::Registered
5980 ));
5981 let _ = supervisor_b.register_authenticated(peer.clone(), b_outbound);
5982
5983 assert!(matches!(
5984 supervisor_a.register_authenticated(peer.clone(), winning_key),
5985 AuthenticatedPeerRegistration::Duplicate
5986 ));
5987 assert!(matches!(
5988 supervisor_b.register_authenticated(peer.clone(), winning_key),
5989 AuthenticatedPeerRegistration::Duplicate
5990 ));
5991 assert!(matches!(
5992 supervisor_a.register_authenticated(peer.clone(), losing_key),
5993 AuthenticatedPeerRegistration::Duplicate
5994 ));
5995 assert!(matches!(
5996 supervisor_b.register_authenticated(peer, losing_key),
5997 AuthenticatedPeerRegistration::Duplicate
5998 ));
5999 }
6000
6001 #[tokio::test]
6002 async fn upgraded_winner_registers_second_loser_cleanup_is_generation_guarded() {
6003 let supervisor = ZakuraSupervisorHandle::new(1);
6004 let peer = test_peer(11);
6005 let remote_ip: IpAddr = "203.0.113.11".parse().expect("test ip parses");
6006 let losing_trace_conn = 100;
6007 let winning_trace_conn = 1;
6008 let losing_hash = [0x80; TRANSCRIPT_HASH_BYTES];
6009 let winning_hash = [0x10; TRANSCRIPT_HASH_BYTES];
6010 let losing_token = CancellationToken::new();
6011 let winning_token = CancellationToken::new();
6012
6013 async fn register_with_hash(
6014 supervisor: &ZakuraSupervisorHandle,
6015 conn_id: ZakuraConnId,
6016 peer: &ZakuraPeerId,
6017 remote_ip: IpAddr,
6018 transcript_hash: [u8; TRANSCRIPT_HASH_BYTES],
6019 token: CancellationToken,
6020 ) -> ZakuraRegistration {
6021 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
6022 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
6023 supervisor
6024 .register(
6025 conn_id,
6026 peer.clone(),
6027 Some(remote_ip),
6028 transcript_hash,
6029 outbound_handle,
6030 token,
6031 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
6032 )
6033 .await
6034 }
6035
6036 let losing_conn = registered_conn_id(
6037 register_with_hash(
6038 &supervisor,
6039 losing_trace_conn,
6040 &peer,
6041 remote_ip,
6042 losing_hash,
6043 losing_token.clone(),
6044 )
6045 .await,
6046 );
6047 let winning_conn = registered_conn_id(
6048 register_with_hash(
6049 &supervisor,
6050 winning_trace_conn,
6051 &peer,
6052 remote_ip,
6053 winning_hash,
6054 winning_token.clone(),
6055 )
6056 .await,
6057 );
6058 assert!(
6059 winning_trace_conn < losing_trace_conn,
6060 "the test deliberately inverts raw trace-id order",
6061 );
6062 assert!(
6063 winning_conn > losing_conn,
6064 "service generations follow successful registration order",
6065 );
6066
6067 assert!(
6068 losing_token.is_cancelled(),
6069 "upgraded winner explicitly cancels the incumbent loser",
6070 );
6071 assert!(
6072 !winning_token.is_cancelled(),
6073 "the winning replacement remains live",
6074 );
6075
6076 supervisor.deregister(&peer, losing_conn).await;
6077
6078 {
6079 let state = supervisor.inner.lock().await;
6080 let entry = state
6081 .active_by_peer
6082 .get(&peer)
6083 .expect("winner remains registered after loser cleanup");
6084 assert_eq!(entry.conn_id, winning_conn);
6085 assert_eq!(entry.remote_ip, Some(remote_ip));
6086 assert_eq!(state.active_by_ip.get(&remote_ip), Some(&1));
6087 state.debug_assert_accounting();
6088 }
6089 assert_eq!(supervisor.registered_ids().await, vec![peer.clone()]);
6090 assert_eq!(supervisor.outbound_peer_handles().await.len(), 1);
6091
6092 supervisor.deregister(&peer, winning_conn).await;
6093 let state = supervisor.inner.lock().await;
6094 assert!(state.active_by_peer.is_empty());
6095 assert!(state.active_by_ip.is_empty());
6096 state.debug_assert_accounting();
6097 }
6098
6099 #[tokio::test]
6100 async fn upgraded_replacement_from_full_different_ip_is_rejected() {
6101 let supervisor = ZakuraSupervisorHandle::new(1);
6102 let peer_a = test_peer(21);
6103 let peer_b = test_peer(22);
6104 let ip1: IpAddr = "203.0.113.21".parse().expect("test ip parses");
6105 let ip2: IpAddr = "203.0.113.22".parse().expect("test ip parses");
6106 let peer_a_losing_conn = test_conn_id();
6107 let peer_b_conn = test_conn_id();
6108 let peer_a_winning_conn = test_conn_id();
6109 let peer_a_losing_hash = [0x80; TRANSCRIPT_HASH_BYTES];
6110 let peer_a_winning_hash = [0x10; TRANSCRIPT_HASH_BYTES];
6111 let peer_b_hash = [0x40; TRANSCRIPT_HASH_BYTES];
6112 let peer_a_losing_token = CancellationToken::new();
6113 let peer_b_token = CancellationToken::new();
6114 let peer_a_winning_token = CancellationToken::new();
6115
6116 let peer_a_losing_generation = registered_conn_id(
6117 register_test_peer_with_hash_and_ip(
6118 &supervisor,
6119 peer_a_losing_conn,
6120 &peer_a,
6121 Some(ip1),
6122 peer_a_losing_hash,
6123 peer_a_losing_token.clone(),
6124 )
6125 .await,
6126 );
6127 let peer_b_generation = registered_conn_id(
6128 register_test_peer_with_hash_and_ip(
6129 &supervisor,
6130 peer_b_conn,
6131 &peer_b,
6132 Some(ip2),
6133 peer_b_hash,
6134 peer_b_token.clone(),
6135 )
6136 .await,
6137 );
6138
6139 let registration = register_test_peer_with_hash_and_ip(
6140 &supervisor,
6141 peer_a_winning_conn,
6142 &peer_a,
6143 Some(ip2),
6144 peer_a_winning_hash,
6145 peer_a_winning_token.clone(),
6146 )
6147 .await;
6148 assert!(
6149 matches!(
6150 registration,
6151 ZakuraRegistration::Rejected(ZakuraRejectReason::ResourceLimit)
6152 ),
6153 "an upgraded same-peer replacement from a different full IP bucket must be rejected",
6154 );
6155 assert!(
6156 !peer_a_losing_token.is_cancelled(),
6157 "the incumbent stays live because the target IP bucket has no room",
6158 );
6159 assert!(
6160 !peer_b_token.is_cancelled(),
6161 "the unrelated peer that fills the target IP bucket stays live",
6162 );
6163 assert!(
6164 !peer_a_winning_token.is_cancelled(),
6165 "the rejected replacement was never registered",
6166 );
6167
6168 let state = supervisor.inner.lock().await;
6169 assert_eq!(
6170 state
6171 .active_by_peer
6172 .get(&peer_a)
6173 .expect("peer A incumbent remains registered")
6174 .conn_id,
6175 peer_a_losing_generation
6176 );
6177 assert_eq!(
6178 state
6179 .active_by_peer
6180 .get(&peer_b)
6181 .expect("peer B remains registered")
6182 .conn_id,
6183 peer_b_generation
6184 );
6185 assert_eq!(state.active_by_ip.get(&ip1), Some(&1));
6186 assert_eq!(state.active_by_ip.get(&ip2), Some(&1));
6187 state.debug_assert_accounting();
6188 }
6189
6190 #[tokio::test]
6191 async fn loser_cleanup_does_not_disconnect_upgraded_winner_service_session() {
6192 let supervisor = ZakuraSupervisorHandle::new(2);
6193 let peer = test_peer(23);
6194 let remote_ip: IpAddr = "203.0.113.23".parse().expect("test ip parses");
6195 let losing_conn = test_conn_id();
6196 let winning_conn = test_conn_id();
6197 let losing_hash = [0x80; TRANSCRIPT_HASH_BYTES];
6198 let winning_hash = [0x10; TRANSCRIPT_HASH_BYTES];
6199 let losing_token = CancellationToken::new();
6200 let winning_token = CancellationToken::new();
6201
6202 let losing_conn = registered_conn_id(
6203 register_test_peer_with_hash_and_ip(
6204 &supervisor,
6205 losing_conn,
6206 &peer,
6207 Some(remote_ip),
6208 losing_hash,
6209 losing_token.clone(),
6210 )
6211 .await,
6212 );
6213 let winning_conn = registered_conn_id(
6214 register_test_peer_with_hash_and_ip(
6215 &supervisor,
6216 winning_conn,
6217 &peer,
6218 Some(remote_ip),
6219 winning_hash,
6220 winning_token.clone(),
6221 )
6222 .await,
6223 );
6224 assert!(
6225 winning_conn > losing_conn,
6226 "winner must have a later service generation",
6227 );
6228
6229 let stream = Stream {
6230 kind: 61,
6231 version: 1,
6232 frame_cap: 1024,
6233 capability: ZAKURA_CAP_LEGACY_GOSSIP,
6234 mode: StreamMode::Ordered,
6235 };
6236 let service = GenerationGuardedRecordingService::new(vec![stream]);
6237 let registry = Arc::new(
6238 ServiceRegistry::new(vec![service.clone()])
6239 .expect("test service declares a valid stream"),
6240 );
6241 let (_inbound_tx, inbound_rx) = crate::zakura::framed_channel(1);
6242 let (outbound_tx, _outbound_rx) = crate::zakura::framed_channel(1);
6243 let streams = HashMap::from([(stream.kind, (inbound_rx, outbound_tx))]);
6244 let winning_service_token = CancellationToken::new();
6245
6246 registry.add_peer(Peer::new_with_conn_id_and_direction(
6247 winning_conn,
6248 peer.clone(),
6249 Some(remote_ip),
6250 ZAKURA_CAP_LEGACY_GOSSIP,
6251 ServicePeerDirection::Inbound,
6252 streams,
6253 winning_service_token.clone(),
6254 ));
6255 assert_eq!(service.active_conn(&peer), Some(winning_conn));
6256
6257 let (_stale_inbound_tx, stale_inbound_rx) = crate::zakura::framed_channel(1);
6258 let (stale_outbound_tx, _stale_outbound_rx) = crate::zakura::framed_channel(1);
6259 let stale_streams = HashMap::from([(stream.kind, (stale_inbound_rx, stale_outbound_tx))]);
6260 let stale_service_token = CancellationToken::new();
6261 registry.add_peer(Peer::new_with_conn_id_and_direction(
6262 losing_conn,
6263 peer.clone(),
6264 Some(remote_ip),
6265 ZAKURA_CAP_LEGACY_GOSSIP,
6266 ServicePeerDirection::Inbound,
6267 stale_streams,
6268 stale_service_token.clone(),
6269 ));
6270 assert_eq!(
6271 service.active_conn(&peer),
6272 Some(winning_conn),
6273 "stale loser add must not overwrite the winner's service session",
6274 );
6275 assert!(
6276 matches!(service.active_token_cancelled(&peer), Some(false)),
6277 "stale loser add must not cancel the winner's service token",
6278 );
6279
6280 registry.remove_peer(&peer, losing_conn, ZAKURA_CAP_LEGACY_GOSSIP);
6281 supervisor.deregister(&peer, losing_conn).await;
6282
6283 assert_eq!(
6284 service.active_conn(&peer),
6285 Some(winning_conn),
6286 "stale loser cleanup must not remove the winner's service session",
6287 );
6288 assert!(
6289 service.disconnected().is_empty(),
6290 "no service disconnect should be emitted for the winner generation",
6291 );
6292 assert!(
6293 matches!(service.active_token_cancelled(&peer), Some(false)),
6294 "the winner's service session token must remain live",
6295 );
6296
6297 {
6298 let state = supervisor.inner.lock().await;
6299 let entry = state
6300 .active_by_peer
6301 .get(&peer)
6302 .expect("winner remains registered after loser cleanup");
6303 assert_eq!(entry.conn_id, winning_conn);
6304 assert_eq!(state.active_by_ip.get(&remote_ip), Some(&1));
6305 state.debug_assert_accounting();
6306 }
6307
6308 registry.remove_peer(&peer, winning_conn, ZAKURA_CAP_LEGACY_GOSSIP);
6309 supervisor.deregister(&peer, winning_conn).await;
6310 }
6311
6312 #[tokio::test]
6313 async fn registered_peer_cleanup_guard_drop_deregisters_and_allows_redial() {
6314 let supervisor = ZakuraSupervisorHandle::new(2);
6315 let peer = test_peer(24);
6316 let transcript_hash = [0x24; TRANSCRIPT_HASH_BYTES];
6317 let disconnect_token = CancellationToken::new();
6318 let conn_id = registered_conn_id(
6319 register_test_peer_with_hash_and_ip(
6320 &supervisor,
6321 test_conn_id(),
6322 &peer,
6323 None,
6324 transcript_hash,
6325 disconnect_token.clone(),
6326 )
6327 .await,
6328 );
6329
6330 let duplicate = register_test_peer_with_hash_and_ip(
6331 &supervisor,
6332 test_conn_id(),
6333 &peer,
6334 None,
6335 transcript_hash,
6336 CancellationToken::new(),
6337 )
6338 .await;
6339 assert!(
6340 matches!(duplicate, ZakuraRegistration::Duplicate { .. }),
6341 "an exact same-direction redial is a duplicate while the incumbent is registered",
6342 );
6343
6344 let stream = Stream {
6345 kind: 62,
6346 version: 1,
6347 frame_cap: 1024,
6348 capability: ZAKURA_CAP_LEGACY_GOSSIP,
6349 mode: StreamMode::Ordered,
6350 };
6351 let service = GenerationGuardedRecordingService::new(vec![stream]);
6352 let registry = Arc::new(
6353 ServiceRegistry::new(vec![service.clone()])
6354 .expect("test service declares a valid stream"),
6355 );
6356 let (_inbound_tx, inbound_rx) = crate::zakura::framed_channel(1);
6357 let (outbound_tx, _outbound_rx) = crate::zakura::framed_channel(1);
6358 registry.add_peer(Peer::new_with_conn_id_and_direction(
6359 conn_id,
6360 peer.clone(),
6361 None,
6362 ZAKURA_CAP_LEGACY_GOSSIP,
6363 ServicePeerDirection::Inbound,
6364 HashMap::from([(stream.kind, (inbound_rx, outbound_tx))]),
6365 disconnect_token.clone(),
6366 ));
6367 assert_eq!(service.active_conn(&peer), Some(conn_id));
6368
6369 let mut cleanup_guard = RegisteredPeerCleanupGuard::new(
6370 supervisor.clone(),
6371 registry,
6372 peer.clone(),
6373 conn_id,
6374 disconnect_token.clone(),
6375 );
6376 cleanup_guard.add_admitted_capabilities(ZAKURA_CAP_LEGACY_GOSSIP);
6377 drop(cleanup_guard);
6378
6379 assert!(
6380 disconnect_token.is_cancelled(),
6381 "drop cleanup cancels the registered connection token",
6382 );
6383 assert_eq!(
6384 service.active_conn(&peer),
6385 None,
6386 "drop cleanup removes admitted service state synchronously",
6387 );
6388 assert_eq!(service.disconnected(), vec![(peer.clone(), conn_id)]);
6389
6390 tokio::time::timeout(Duration::from_secs(1), async {
6391 loop {
6392 if supervisor.registered_ids().await.is_empty() {
6393 break;
6394 }
6395 tokio::task::yield_now().await;
6396 }
6397 })
6398 .await
6399 .expect("drop cleanup deregisters the supervisor entry promptly");
6400
6401 let redial = register_test_peer_with_hash_and_ip(
6402 &supervisor,
6403 test_conn_id(),
6404 &peer,
6405 None,
6406 transcript_hash,
6407 CancellationToken::new(),
6408 )
6409 .await;
6410 assert!(
6411 matches!(redial, ZakuraRegistration::Registered { .. }),
6412 "after drop cleanup, the deterministic same-direction redial can register again",
6413 );
6414 }
6415
6416 #[tokio::test]
6417 async fn registration_generation_exhaustion_rejects_without_mutating_state() {
6418 let supervisor = ZakuraSupervisorHandle::new(2);
6419 let peer = test_peer(25);
6420
6421 {
6422 let mut state = supervisor.inner.lock().await;
6423 state.next_registration_id = u64::MAX;
6424 }
6425
6426 let registration = register_test_peer_with_hash_and_ip(
6427 &supervisor,
6428 test_conn_id(),
6429 &peer,
6430 None,
6431 [0x25; TRANSCRIPT_HASH_BYTES],
6432 CancellationToken::new(),
6433 )
6434 .await;
6435
6436 assert!(
6437 matches!(
6438 registration,
6439 ZakuraRegistration::Rejected(ZakuraRejectReason::TemporaryUnavailable)
6440 ),
6441 "exhausted registration generations reject the new connection instead of panicking",
6442 );
6443 assert!(
6444 supervisor.registered_ids().await.is_empty(),
6445 "failed registration must not mutate active supervisor state",
6446 );
6447 }
6448
6449 #[test]
6450 fn ordered_stream_collision_winner_is_mirror_stable() {
6451 let node_a = LocalEndpointFactory::secret_key(1).public();
6456 let node_b = LocalEndpointFactory::secret_key(2).public();
6457 assert_ne!(node_a, node_b);
6458 assert_ne!(
6459 i_open_collision_winner(&node_a, &node_b),
6460 i_open_collision_winner(&node_b, &node_a),
6461 "exactly one side must win a same-kind ordered-stream collision",
6462 );
6463 let a_wins = node_a.as_bytes() < node_b.as_bytes();
6465 assert_eq!(i_open_collision_winner(&node_a, &node_b), a_wins);
6466 }
6467
6468 fn header_sync_test_session(peer: ZakuraPeerId) -> (PeerSession, crate::zakura::FramedRecv) {
6469 let (send, recv) = crate::zakura::framed_channel(32);
6470 (
6471 PeerSession::from_parts(peer, send, CancellationToken::new()),
6472 recv,
6473 )
6474 }
6475
6476 fn header_sync_startup(shutdown: CancellationToken) -> HeaderSyncStartup {
6477 let network = Network::Mainnet;
6478 let anchor = (block::Height(0), network.genesis_hash());
6479 let mut startup = HeaderSyncStartup::new(
6480 network,
6481 anchor,
6482 FullStateFrontiers {
6483 finalized_height: anchor.0,
6484 verified_block_tip: anchor.0,
6485 verified_block_hash: anchor.1,
6486 },
6487 Some(anchor),
6488 ZakuraHeaderSyncConfig::default(),
6489 LOCAL_MAX_MESSAGE_BYTES,
6490 );
6491 startup.shutdown = shutdown;
6492 startup
6493 }
6494
6495 async fn register_test_peer(
6496 supervisor: &ZakuraSupervisorHandle,
6497 peer: ZakuraPeerId,
6498 disconnect_token: CancellationToken,
6499 ) {
6500 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
6501 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
6502 let registration = supervisor
6503 .register(
6504 test_conn_id(),
6505 peer.clone(),
6506 None,
6507 [peer.as_bytes()[0]; TRANSCRIPT_HASH_BYTES],
6508 outbound_handle,
6509 disconnect_token,
6510 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
6511 )
6512 .await;
6513
6514 assert!(
6515 matches!(registration, ZakuraRegistration::Registered { .. }),
6516 "test peer should register once"
6517 );
6518 }
6519
6520 async fn register_test_peer_with_hash_and_ip(
6521 supervisor: &ZakuraSupervisorHandle,
6522 conn_id: ZakuraConnId,
6523 peer: &ZakuraPeerId,
6524 remote_ip: Option<IpAddr>,
6525 transcript_hash: [u8; TRANSCRIPT_HASH_BYTES],
6526 disconnect_token: CancellationToken,
6527 ) -> ZakuraRegistration {
6528 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
6529 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
6530 supervisor
6531 .register(
6532 conn_id,
6533 peer.clone(),
6534 remote_ip,
6535 transcript_hash,
6536 outbound_handle,
6537 disconnect_token,
6538 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
6539 )
6540 .await
6541 }
6542
6543 #[test]
6544 fn local_limits_clamp_negotiated_values_down() {
6545 let config = Config::default();
6546 let limits = ZakuraLocalLimits::from_config(&config);
6547 let negotiated = ZakuraAcceptedLimits {
6548 max_frame_bytes: u32::MAX,
6549 max_message_bytes: u32::MAX,
6550 max_open_streams: u16::MAX,
6551 max_inbound_queue_depth: u16::MAX,
6552 idle_timeout_millis: u32::MAX,
6553 };
6554
6555 let clamped = limits.clamp(&negotiated);
6556
6557 assert_eq!(clamped.max_frame_bytes, limits.max_frame_bytes);
6558 assert_eq!(clamped.max_message_bytes, limits.max_message_bytes);
6559 assert_eq!(clamped.max_open_streams, limits.max_open_streams);
6560 assert_eq!(
6561 clamped.max_inbound_queue_depth,
6562 limits.max_inbound_queue_depth
6563 );
6564 assert!(clamped.idle_timeout < limits.quic_idle_timeout);
6565 }
6566
6567 #[test]
6568 fn inbound_queue_depth_is_split_across_ordered_streams() {
6569 assert_eq!(per_stream_inbound_queue_depth(64, 2), 32);
6570 assert_eq!(per_stream_inbound_queue_depth(63, 2), 31);
6571 assert_eq!(per_stream_inbound_queue_depth(64, 0), 64);
6572 assert!(per_stream_inbound_queue_depth(63, 2) * 2 <= 63);
6573 }
6574
6575 #[test]
6576 fn request_response_only_peers_do_not_use_ordered_stream_freshness_reaper() {
6577 assert!(!should_run_freshness_reaper(0, 1));
6578 assert!(should_run_freshness_reaper(1, 1));
6579 assert!(should_run_freshness_reaper(1, 0));
6580 assert!(should_run_freshness_reaper(0, 0));
6581 }
6582
6583 #[tokio::test]
6584 async fn v2_p2p_false_leaves_header_sync_disabled() -> Result<(), BoxError> {
6585 let _guard = zakura_test::init();
6586 let config = Config::for_test(P2pStack::Legacy);
6587
6588 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
6589 Arc::new(NoopService) as Arc<dyn Service>
6590 })
6591 .await?;
6592
6593 assert!(endpoint.is_none());
6594 Ok(())
6595 }
6596
6597 #[tokio::test]
6598 async fn v2_p2p_true_starts_header_sync_handle() -> Result<(), BoxError> {
6599 let _guard = zakura_test::init();
6600 let config = Config::for_test(P2pStack::Dual);
6601
6602 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
6603 Arc::new(NoopService) as Arc<dyn Service>
6604 })
6605 .await?
6606 .expect("v2_p2p is enabled in test config");
6607
6608 assert!(endpoint.header_sync().is_some());
6609 endpoint.shutdown().await;
6610 Ok(())
6611 }
6612
6613 #[tokio::test]
6614 async fn endpoint_shutdown_stops_header_sync_task() -> Result<(), BoxError> {
6615 let _guard = zakura_test::init();
6616 let config = Config::for_test(P2pStack::Dual);
6617 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
6618 Arc::new(NoopService) as Arc<dyn Service>
6619 })
6620 .await?
6621 .expect("v2_p2p is enabled in test config");
6622 let header_sync = endpoint
6623 .header_sync()
6624 .expect("header sync handle exists for a v2 endpoint");
6625
6626 endpoint.shutdown().await;
6627
6628 let (session, _recv) = header_sync_test_session(
6629 ZakuraPeerId::new(vec![5u8; 32]).expect("32-byte node id is valid"),
6630 );
6631 let send_result = tokio::time::timeout(
6632 Duration::from_secs(1),
6633 header_sync.send(Event::PeerConnected(session)),
6634 )
6635 .await
6636 .expect("send returns promptly after header-sync shutdown");
6637 assert!(send_result.is_err());
6638
6639 Ok(())
6640 }
6641
6642 #[tokio::test]
6650 async fn endpoint_shutdown_stops_maintained_native_dial_loop() -> Result<(), BoxError> {
6651 let _guard = zakura_test::init();
6652 let config = Config::for_test(P2pStack::Dual);
6653 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
6654 Arc::new(NoopService) as Arc<dyn Service>
6655 })
6656 .await?
6657 .expect("v2_p2p is enabled in test config");
6658
6659 let unreachable_addr: SocketAddr = "192.0.2.1:65535".parse().expect("valid test address");
6662 let unreachable = NodeAddr::new(LocalEndpointFactory::secret_key(987_654).public())
6663 .with_direct_addresses([unreachable_addr]);
6664 let dial = endpoint.spawn_native_dial(unreachable);
6665
6666 tokio::time::sleep(Duration::from_millis(50)).await;
6668 assert!(
6669 !dial.is_finished(),
6670 "maintained dial loop to an unreachable peer should still be running"
6671 );
6672
6673 endpoint.shutdown().await;
6674
6675 tokio::time::timeout(Duration::from_secs(5), dial)
6676 .await
6677 .expect("maintained native dial loop must exit after endpoint shutdown")
6678 .expect("native dial task must not panic");
6679 Ok(())
6680 }
6681
6682 #[tokio::test]
6688 async fn endpoint_shutdown_stops_discovery_candidate_dialer() -> Result<(), BoxError> {
6689 let _guard = zakura_test::init();
6690 let config = Config::for_test(P2pStack::Dual);
6691 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
6692 Arc::new(NoopService) as Arc<dyn Service>
6693 })
6694 .await?
6695 .expect("v2_p2p is enabled in test config");
6696
6697 let limits = ZakuraLocalLimits::from_config(&config);
6698 let handshake = ZakuraHandshakeConfig::for_network(&config.network);
6699 let discovery = crate::zakura::discovery::build_discovery_handle(
6700 SecretKey::generate(OsRng),
6701 Vec::new(),
6702 crate::zakura::discovery::default_advertised_services(),
6703 &handshake,
6704 limits.max_connections,
6705 0,
6706 endpoint.supervisor().subscribe(),
6707 )?;
6708 let dialer = crate::zakura::discovery::spawn_native_discovery_dialer(
6709 endpoint.clone(),
6710 discovery,
6711 limits,
6712 Vec::new(),
6713 );
6714
6715 tokio::time::sleep(Duration::from_millis(50)).await;
6716 assert!(
6717 !dialer.is_finished(),
6718 "discovery candidate dialer should still be running"
6719 );
6720
6721 endpoint.shutdown().await;
6722
6723 tokio::time::timeout(Duration::from_secs(5), dialer)
6724 .await
6725 .expect("discovery candidate dialer must exit after endpoint shutdown")
6726 .expect("discovery dialer task must not panic");
6727 Ok(())
6728 }
6729
6730 #[tokio::test(start_paused = true)]
6742 async fn failed_legacy_upgrade_does_not_leak_maintained_dial() -> Result<(), BoxError> {
6743 let _guard = zakura_test::init();
6744 let config = Config::for_test(P2pStack::Dual);
6745 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
6746 Arc::new(NoopService) as Arc<dyn Service>
6747 })
6748 .await?
6749 .expect("v2_p2p is enabled in test config");
6750
6751 let node_id = LocalEndpointFactory::secret_key(0x0BAD_C0DE)
6755 .public()
6756 .as_bytes()
6757 .to_vec();
6758 let peer_id = ZakuraPeerId::new(node_id.clone()).expect("32-byte node id is valid");
6759 let direct_addresses = vec![b"192.0.2.1:1".to_vec()];
6760
6761 let connector =
6762 crate::zakura::ZakuraHandshakeConnector::new_with_endpoint(endpoint.clone());
6763 let handoff = connector
6764 .spawn_zakura_dial_to_hints_and_wait(&peer_id, &node_id, &direct_addresses)
6765 .await;
6766
6767 assert!(
6768 matches!(handoff, crate::zakura::ZakuraNativeHandoff::Failed),
6769 "an unreachable upgrade peer must not report a completed hand-off",
6770 );
6771 assert!(
6772 endpoint
6773 .upgrade_dials
6774 .lock()
6775 .expect("Zakura upgrade dial registry mutex is never poisoned")
6776 .is_empty(),
6777 "a failed legacy upgrade leaked a maintained native dial / upgrade_dials entry",
6778 );
6779
6780 endpoint.shutdown().await;
6781 Ok(())
6782 }
6783
6784 #[tokio::test]
6785 async fn maintained_upgrade_dial_keeps_its_original_handoff_owner() -> Result<(), BoxError> {
6786 let _guard = zakura_test::init();
6787 let config = Config::for_test(P2pStack::Dual);
6788 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
6789 Arc::new(NoopService) as Arc<dyn Service>
6790 })
6791 .await?
6792 .expect("v2_p2p is enabled in test config");
6793 let node_id = LocalEndpointFactory::secret_key(0x0BAD_CAFE).public();
6794 let peer_id = ZakuraPeerId::new(node_id.as_bytes().to_vec())?;
6795 let node_addr = NodeAddr::new(node_id)
6796 .with_direct_addresses(["192.0.2.2:1".parse().expect("test direct address parses")]);
6797
6798 assert_eq!(
6799 endpoint.start_upgrade_native_dial(node_addr.clone()),
6800 crate::zakura::ZakuraUpgradeDialStart::Started,
6801 );
6802 assert_eq!(
6803 endpoint.start_upgrade_native_dial(node_addr),
6804 crate::zakura::ZakuraUpgradeDialStart::AlreadyRunning,
6805 );
6806
6807 endpoint.cancel_upgrade_native_dial(&peer_id);
6808 endpoint.shutdown().await;
6809 Ok(())
6810 }
6811
6812 #[tokio::test]
6813 async fn custom_ordered_service_round_trips_over_zakura() -> Result<(), BoxError> {
6814 const CUSTOM_STREAM: Stream = Stream {
6815 kind: 64,
6816 version: 1,
6817 frame_cap: 64 * 1024,
6818 capability: 1 << 16,
6819 mode: StreamMode::Ordered,
6820 };
6821
6822 let _guard = zakura_test::init();
6823 let service_id = ZakuraServiceId::new("zakura.test.ordered.v1")?;
6824 let listen_addr = "127.0.0.1:0".parse().expect("valid loopback address");
6825 let server_identity = tempfile::tempdir()?;
6826 let client_identity = tempfile::tempdir()?;
6827 let (server_sessions, mut server_session_rx) = mpsc::unbounded_channel();
6828 let (client_sessions, mut client_session_rx) = mpsc::unbounded_channel();
6829
6830 let mut server_config = Config::for_test(P2pStack::Dual);
6831 server_config.identity_dir = server_identity.path().to_owned();
6832 server_config.zakura.listen_addr = Some(listen_addr);
6833 server_config.zakura.bootstrap_peers.clear();
6834 let server = spawn_zakura_endpoint_with_services(
6835 &server_config,
6836 |_supervisor, _trace| Arc::new(NoopService),
6837 None,
6838 vec![CustomService {
6839 service: Arc::new(OrderedStreamService {
6840 stream: CUSTOM_STREAM,
6841 sessions: server_sessions,
6842 }),
6843 provides: vec![service_id.clone()],
6844 seeks: Vec::new(),
6845 }],
6846 )
6847 .await?
6848 .expect("test server uses Zakura");
6849 let server_addr = server.node_addr().await;
6850 let server_direct = server_addr
6851 .direct_addresses()
6852 .copied()
6853 .find(|addr| addr.ip().is_loopback())
6854 .ok_or("test server has no loopback address")?;
6855
6856 let mut client_config = Config::for_test(P2pStack::Dual);
6857 client_config.identity_dir = client_identity.path().to_owned();
6858 client_config.zakura.listen_addr = Some(listen_addr);
6859 client_config.zakura.bootstrap_peers =
6860 vec![format!("{}@{server_direct}", server_addr.node_id)];
6861 let client = spawn_zakura_endpoint_with_services(
6862 &client_config,
6863 |_supervisor, _trace| Arc::new(NoopService),
6864 None,
6865 vec![CustomService {
6866 service: Arc::new(OrderedStreamService {
6867 stream: CUSTOM_STREAM,
6868 sessions: client_sessions,
6869 }),
6870 provides: Vec::new(),
6871 seeks: vec![service_id],
6872 }],
6873 )
6874 .await?
6875 .expect("test client uses Zakura");
6876
6877 let round_trip = timeout(Duration::from_secs(10), async {
6878 let (mut server_recv, server_send) = server_session_rx
6879 .recv()
6880 .await
6881 .ok_or_else(|| -> BoxError { "server custom stream did not open".into() })?;
6882 let (mut client_recv, client_send) = client_session_rx
6883 .recv()
6884 .await
6885 .ok_or_else(|| -> BoxError { "client custom stream did not open".into() })?;
6886 let outbound = Frame {
6887 message_type: 1_001,
6888 flags: 3,
6889 payload: b"custom client frame".to_vec(),
6890 };
6891 client_send.send(outbound.clone()).await?;
6892 let received = server_recv
6893 .recv()
6894 .await
6895 .ok_or_else(|| -> BoxError { "server custom stream closed".into() })?;
6896 assert_eq!(received, outbound);
6897
6898 let response = Frame {
6899 message_type: 2_002,
6900 flags: 5,
6901 payload: b"custom server frame".to_vec(),
6902 };
6903 server_send.send(response.clone()).await?;
6904 let received = client_recv
6905 .recv()
6906 .await
6907 .ok_or_else(|| -> BoxError { "client custom stream closed".into() })?;
6908 assert_eq!(received, response);
6909 Ok::<_, BoxError>(())
6910 })
6911 .await
6912 .map_err(|_| -> BoxError { "custom ordered stream round trip timed out".into() })?;
6913 client.shutdown().await;
6914 server.shutdown().await;
6915 round_trip
6916 }
6917
6918 #[tokio::test]
6919 async fn lazy_escalation_service_demand_respects_directional_caps() -> Result<(), BoxError> {
6920 let _guard = zakura_test::init();
6921 let shutdown = CancellationToken::new();
6922 let mut startup = header_sync_startup(shutdown.clone());
6923 startup.config.peer_limits = ServicePeerLimits {
6924 max_inbound_peers: 0,
6925 max_outbound_peers: 1,
6926 ..ServicePeerLimits::default()
6927 };
6928 let (header_sync, _actions, reactor_task) = spawn_header_sync_reactor(startup)?;
6929 let service = HeaderSyncService::new(header_sync);
6930 let peer = test_peer(17);
6931
6932 assert!(
6933 !service.wants_peer(&peer, ZAKURA_CAP_HEADER_SYNC, ServicePeerDirection::Inbound,),
6934 "a full inbound cap must not open a new header-sync stream"
6935 );
6936 assert!(
6937 service.wants_peer(
6938 &peer,
6939 ZAKURA_CAP_HEADER_SYNC,
6940 ServicePeerDirection::Outbound,
6941 ),
6942 "an outbound slot should still allow lazy header-sync escalation"
6943 );
6944
6945 shutdown.cancel();
6946 reactor_task.await?;
6947 Ok(())
6948 }
6949
6950 #[tokio::test]
6951 async fn header_ordered_session_waits_for_coordinator_capability_epoch() -> Result<(), BoxError>
6952 {
6953 use zakura_node_services::sync_lifecycle::{
6954 BlockServiceDemand, HeaderServiceDemand, LifecycleEpoch, SyncServiceDemand,
6955 };
6956
6957 let shutdown = CancellationToken::new();
6958 let (header_sync, _actions, reactor_task) =
6959 spawn_header_sync_reactor(header_sync_startup(shutdown.clone()))?;
6960 let disabled = SyncServiceDemand {
6961 header: HeaderServiceDemand::Disabled {
6962 runtime_epoch: LifecycleEpoch::INITIAL,
6963 },
6964 block: BlockServiceDemand::ServingOnly {
6965 apply_epoch: LifecycleEpoch::INITIAL,
6966 },
6967 };
6968 let (demand_tx, demand_rx) = watch::channel(disabled);
6969 let service = HeaderSyncService::new(header_sync).with_service_demand(Some(demand_rx));
6970 let peer = test_peer(171);
6971 assert!(!service.wants_peer(
6972 &peer,
6973 ZAKURA_CAP_HEADER_SYNC,
6974 ServicePeerDirection::Outbound,
6975 ));
6976 let OrderedSessionDemand::WaitForChange(changed) = service.ordered_session_demand(
6977 test_conn_id(),
6978 &peer,
6979 ZAKURA_CAP_HEADER_SYNC,
6980 ServicePeerDirection::Outbound,
6981 ) else {
6982 panic!("disabled coordinator demand must park the ordered header session");
6983 };
6984
6985 demand_tx
6986 .send(SyncServiceDemand {
6987 header: HeaderServiceDemand::Enabled {
6988 capability_epoch: LifecycleEpoch::new(1),
6989 },
6990 ..disabled
6991 })
6992 .expect("the header service retains its demand receiver");
6993 tokio::time::timeout(Duration::from_secs(1), changed)
6994 .await
6995 .expect("the capability epoch wakes ordered-session demand");
6996 assert!(service.wants_peer(
6997 &peer,
6998 ZAKURA_CAP_HEADER_SYNC,
6999 ServicePeerDirection::Outbound,
7000 ));
7001
7002 shutdown.cancel();
7003 reactor_task.await?;
7004 Ok(())
7005 }
7006
7007 #[tokio::test]
7008 async fn header_sync_peer_connected_has_immediately_ready_outbound_source(
7009 ) -> Result<(), BoxError> {
7010 let shutdown = CancellationToken::new();
7011 let mut startup = header_sync_startup(shutdown.clone());
7012 let anchor = zakura_header_chain::Frontier::new(startup.anchor.0, startup.anchor.1);
7013 let snapshot = zakura_header_chain::EngineSnapshot {
7014 mode: zakura_header_chain::EngineMode::Integrated,
7015 state_version: zakura_header_chain::StateVersion::new(1),
7016 header_generation: zakura_header_chain::HeaderGeneration::new(1),
7017 verified_generation: zakura_header_chain::VerifiedGeneration::new(1),
7018 frontiers: zakura_header_chain::FrontierSet {
7019 finalized: anchor,
7020 header_best: anchor,
7021 verified_best: anchor,
7022 },
7023 header_best_score: zakura_header_chain::ChainScore::new(
7024 zakura_header_chain::SuffixWork::zero(),
7025 anchor.hash,
7026 ),
7027 oldest_retained_height: anchor.height,
7028 alarms: Default::default(),
7029 };
7030 let (_snapshot_tx, snapshot_rx) = watch::channel(Some(snapshot));
7031 startup.committed_snapshots = Some(snapshot_rx);
7032 let (header_sync, _actions, reactor_task) = spawn_header_sync_reactor(startup)?;
7033 let codec = header_sync.codec();
7034 let service = HeaderSyncService::new(header_sync);
7035 let peer = test_peer(18);
7036 let connection_cancel = CancellationToken::new();
7037 let service_cancel = connection_cancel.child_token();
7038 let (_inbound_tx, inbound_rx) = crate::zakura::framed_channel(8);
7039 let (outbound_tx, mut outbound_rx) = crate::zakura::framed_channel(8);
7040 let session_id = 7;
7041 let streams = HashMap::from([(
7042 ZAKURA_STREAM_HEADER_SYNC,
7043 ServiceStream::new(
7044 session_id,
7045 ZAKURA_HEADER_SYNC_STREAM_VERSION,
7046 inbound_rx,
7047 outbound_tx,
7048 service_cancel,
7049 ),
7050 )]);
7051
7052 service.add_peer(Peer::new_with_service_streams(
7053 test_conn_id(),
7054 peer,
7055 None,
7056 ZAKURA_CAP_HEADER_SYNC,
7057 ServicePeerDirection::Inbound,
7058 streams,
7059 connection_cancel.clone(),
7060 CloseCause::new(),
7061 ));
7062
7063 let frame = tokio::time::timeout(Duration::from_secs(1), outbound_rx.recv())
7064 .await
7065 .expect("the first outbound source is immediately ready")
7066 .expect("the admitted source remains open");
7067 assert!(matches!(
7068 codec
7069 .decode_frame(frame, None)
7070 .expect("the first outbound frame decodes"),
7071 crate::zakura::HeaderSyncMessage::Status(_)
7072 ));
7073 assert!(!connection_cancel.is_cancelled());
7074
7075 shutdown.cancel();
7076 reactor_task.await?;
7077 Ok(())
7078 }
7079
7080 #[tokio::test]
7081 async fn supervisor_disconnect_peer_cancels_registered_token() -> Result<(), BoxError> {
7082 let supervisor = ZakuraSupervisorHandle::new(4);
7083 let peer = test_peer(8);
7084 let disconnect_token = CancellationToken::new();
7085 register_test_peer(&supervisor, peer.clone(), disconnect_token.clone()).await;
7086
7087 assert!(supervisor.disconnect_peer(&peer).await);
7088 tokio::time::timeout(Duration::from_secs(1), disconnect_token.cancelled())
7089 .await
7090 .expect("disconnect token is cancelled promptly");
7091 assert!(!supervisor.disconnect_peer(&test_peer(9)).await);
7092
7093 Ok(())
7094 }
7095
7096 #[tokio::test(start_paused = true)]
7097 async fn duplicate_evicts_stale_incumbent_but_keeps_fresh_one() -> Result<(), BoxError> {
7098 let supervisor = ZakuraSupervisorHandle::new(4);
7104
7105 async fn register_duplicate(
7106 supervisor: &ZakuraSupervisorHandle,
7107 peer: &ZakuraPeerId,
7108 token: CancellationToken,
7109 ) -> ZakuraRegistration {
7110 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
7111 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
7112 supervisor
7113 .register(
7114 test_conn_id(),
7115 peer.clone(),
7116 None,
7117 [peer.as_bytes()[0]; TRANSCRIPT_HASH_BYTES],
7118 outbound_handle,
7119 token,
7120 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
7121 )
7122 .await
7123 }
7124
7125 let fresh_peer = test_peer(8);
7128 let fresh_incumbent = CancellationToken::new();
7129 register_test_peer(&supervisor, fresh_peer.clone(), fresh_incumbent.clone()).await;
7130 let registration =
7131 register_duplicate(&supervisor, &fresh_peer, CancellationToken::new()).await;
7132 assert!(matches!(registration, ZakuraRegistration::Duplicate { .. }));
7133 assert!(
7134 !fresh_incumbent.is_cancelled(),
7135 "a young incumbent is kept so simultaneous-open races do not flap",
7136 );
7137
7138 let stale_peer = test_peer(9);
7141 let stale_incumbent = CancellationToken::new();
7142 register_test_peer(&supervisor, stale_peer.clone(), stale_incumbent.clone()).await;
7143 tokio::time::advance(ZAKURA_DUPLICATE_EVICT_MIN_AGE + Duration::from_secs(1)).await;
7144 let newcomer = CancellationToken::new();
7145 let registration = register_duplicate(&supervisor, &stale_peer, newcomer.clone()).await;
7146 assert!(matches!(registration, ZakuraRegistration::Duplicate { .. }));
7147 assert!(
7148 stale_incumbent.is_cancelled(),
7149 "a stale incumbent is evicted so the restarted peer's redial reclaims the slot",
7150 );
7151 assert!(
7152 !newcomer.is_cancelled(),
7153 "the rejected newcomer's token is never registered, so it is left to redial",
7154 );
7155
7156 Ok(())
7157 }
7158
7159 #[tokio::test(start_paused = true)]
7172 async fn same_peer_duplicate_at_per_ip_cap_still_evicts_stale_incumbent() -> Result<(), BoxError>
7173 {
7174 async fn register_from_ip(
7175 supervisor: &ZakuraSupervisorHandle,
7176 peer: &ZakuraPeerId,
7177 ip: IpAddr,
7178 token: CancellationToken,
7179 ) -> ZakuraRegistration {
7180 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
7181 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
7182 supervisor
7183 .register(
7184 test_conn_id(),
7185 peer.clone(),
7186 Some(ip),
7187 [peer.as_bytes()[0]; TRANSCRIPT_HASH_BYTES],
7188 outbound_handle,
7189 token,
7190 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
7191 )
7192 .await
7193 }
7194
7195 let supervisor = ZakuraSupervisorHandle::new(1);
7196 let ip: IpAddr = "203.0.113.9".parse().expect("test ip parses");
7197 let peer = test_peer(42);
7198
7199 let incumbent = CancellationToken::new();
7202 let registration = register_from_ip(&supervisor, &peer, ip, incumbent.clone()).await;
7203 assert!(
7204 matches!(registration, ZakuraRegistration::Registered { .. }),
7205 "the first connection from the IP registers",
7206 );
7207
7208 tokio::time::advance(ZAKURA_DUPLICATE_EVICT_MIN_AGE + Duration::from_secs(1)).await;
7211
7212 let newcomer = CancellationToken::new();
7216 let registration = register_from_ip(&supervisor, &peer, ip, newcomer.clone()).await;
7217 assert!(
7218 matches!(registration, ZakuraRegistration::Duplicate { .. }),
7219 "a same-peer redial from a capped IP must reach duplicate handling, not be \
7220 rejected as a resource limit before stale eviction can run",
7221 );
7222 assert!(
7223 incumbent.is_cancelled(),
7224 "the stale incumbent must be evicted so the restarted peer's redial reclaims \
7225 the slot in milliseconds instead of waiting for the QUIC idle timeout",
7226 );
7227 assert!(
7228 !newcomer.is_cancelled(),
7229 "the rejected newcomer's token is never registered, so it is left to redial",
7230 );
7231
7232 Ok(())
7233 }
7234
7235 #[tokio::test]
7236 async fn ipv4_embedded_ipv6_addresses_share_the_supervisor_ip_bucket() {
7237 async fn register_from_ip(
7238 supervisor: &ZakuraSupervisorHandle,
7239 peer: ZakuraPeerId,
7240 ip: IpAddr,
7241 ) -> ZakuraRegistration {
7242 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
7243 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
7244 supervisor
7245 .register(
7246 test_conn_id(),
7247 peer.clone(),
7248 Some(ip),
7249 [peer.as_bytes()[0]; TRANSCRIPT_HASH_BYTES],
7250 outbound_handle,
7251 CancellationToken::new(),
7252 ZAKURA_CAP_DISCOVERY,
7253 )
7254 .await
7255 }
7256
7257 let ipv4: IpAddr = "93.184.216.34".parse().expect("IPv4 test address parses");
7258 let mapped: IpAddr = "::ffff:93.184.216.34"
7259 .parse()
7260 .expect("mapped test address parses");
7261 let compatible: IpAddr = "::93.184.216.34"
7262 .parse()
7263 .expect("compatible test address parses");
7264 let six_to_four: IpAddr = "2002:5db8:d822::"
7265 .parse()
7266 .expect("6to4 test address parses");
7267 let teredo: IpAddr = "2001:0:c000:22d::a247:27dd"
7268 .parse()
7269 .expect("Teredo test address parses");
7270 let nat64: IpAddr = "64:ff9b::5db8:d822"
7271 .parse()
7272 .expect("NAT64 test address parses");
7273 let supervisor = ZakuraSupervisorHandle::new(1);
7274
7275 assert!(matches!(
7276 register_from_ip(&supervisor, test_peer(60), six_to_four).await,
7277 ZakuraRegistration::Registered { .. }
7278 ));
7279
7280 for alias in [ipv4, mapped, compatible, six_to_four, teredo, nat64] {
7281 assert!(
7282 !supervisor
7283 .can_accept_remote_ip_with_in_flight(alias, 0)
7284 .await,
7285 "{alias} must share the occupied IPv4 identity bucket"
7286 );
7287 }
7288
7289 assert!(matches!(
7290 register_from_ip(&supervisor, test_peer(61), ipv4).await,
7291 ZakuraRegistration::Rejected(ZakuraRejectReason::ResourceLimit)
7292 ));
7293 }
7294
7295 #[tokio::test(start_paused = true)]
7304 async fn same_ip_duplicate_uses_short_eviction_gate() -> Result<(), BoxError> {
7305 async fn register_from_ip(
7306 supervisor: &ZakuraSupervisorHandle,
7307 peer: &ZakuraPeerId,
7308 ip: Option<IpAddr>,
7309 token: CancellationToken,
7310 ) -> ZakuraRegistration {
7311 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
7312 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
7313 supervisor
7314 .register(
7315 test_conn_id(),
7316 peer.clone(),
7317 ip,
7318 [peer.as_bytes()[0]; TRANSCRIPT_HASH_BYTES],
7319 outbound_handle,
7320 token,
7321 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
7322 )
7323 .await
7324 }
7325
7326 let supervisor = ZakuraSupervisorHandle::new(4);
7327 let ip_a: IpAddr = "203.0.113.10".parse().expect("test ip parses");
7328 let peer = test_peer(50);
7329
7330 let incumbent = CancellationToken::new();
7332 let registration =
7333 register_from_ip(&supervisor, &peer, Some(ip_a), incumbent.clone()).await;
7334 assert!(matches!(
7335 registration,
7336 ZakuraRegistration::Registered { .. }
7337 ));
7338
7339 tokio::time::advance(ZAKURA_SAME_IP_DUPLICATE_EVICT_MIN_AGE - Duration::from_secs(1)).await;
7342 let early = CancellationToken::new();
7343 let registration = register_from_ip(&supervisor, &peer, Some(ip_a), early.clone()).await;
7344 assert!(matches!(registration, ZakuraRegistration::Duplicate { .. }));
7345 assert!(
7346 !incumbent.is_cancelled(),
7347 "a same-IP duplicate below the short gate keeps the incumbent",
7348 );
7349
7350 tokio::time::advance(Duration::from_secs(2)).await;
7354 let redial = CancellationToken::new();
7355 let registration = register_from_ip(&supervisor, &peer, Some(ip_a), redial.clone()).await;
7356 assert!(matches!(registration, ZakuraRegistration::Duplicate { .. }));
7357 assert!(
7358 incumbent.is_cancelled(),
7359 "a same-IP duplicate past the short gate evicts the stale incumbent",
7360 );
7361 assert!(
7362 !redial.is_cancelled(),
7363 "the rejected newcomer's token is never registered, so it is left to redial",
7364 );
7365
7366 let peer2 = test_peer(51);
7369 let ip_b: IpAddr = "203.0.113.11".parse().expect("test ip parses");
7370 let ip_c: IpAddr = "203.0.113.12".parse().expect("test ip parses");
7371 let incumbent2 = CancellationToken::new();
7372 let registration =
7373 register_from_ip(&supervisor, &peer2, Some(ip_b), incumbent2.clone()).await;
7374 assert!(matches!(
7375 registration,
7376 ZakuraRegistration::Registered { .. }
7377 ));
7378 tokio::time::advance(ZAKURA_SAME_IP_DUPLICATE_EVICT_MIN_AGE + Duration::from_secs(1)).await;
7379 let diff_ip = CancellationToken::new();
7380 let registration = register_from_ip(&supervisor, &peer2, Some(ip_c), diff_ip.clone()).await;
7381 assert!(matches!(registration, ZakuraRegistration::Duplicate { .. }));
7382 assert!(
7383 !incumbent2.is_cancelled(),
7384 "a different-IP duplicate keeps the long gate and does not evict a young incumbent",
7385 );
7386
7387 Ok(())
7388 }
7389
7390 #[tokio::test]
7391 async fn header_sync_misbehavior_action_does_not_disconnect_peer() -> Result<(), BoxError> {
7392 let _guard = zakura_test::init();
7393 let reactor_shutdown = CancellationToken::new();
7394 let startup = header_sync_startup(reactor_shutdown.clone());
7395 let (header_sync, _reactor_actions, reactor_task) = spawn_header_sync_reactor(startup)?;
7396 let supervisor = ZakuraSupervisorHandle::new(4);
7397 let peer = test_peer(10);
7398 let disconnect_token = CancellationToken::new();
7399 register_test_peer(&supervisor, peer.clone(), disconnect_token.clone()).await;
7400
7401 let (actions_tx, actions_rx) = mpsc::channel(4);
7402 let driver_shutdown = CancellationToken::new();
7403 let driver_task = tokio::spawn(drive_header_sync_actions(
7404 actions_rx,
7405 header_sync,
7406 supervisor,
7407 driver_shutdown.clone(),
7408 ));
7409 actions_tx
7410 .send(HeaderSyncAction::Misbehavior {
7411 peer,
7412 reason: HeaderSyncMisbehavior::MalformedMessage,
7413 })
7414 .await?;
7415
7416 assert!(
7417 tokio::time::timeout(Duration::from_millis(100), disconnect_token.cancelled())
7418 .await
7419 .is_err(),
7420 "misbehavior is record-only and must not cancel the registered connection",
7421 );
7422
7423 driver_shutdown.cancel();
7424 driver_task.await?;
7425 reactor_shutdown.cancel();
7426 reactor_task.await?;
7427 Ok(())
7428 }
7429
7430 #[test]
7433 fn ordered_session_reopen_follows_transport_opening_policy() {
7434 let initiator_opened = Stream {
7435 kind: ZAKURA_STREAM_HEADER_SYNC,
7436 version: ZAKURA_HEADER_SYNC_STREAM_VERSION,
7437 frame_cap: 1,
7438 capability: ZAKURA_CAP_HEADER_SYNC,
7439 mode: StreamMode::Ordered,
7440 };
7441 let exit = OrderedSessionExit {
7442 stream: initiator_opened,
7443 session_id: 1,
7444 opened_locally: true,
7445 };
7446 let initiator_policy = OrderedStreamPolicy {
7447 opening: OrderedStreamOpening::InitiatorOnly,
7448 reopen: true,
7449 };
7450 assert!(should_reopen_ordered_session(
7451 exit,
7452 initiator_policy,
7453 true,
7454 false,
7455 false
7456 ));
7457 assert!(!should_reopen_ordered_session(
7458 exit,
7459 initiator_policy,
7460 false,
7461 true,
7462 false
7463 ));
7464 assert!(!should_reopen_ordered_session(
7465 exit,
7466 initiator_policy,
7467 true,
7468 false,
7469 true
7470 ));
7471 assert!(!should_reopen_ordered_session(
7472 exit,
7473 OrderedStreamPolicy::default(),
7474 true,
7475 false,
7476 false,
7477 ));
7478
7479 assert!(should_reopen_ordered_session(
7482 OrderedSessionExit {
7483 opened_locally: false,
7484 ..exit
7485 },
7486 initiator_policy,
7487 true,
7488 false,
7489 false,
7490 ));
7491
7492 let either_peer = OrderedSessionExit {
7493 stream: Stream {
7494 kind: ZAKURA_STREAM_BLOCK_SYNC,
7495 ..initiator_opened
7496 },
7497 ..exit
7498 };
7499 let either_policy = OrderedStreamPolicy {
7500 opening: OrderedStreamOpening::EitherSide,
7501 reopen: true,
7502 };
7503 assert!(should_reopen_ordered_session(
7504 either_peer,
7505 either_policy,
7506 true,
7507 true,
7508 false
7509 ));
7510 assert!(!should_reopen_ordered_session(
7511 either_peer,
7512 either_policy,
7513 false,
7514 false,
7515 false
7516 ));
7517 assert!(opens_ordered_stream_locally(either_policy, false, true));
7518 assert!(!opens_ordered_stream_locally(either_policy, true, false));
7519
7520 let request_response = OrderedSessionExit {
7521 stream: Stream {
7522 kind: 102,
7523 mode: StreamMode::RequestResponse,
7524 ..initiator_opened
7525 },
7526 ..exit
7527 };
7528 assert!(!should_reopen_ordered_session(
7529 request_response,
7530 either_policy,
7531 true,
7532 true,
7533 false,
7534 ));
7535 }
7536
7537 #[test]
7538 fn either_side_session_has_one_proactive_opener_across_connection_roles() {
7539 let policy = OrderedStreamPolicy {
7540 opening: OrderedStreamOpening::EitherSide,
7541 reopen: true,
7542 };
7543 let exit = OrderedSessionExit {
7544 stream: Stream {
7545 kind: ZAKURA_STREAM_BLOCK_SYNC,
7546 version: ZAKURA_BLOCK_SYNC_STREAM_VERSION,
7547 frame_cap: 1,
7548 capability: ZAKURA_CAP_BLOCK_SYNC,
7549 mode: StreamMode::Ordered,
7550 },
7551 session_id: 1,
7552 opened_locally: false,
7553 };
7554
7555 for winner_is_initiator in [false, true] {
7556 let opens = [
7557 opens_ordered_stream_locally(policy, winner_is_initiator, true),
7558 opens_ordered_stream_locally(policy, !winner_is_initiator, false),
7559 ];
7560 assert_eq!(opens, [true, false]);
7561 assert_eq!(opens.into_iter().filter(|opens| *opens).count(), 1);
7562
7563 let reopens = [
7564 should_reopen_ordered_session(exit, policy, winner_is_initiator, true, false),
7565 should_reopen_ordered_session(exit, policy, !winner_is_initiator, false, false),
7566 ];
7567 assert_eq!(reopens, [true, false]);
7568 assert_eq!(reopens.into_iter().filter(|reopens| *reopens).count(), 1);
7569 }
7570
7571 assert!(may_open_ordered_stream(policy, true));
7574 assert!(may_open_ordered_stream(policy, false));
7575 }
7576
7577 #[test]
7578 fn stale_ordered_session_exit_keeps_newer_active_generation() {
7579 let stream = Stream {
7580 kind: ZAKURA_STREAM_HEADER_SYNC,
7581 version: ZAKURA_HEADER_SYNC_STREAM_VERSION,
7582 frame_cap: 1,
7583 capability: ZAKURA_CAP_HEADER_SYNC,
7584 mode: StreamMode::Ordered,
7585 };
7586 let mut session = OrderedSessionState::new(stream);
7587 session.remote_session_id = Some(2);
7588
7589 assert!(!session.remove_active_session(false, 1));
7590 assert_eq!(session.remote_session_id, Some(2));
7591 assert!(session.remove_active_session(false, 2));
7592 assert_eq!(session.remote_session_id, None);
7593 }
7594
7595 fn test_ordered_stream() -> Stream {
7596 Stream {
7597 kind: ZAKURA_STREAM_HEADER_SYNC,
7598 version: ZAKURA_HEADER_SYNC_STREAM_VERSION,
7599 frame_cap: 1,
7600 capability: ZAKURA_CAP_HEADER_SYNC,
7601 mode: StreamMode::Ordered,
7602 }
7603 }
7604
7605 #[tokio::test(start_paused = true)]
7606 async fn ordered_session_waits_deduplicate_demand() {
7607 let stream = test_ordered_stream();
7608 let mut session = OrderedSessionState::new(stream);
7609 let mut waits = OrderedSessionWaits::new();
7610 let (change_tx, change_rx) = watch::channel(());
7611
7612 session.schedule_demand(
7613 &mut waits,
7614 OrderedSessionDemand::WaitForChange(Box::pin(async move {
7615 let mut change_rx = change_rx;
7616 let _ = change_rx.changed().await;
7617 })),
7618 );
7619 session.schedule_demand(
7620 &mut waits,
7621 OrderedSessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)),
7622 );
7623
7624 assert_eq!(waits.len(), 1);
7625 assert_eq!(
7626 session.reopen_state,
7627 OrderedSessionReopenState::Waiting(OrderedSessionWaitReason::Demand)
7628 );
7629 assert_eq!(
7630 change_tx.receiver_count(),
7631 1,
7632 "repeated demand must retain the original reactor wait"
7633 );
7634 }
7635
7636 #[tokio::test(start_paused = true)]
7637 async fn ordered_session_demand_replaces_transport_backoff() {
7638 let stream = test_ordered_stream();
7639 let mut session = OrderedSessionState::new(stream);
7640 let mut waits = OrderedSessionWaits::new();
7641
7642 assert_eq!(
7643 session.schedule_transport_backoff(&mut waits),
7644 Some(ORDERED_STREAM_REOPEN_BACKOFF)
7645 );
7646 session.schedule_demand(
7647 &mut waits,
7648 OrderedSessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)),
7649 );
7650
7651 assert_eq!(session.reopen_attempts, 1);
7652 assert_eq!(waits.len(), 1);
7653 assert_eq!(
7654 session.reopen_state,
7655 OrderedSessionReopenState::Waiting(OrderedSessionWaitReason::Demand)
7656 );
7657 }
7658
7659 #[tokio::test(start_paused = true)]
7660 async fn ordered_session_exit_keeps_existing_demand_wait() {
7661 let stream = test_ordered_stream();
7662 let mut session = OrderedSessionState::new(stream);
7663 let mut waits = OrderedSessionWaits::new();
7664
7665 session.schedule_demand(
7666 &mut waits,
7667 OrderedSessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)),
7668 );
7669
7670 assert_eq!(session.schedule_transport_backoff(&mut waits), None);
7671 assert_eq!(session.reopen_attempts, 0);
7672 assert_eq!(waits.len(), 1);
7673 assert_eq!(
7674 session.reopen_state,
7675 OrderedSessionReopenState::Waiting(OrderedSessionWaitReason::Demand)
7676 );
7677 }
7678
7679 #[tokio::test]
7680 async fn ordered_session_retirement_cancels_and_blocks_transport_waits() {
7681 let stream = test_ordered_stream();
7682 let mut session = OrderedSessionState::new(stream);
7683 let mut waits = OrderedSessionWaits::new();
7684
7685 session.schedule_demand(
7686 &mut waits,
7687 OrderedSessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)),
7688 );
7689 session.schedule_demand(&mut waits, OrderedSessionDemand::Retire);
7690
7691 assert!(waits.is_empty());
7692 assert_eq!(session.reopen_state, OrderedSessionReopenState::Retired);
7693 assert_eq!(session.schedule_transport_backoff(&mut waits), None);
7694 assert_eq!(session.reopen_attempts, 0);
7695 }
7696
7697 #[tokio::test]
7698 async fn ordered_session_adoption_drops_pending_wait() {
7699 let stream = test_ordered_stream();
7700 let mut session = OrderedSessionState::new(stream);
7701 let mut waits = OrderedSessionWaits::new();
7702
7703 session.schedule_demand(&mut waits, OrderedSessionDemand::OpenNow);
7704 session.cancel_wait(&mut waits);
7705
7706 assert!(waits.is_empty());
7707 assert_eq!(session.reopen_state, OrderedSessionReopenState::Idle);
7708 }
7709
7710 #[tokio::test]
7711 async fn ordered_session_demand_change_yields_exactly_one_reopen() {
7712 let stream = test_ordered_stream();
7713 let mut session = OrderedSessionState::new(stream);
7714 let mut waits = OrderedSessionWaits::new();
7715 let (change_tx, mut change_rx) = watch::channel(());
7716
7717 session.schedule_demand(
7718 &mut waits,
7719 OrderedSessionDemand::WaitForChange(Box::pin(async move {
7720 let _ = change_rx.changed().await;
7721 })),
7722 );
7723 change_tx.send_replace(());
7724 let (kind, ()) = timeout(Duration::from_secs(1), waits.next())
7725 .await
7726 .expect("reactor change is observed before timeout")
7727 .expect("connection still owns a pending session wait");
7728
7729 assert_eq!(kind, stream.kind);
7730 session.finish_wait(&mut waits);
7731 assert!(
7732 waits.is_empty(),
7733 "one state change yields one reopen attempt"
7734 );
7735 }
7736
7737 #[tokio::test(start_paused = true)]
7738 async fn ordered_session_demand_deadline_yields_exactly_one_reopen() {
7739 let stream = test_ordered_stream();
7740 let mut session = OrderedSessionState::new(stream);
7741 let mut waits = OrderedSessionWaits::new();
7742
7743 session.schedule_demand(
7744 &mut waits,
7745 OrderedSessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)),
7746 );
7747 tokio::time::advance(Duration::from_secs(60)).await;
7748 let (kind, ()) = waits
7749 .next()
7750 .await
7751 .expect("reactor deadline keeps a pending session wait");
7752
7753 assert_eq!(kind, stream.kind);
7754 session.finish_wait(&mut waits);
7755 assert!(waits.is_empty(), "one deadline yields one reopen attempt");
7756 }
7757
7758 #[test]
7759 fn ordered_session_connection_teardown_drops_pending_waits() {
7760 let stream = test_ordered_stream();
7761 let mut session = OrderedSessionState::new(stream);
7762 let mut waits = OrderedSessionWaits::new();
7763 let (change_tx, mut change_rx) = watch::channel(());
7764
7765 session.schedule_demand(
7766 &mut waits,
7767 OrderedSessionDemand::WaitForChange(Box::pin(async move {
7768 let _ = change_rx.changed().await;
7769 })),
7770 );
7771 assert_eq!(change_tx.receiver_count(), 1);
7772 drop(waits);
7773
7774 assert_eq!(
7775 change_tx.receiver_count(),
7776 0,
7777 "dropping the connection-owned wait map must drop every pending wait"
7778 );
7779 }
7780
7781 #[test]
7782 fn ordered_session_reopen_backoff_grows_and_caps() {
7783 assert_eq!(
7784 ordered_session_reopen_backoff(0),
7785 ORDERED_STREAM_REOPEN_BACKOFF
7786 );
7787 assert_eq!(
7788 ordered_session_reopen_backoff(1),
7789 ORDERED_STREAM_REOPEN_BACKOFF * 2
7790 );
7791 assert_eq!(
7792 ordered_session_reopen_backoff(5),
7793 ORDERED_STREAM_REOPEN_BACKOFF_CAP
7794 );
7795 assert_eq!(
7796 ordered_session_reopen_backoff(u32::MAX),
7797 ORDERED_STREAM_REOPEN_BACKOFF_CAP
7798 );
7799 }
7800
7801 #[tokio::test]
7802 async fn stream_cancel_closes_ordered_worker_without_connection_cancel() -> Result<(), BoxError>
7803 {
7804 const ALPN: &[u8] = b"/zakura/testkit/stream-cancel/0";
7805
7806 let _guard = zakura_test::init();
7807 let server = LocalEndpointFactory::new().endpoint(50).await?;
7808 let (conn_tx, mut conn_rx) = mpsc::channel(1);
7809 let (stream_tx, mut stream_rx) = mpsc::channel(2);
7810 let router = Router::builder(server)
7811 .accept(
7812 ALPN,
7813 CaptureConnection {
7814 connection_tx: conn_tx,
7815 stream_tx,
7816 },
7817 )
7818 .spawn();
7819 let client = LocalEndpointFactory::new().endpoint(51).await?;
7820 let server_addr = router.endpoint().node_addr().initialized().await;
7821 client.add_node_addr(server_addr.clone())?;
7822
7823 let client_conn = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
7824 .await
7825 .expect("client connects to loopback capture endpoint")?;
7826 let server_conn = timeout(Duration::from_secs(10), conn_rx.recv())
7827 .await
7828 .expect("server connection is captured")
7829 .expect("capture handler sends the accepted connection");
7830 drop(server_conn);
7831 let (mut client_send, _client_recv) =
7832 timeout(Duration::from_secs(1), client_conn.open_bi())
7833 .await
7834 .expect("client opens the worker stream")?;
7835 let test_frame = Frame {
7836 message_type: 1,
7837 flags: 0,
7838 payload: Vec::new(),
7839 }
7840 .encode(LOCAL_MAX_CONTROL_FRAME_BYTES)?;
7841 timeout(Duration::from_secs(1), client_send.write_all(&test_frame))
7842 .await
7843 .expect("client writes the worker stream frame")?;
7844 let (server_send, server_recv) = timeout(Duration::from_secs(1), stream_rx.recv())
7845 .await
7846 .expect("server accepts the worker stream")
7847 .expect("capture handler sends the worker stream");
7848
7849 let mut limits = test_connection_limits();
7850 limits.idle_timeout = Duration::from_millis(50);
7851 let stream_kind = DISCOVERY_STREAM_KIND;
7852 let connection_token = CancellationToken::new();
7853 let stream_token = connection_token.child_token();
7854 let (freshness_tx, _freshness_rx) = watch::channel(Instant::now());
7855 let permit = Arc::new(Semaphore::new(1))
7856 .try_acquire_owned()
7857 .expect("test semaphore starts with one permit");
7858 let stream = Stream {
7859 kind: stream_kind,
7860 version: ZAKURA_STREAM_VERSION_1,
7861 frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES,
7862 capability: ZAKURA_CAP_DISCOVERY,
7863 mode: StreamMode::Ordered,
7864 };
7865 let context = StreamWorkerContext {
7866 conn: ZakuraConnTrace::without_peer(1),
7867 peer_id: test_peer(55),
7868 stream_id: 1,
7869 _permit: permit,
7870 limits,
7871 inbound_frame_cap: inbound_frame_cap_for_stream(&limits, stream),
7872 outbound_frame_cap: application_frame_cap(&limits, stream),
7873 message_bucket: Arc::new(std::sync::Mutex::new(TokenBucket::new(128))),
7874 connection_token: connection_token.clone(),
7875 stream_token: stream_token.clone(),
7876 close_cause: CloseCause::new(),
7877 freshness_tx,
7878 };
7879 let prelude = StreamPrelude {
7880 magic: STREAM_PRELUDE_MAGIC,
7881 stream_kind,
7882 stream_version: ZAKURA_STREAM_VERSION_1,
7883 request_id: None,
7884 max_frame_bytes: inbound_frame_cap_for_stream(&limits, stream),
7885 };
7886 let mut workers = JoinSet::new();
7887 let (ordered_session_exit_tx, mut ordered_session_exit_rx) = mpsc::unbounded_channel();
7888 let admitted = spawn_persistent_stream_worker(
7889 &mut workers,
7890 server_send,
7891 server_recv,
7892 stream,
7893 prelude,
7894 context,
7895 1,
7896 true,
7897 ordered_session_exit_tx,
7898 );
7899
7900 admitted.cancel_token.cancel();
7901 let exited = timeout(Duration::from_secs(1), ordered_session_exit_rx.recv())
7904 .await
7905 .expect("stream cancellation reports worker exit")
7906 .expect("exit channel stays open");
7907 assert_eq!(exited.stream, stream);
7908 assert_eq!(exited.session_id, 1);
7909 assert!(exited.opened_locally);
7910 timeout(Duration::from_secs(1), workers.join_next())
7911 .await
7912 .expect("stream cancellation stops the worker promptly")
7913 .expect("worker exit is observed")?;
7914 tokio::time::sleep(Duration::from_millis(100)).await;
7915 assert!(
7916 !connection_token.is_cancelled(),
7917 "stream-local cancellation must not cancel the shared connection"
7918 );
7919
7920 let (mut sibling_send, _sibling_recv) =
7921 timeout(Duration::from_secs(1), client_conn.open_bi())
7922 .await
7923 .expect("connection survives to open a sibling stream")?;
7924 timeout(Duration::from_secs(1), sibling_send.write_all(&test_frame))
7925 .await
7926 .expect("client writes the sibling stream frame")?;
7927 timeout(Duration::from_secs(1), stream_rx.recv())
7928 .await
7929 .expect("server accepts the sibling stream")
7930 .expect("capture handler sends the sibling stream");
7931
7932 client_conn.close(0u32.into(), b"done");
7933 client.close().await;
7934 router.shutdown().await?;
7935 Ok(())
7936 }
7937
7938 #[tokio::test]
7963 async fn quiet_ordered_stream_does_not_cancel_connection_with_active_sibling(
7964 ) -> Result<(), BoxError> {
7965 const ALPN: &[u8] = b"/zakura/testkit/quiet-ordered-stream/0";
7966
7967 let _guard = zakura_test::init();
7968 let server = LocalEndpointFactory::new().endpoint(80).await?;
7969 let (conn_tx, _conn_rx) = mpsc::channel(1);
7970 let (stream_tx, mut stream_rx) = mpsc::channel(2);
7971 let router = Router::builder(server)
7972 .accept(
7973 ALPN,
7974 CaptureConnection {
7975 connection_tx: conn_tx,
7976 stream_tx,
7977 },
7978 )
7979 .spawn();
7980 let client = LocalEndpointFactory::new().endpoint(81).await?;
7981 let server_addr = router.endpoint().node_addr().initialized().await;
7982 client.add_node_addr(server_addr.clone())?;
7983
7984 let client_conn = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
7985 .await
7986 .expect("client connects to loopback capture endpoint")?;
7987
7988 let frame = Frame {
7989 message_type: 1,
7990 flags: 0,
7991 payload: Vec::new(),
7992 }
7993 .encode(LOCAL_MAX_CONTROL_FRAME_BYTES)?;
7994
7995 let (mut active_send, _active_recv) =
8001 timeout(Duration::from_secs(1), client_conn.open_bi())
8002 .await
8003 .expect("client opens the active stream")?;
8004 timeout(Duration::from_secs(1), active_send.write_all(&frame))
8005 .await
8006 .expect("client writes the active stream's first frame")?;
8007 let (active_server_send, active_server_recv) =
8008 timeout(Duration::from_secs(1), stream_rx.recv())
8009 .await
8010 .expect("server accepts the active stream")
8011 .expect("capture handler forwards the active stream");
8012
8013 let (mut quiet_send, _quiet_recv) = timeout(Duration::from_secs(1), client_conn.open_bi())
8014 .await
8015 .expect("client opens the quiet stream")?;
8016 timeout(Duration::from_secs(1), quiet_send.write_all(&frame))
8017 .await
8018 .expect("client writes the quiet stream's only frame")?;
8019 let (quiet_server_send, quiet_server_recv) =
8020 timeout(Duration::from_secs(1), stream_rx.recv())
8021 .await
8022 .expect("server accepts the quiet stream")
8023 .expect("capture handler forwards the quiet stream");
8024
8025 let idle_timeout = Duration::from_millis(250);
8029 let mut limits = test_connection_limits();
8030 limits.idle_timeout = idle_timeout;
8031 let stream_kind = LEGACY_GOSSIP_STREAM_KIND;
8032 let connection_token = CancellationToken::new();
8033 let (freshness_tx, freshness_rx) = watch::channel(Instant::now());
8034
8035 let spawn_ordered_worker = |send: SendStream, recv: RecvStream, stream_id: u64| {
8036 let permit = Arc::new(Semaphore::new(1))
8037 .try_acquire_owned()
8038 .expect("test semaphore starts with one permit");
8039 let stream = Stream {
8040 kind: stream_kind,
8041 version: ZAKURA_STREAM_VERSION_1,
8042 frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES,
8043 capability: ZAKURA_CAP_LEGACY_GOSSIP,
8044 mode: StreamMode::Ordered,
8045 };
8046 let context = StreamWorkerContext {
8047 conn: ZakuraConnTrace::without_peer(1),
8048 peer_id: test_peer(80),
8049 stream_id,
8050 _permit: permit,
8051 limits,
8052 inbound_frame_cap: inbound_frame_cap_for_stream(&limits, stream),
8053 outbound_frame_cap: application_frame_cap(&limits, stream),
8054 message_bucket: Arc::new(std::sync::Mutex::new(TokenBucket::new(
8055 limits.message_rate_per_second,
8056 ))),
8057 connection_token: connection_token.clone(),
8058 stream_token: connection_token.child_token(),
8059 close_cause: CloseCause::new(),
8060 freshness_tx: freshness_tx.clone(),
8061 };
8062 let prelude = StreamPrelude {
8063 magic: STREAM_PRELUDE_MAGIC,
8064 stream_kind,
8065 stream_version: ZAKURA_STREAM_VERSION_1,
8066 request_id: None,
8067 max_frame_bytes: inbound_frame_cap_for_stream(&limits, stream),
8068 };
8069 let (inbound_tx, mut inbound_rx) = mpsc::channel(16);
8073 let (_outbound_tx, outbound_rx) = mpsc::channel(1);
8074 tokio::spawn(async move { while inbound_rx.recv().await.is_some() {} });
8075 tokio::spawn(persistent_stream_worker(
8076 send,
8077 recv,
8078 prelude,
8079 context,
8080 inbound_tx,
8081 outbound_rx,
8082 16,
8083 ))
8084 };
8085
8086 let _active_worker = spawn_ordered_worker(active_server_send, active_server_recv, 1);
8087 let _quiet_worker = spawn_ordered_worker(quiet_server_send, quiet_server_recv, 2);
8088
8089 let serve = tokio::spawn({
8094 let connection_token = connection_token.clone();
8095 async move {
8096 tokio::select! {
8097 biased;
8098 _ = connection_token.cancelled() => {}
8099 _ = freshness_reaper(freshness_rx, idle_timeout) => connection_token.cancel(),
8100 }
8101 }
8102 });
8103
8104 let keepalive_frame = frame.clone();
8107 let keepalive = tokio::spawn(async move {
8108 for _ in 0..16 {
8109 tokio::time::sleep(Duration::from_millis(80)).await;
8110 if active_send.write_all(&keepalive_frame).await.is_err() {
8111 break;
8112 }
8113 }
8114 });
8115
8116 tokio::time::sleep(idle_timeout * 4).await;
8119 assert!(
8120 !connection_token.is_cancelled(),
8121 "a quiet ordered stream must not cancel a connection whose sibling stream is \
8122 actively transferring; the shared connection_token was cancelled within {:?}",
8123 idle_timeout * 4,
8124 );
8125
8126 connection_token.cancel();
8127 keepalive.abort();
8128 let _ = serve.await;
8129 client_conn.close(0u32.into(), b"done");
8130 client.close().await;
8131 router.shutdown().await?;
8132 Ok(())
8133 }
8134
8135 #[tokio::test]
8148 async fn write_ordered_frame_rejects_payload_over_message_cap() -> Result<(), BoxError> {
8149 const ALPN: &[u8] = b"/zakura/testkit/ordered-message-cap/0";
8150
8151 let _guard = zakura_test::init();
8152 let server = LocalEndpointFactory::new().endpoint(74).await?;
8153 let (conn_tx, _conn_rx) = mpsc::channel(1);
8154 let (stream_tx, _stream_rx) = mpsc::channel(2);
8155 let router = Router::builder(server)
8156 .accept(
8157 ALPN,
8158 CaptureConnection {
8159 connection_tx: conn_tx,
8160 stream_tx,
8161 },
8162 )
8163 .spawn();
8164 let client = LocalEndpointFactory::new().endpoint(75).await?;
8165 let server_addr = router.endpoint().node_addr().initialized().await;
8166 client.add_node_addr(server_addr.clone())?;
8167
8168 let client_conn = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
8169 .await
8170 .expect("client connects to the ordered-message-cap endpoint")?;
8171 let (mut send, _recv) = timeout(Duration::from_secs(1), client_conn.open_bi())
8172 .await
8173 .expect("client opens an ordered stream")?;
8174
8175 let mut limits = test_connection_limits();
8178 limits.max_message_bytes = 256;
8179 let stream = Stream {
8180 kind: DISCOVERY_STREAM_KIND,
8181 version: ZAKURA_STREAM_VERSION_1,
8182 frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES,
8183 capability: ZAKURA_CAP_DISCOVERY,
8184 mode: StreamMode::Ordered,
8185 };
8186 let frame_cap = application_frame_cap(&limits, stream);
8187
8188 let oversized = Frame {
8191 message_type: 1,
8192 flags: 0,
8193 payload: vec![0xab; 4096],
8194 };
8195 assert!(
8196 oversized.payload.len() <= frame_cap as usize,
8197 "test payload must fit the frame cap so only the message cap can reject it"
8198 );
8199
8200 let result = write_ordered_frame(&mut send, oversized, limits, frame_cap).await;
8201 assert!(
8202 result.is_err(),
8203 "write_ordered_frame must reject a payload over the negotiated max_message_bytes \
8204 before encoding/writing it, mirroring write_response_frame; got {result:?}"
8205 );
8206
8207 let within_cap = Frame {
8209 message_type: 1,
8210 flags: 0,
8211 payload: vec![0xcd; 128],
8212 };
8213 write_ordered_frame(&mut send, within_cap, limits, frame_cap)
8214 .await
8215 .expect("a frame within the negotiated message cap must still be written");
8216
8217 client_conn.close(0u32.into(), b"done");
8218 client.close().await;
8219 router.shutdown().await?;
8220 Ok(())
8221 }
8222
8223 #[tokio::test]
8231 async fn read_control_payload_enforces_hard_cap_before_reading_body() -> Result<(), BoxError> {
8232 const ALPN: &[u8] = b"/zakura/testkit/control-hard-cap/0";
8233
8234 let _guard = zakura_test::init();
8235 let server = LocalEndpointFactory::new().endpoint(70).await?;
8236 let (conn_tx, _conn_rx) = mpsc::channel(1);
8237 let (stream_tx, mut stream_rx) = mpsc::channel(2);
8238 let router = Router::builder(server)
8239 .accept(
8240 ALPN,
8241 CaptureConnection {
8242 connection_tx: conn_tx,
8243 stream_tx,
8244 },
8245 )
8246 .spawn();
8247 let client = LocalEndpointFactory::new().endpoint(71).await?;
8248 let server_addr = router.endpoint().node_addr().initialized().await;
8249 client.add_node_addr(server_addr.clone())?;
8250
8251 let client_conn = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
8252 .await
8253 .expect("client connects to the control-hard-cap endpoint")?;
8254
8255 let oversized_len =
8259 u32::try_from(MAX_CONTROL_PAYLOAD_BYTES + 1).expect("control hard cap + 1 fits in u32");
8260 let (mut over_send, _over_recv) = timeout(Duration::from_secs(1), client_conn.open_bi())
8261 .await
8262 .expect("client opens the oversized control stream")?;
8263 timeout(
8264 Duration::from_secs(1),
8265 over_send.write_all(&oversized_len.to_le_bytes()),
8266 )
8267 .await
8268 .expect("client writes the oversized control length")?;
8269 let _ = over_send.finish();
8270
8271 let (_over_server_send, mut over_server_recv) =
8272 timeout(Duration::from_secs(1), stream_rx.recv())
8273 .await
8274 .expect("server accepts the oversized control stream")
8275 .expect("capture handler forwards the oversized control stream");
8276
8277 let oversized = read_control_payload(
8280 &mut over_server_recv,
8281 LOCAL_MAX_CONTROL_FRAME_BYTES,
8282 Duration::from_secs(2),
8283 )
8284 .await;
8285 assert!(
8286 matches!(oversized, Err(ZakuraHandlerError::Oversize)),
8287 "a control length over the 16 KiB hard cap must be rejected as Oversize \
8288 on the length prefix, before the body is read; got {oversized:?}"
8289 );
8290
8291 let valid_body = vec![0xa5u8; 128];
8294 let valid_len = u32::try_from(valid_body.len()).expect("128 fits in u32");
8295 let (mut ok_send, _ok_recv) = timeout(Duration::from_secs(1), client_conn.open_bi())
8296 .await
8297 .expect("client opens the in-cap control stream")?;
8298 timeout(
8299 Duration::from_secs(1),
8300 ok_send.write_all(&valid_len.to_le_bytes()),
8301 )
8302 .await
8303 .expect("client writes the in-cap control length")?;
8304 timeout(Duration::from_secs(1), ok_send.write_all(&valid_body))
8305 .await
8306 .expect("client writes the in-cap control body")?;
8307 let _ = ok_send.finish();
8308
8309 let (_ok_server_send, mut ok_server_recv) =
8310 timeout(Duration::from_secs(1), stream_rx.recv())
8311 .await
8312 .expect("server accepts the in-cap control stream")
8313 .expect("capture handler forwards the in-cap control stream");
8314 let read_back = read_control_payload(
8315 &mut ok_server_recv,
8316 LOCAL_MAX_CONTROL_FRAME_BYTES,
8317 Duration::from_secs(2),
8318 )
8319 .await
8320 .expect("a control payload within the hard cap is read");
8321 assert_eq!(
8322 read_back, valid_body,
8323 "an in-cap control payload must round-trip unchanged"
8324 );
8325
8326 client_conn.close(0u32.into(), b"done");
8327 client.close().await;
8328 router.shutdown().await?;
8329 Ok(())
8330 }
8331
8332 #[tokio::test]
8342 async fn inbound_frame_cap_rejects_over_message_frame_before_reading_payload(
8343 ) -> Result<(), BoxError> {
8344 const ALPN: &[u8] = b"/zakura/testkit/late-message-cap/0";
8345 const MAX_FRAME_BYTES: u32 = 64 * 1024;
8348 const MAX_MESSAGE_BYTES: u32 = 1024;
8349 const OVER_MESSAGE_PAYLOAD_LEN: u32 = 2048;
8352 let stream = Stream {
8353 kind: LEGACY_GOSSIP_STREAM_KIND,
8354 version: ZAKURA_STREAM_VERSION_1,
8355 frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES,
8356 capability: ZAKURA_CAP_LEGACY_GOSSIP,
8357 mode: StreamMode::Ordered,
8358 };
8359
8360 let limits = ZakuraConnectionLimits {
8361 max_frame_bytes: MAX_FRAME_BYTES,
8362 max_message_bytes: MAX_MESSAGE_BYTES,
8363 ..test_connection_limits()
8364 };
8365 let inbound_cap = inbound_frame_cap_for_stream(&limits, stream);
8368 let raw_cap = application_frame_cap(&limits, stream);
8369 assert!(
8370 inbound_cap < raw_cap,
8371 "the inbound cap must be tighter than the raw frame cap when the caps diverge \
8372 (inbound_cap={inbound_cap}, raw_cap={raw_cap})"
8373 );
8374
8375 let _guard = zakura_test::init();
8376 let server = LocalEndpointFactory::new().endpoint(72).await?;
8377 let (conn_tx, _conn_rx) = mpsc::channel(2);
8378 let (stream_tx, mut stream_rx) = mpsc::channel(4);
8379 let router = Router::builder(server)
8380 .accept(
8381 ALPN,
8382 CaptureConnection {
8383 connection_tx: conn_tx,
8384 stream_tx,
8385 },
8386 )
8387 .spawn();
8388 let client = LocalEndpointFactory::new().endpoint(73).await?;
8389 let server_addr = router.endpoint().node_addr().initialized().await;
8390 client.add_node_addr(server_addr.clone())?;
8391
8392 let frame_header = |payload_len: u32| -> Vec<u8> {
8394 let mut header = Vec::with_capacity(FRAME_HEADER_BYTES);
8395 header.extend_from_slice(&0u16.to_le_bytes());
8396 header.extend_from_slice(&0u16.to_le_bytes());
8397 header.extend_from_slice(&payload_len.to_le_bytes());
8398 header
8399 };
8400
8401 let conn_a = timeout(
8405 Duration::from_secs(10),
8406 client.connect(server_addr.clone(), ALPN),
8407 )
8408 .await
8409 .expect("client connects for the oversized streams")?;
8410
8411 let (mut over_send, _over_recv) =
8414 timeout(Duration::from_secs(1), conn_a.open_bi())
8415 .await
8416 .expect("client opens the oversized inbound-cap stream")?;
8417 timeout(
8418 Duration::from_secs(1),
8419 over_send.write_all(&frame_header(OVER_MESSAGE_PAYLOAD_LEN)),
8420 )
8421 .await
8422 .expect("client writes the oversized frame header")?;
8423 let _ = over_send.finish();
8424 let (_s1_send, mut s1_recv) = timeout(Duration::from_secs(1), stream_rx.recv())
8425 .await
8426 .expect("server accepts the oversized inbound-cap stream")
8427 .expect("capture handler forwards the oversized inbound-cap stream");
8428 let rejected = read_frame(
8429 &mut s1_recv,
8430 inbound_cap,
8431 Duration::from_secs(2),
8432 Some(Duration::from_secs(2)),
8433 )
8434 .await;
8435 assert!(
8436 matches!(rejected, Err(ZakuraHandlerError::OversizeFrame { .. })),
8437 "a frame whose payload exceeds max_message_bytes must be rejected as Oversize \
8438 on the header alone with the inbound (message-limited) cap, before the payload \
8439 is allocated and read; got {rejected:?}"
8440 );
8441
8442 let (mut raw_send, _raw_recv) = timeout(Duration::from_secs(1), conn_a.open_bi())
8448 .await
8449 .expect("client opens the oversized raw-cap stream")?;
8450 timeout(
8451 Duration::from_secs(1),
8452 raw_send.write_all(&frame_header(OVER_MESSAGE_PAYLOAD_LEN)),
8453 )
8454 .await
8455 .expect("client writes the oversized frame header again")?;
8456 let _ = raw_send.finish();
8457 let (_s2_send, mut s2_recv) = timeout(Duration::from_secs(1), stream_rx.recv())
8458 .await
8459 .expect("server accepts the oversized raw-cap stream")
8460 .expect("capture handler forwards the oversized raw-cap stream");
8461 let allocated = read_frame(
8462 &mut s2_recv,
8463 raw_cap,
8464 Duration::from_secs(2),
8465 Some(Duration::from_secs(2)),
8466 )
8467 .await;
8468 assert!(
8469 allocated.is_err()
8470 && !matches!(
8471 allocated,
8472 Err(ZakuraHandlerError::Oversize | ZakuraHandlerError::OversizeFrame { .. })
8473 ),
8474 "with the raw frame cap the same oversized frame passes the size check and \
8475 read_frame proceeds to allocate/read the payload (it is not rejected as \
8476 Oversize), proving the message cap is enforced too late; got {allocated:?}"
8477 );
8478
8479 let conn_b = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
8482 .await
8483 .expect("client connects for the in-cap stream")?;
8484 let valid_payload = vec![0x5au8; (MAX_MESSAGE_BYTES / 2) as usize];
8485 let mut valid_bytes =
8486 frame_header(u32::try_from(valid_payload.len()).expect("in-cap payload len fits u32"));
8487 valid_bytes.extend_from_slice(&valid_payload);
8488 let (mut ok_send, _ok_recv) = timeout(Duration::from_secs(1), conn_b.open_bi())
8489 .await
8490 .expect("client opens the in-cap stream")?;
8491 timeout(Duration::from_secs(1), ok_send.write_all(&valid_bytes))
8492 .await
8493 .expect("client writes the in-cap frame")?;
8494 let _ = ok_send.finish();
8495 let (_s3_send, mut s3_recv) = timeout(Duration::from_secs(1), stream_rx.recv())
8496 .await
8497 .expect("server accepts the in-cap stream")
8498 .expect("capture handler forwards the in-cap stream");
8499 let frame = read_frame(
8500 &mut s3_recv,
8501 inbound_cap,
8502 Duration::from_secs(2),
8503 Some(Duration::from_secs(2)),
8504 )
8505 .await
8506 .expect("a frame within the message cap is read with the inbound cap");
8507 assert_eq!(
8508 frame.payload, valid_payload,
8509 "an in-cap frame payload must round-trip unchanged"
8510 );
8511
8512 conn_a.close(0u32.into(), b"done");
8513 conn_b.close(0u32.into(), b"done");
8514 client.close().await;
8515 router.shutdown().await?;
8516 Ok(())
8517 }
8518
8519 #[tokio::test]
8520 async fn invalid_prelude_stream_churn_charges_open_rate_token() -> Result<(), BoxError> {
8521 const ALPN: &[u8] = b"/zakura/testkit/open-rate-churn/0";
8530
8531 let _guard = zakura_test::init();
8532 let server = LocalEndpointFactory::new().endpoint(90).await?;
8533 let (conn_tx, _conn_rx) = mpsc::channel(1);
8534 let (stream_tx, mut stream_rx) = mpsc::channel(2);
8535 let router = Router::builder(server)
8536 .accept(
8537 ALPN,
8538 CaptureConnection {
8539 connection_tx: conn_tx,
8540 stream_tx,
8541 },
8542 )
8543 .spawn();
8544 let client = LocalEndpointFactory::new().endpoint(91).await?;
8545 let server_addr = router.endpoint().node_addr().initialized().await;
8546 client.add_node_addr(server_addr.clone())?;
8547 let client_conn = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
8548 .await
8549 .expect("client connects to the open-rate-churn endpoint")?;
8550
8551 let prelude = StreamPrelude {
8555 magic: STREAM_PRELUDE_MAGIC,
8556 stream_kind: 9,
8557 stream_version: 1,
8558 request_id: None,
8559 max_frame_bytes: 1024,
8560 };
8561 let (mut client_send, _client_recv) =
8562 timeout(Duration::from_secs(1), client_conn.open_bi())
8563 .await
8564 .expect("client opens the unknown-kind stream")?;
8565 timeout(
8566 Duration::from_secs(1),
8567 client_send.write_all(&prelude.encode()?),
8568 )
8569 .await
8570 .expect("client writes the unknown-kind prelude")?;
8571 let _ = client_send.finish();
8572
8573 let (server_send, server_recv) = timeout(Duration::from_secs(1), stream_rx.recv())
8574 .await
8575 .expect("server accepts the unknown-kind stream")
8576 .expect("capture handler forwards the unknown-kind stream");
8577
8578 let supervisor = ZakuraSupervisorHandle::new(16);
8579 let handler = ZakuraProtocolHandler::new(
8580 supervisor,
8581 Network::Mainnet,
8582 ZakuraHandshakeConfig::for_network(&Network::Mainnet),
8583 ZakuraLocalLimits::from_config(&Config::default()),
8584 );
8585
8586 let peer_id = test_peer(9);
8587 let stream_sem = Arc::new(Semaphore::new(16));
8588 let mut open_limiter = TokenBucket::new(4);
8591 let mut message_buckets = MessageRateBuckets::new();
8592 let mut workers = JoinSet::new();
8593 let connection_token = CancellationToken::new();
8594 let (freshness_tx, _freshness_rx) = watch::channel(Instant::now());
8595
8596 let mut admission = StreamAdmission {
8597 conn: ZakuraConnTrace::placeholder(),
8598 peer_id: &peer_id,
8599 stream_sem: &stream_sem,
8600 open_limiter: &mut open_limiter,
8601 message_buckets: &mut message_buckets,
8602 workers: &mut workers,
8603 limits: test_connection_limits(),
8604 accepted_capabilities: 0,
8605 connection_token: connection_token.clone(),
8606 close_cause: CloseCause::new(),
8607 freshness_tx,
8608 };
8609 let (ordered_session_exit_tx, _ordered_session_exit_rx) = mpsc::unbounded_channel();
8610 let admitted = handler
8611 .admit_bi_stream(
8612 server_send,
8613 server_recv,
8614 &mut admission,
8615 16,
8616 ordered_session_exit_tx,
8617 )
8618 .await;
8619
8620 assert!(
8621 admitted.is_none(),
8622 "an unknown-kind stream must be rejected, not admitted"
8623 );
8624 assert!(
8625 !connection_token.is_cancelled(),
8626 "an unknown-kind stream is reset stream-only and must keep the connection alive"
8627 );
8628 assert_eq!(
8629 admission.open_limiter.tokens, 3,
8630 "the protocol-invalid stream open must spend exactly one open-rate token \
8631 (capacity 4 -> 3); before the fix the unknown-kind reject returned before \
8632 reaching the limiter, leaving the bucket full at 4"
8633 );
8634
8635 client.close().await;
8636 router.shutdown().await?;
8637 Ok(())
8638 }
8639
8640 #[test]
8641 fn local_application_frame_cap_admits_default_header_sync_response() {
8642 let limits = ZakuraLocalLimits::from_config(&Config::default());
8643 let default_header_sync_frame_bytes =
8644 u32::try_from(MAX_HS_MESSAGE_BYTES.saturating_add(FRAME_HEADER_BYTES))
8645 .expect("header-sync frame cap fits in u32");
8646
8647 assert!(limits.max_frame_bytes >= default_header_sync_frame_bytes);
8648 assert!(limits.initial_limits().max_frame_bytes >= default_header_sync_frame_bytes);
8649 }
8650
8651 #[test]
8652 fn ordered_stream_stopped_write_is_stream_local() {
8653 let error: BoxError = iroh::endpoint::WriteError::Stopped(VarInt::from_u32(0)).into();
8654
8655 assert!(ordered_stream_write_was_stopped(&error));
8656 }
8657
8658 #[test]
8659 fn declared_stream_frame_caps_are_authoritative() {
8660 let limits = ZakuraLocalLimits::from_config(&Config::default());
8661 let negotiated = limits.clamp(&ZakuraAcceptedLimits {
8662 max_frame_bytes: u32::MAX,
8663 max_message_bytes: u32::MAX,
8664 max_open_streams: u16::MAX,
8665 max_inbound_queue_depth: u16::MAX,
8666 idle_timeout_millis: u32::MAX,
8667 });
8668 const CUSTOM_FRAME_CAP: u32 = 64 * 1024;
8669 let custom_stream = Stream {
8670 kind: 42,
8671 version: ZAKURA_STREAM_VERSION_1,
8672 frame_cap: CUSTOM_FRAME_CAP,
8673 capability: 1 << 20,
8674 mode: StreamMode::RequestResponse,
8675 };
8676
8677 assert_eq!(
8678 application_frame_cap(&negotiated, custom_stream),
8679 CUSTOM_FRAME_CAP,
8680 "a custom declaration below the negotiated cap must remain authoritative"
8681 );
8682 assert_eq!(
8683 peer_accepted_frame_cap(&negotiated, custom_stream, CUSTOM_FRAME_CAP / 2),
8684 CUSTOM_FRAME_CAP / 2,
8685 "writes on peer-opened streams must also honor the peer's advertised receive cap"
8686 );
8687 assert_eq!(
8688 peer_accepted_frame_cap(&negotiated, custom_stream, CUSTOM_FRAME_CAP * 2),
8689 CUSTOM_FRAME_CAP,
8690 "a peer cannot raise the locally declared stream cap"
8691 );
8692
8693 let at_cap = Frame {
8694 message_type: 1,
8695 flags: 0,
8696 payload: vec![
8697 0;
8698 usize::try_from(CUSTOM_FRAME_CAP)
8699 .expect("custom test cap fits usize")
8700 .saturating_sub(FRAME_HEADER_BYTES)
8701 ],
8702 };
8703 let over_cap = Frame {
8704 message_type: 1,
8705 flags: 0,
8706 payload: vec![0; at_cap.payload.len() + 1],
8707 };
8708
8709 assert!(
8710 at_cap
8711 .encode(application_frame_cap(&negotiated, custom_stream))
8712 .is_ok(),
8713 "a frame exactly at the declared custom cap must be accepted"
8714 );
8715 assert!(
8716 over_cap
8717 .encode(application_frame_cap(&negotiated, custom_stream))
8718 .is_err(),
8719 "a frame one byte over the declared custom cap must be rejected"
8720 );
8721 }
8722
8723 #[test]
8724 fn token_bucket_rejects_churn_until_refill() {
8725 let clock = crate::zakura::testkit::TestClock::new();
8726 let mut bucket = TokenBucket::with_clock(2, clock.clone());
8727
8728 assert!(bucket.try_take());
8729 assert!(bucket.try_take());
8730 assert!(!bucket.try_take());
8731 clock.advance(Duration::from_millis(500));
8732 assert!(bucket.try_take());
8733 assert!(!bucket.try_take());
8734 clock.advance(Duration::from_millis(500));
8735 assert!(bucket.try_take());
8736 }
8737
8738 #[test]
8739 fn supported_stream_accepts_registered_kinds_at_declared_version_only() {
8740 let registry = ServiceRegistry::new(vec![Arc::new(DeclaredStreamService {
8741 streams: vec![
8742 Stream {
8743 kind: LEGACY_GOSSIP_STREAM_KIND,
8744 version: ZAKURA_STREAM_VERSION_1,
8745 frame_cap: 1024,
8746 capability: ZAKURA_CAP_LEGACY_GOSSIP,
8747 mode: StreamMode::Ordered,
8748 },
8749 Stream {
8750 kind: LEGACY_REQUEST_STREAM_KIND,
8751 version: ZAKURA_STREAM_VERSION_1,
8752 frame_cap: 1024,
8753 capability: ZAKURA_CAP_LEGACY_GOSSIP,
8754 mode: StreamMode::RequestResponse,
8755 },
8756 Stream {
8757 kind: DISCOVERY_STREAM_KIND,
8758 version: ZAKURA_STREAM_VERSION_1,
8759 frame_cap: 1024,
8760 capability: ZAKURA_CAP_DISCOVERY,
8761 mode: StreamMode::Ordered,
8762 },
8763 Stream {
8764 kind: HEADER_SYNC_STREAM_KIND,
8765 version: ZAKURA_HEADER_SYNC_STREAM_VERSION,
8766 frame_cap: 1024,
8767 capability: ZAKURA_CAP_HEADER_SYNC,
8768 mode: StreamMode::Ordered,
8769 },
8770 Stream {
8771 kind: ZAKURA_STREAM_BLOCK_SYNC,
8772 version: ZAKURA_STREAM_VERSION_1,
8773 frame_cap: MAX_BS_FRAME_BYTES,
8774 capability: crate::zakura::ZAKURA_CAP_BLOCK_SYNC,
8775 mode: StreamMode::Ordered,
8776 },
8777 ],
8778 }) as Arc<dyn Service>])
8779 .expect("test registry declares unique stream kinds");
8780
8781 for (kind, version) in [
8782 (LEGACY_GOSSIP_STREAM_KIND, ZAKURA_STREAM_VERSION_1),
8783 (LEGACY_REQUEST_STREAM_KIND, ZAKURA_STREAM_VERSION_1),
8784 (DISCOVERY_STREAM_KIND, ZAKURA_STREAM_VERSION_1),
8785 (HEADER_SYNC_STREAM_KIND, ZAKURA_HEADER_SYNC_STREAM_VERSION),
8786 (ZAKURA_STREAM_BLOCK_SYNC, ZAKURA_STREAM_VERSION_1),
8787 ] {
8788 assert!(
8789 is_supported_stream(®istry, kind, version),
8790 "registered kind {kind} at declared version {version} must be supported"
8791 );
8792 assert!(
8793 !is_supported_stream(®istry, kind, 0),
8794 "registered kind {kind} at version 0 must be rejected"
8795 );
8796 assert!(
8797 !is_supported_stream(®istry, kind, u16::MAX),
8798 "registered kind {kind} at an unsupported version must be rejected"
8799 );
8800 }
8801
8802 assert!(
8803 !is_supported_stream(®istry, HEADER_SYNC_STREAM_KIND, ZAKURA_STREAM_VERSION_1),
8804 "header-sync v1 is rejected because native header sync uses exact version matching"
8805 );
8806 assert!(
8807 !is_supported_stream(®istry, HEADER_SYNC_STREAM_KIND, 5),
8808 "header-sync v5 is rejected after the expanded Ironwood root-record wire break"
8809 );
8810 assert!(
8811 !is_supported_stream(®istry, HEADER_SYNC_STREAM_KIND, 7),
8812 "the predecessor header-sync stream version is rejected"
8813 );
8814 assert!(
8815 registry.ordered_streams_for_negotiated(1 << 4).is_empty(),
8816 "the retired predecessor capability opens no header-sync stream"
8817 );
8818
8819 assert_eq!(stream_kind_label(2), "gossip");
8820 assert_eq!(stream_kind_label(3), "legacy_request");
8821 assert_eq!(stream_kind_label(4), "discovery");
8822 assert_eq!(stream_kind_label(5), "header_sync");
8823 assert_eq!(stream_kind_label(6), "block_sync");
8824
8825 for kind in [0u16, 1, 7, 255, u16::MAX] {
8826 assert!(
8827 !is_supported_stream(®istry, kind, ZAKURA_STREAM_VERSION_1),
8828 "unknown kind {kind} must be rejected even at version 1"
8829 );
8830 }
8831
8832 for kind in [7u16, 255, u16::MAX] {
8833 assert_eq!(stream_kind_label(kind), "unknown");
8834 }
8835 }
8836
8837 #[test]
8838 fn legacy_transport_response_validator_accepts_codec_frames() -> Result<(), BoxError> {
8839 let block = Arc::new(Block::zcash_deserialize(
8840 BLOCK_TESTNET_141042_BYTES.as_slice(),
8841 )?);
8842 let header = block::CountedHeader {
8843 header: block.header.clone(),
8844 };
8845 let tx_id = legacy_tx_id(7);
8846
8847 assert_codec_frames_validate_at_transport(
8848 LegacyRequestFrame::BlocksByHash(vec![block.hash()]),
8849 LegacyRequestKind::Blocks,
8850 Response::Blocks(vec![InventoryResponse::Available((block, None))]),
8851 )?;
8852 assert_codec_frames_validate_at_transport(
8853 LegacyRequestFrame::TransactionsById(vec![tx_id]),
8854 LegacyRequestKind::Transactions,
8855 Response::Transactions(vec![InventoryResponse::Missing(tx_id)]),
8856 )?;
8857 assert_codec_frames_validate_at_transport(
8858 LegacyRequestFrame::FindBlocks {
8859 known_blocks: vec![block_hash(1)],
8860 stop: None,
8861 },
8862 LegacyRequestKind::FindBlocks,
8863 Response::BlockHashes(vec![block_hash(2)]),
8864 )?;
8865 assert_codec_frames_validate_at_transport(
8866 LegacyRequestFrame::FindHeaders {
8867 known_blocks: vec![block_hash(1)],
8868 stop: None,
8869 },
8870 LegacyRequestKind::FindHeaders,
8871 Response::BlockHeaders(vec![header]),
8872 )?;
8873 assert_codec_frames_validate_at_transport(
8874 LegacyRequestFrame::MempoolTransactionIds,
8875 LegacyRequestKind::MempoolTransactionIds,
8876 Response::TransactionIds(vec![tx_id]),
8877 )?;
8878 assert_codec_frames_validate_at_transport(
8879 LegacyRequestFrame::Ping,
8880 LegacyRequestKind::Ping,
8881 Response::Pong(Duration::ZERO),
8882 )?;
8883 assert_codec_frames_validate_at_transport(
8884 LegacyRequestFrame::PushTransaction(empty_v5_transaction(8).into()),
8885 LegacyRequestKind::PushTransaction,
8886 Response::Nil,
8887 )?;
8888
8889 assert_codec_frames_validate_at_transport(
8894 LegacyRequestFrame::FindBlocks {
8895 known_blocks: vec![block_hash(1)],
8896 stop: None,
8897 },
8898 LegacyRequestKind::FindBlocks,
8899 Response::Nil,
8900 )?;
8901 assert_codec_frames_validate_at_transport(
8902 LegacyRequestFrame::FindHeaders {
8903 known_blocks: vec![block_hash(1)],
8904 stop: None,
8905 },
8906 LegacyRequestKind::FindHeaders,
8907 Response::Nil,
8908 )?;
8909 assert_codec_frames_validate_at_transport(
8910 LegacyRequestFrame::MempoolTransactionIds,
8911 LegacyRequestKind::MempoolTransactionIds,
8912 Response::Nil,
8913 )?;
8914
8915 Ok(())
8916 }
8917
8918 fn assert_codec_frames_validate_at_transport(
8919 request: LegacyRequestFrame,
8920 request_kind: LegacyRequestKind,
8921 response: Response,
8922 ) -> Result<(), BoxError> {
8923 let limits = test_connection_limits();
8924 let request_id = 99;
8925 let request_frame = request.encode_frame()?;
8926 let budget = LegacyResponseBudget::from_request(
8927 request_frame.message_type,
8928 &request_frame.payload,
8929 limits,
8930 )
8931 .map_err(|error| -> BoxError { format!("{error:?}").into() })?;
8932 let frames = LegacyResponseCodec::encode_response(
8933 request_id,
8934 response,
8935 limits.max_frame_bytes,
8936 limits.max_message_bytes,
8937 )?;
8938 LegacyResponseCodec::decode_response(request_id, request_kind, frames.clone(), None)?;
8939
8940 let mut state = LegacyResponseReadState::new(budget);
8941 for frame in &frames {
8942 state
8943 .validate_frame(request_id, frame)
8944 .map_err(|error| -> BoxError { format!("{error:?}").into() })?;
8945 }
8946 state
8947 .finish()
8948 .map_err(|error| -> BoxError { format!("{error:?}").into() })?;
8949
8950 Ok(())
8951 }
8952
8953 fn test_connection_limits() -> ZakuraConnectionLimits {
8954 let max_protocol_message_len =
8955 u32::try_from(MAX_PROTOCOL_MESSAGE_LEN).expect("protocol message length fits in u32");
8956
8957 ZakuraConnectionLimits {
8958 max_frame_bytes: max_protocol_message_len,
8959 max_message_bytes: max_protocol_message_len,
8960 max_open_streams: 16,
8961 max_inbound_queue_depth: 16,
8962 idle_timeout: Duration::from_secs(1),
8963 prelude_timeout: Duration::from_secs(1),
8964 control_timeout: Duration::from_secs(1),
8965 stream_open_rate_per_second: 16,
8966 message_rate_per_second: 128,
8967 }
8968 }
8969
8970 fn block_hash(byte: u8) -> block::Hash {
8971 block::Hash([byte; 32])
8972 }
8973
8974 fn legacy_tx_id(byte: u8) -> UnminedTxId {
8975 UnminedTxId::from_legacy_id(transaction::Hash([byte; 32]))
8976 }
8977
8978 fn empty_v5_transaction(byte: u8) -> transaction::Transaction {
8979 transaction::Transaction::V5 {
8980 network_upgrade: zakura_chain::parameters::NetworkUpgrade::Nu5,
8981 lock_time: transaction::LockTime::min_lock_time_timestamp(),
8982 expiry_height: block::Height(u32::from(byte)),
8983 inputs: Vec::new(),
8984 outputs: Vec::new(),
8985 sapling_shielded_data: None,
8986 orchard_shielded_data: None,
8987 }
8988 }
8989
8990 #[test]
8991 fn same_kind_streams_share_one_connection_message_budget() {
8992 let clock = crate::zakura::testkit::TestClock::new();
8997 let mut buckets: MessageRateBuckets<crate::zakura::testkit::TestClock> =
8998 MessageRateBuckets::new();
8999 let rate = 4;
9000
9001 let stream_a = message_bucket_for(&mut buckets, 1, rate, clock.clone());
9003 let stream_b = message_bucket_for(&mut buckets, 1, rate, clock.clone());
9004 assert!(
9005 Arc::ptr_eq(&stream_a, &stream_b),
9006 "same kind must reuse one shared bucket"
9007 );
9008 assert_eq!(
9009 buckets.len(),
9010 1,
9011 "no extra bucket created for the same kind"
9012 );
9013
9014 let take = |bucket: &SharedMessageBucket<crate::zakura::testkit::TestClock>| {
9016 bucket
9017 .lock()
9018 .expect("test bucket mutex is never poisoned")
9019 .try_take()
9020 };
9021 let mut accepted = 0;
9022 for _ in 0..rate as usize * 4 {
9023 if take(&stream_a) {
9024 accepted += 1;
9025 }
9026 if take(&stream_b) {
9027 accepted += 1;
9028 }
9029 }
9030 assert_eq!(
9031 accepted, rate as usize,
9032 "the second same-kind stream must NOT get a fresh full budget"
9033 );
9034
9035 clock.advance(Duration::from_secs(1));
9038 let mut refilled = 0;
9039 for _ in 0..rate as usize * 4 {
9040 if take(&stream_a) {
9041 refilled += 1;
9042 }
9043 if take(&stream_b) {
9044 refilled += 1;
9045 }
9046 }
9047 assert_eq!(
9048 refilled, rate as usize,
9049 "refill restores one shared budget, not one-per-stream"
9050 );
9051 }
9052
9053 #[test]
9054 fn different_kinds_get_independent_message_budgets() {
9055 let clock = crate::zakura::testkit::TestClock::new();
9058 let mut buckets: MessageRateBuckets<crate::zakura::testkit::TestClock> =
9059 MessageRateBuckets::new();
9060
9061 let request = message_bucket_for(&mut buckets, 1, 2, clock.clone());
9062 let gossip = message_bucket_for(&mut buckets, 2, 2, clock.clone());
9063 assert!(
9064 !Arc::ptr_eq(&request, &gossip),
9065 "distinct kinds must not share a bucket"
9066 );
9067 assert_eq!(buckets.len(), 2);
9068
9069 let take = |bucket: &SharedMessageBucket<crate::zakura::testkit::TestClock>| {
9070 bucket
9071 .lock()
9072 .expect("test bucket mutex is never poisoned")
9073 .try_take()
9074 };
9075 assert!(take(&request));
9077 assert!(take(&request));
9078 assert!(!take(&request));
9079 assert!(take(&gossip));
9080 assert!(take(&gossip));
9081 assert!(!take(&gossip));
9082 }
9083
9084 #[test]
9085 fn idle_invariant_keeps_app_timeout_below_quic_timeout() {
9086 let limits = ZakuraLocalLimits::from_config(&Config::default());
9087
9088 validate_idle_invariant(&limits).expect("default Zakura limits satisfy idle invariant");
9089 assert!(
9090 (limits.initial_limits().idle_timeout_millis as u128)
9091 < limits.quic_idle_timeout.as_millis()
9092 );
9093 }
9094
9095 #[test]
9114 fn requester_response_budget_is_capped_for_large_inventory_request() {
9115 let limits = test_connection_limits();
9116 assert_eq!(
9117 limits.max_message_bytes as usize, MAX_PROTOCOL_MESSAGE_LEN,
9118 "fixture should negotiate the protocol-max message cap so the unclamped \
9119 budget is maximal",
9120 );
9121
9122 let max_items =
9123 usize::try_from(MAX_TX_INV_IN_SENT_MESSAGE).expect("inventory cap fits in usize");
9124
9125 let blocks = LegacyRequestFrame::BlocksByHash(vec![block_hash(7); max_items])
9128 .encode_frame()
9129 .expect("max-inventory blocks request encodes");
9130 let blocks_budget =
9131 LegacyResponseBudget::from_request(blocks.message_type, &blocks.payload, limits)
9132 .expect("budget derives from a max-inventory blocks request");
9133 assert!(
9134 blocks_budget.max_bytes <= LEGACY_RESPONSE_MAX_AGGREGATE_BYTES,
9135 "BlocksByHash retained-frame budget {} must be clamped to the aggregate cap {}",
9136 blocks_budget.max_bytes,
9137 LEGACY_RESPONSE_MAX_AGGREGATE_BYTES,
9138 );
9139
9140 let txs = LegacyRequestFrame::TransactionsById(vec![legacy_tx_id(7); max_items])
9142 .encode_frame()
9143 .expect("max-inventory transactions request encodes");
9144 let txs_budget = LegacyResponseBudget::from_request(txs.message_type, &txs.payload, limits)
9145 .expect("budget derives from a max-inventory transactions request");
9146 assert!(
9147 txs_budget.max_bytes <= LEGACY_RESPONSE_MAX_AGGREGATE_BYTES,
9148 "TransactionsById retained-frame budget {} must be clamped to the aggregate cap {}",
9149 txs_budget.max_bytes,
9150 LEGACY_RESPONSE_MAX_AGGREGATE_BYTES,
9151 );
9152
9153 let small = LegacyRequestFrame::BlocksByHash(vec![block_hash(1)])
9157 .encode_frame()
9158 .expect("single-item blocks request encodes");
9159 let small_budget =
9160 LegacyResponseBudget::from_request(small.message_type, &small.payload, limits)
9161 .expect("budget derives from a single-item blocks request");
9162 assert!(
9163 small_budget.max_bytes >= limits.max_message_bytes as usize,
9164 "a single-item request must still permit one full response message; budget was {}",
9165 small_budget.max_bytes,
9166 );
9167 assert!(
9168 small_budget.max_bytes <= LEGACY_RESPONSE_MAX_AGGREGATE_BYTES,
9169 "even a single-item budget stays within the aggregate cap",
9170 );
9171 }
9172
9173 #[test]
9192 fn nil_response_to_inventory_or_ping_request_is_not_fail_closed() {
9193 let limits = test_connection_limits();
9194 let request_id = 99;
9195
9196 let cases: [(LegacyRequestFrame, LegacyRequestKind); 3] = [
9197 (
9198 LegacyRequestFrame::BlocksByHash(vec![block_hash(1)]),
9199 LegacyRequestKind::Blocks,
9200 ),
9201 (
9202 LegacyRequestFrame::TransactionsById(vec![legacy_tx_id(2)]),
9203 LegacyRequestKind::Transactions,
9204 ),
9205 (LegacyRequestFrame::Ping, LegacyRequestKind::Ping),
9206 ];
9207
9208 for (request, request_kind) in cases {
9209 let request_frame = request.encode_frame().expect("request frame encodes");
9210 let budget = LegacyResponseBudget::from_request(
9211 request_frame.message_type,
9212 &request_frame.payload,
9213 limits,
9214 )
9215 .expect("budget derives from request");
9216
9217 let nil_frames = LegacyResponseCodec::encode_response(
9220 request_id,
9221 Response::Nil,
9222 limits.max_frame_bytes,
9223 limits.max_message_bytes,
9224 )
9225 .expect("nil response encodes");
9226
9227 let decoded = LegacyResponseCodec::decode_response(
9229 request_id,
9230 request_kind,
9231 nil_frames.clone(),
9232 None,
9233 );
9234 assert!(
9235 decoded.is_err(),
9236 "decode_response must reject a bare NIL for {request_kind:?}",
9237 );
9238
9239 let mut state = LegacyResponseReadState::new(budget);
9243 let mut validate = Ok(());
9244 for frame in &nil_frames {
9245 validate = state.validate_frame(request_id, frame);
9246 if validate.is_err() {
9247 break;
9248 }
9249 }
9250 let validate = validate.and_then(|()| state.finish());
9251 assert!(
9252 matches!(validate, Err(OutboundRequestError::Fatal(_))),
9253 "transport must fail closed (Fatal) on a NIL answer to {request_kind:?} so the \
9254 request stream worker disconnects the peer; got {validate:?}",
9255 );
9256 }
9257 }
9258
9259 #[tokio::test]
9280 async fn inbound_accept_enforces_per_ip_cap() -> Result<(), BoxError> {
9281 let _guard = zakura_test::init();
9282
9283 async fn try_register(
9287 supervisor: &ZakuraSupervisorHandle,
9288 peer: &ZakuraPeerId,
9289 remote_ip: Option<IpAddr>,
9290 ) -> ZakuraRegistration {
9291 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
9292 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
9293 supervisor
9294 .register(
9295 test_conn_id(),
9296 peer.clone(),
9297 remote_ip,
9298 [peer.as_bytes()[0]; TRANSCRIPT_HASH_BYTES],
9299 outbound_handle,
9300 CancellationToken::new(),
9301 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
9302 )
9303 .await
9304 }
9305 let ip: IpAddr = "203.0.113.7".parse().expect("test ip parses");
9306 let supervisor = ZakuraSupervisorHandle::new(1);
9307 assert!(
9308 matches!(
9309 try_register(&supervisor, &test_peer(1), Some(ip)).await,
9310 ZakuraRegistration::Registered { .. }
9311 ),
9312 "first identity from the IP registers",
9313 );
9314 assert!(
9315 matches!(
9316 try_register(&supervisor, &test_peer(2), Some(ip)).await,
9317 ZakuraRegistration::Rejected(ZakuraRejectReason::ResourceLimit)
9318 ),
9319 "a second distinct identity from the same IP must be Rejected(ResourceLimit) at cap 1",
9320 );
9321
9322 let limits = ZakuraLocalLimits::from_config(&Config::default());
9327 assert!(
9328 limits.max_connections > 1,
9329 "global admission cap must exceed the per-IP cap so the second same-IP identity is \
9330 turned away by the per-IP cap rather than the global gate",
9331 );
9332 let server_ep = LocalEndpointFactory::with_transport_config(limits.transport_config())
9333 .endpoint(880)
9334 .await?;
9335 let supervisor = ZakuraSupervisorHandle::new(1);
9336 let handler = ZakuraProtocolHandler::new(
9337 supervisor.clone(),
9338 Network::Mainnet,
9339 ZakuraHandshakeConfig::for_network(&Network::Mainnet),
9340 limits.clone(),
9341 )
9342 .with_endpoint(server_ep.clone());
9343 let router = Router::builder(server_ep)
9344 .accept(P2P_V2_ALPN, handler)
9345 .spawn();
9346 let full_addr = router.endpoint().node_addr().initialized().await;
9352 let loopback_addr = NodeAddr::new(full_addr.node_id).with_direct_addresses(
9353 full_addr
9354 .direct_addresses()
9355 .copied()
9356 .filter(|addr| addr.is_ipv4() && addr.ip().is_loopback()),
9357 );
9358 assert!(
9359 loopback_addr.direct_addresses().next().is_some(),
9360 "server must advertise an IPv4 loopback direct address",
9361 );
9362 let server_addr = loopback_addr;
9363
9364 async fn connect_native(
9368 server_addr: &NodeAddr,
9369 seed: u64,
9370 limits: &ZakuraLocalLimits,
9371 ) -> Result<(Endpoint, Connection), BoxError> {
9372 let endpoint = LocalEndpointFactory::with_transport_config(limits.transport_config())
9373 .endpoint(seed)
9374 .await?;
9375 endpoint.add_node_addr(server_addr.clone())?;
9376 let connection = endpoint.connect(server_addr.clone(), P2P_V2_ALPN).await?;
9377 Ok((endpoint, connection))
9378 }
9379 async fn run_handshake(
9380 endpoint: &Endpoint,
9381 connection: &Connection,
9382 limits: &ZakuraLocalLimits,
9383 ) -> Result<(), BoxError> {
9384 let config = ZakuraHandshakeConfig::for_network(&Config::default().network);
9385 let local_peer_id = ZakuraPeerId::new(endpoint.node_id().as_bytes().to_vec())?;
9386 run_native_initiator_handshake_without_trace(
9387 connection,
9388 limits,
9389 &config,
9390 &local_peer_id,
9391 )
9392 .await?;
9393 Ok(())
9394 }
9395
9396 let (_ep1, _conn1) = connect_native(&server_addr, 881, &limits).await?;
9399 run_handshake(&_ep1, &_conn1, &limits).await?;
9400 let mut first_registered = 0;
9401 for _ in 0..200 {
9402 first_registered = supervisor.registered_ids().await.len();
9403 if first_registered >= 1 {
9404 break;
9405 }
9406 tokio::time::sleep(Duration::from_millis(25)).await;
9407 }
9408 assert_eq!(
9409 first_registered, 1,
9410 "the first inbound identity from the source IP must register (a real source IP was \
9411 resolved from the endpoint and counted against the per-IP cap)",
9412 );
9413
9414 let (_ep2, conn2) = connect_native(&server_addr, 882, &limits).await?;
9421 let _ = run_handshake(&_ep2, &conn2, &limits).await;
9422 let mut rejected_close = false;
9423 for _ in 0..400 {
9424 if supervisor.registered_ids().await.len() >= 2 {
9425 break;
9426 }
9427 if matches!(
9428 conn2.close_reason(),
9429 Some(iroh::endpoint::ConnectionError::ApplicationClosed(ref close))
9430 if close.error_code == VarInt::from_u32(ZAKURA_CLOSE_RESOURCE)
9431 ) {
9432 rejected_close = true;
9433 break;
9434 }
9435 tokio::time::sleep(Duration::from_millis(25)).await;
9436 }
9437 let registered = supervisor.registered_ids().await.len();
9438 assert!(
9439 rejected_close && registered == 1,
9440 "a second distinct identity from the same source IP must be rejected by the per-IP \
9441 cap with a resource-limit close (resource_close={rejected_close}, \
9442 registered={registered}); before the fix the inbound accept passed remote_ip = None, \
9443 so both identities registered and one source IP could exhaust the connection budget",
9444 );
9445
9446 router.shutdown().await?;
9447 Ok(())
9448 }
9449}