1use std::{
4 collections::{HashMap, HashSet},
5 future,
6 io::{Cursor, Read},
7 net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
8 path::PathBuf,
9 sync::{
10 atomic::{AtomicU64, Ordering},
11 Arc, Mutex as StdMutex,
12 },
13 time::Duration,
14};
15
16use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
17use iroh::Watcher as _;
18use iroh::{
19 endpoint::{
20 Connection, ConnectionType, Endpoint, RecvStream, SendStream, TransportConfig, VarInt,
21 },
22 protocol::{AcceptError, ProtocolHandler, Router},
23 NodeAddr, NodeId, SecretKey,
24};
25use rand::{rngs::OsRng, RngCore};
26use thiserror::Error;
27use tokio::{
28 sync::{mpsc, oneshot, watch, Mutex, OwnedSemaphorePermit, Semaphore},
29 task::{AbortHandle, JoinHandle, JoinSet},
30 time::{timeout, Instant},
31};
32use tokio_util::sync::CancellationToken;
33use zakura_chain::{
34 block::{self, Block, CountedHeader},
35 parameters::{Network, NetworkKind},
36 serialization::{CompactSizeMessage, ZcashDeserialize, MAX_HEADERS_PER_MESSAGE},
37 transaction::Transaction,
38};
39
40use super::discovery::{native_dial_supervised, spawn_native_bootstrap_dialer, RedialPolicy};
41use super::{
42 trace::{
43 peer_label as trace_peer_label, reject_reason_label, ZakuraTrace, CONN_TABLE,
44 HANDSHAKE_TABLE, RATELIMIT_TABLE, STREAM_TABLE,
45 },
46 ZakuraTraceEvent,
47};
48use crate::{
49 protocol::external::InventoryHash,
50 zakura::{
51 direct_endpoint_builder, drive_header_sync_actions, spawn_block_sync_reactor,
52 spawn_header_sync_reactor, BlockSyncAction, BlockSyncFrontiers, BlockSyncHandle,
53 BlockSyncService, BlockSyncStartup, Clock, CloseCause, Frame, FramedRecv, FramedSend,
54 Frontier, FrontierChange, FrontierUpdate, HeaderSyncAction, HeaderSyncFrontiers,
55 HeaderSyncPassthroughService, HeaderSyncService, HeaderSyncStartup, Peer, RealClock,
56 Service, ServicePeerDirection, ServiceRegistry, ServiceStream, SinkReject, Stream,
57 StreamMode, StreamPrelude, ZakuraAcceptedLimits, ZakuraBlockSyncConfig, ZakuraConnId,
58 ZakuraControlAck, ZakuraControlHello, ZakuraControlRole, ZakuraControlValidation,
59 ZakuraHandshakeConfig, ZakuraHandshakePath, ZakuraHeaderSyncConfig, ZakuraInitialLimits,
60 ZakuraLimits, ZakuraPeerId, ZakuraPeerSupervisor, ZakuraProtocolError, ZakuraRejectReason,
61 ZakuraSyncExchange, ZakuraUpgradeOutcome, CONTROL_ACK_MAGIC, CONTROL_HELLO_MAGIC,
62 CONTROL_VERSION, FRAME_HEADER_BYTES, LOCAL_MAX_CONTROL_FRAME_BYTES, MAX_BS_FRAME_BYTES,
63 MAX_CONTROL_PAYLOAD_BYTES, MAX_HS_MESSAGE_BYTES, P2P_V2_ALPN, STREAM_PRELUDE_MAGIC,
64 TRANSCRIPT_HASH_BYTES, ZAKURA_HEADER_SYNC_STREAM_VERSION, ZAKURA_PROTOCOL_VERSION_1,
65 ZAKURA_STREAM_BLOCK_SYNC, ZAKURA_STREAM_HEADER_SYNC,
66 },
67};
68use crate::{BoxError, Config, MAX_TX_INV_IN_SENT_MESSAGE};
69
70pub const DEFAULT_ZAKURA_MAX_CONNECTIONS: usize = 256;
72pub const DEFAULT_ZAKURA_MAX_CONNS_PER_IP: usize = 16;
74pub const DEFAULT_ZAKURA_MAX_PENDING_HANDSHAKES: usize = 32;
76pub const DEFAULT_ZAKURA_STREAM_OPEN_RATE_PER_SECOND: u32 = 32;
78pub const DEFAULT_ZAKURA_MESSAGE_RATE_PER_SECOND: u32 = 2048;
88pub const DEFAULT_ZAKURA_LISTEN_ADDR: SocketAddr =
90 SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 8234));
91pub const DEFAULT_MAINNET_ZAKURA_BOOTSTRAP_PEERS: &[&str] = &[
93 "1398f62c6d1a457c51ba6a4b5f3dbd2f69fca93216218dc8997e416bd17d93ca@165.22.54.66:8234",
94 "fd1724385aa0c75b64fb78cd602fa1d991fdebf76b13c58ed702eac835e9f618@104.131.184.123:8234",
95 "9ec67ad6834bc2ca0d659c240e042d3446c37cabcc092b527d459c87d938b4a4@159.65.183.89:8234",
96 "bd3dc5d2a3d44c6bf90e364bf446231dbf9737e38a562ccf9e91ea631ea59b22@143.244.184.176:8234",
97 "14ab98fa0c4b07d40119e1dbc9f3c36d20c8f226ae5ba4216218a2034f148e57@159.203.38.10:8234",
98 "681d21b18644cd82ec13256a97f92bec1fff815683ef6f65dc7c993f098a4fe5@64.227.44.93:8234",
99 "058b3f20dc9bef7bb447f94d7663d793cfbc036720f97e52d7f13661b21818e1@161.35.156.226:8234",
100 "291323d78eb7186c3fa225ef5e305e95363e0ef06d42dca91bd4ef0254aed1ae@139.59.64.115:8234",
101 "85e425233a68697d4be91dd5d542305a8a327cd06d992d53c0913cef2fa75084@168.144.173.250:8234",
102];
103pub const DEFAULT_TESTNET_ZAKURA_BOOTSTRAP_PEERS: &[&str] = &[
105 "57ad39fad4f0bca46cf1ea831772a99d5027b372fef2be5a0ea68e1b5bb4da49@167.99.103.111:8234",
106 "4faac8f988a7820690d63b57a385cd6f833638b068e774550712c05e4b692426@167.99.110.145:8234",
107 "9ce6b95aa197d169399788fe01dd8a88140e81d23d00b4739aeeb1113c6247a2@138.68.229.254:8234",
108 "2bbb907b5d90598ef49f2e637066586b311a64587479be6ed43e8388587fcd2a@164.92.209.78:8234",
109 "50999835f48f4a048c0e9042e5332844c9673943d7fab1f7e993bae698c27ea3@206.189.148.0:8234",
110];
111pub const DEFAULT_ZAKURA_BOOTSTRAP_PEERS: &[&str] = DEFAULT_MAINNET_ZAKURA_BOOTSTRAP_PEERS;
113pub const DEFAULT_ZAKURA_PRELUDE_TIMEOUT: Duration = Duration::from_secs(3);
115pub const DEFAULT_ZAKURA_CONTROL_TIMEOUT: Duration = Duration::from_secs(10);
117pub const DEFAULT_ZAKURA_QUIC_IDLE_TIMEOUT: Duration = Duration::from_secs(150);
125pub const DEFAULT_ZAKURA_KEEP_ALIVE_INTERVAL: Duration = Duration::from_secs(10);
127pub const ZAKURA_DUPLICATE_EVICT_MIN_AGE: Duration = Duration::from_secs(300);
137pub const ZAKURA_SAME_IP_DUPLICATE_EVICT_MIN_AGE: Duration = Duration::from_secs(5);
149pub const DEFAULT_ZAKURA_STREAM_RECEIVE_WINDOW: u32 = 32 * 1024 * 1024;
151pub const DEFAULT_ZAKURA_RECEIVE_WINDOW: u32 = 32 * 1024 * 1024;
153pub const DEFAULT_ZAKURA_SEND_WINDOW: u64 = 32 * 1024 * 1024;
155pub const DEFAULT_ZAKURA_REDIAL_INITIAL_BACKOFF: Duration = Duration::from_secs(1);
157pub const DEFAULT_ZAKURA_REDIAL_MAX_BACKOFF: Duration = Duration::from_secs(30);
159const CONTROL_LENGTH_BYTES: usize = 4;
160const STREAM_PRELUDE_FIXED_BYTES: usize = 4 + 2 + 2 + 1;
161const STREAM_PRELUDE_REQUEST_ID_FLAG_OFFSET: usize = STREAM_PRELUDE_FIXED_BYTES - 1;
162const STREAM_PRELUDE_REQUEST_ID_BYTES: usize = 8;
163const STREAM_PRELUDE_CAP_BYTES: usize = 4;
164const STREAM_WORKER_DRAIN_TIMEOUT: Duration = Duration::from_secs(1);
165const ORDERED_STREAM_REOPEN_BACKOFF: Duration = Duration::from_millis(250);
166const ORDERED_STREAM_REOPEN_BACKOFF_CAP: Duration = Duration::from_secs(8);
167const OUTBOUND_STREAM_WRITE_TIMEOUT: Duration = Duration::from_secs(10);
168const OUTBOUND_REQUEST_RESPONSE_TIMEOUT: Duration = Duration::from_secs(30);
169const LEGACY_GOSSIP_STREAM_KIND: u16 = 2;
172const LEGACY_REQUEST_STREAM_KIND: u16 = 3;
173const DISCOVERY_STREAM_KIND: u16 = 4;
174const HEADER_SYNC_STREAM_KIND: u16 = ZAKURA_STREAM_HEADER_SYNC;
175const LEGACY_REQUEST_BLOCKS_BY_HASH: u16 = 3;
176const LEGACY_REQUEST_TRANSACTIONS_BY_ID: u16 = 4;
177const LEGACY_RESPONSE_BLOCK: u16 = 5;
178const LEGACY_RESPONSE_TRANSACTION: u16 = 6;
179const LEGACY_RESPONSE_MISSING_BLOCKS: u16 = 7;
180const LEGACY_RESPONSE_MISSING_TRANSACTIONS: u16 = 8;
181const LEGACY_REQUEST_FIND_BLOCKS: u16 = 9;
182const LEGACY_REQUEST_FIND_HEADERS: u16 = 10;
183const LEGACY_REQUEST_MEMPOOL_TRANSACTION_IDS: u16 = 11;
184const LEGACY_REQUEST_PING: u16 = 12;
185const LEGACY_REQUEST_PUSH_TRANSACTION: u16 = 13;
186const LEGACY_RESPONSE_BLOCK_HASHES: u16 = 14;
187const LEGACY_RESPONSE_BLOCK_HEADERS: u16 = 15;
188const LEGACY_RESPONSE_TRANSACTION_IDS: u16 = 16;
189const LEGACY_RESPONSE_PONG: u16 = 17;
190const LEGACY_RESPONSE_NIL: u16 = 18;
191const LEGACY_RESPONSE_REQUEST_ID_BYTES: usize = 8;
192const LEGACY_RESPONSE_CHUNK_HEADER_BYTES: usize = LEGACY_RESPONSE_REQUEST_ID_BYTES + 1;
193const LEGACY_COMPACT_SIZE_PREFIX_BYTES: usize = 9;
194const LEGACY_BLOCK_HASH_BYTES: usize = 32;
195const LEGACY_INVENTORY_HASH_BYTES: usize = 36;
196const LEGACY_RESPONSE_MAX_FRAMES_PER_ITEM: usize = 8;
197const LEGACY_RESPONSE_MAX_AGGREGATE_BYTES: usize =
210 8 * zakura_chain::serialization::MAX_PROTOCOL_MESSAGE_LEN;
211const _: () = assert!(LEGACY_GOSSIP_STREAM_KIND == super::legacy_gossip::ZAKURA_STREAM_GOSSIP);
212const _: () =
213 assert!(LEGACY_REQUEST_STREAM_KIND == super::legacy_gossip::ZAKURA_STREAM_LEGACY_REQUESTS);
214const _: () = assert!(DISCOVERY_STREAM_KIND == super::discovery::ZAKURA_STREAM_DISCOVERY);
215const _: () = assert!(HEADER_SYNC_STREAM_KIND == super::header_sync::ZAKURA_STREAM_HEADER_SYNC);
216const _: () = assert!(ZAKURA_STREAM_VERSION_7 == ZAKURA_HEADER_SYNC_STREAM_VERSION);
217const _: () =
218 assert!(LEGACY_REQUEST_BLOCKS_BY_HASH == super::legacy_gossip::MSG_REQUEST_BLOCKS_BY_HASH);
219const _: () = assert!(
220 LEGACY_REQUEST_TRANSACTIONS_BY_ID == super::legacy_gossip::MSG_REQUEST_TRANSACTIONS_BY_ID
221);
222const _: () = assert!(LEGACY_RESPONSE_BLOCK == super::legacy_gossip::MSG_RESPONSE_BLOCK);
223const _: () =
224 assert!(LEGACY_RESPONSE_TRANSACTION == super::legacy_gossip::MSG_RESPONSE_TRANSACTION);
225const _: () =
226 assert!(LEGACY_RESPONSE_MISSING_BLOCKS == super::legacy_gossip::MSG_RESPONSE_MISSING_BLOCKS);
227const _: () = assert!(
228 LEGACY_RESPONSE_MISSING_TRANSACTIONS == super::legacy_gossip::MSG_RESPONSE_MISSING_TRANSACTIONS
229);
230const _: () = assert!(LEGACY_REQUEST_FIND_BLOCKS == super::legacy_gossip::MSG_REQUEST_FIND_BLOCKS);
231const _: () =
232 assert!(LEGACY_REQUEST_FIND_HEADERS == super::legacy_gossip::MSG_REQUEST_FIND_HEADERS);
233const _: () = assert!(
234 LEGACY_REQUEST_MEMPOOL_TRANSACTION_IDS
235 == super::legacy_gossip::MSG_REQUEST_MEMPOOL_TRANSACTION_IDS
236);
237const _: () = assert!(LEGACY_REQUEST_PING == super::legacy_gossip::MSG_REQUEST_PING);
238const _: () =
239 assert!(LEGACY_REQUEST_PUSH_TRANSACTION == super::legacy_gossip::MSG_REQUEST_PUSH_TRANSACTION);
240const _: () =
241 assert!(LEGACY_RESPONSE_BLOCK_HASHES == super::legacy_gossip::MSG_RESPONSE_BLOCK_HASHES);
242const _: () =
243 assert!(LEGACY_RESPONSE_BLOCK_HEADERS == super::legacy_gossip::MSG_RESPONSE_BLOCK_HEADERS);
244const _: () =
245 assert!(LEGACY_RESPONSE_TRANSACTION_IDS == super::legacy_gossip::MSG_RESPONSE_TRANSACTION_IDS);
246const _: () = assert!(LEGACY_RESPONSE_PONG == super::legacy_gossip::MSG_RESPONSE_PONG);
247const _: () = assert!(LEGACY_RESPONSE_NIL == super::legacy_gossip::MSG_RESPONSE_NIL);
248const ZAKURA_CLOSE_NEUTRAL: u32 = 0;
249const ZAKURA_CLOSE_RESOURCE: u32 = 1;
250const ZAKURA_CLOSE_BAD_PRELUDE: u32 = 2;
251const ZAKURA_CLOSE_RATE_LIMIT: u32 = 3;
252const ZAKURA_CLOSE_OVERSIZE: u32 = 4;
253const ZAKURA_CLOSE_UNKNOWN_STREAM: u32 = 5;
258
259#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
261#[serde(deny_unknown_fields, default)]
262pub struct ZakuraConfig {
263 pub bootstrap_peers: Vec<String>,
269 pub listen_addr: Option<SocketAddr>,
279 pub max_connections: usize,
281 pub max_connections_per_ip: usize,
288 pub max_pending_handshakes: usize,
290 pub stream_open_rate_per_second: u32,
292 pub message_rate_per_second: u32,
294 pub trace_dir: Option<PathBuf>,
299 pub header_sync: ZakuraHeaderSyncConfig,
301 pub block_sync: ZakuraBlockSyncConfig,
303 pub dev_network: Option<String>,
316}
317
318impl Default for ZakuraConfig {
319 fn default() -> Self {
320 Self {
321 bootstrap_peers: DEFAULT_ZAKURA_BOOTSTRAP_PEERS
322 .iter()
323 .map(ToString::to_string)
324 .collect(),
325 listen_addr: Some(DEFAULT_ZAKURA_LISTEN_ADDR),
326 max_connections: DEFAULT_ZAKURA_MAX_CONNECTIONS,
327 max_connections_per_ip: DEFAULT_ZAKURA_MAX_CONNS_PER_IP,
328 max_pending_handshakes: DEFAULT_ZAKURA_MAX_PENDING_HANDSHAKES,
329 stream_open_rate_per_second: DEFAULT_ZAKURA_STREAM_OPEN_RATE_PER_SECOND,
330 message_rate_per_second: DEFAULT_ZAKURA_MESSAGE_RATE_PER_SECOND,
331 trace_dir: None,
332 header_sync: ZakuraHeaderSyncConfig::default(),
333 block_sync: ZakuraBlockSyncConfig::default(),
334 dev_network: None,
335 }
336 }
337}
338
339impl ZakuraConfig {
340 pub fn max_connections_per_ip(&self) -> usize {
342 if self.max_connections_per_ip == 0 {
343 DEFAULT_ZAKURA_MAX_CONNS_PER_IP
344 } else {
345 self.max_connections_per_ip
346 }
347 }
348
349 pub fn default_bootstrap_peers_for_network(network: &Network) -> Vec<String> {
351 match network.kind() {
352 NetworkKind::Mainnet => {
353 bootstrap_peers_to_strings(DEFAULT_MAINNET_ZAKURA_BOOTSTRAP_PEERS)
354 }
355 NetworkKind::Testnet => {
356 bootstrap_peers_to_strings(DEFAULT_TESTNET_ZAKURA_BOOTSTRAP_PEERS)
357 }
358 NetworkKind::Regtest => Vec::new(),
359 }
360 }
361
362 pub fn apply_network_defaults(&mut self, network: &Network) {
367 if self.bootstrap_peers
368 != bootstrap_peers_to_strings(DEFAULT_MAINNET_ZAKURA_BOOTSTRAP_PEERS)
369 {
370 return;
371 }
372
373 self.bootstrap_peers = Self::default_bootstrap_peers_for_network(network);
374 }
375}
376
377fn bootstrap_peers_to_strings(peers: &[&str]) -> Vec<String> {
378 peers.iter().map(ToString::to_string).collect()
379}
380
381#[derive(Clone, Debug)]
383pub struct ZakuraLocalLimits {
384 pub max_connections: usize,
386 pub max_pending_handshakes: usize,
388 pub quic_idle_timeout: Duration,
390 pub keep_alive_interval: Duration,
392 pub prelude_timeout: Duration,
394 pub control_timeout: Duration,
396 pub stream_open_rate_per_second: u32,
398 pub message_rate_per_second: u32,
400 pub max_frame_bytes: u32,
402 pub max_message_bytes: u32,
404 pub max_open_streams: u16,
406 pub max_inbound_queue_depth: u16,
408}
409
410impl ZakuraLocalLimits {
411 pub fn from_config(config: &Config) -> Self {
413 let handshake = ZakuraHandshakeConfig::for_network(&config.network);
414 Self {
415 max_connections: config.zakura.max_connections.max(1),
416 max_pending_handshakes: config.zakura.max_pending_handshakes.max(1),
417 quic_idle_timeout: DEFAULT_ZAKURA_QUIC_IDLE_TIMEOUT,
418 keep_alive_interval: DEFAULT_ZAKURA_KEEP_ALIVE_INTERVAL,
419 prelude_timeout: DEFAULT_ZAKURA_PRELUDE_TIMEOUT,
420 control_timeout: DEFAULT_ZAKURA_CONTROL_TIMEOUT,
421 stream_open_rate_per_second: config.zakura.stream_open_rate_per_second.max(1),
422 message_rate_per_second: config.zakura.message_rate_per_second.max(1),
423 max_frame_bytes: handshake.max_message_bytes,
424 max_message_bytes: handshake.max_message_bytes,
425 max_open_streams: handshake.max_open_streams,
426 max_inbound_queue_depth: handshake.max_inbound_queue_depth,
427 }
428 }
429
430 pub fn clamp(&self, negotiated: &ZakuraAcceptedLimits) -> ZakuraConnectionLimits {
432 let max_open_streams = negotiated
433 .max_open_streams
434 .min(self.max_open_streams)
435 .max(1);
436 let idle_timeout = Duration::from_millis(
437 u64::from(negotiated.idle_timeout_millis)
438 .min(self.quic_idle_timeout.as_millis().saturating_sub(1) as u64)
439 .max(1),
440 );
441
442 ZakuraConnectionLimits {
443 max_frame_bytes: negotiated.max_frame_bytes.min(self.max_frame_bytes).max(1),
444 max_message_bytes: negotiated
445 .max_message_bytes
446 .min(self.max_message_bytes)
447 .max(1),
448 max_open_streams,
449 max_inbound_queue_depth: negotiated
450 .max_inbound_queue_depth
451 .min(self.max_inbound_queue_depth)
452 .max(1),
453 idle_timeout,
454 prelude_timeout: self.prelude_timeout,
455 control_timeout: self.control_timeout,
456 stream_open_rate_per_second: self.stream_open_rate_per_second,
457 message_rate_per_second: self.message_rate_per_second,
458 }
459 }
460
461 pub fn initial_limits(&self) -> ZakuraInitialLimits {
463 ZakuraLimits {
464 max_frame_bytes: self.max_frame_bytes,
465 max_message_bytes: self.max_message_bytes,
466 max_open_streams: self.max_open_streams,
467 max_inbound_queue_depth: self.max_inbound_queue_depth,
468 idle_timeout_millis: self.quic_idle_timeout.as_millis().saturating_sub(1) as u32,
469 }
470 }
471
472 pub fn transport_config(&self) -> TransportConfig {
474 let mut transport = TransportConfig::default();
475 transport
476 .max_concurrent_bidi_streams(VarInt::from_u32(u32::from(self.max_open_streams)))
477 .max_concurrent_uni_streams(VarInt::from_u32(0))
478 .stream_receive_window(VarInt::from_u32(DEFAULT_ZAKURA_STREAM_RECEIVE_WINDOW))
479 .receive_window(VarInt::from_u32(DEFAULT_ZAKURA_RECEIVE_WINDOW))
480 .send_window(DEFAULT_ZAKURA_SEND_WINDOW)
481 .max_idle_timeout(Some(
482 self.quic_idle_timeout
483 .try_into()
484 .expect("default Zakura idle timeout is a valid QUIC idle timeout"),
485 ))
486 .keep_alive_interval(Some(self.keep_alive_interval))
487 .datagram_receive_buffer_size(None)
488 .datagram_send_buffer_size(0);
489 transport
490 }
491}
492
493#[derive(Copy, Clone, Debug, Eq, PartialEq)]
495pub struct ZakuraConnectionLimits {
496 pub max_frame_bytes: u32,
498 pub max_message_bytes: u32,
500 pub max_open_streams: u16,
502 pub max_inbound_queue_depth: u16,
504 pub idle_timeout: Duration,
506 pub prelude_timeout: Duration,
508 pub control_timeout: Duration,
510 pub stream_open_rate_per_second: u32,
512 pub message_rate_per_second: u32,
514}
515
516#[derive(Debug, Clone)]
518pub struct ZakuraEndpoint {
519 router: Router,
520 supervisor: ZakuraSupervisorHandle,
521 handler: ZakuraProtocolHandler,
522 header_sync: Option<super::HeaderSyncHandle>,
523 block_sync: Option<BlockSyncHandle>,
524 sync_frontier: Option<ZakuraSyncExchange>,
525 header_sync_tasks: Option<Arc<HeaderSyncBackgroundTasks>>,
526 header_sync_actions: Option<Arc<Mutex<Option<mpsc::Receiver<HeaderSyncAction>>>>>,
527 block_sync_actions: Option<Arc<Mutex<Option<mpsc::Receiver<BlockSyncAction>>>>>,
528 upgrade_dials: Arc<StdMutex<HashMap<ZakuraPeerId, AbortHandle>>>,
533}
534
535#[derive(Debug)]
536struct HeaderSyncBackgroundTasks {
537 shutdown: CancellationToken,
538 tasks: Mutex<Vec<JoinHandle<()>>>,
539}
540
541#[derive(Clone, Debug)]
543pub struct ZakuraHeaderSyncDriverStartup {
544 pub frontiers: HeaderSyncFrontiers,
546 pub best_header_tip: Option<(block::Height, block::Hash)>,
548 pub verified_block_tip_hash: block::Hash,
550}
551
552impl ZakuraEndpoint {
553 pub fn connector(&self) -> super::ZakuraHandshakeConnector {
555 super::ZakuraHandshakeConnector::new_with_endpoint(self.clone())
556 }
557
558 pub(crate) async fn local_upgrade_hints(&self) -> (Vec<u8>, Vec<Vec<u8>>) {
564 let endpoint = self.router.endpoint();
565 let node_id = endpoint.node_id().as_bytes().to_vec();
566 let node_addr = endpoint.node_addr().initialized().await;
567 let direct_addresses = node_addr
568 .direct_addresses()
569 .take(super::MAX_IROH_DIRECT_ADDRESSES)
570 .map(|addr| addr.to_string().into_bytes())
571 .collect();
572 (node_id, direct_addresses)
573 }
574
575 pub fn supervisor(&self) -> ZakuraSupervisorHandle {
577 self.supervisor.clone()
578 }
579
580 pub fn trace(&self) -> ZakuraTrace {
582 self.handler.trace.clone()
583 }
584
585 pub fn header_sync(&self) -> Option<super::HeaderSyncHandle> {
587 self.header_sync.clone()
588 }
589
590 pub fn block_sync(&self) -> Option<BlockSyncHandle> {
592 self.block_sync.clone()
593 }
594
595 pub fn subscribe_sync_frontier(&self) -> Option<watch::Receiver<FrontierUpdate>> {
597 self.sync_frontier
598 .as_ref()
599 .map(ZakuraSyncExchange::subscribe_frontier)
600 }
601
602 pub fn current_sync_frontier(&self) -> Option<FrontierUpdate> {
604 self.sync_frontier
605 .as_ref()
606 .map(ZakuraSyncExchange::current_frontier)
607 }
608
609 pub fn publish_sync_frontier(&self, update: FrontierUpdate) {
611 self.publish_sync_frontier_from(update, "unknown");
612 }
613
614 pub fn publish_sync_frontier_from(&self, update: FrontierUpdate, source: &'static str) {
616 if let Some(sync_frontier) = &self.sync_frontier {
617 sync_frontier.publish_frontier(update, source);
618 }
619 }
620
621 pub async fn take_header_sync_actions(&self) -> Option<mpsc::Receiver<HeaderSyncAction>> {
623 let actions = self.header_sync_actions.as_ref()?;
624 actions.lock().await.take()
625 }
626
627 pub async fn take_block_sync_actions(&self) -> Option<mpsc::Receiver<BlockSyncAction>> {
629 let actions = self.block_sync_actions.as_ref()?;
630 actions.lock().await.take()
631 }
632
633 pub fn header_sync_shutdown(&self) -> Option<CancellationToken> {
635 self.header_sync_tasks
636 .as_ref()
637 .map(|tasks| tasks.shutdown.clone())
638 }
639
640 pub async fn shutdown_sync_tasks(&self) {
646 if let Some(tasks) = &self.header_sync_tasks {
647 tasks.shutdown.cancel();
648 let mut tasks = tasks.tasks.lock().await;
649 for task in tasks.drain(..) {
650 let _ = task.await;
651 }
652 }
653 }
654
655 pub(crate) fn background_shutdown_token(&self) -> CancellationToken {
665 self.header_sync_tasks
666 .as_ref()
667 .map(|tasks| tasks.shutdown.clone())
668 .unwrap_or_default()
669 }
670
671 pub async fn push_header_sync_task(&self, task: JoinHandle<()>) {
673 if let Some(tasks) = self.header_sync_tasks.as_ref() {
674 tasks.tasks.lock().await.push(task);
675 }
676 }
677
678 pub async fn push_block_sync_task(&self, task: JoinHandle<()>) {
680 self.push_header_sync_task(task).await;
681 }
682
683 pub async fn node_addr(&self) -> NodeAddr {
685 self.router.endpoint().node_addr().initialized().await
686 }
687
688 pub fn add_node_addr(
690 &self,
691 node_addr: NodeAddr,
692 ) -> Result<(), iroh::endpoint::AddNodeAddrError> {
693 self.router.endpoint().add_node_addr(node_addr)
694 }
695
696 pub fn spawn_native_dial(&self, node_addr: NodeAddr) -> tokio::task::JoinHandle<()> {
708 let endpoint = self.clone();
709 let limits = self.handler.limits.clone();
710 let policy = RedialPolicy::maintain(
711 DEFAULT_ZAKURA_REDIAL_INITIAL_BACKOFF,
712 DEFAULT_ZAKURA_REDIAL_MAX_BACKOFF,
713 );
714 tokio::spawn(native_dial_supervised(endpoint, node_addr, limits, policy))
715 }
716
717 pub(crate) fn ensure_upgrade_native_dial(&self, node_addr: NodeAddr) -> bool {
724 let Ok(peer_id) = ZakuraPeerId::new(node_addr.node_id.as_bytes().to_vec()) else {
725 return false;
726 };
727
728 let mut upgrade_dials = self
734 .upgrade_dials
735 .lock()
736 .expect("Zakura upgrade dial registry mutex is never poisoned");
737 if upgrade_dials.contains_key(&peer_id) {
738 return true;
739 }
740
741 let endpoint = self.clone();
742 let limits = self.handler.limits.clone();
743 let policy = RedialPolicy::maintain(
744 DEFAULT_ZAKURA_REDIAL_INITIAL_BACKOFF,
745 DEFAULT_ZAKURA_REDIAL_MAX_BACKOFF,
746 );
747 let task_peer_id = peer_id.clone();
748 let dial = tokio::spawn(async move {
749 native_dial_supervised(endpoint.clone(), node_addr, limits, policy).await;
750 endpoint
751 .upgrade_dials
752 .lock()
753 .expect("Zakura upgrade dial registry mutex is never poisoned")
754 .remove(&task_peer_id);
755 });
756 upgrade_dials.insert(peer_id, dial.abort_handle());
757 true
758 }
759
760 pub(crate) fn cancel_upgrade_native_dial(&self, peer_id: &ZakuraPeerId) {
772 let handle = self
773 .upgrade_dials
774 .lock()
775 .expect("Zakura upgrade dial registry mutex is never poisoned")
776 .remove(peer_id);
777 if let Some(handle) = handle {
778 handle.abort();
779 }
780 }
781
782 pub(crate) fn has_native_admission_capacity(&self) -> bool {
787 self.handler.admission.available_permits() > 0
788 }
789
790 pub async fn shutdown(&self) {
792 if let Some(tasks) = &self.header_sync_tasks {
793 tasks.shutdown.cancel();
794 let mut tasks = tasks.tasks.lock().await;
795 for mut task in tasks.drain(..) {
796 if timeout(Duration::from_secs(1), &mut task).await.is_err() {
797 task.abort();
798 let _ = task.await;
799 }
800 }
801 }
802 self.supervisor.shutdown();
803 let _ = self.router.shutdown().await;
804 }
805
806 #[cfg(any(test, feature = "zakura-testkit"))]
807 pub(crate) fn from_parts(
808 router: Router,
809 supervisor: ZakuraSupervisorHandle,
810 handler: ZakuraProtocolHandler,
811 ) -> Self {
812 Self {
813 router,
814 supervisor,
815 handler,
816 header_sync: None,
817 block_sync: None,
818 sync_frontier: None,
819 header_sync_tasks: None,
820 header_sync_actions: None,
821 block_sync_actions: None,
822 upgrade_dials: Arc::new(StdMutex::new(HashMap::new())),
823 }
824 }
825
826 #[cfg(any(test, feature = "zakura-testkit"))]
827 #[allow(dead_code)]
828 pub(crate) fn from_parts_with_header_sync(
829 router: Router,
830 supervisor: ZakuraSupervisorHandle,
831 handler: ZakuraProtocolHandler,
832 header_sync: super::HeaderSyncHandle,
833 shutdown: CancellationToken,
834 tasks: Vec<JoinHandle<()>>,
835 actions: Option<mpsc::Receiver<HeaderSyncAction>>,
836 ) -> Self {
837 Self {
838 router,
839 supervisor,
840 handler,
841 header_sync: Some(header_sync),
842 block_sync: None,
843 sync_frontier: None,
844 header_sync_tasks: Some(Arc::new(HeaderSyncBackgroundTasks {
845 shutdown,
846 tasks: Mutex::new(tasks),
847 })),
848 header_sync_actions: actions.map(|actions| Arc::new(Mutex::new(Some(actions)))),
849 block_sync_actions: None,
850 upgrade_dials: Arc::new(StdMutex::new(HashMap::new())),
851 }
852 }
853
854 #[cfg(any(test, feature = "zakura-testkit"))]
855 #[allow(clippy::too_many_arguments)]
856 pub(crate) fn from_parts_with_sync_services(
857 router: Router,
858 supervisor: ZakuraSupervisorHandle,
859 handler: ZakuraProtocolHandler,
860 header_sync: super::HeaderSyncHandle,
861 block_sync: BlockSyncHandle,
862 shutdown: CancellationToken,
863 tasks: Vec<JoinHandle<()>>,
864 header_sync_actions: Option<mpsc::Receiver<HeaderSyncAction>>,
865 block_sync_actions: Option<mpsc::Receiver<BlockSyncAction>>,
866 ) -> Self {
867 Self {
868 router,
869 supervisor,
870 handler,
871 header_sync: Some(header_sync),
872 block_sync: Some(block_sync),
873 sync_frontier: None,
874 header_sync_tasks: Some(Arc::new(HeaderSyncBackgroundTasks {
875 shutdown,
876 tasks: Mutex::new(tasks),
877 })),
878 header_sync_actions: header_sync_actions
879 .map(|actions| Arc::new(Mutex::new(Some(actions)))),
880 block_sync_actions: block_sync_actions
881 .map(|actions| Arc::new(Mutex::new(Some(actions)))),
882 upgrade_dials: Arc::new(StdMutex::new(HashMap::new())),
883 }
884 }
885}
886
887#[derive(Clone, Debug)]
889pub struct ZakuraSupervisorHandle {
890 id: u64,
891 inner: Arc<Mutex<ZakuraSupervisorState>>,
892 shutdown: CancellationToken,
893 peer_set_tx: watch::Sender<Vec<ZakuraPeerId>>,
894}
895
896static NEXT_SUPERVISOR_ID: AtomicU64 = AtomicU64::new(1);
897
898#[derive(Debug)]
899struct ZakuraSupervisorState {
900 supervisor: ZakuraPeerSupervisor,
901 active_by_peer: HashMap<ZakuraPeerId, ZakuraPeerConnectionEntry>,
902 active_by_ip: HashMap<IpAddr, usize>,
903 max_connections_per_ip: usize,
904 next_registration_id: ZakuraConnId,
905}
906
907#[derive(Debug)]
908struct ZakuraPeerConnectionEntry {
909 conn_id: ZakuraConnId,
912 outbound_handle: ZakuraPeerHandle,
913 disconnect_token: CancellationToken,
914 registered_at: Instant,
915 remote_ip: Option<IpAddr>,
916}
917
918impl ZakuraSupervisorState {
919 fn increment_ip(&mut self, remote_ip: Option<IpAddr>) {
920 if let Some(remote_ip) = remote_ip {
921 *self.active_by_ip.entry(remote_ip).or_default() += 1;
922 }
923 }
924
925 fn decrement_ip(&mut self, remote_ip: Option<IpAddr>) {
926 if let Some(remote_ip) = remote_ip {
927 if let Some(count) = self.active_by_ip.get_mut(&remote_ip) {
928 *count = count.saturating_sub(1);
929 if *count == 0 {
930 self.active_by_ip.remove(&remote_ip);
931 }
932 }
933 }
934 }
935
936 #[cfg(debug_assertions)]
937 fn debug_assert_accounting(&self) {
938 let active_by_ip_total: usize = self.active_by_ip.values().sum();
939 debug_assert_eq!(
940 active_by_ip_total,
941 self.active_by_peer
942 .values()
943 .filter(|entry| entry.remote_ip.is_some())
944 .count(),
945 "Zakura active_by_ip totals must match active peer registrations with known IPs",
946 );
947 }
948
949 #[cfg(not(debug_assertions))]
950 fn debug_assert_accounting(&self) {}
951}
952
953#[derive(Clone, Debug)]
955pub struct ZakuraPeerHandle {
956 peer_id: ZakuraPeerId,
957 sender: mpsc::Sender<ZakuraOutboundFrame>,
958}
959
960impl ZakuraPeerHandle {
961 #[cfg(test)]
962 pub(crate) fn new_for_tests(
963 peer_id: ZakuraPeerId,
964 sender: mpsc::Sender<ZakuraOutboundFrame>,
965 ) -> Self {
966 Self { peer_id, sender }
967 }
968
969 pub fn peer_id(&self) -> &ZakuraPeerId {
971 &self.peer_id
972 }
973
974 fn has_outbound_capacity(&self) -> bool {
975 self.sender.capacity() > 0
976 }
977
978 pub async fn request(
980 &self,
981 stream_kind: u16,
982 request_id: u64,
983 message_type: u16,
984 flags: u16,
985 payload: Vec<u8>,
986 ) -> Result<Vec<Frame>, BoxError> {
987 let (completion, completed) = oneshot::channel();
988 let frame = ZakuraOutboundFrame::Request {
989 stream_kind,
990 request_id,
991 message_type,
992 flags,
993 payload,
994 completion,
995 };
996 self.sender
997 .send(frame)
998 .await
999 .map_err(|_| -> BoxError { "Zakura outbound peer queue closed".into() })?;
1000 completed
1001 .await
1002 .map_err(|_| -> BoxError { "Zakura outbound completion dropped".into() })?
1003 }
1004}
1005
1006#[derive(Debug)]
1008pub enum ZakuraOutboundFrame {
1009 Request {
1011 stream_kind: u16,
1013 request_id: u64,
1015 message_type: u16,
1017 flags: u16,
1019 payload: Vec<u8>,
1021 completion: oneshot::Sender<Result<Vec<Frame>, BoxError>>,
1023 },
1024}
1025
1026impl ZakuraSupervisorHandle {
1027 pub fn new(max_connections_per_ip: usize) -> Self {
1029 Self {
1030 id: NEXT_SUPERVISOR_ID.fetch_add(1, Ordering::Relaxed),
1031 inner: Arc::new(Mutex::new(ZakuraSupervisorState {
1032 supervisor: ZakuraPeerSupervisor::default(),
1033 active_by_peer: HashMap::new(),
1034 active_by_ip: HashMap::new(),
1035 max_connections_per_ip: max_connections_per_ip.max(1),
1036 next_registration_id: 1,
1037 })),
1038 shutdown: CancellationToken::new(),
1039 peer_set_tx: watch::channel(Vec::new()).0,
1040 }
1041 }
1042
1043 pub async fn registered_ids(&self) -> Vec<ZakuraPeerId> {
1045 let state = self.inner.lock().await;
1046 state.active_by_peer.keys().cloned().collect()
1047 }
1048
1049 pub(crate) fn id(&self) -> u64 {
1050 self.id
1051 }
1052
1053 pub async fn outbound_peer_handles(&self) -> Vec<ZakuraPeerHandle> {
1055 let state = self.inner.lock().await;
1056 state
1057 .active_by_peer
1058 .values()
1059 .map(|entry| &entry.outbound_handle)
1060 .filter(|handle| handle.has_outbound_capacity())
1061 .cloned()
1062 .collect()
1063 }
1064
1065 pub fn subscribe(&self) -> watch::Receiver<Vec<ZakuraPeerId>> {
1067 self.peer_set_tx.subscribe()
1068 }
1069
1070 pub async fn disconnect_peer(&self, peer_id: &ZakuraPeerId) -> bool {
1072 let token = {
1073 let state = self.inner.lock().await;
1074 state
1075 .active_by_peer
1076 .get(peer_id)
1077 .map(|entry| entry.disconnect_token.clone())
1078 };
1079
1080 if let Some(token) = token {
1081 token.cancel();
1082 true
1083 } else {
1084 false
1085 }
1086 }
1087
1088 #[allow(clippy::too_many_arguments)]
1089 async fn register(
1090 &self,
1091 _trace_conn_id: ZakuraConnId,
1092 peer_id: ZakuraPeerId,
1093 remote_ip: Option<IpAddr>,
1094 transcript_hash: [u8; TRANSCRIPT_HASH_BYTES],
1095 outbound_handle: ZakuraPeerHandle,
1096 disconnect_token: CancellationToken,
1097 _accepted_capabilities: u64,
1098 ) -> ZakuraRegistration {
1099 let mut state = self.inner.lock().await;
1100 let same_ip_duplicate_redial = state
1109 .active_by_peer
1110 .get(&peer_id)
1111 .is_some_and(|entry| entry.remote_ip == remote_ip);
1112 if let Some(remote_ip) = remote_ip {
1113 if !same_ip_duplicate_redial {
1114 let ip_count = state
1115 .active_by_ip
1116 .get(&remote_ip)
1117 .copied()
1118 .unwrap_or_default();
1119 if ip_count >= state.max_connections_per_ip {
1120 metrics::counter!("zakura.p2p.conn.rejected.admission").increment(1);
1121 return ZakuraRegistration::Rejected(ZakuraRejectReason::ResourceLimit);
1122 }
1123 }
1124 }
1125 if state.next_registration_id == u64::MAX {
1126 metrics::counter!("zakura.p2p.conn.rejected.admission").increment(1);
1127 return ZakuraRegistration::Rejected(ZakuraRejectReason::TemporaryUnavailable);
1128 }
1129
1130 match state
1131 .supervisor
1132 .register_authenticated(peer_id.clone(), transcript_hash)
1133 {
1134 ZakuraUpgradeOutcome::Upgraded { .. } => {
1135 let conn_id = state.next_registration_id;
1136 state.next_registration_id += 1;
1137 let entry = ZakuraPeerConnectionEntry {
1138 conn_id,
1139 outbound_handle,
1140 disconnect_token,
1141 registered_at: Instant::now(),
1142 remote_ip,
1143 };
1144 if let Some(old_entry) = state.active_by_peer.insert(peer_id.clone(), entry) {
1145 state.decrement_ip(old_entry.remote_ip);
1146 old_entry.disconnect_token.cancel();
1147 metrics::counter!("zakura.p2p.conn.duplicate.evicted_upgraded").increment(1);
1148 }
1149 state.increment_ip(remote_ip);
1150 state.debug_assert_accounting();
1151 let registered_ids: Vec<_> = state.active_by_peer.keys().cloned().collect();
1152 set_active_connection_gauge(registered_ids.len());
1153 self.peer_set_tx.send_replace(registered_ids);
1154 let disconnect_token = state
1155 .active_by_peer
1156 .get(&peer_id)
1157 .map(|entry| entry.disconnect_token.clone())
1158 .expect("disconnect token exists because this peer was just registered");
1159 ZakuraRegistration::Registered {
1160 conn_id,
1161 peer_id,
1162 remote_ip,
1163 disconnect_token,
1164 }
1165 }
1166 ZakuraUpgradeOutcome::Duplicate { .. } => {
1167 let evict_min_age = if same_ip_duplicate_redial {
1185 ZAKURA_SAME_IP_DUPLICATE_EVICT_MIN_AGE
1186 } else {
1187 ZAKURA_DUPLICATE_EVICT_MIN_AGE
1188 };
1189 if let Some(entry) = state.active_by_peer.get(&peer_id) {
1190 if entry.registered_at.elapsed() >= evict_min_age
1191 && !entry.disconnect_token.is_cancelled()
1192 {
1193 entry.disconnect_token.cancel();
1194 metrics::counter!("zakura.p2p.conn.duplicate.evicted_stale").increment(1);
1195 }
1196 }
1197 ZakuraRegistration::Duplicate { peer_id }
1198 }
1199 ZakuraUpgradeOutcome::Rejected { reason } => ZakuraRegistration::Rejected(reason),
1200 }
1201 }
1202
1203 async fn deregister(&self, peer_id: &ZakuraPeerId, conn_id: ZakuraConnId) {
1204 let mut state = self.inner.lock().await;
1205 let Some(entry) = state.active_by_peer.get(peer_id) else {
1206 state.debug_assert_accounting();
1207 return;
1208 };
1209 if entry.conn_id != conn_id {
1210 state.debug_assert_accounting();
1211 return;
1212 }
1213 let entry = state
1214 .active_by_peer
1215 .remove(peer_id)
1216 .expect("peer entry exists because it was just checked");
1217 state.decrement_ip(entry.remote_ip);
1218 state.supervisor.deregister_authenticated(peer_id);
1219 state.debug_assert_accounting();
1220 let registered_ids: Vec<_> = state.active_by_peer.keys().cloned().collect();
1221 set_active_connection_gauge(registered_ids.len());
1222 self.peer_set_tx.send_replace(registered_ids);
1223 }
1224
1225 fn shutdown(&self) {
1226 self.shutdown.cancel();
1227 }
1228
1229 pub(crate) async fn can_accept_remote_ip_with_in_flight(
1234 &self,
1235 remote_ip: IpAddr,
1236 in_flight_count: usize,
1237 ) -> bool {
1238 let state = self.inner.lock().await;
1239 let active_count = state
1240 .active_by_ip
1241 .get(&remote_ip)
1242 .copied()
1243 .unwrap_or_default();
1244 active_count.saturating_add(in_flight_count) < state.max_connections_per_ip
1245 }
1246}
1247
1248fn set_active_connection_gauge(active_connections: usize) {
1249 metrics::gauge!("zakura.p2p.conn.active").set(active_connections as f64);
1252}
1253
1254#[derive(Debug)]
1255enum ZakuraRegistration {
1256 Registered {
1257 conn_id: ZakuraConnId,
1258 peer_id: ZakuraPeerId,
1259 remote_ip: Option<IpAddr>,
1260 disconnect_token: CancellationToken,
1261 },
1262 Duplicate {
1263 peer_id: ZakuraPeerId,
1264 },
1265 Rejected(ZakuraRejectReason),
1266}
1267
1268#[derive(Debug)]
1269struct RegisteredPeerCleanupGuard {
1270 supervisor: ZakuraSupervisorHandle,
1271 registry: Arc<ServiceRegistry>,
1272 peer_id: ZakuraPeerId,
1273 conn_id: ZakuraConnId,
1274 disconnect_token: CancellationToken,
1275 admitted_capabilities: u64,
1276 armed: bool,
1277}
1278
1279impl RegisteredPeerCleanupGuard {
1280 fn new(
1281 supervisor: ZakuraSupervisorHandle,
1282 registry: Arc<ServiceRegistry>,
1283 peer_id: ZakuraPeerId,
1284 conn_id: ZakuraConnId,
1285 disconnect_token: CancellationToken,
1286 ) -> Self {
1287 Self {
1288 supervisor,
1289 registry,
1290 peer_id,
1291 conn_id,
1292 disconnect_token,
1293 admitted_capabilities: 0,
1294 armed: true,
1295 }
1296 }
1297
1298 fn add_admitted_capabilities(&mut self, capabilities: u64) {
1299 self.admitted_capabilities |= capabilities;
1300 }
1301
1302 fn remove_admitted_services(&mut self) {
1303 if self.admitted_capabilities != 0 {
1304 self.registry
1305 .remove_peer(&self.peer_id, self.conn_id, self.admitted_capabilities);
1306 self.admitted_capabilities = 0;
1307 }
1308 }
1309
1310 async fn cleanup_registered_peer(mut self) {
1311 if self.armed {
1312 self.disconnect_token.cancel();
1313 self.remove_admitted_services();
1314 self.supervisor
1315 .deregister(&self.peer_id, self.conn_id)
1316 .await;
1317 self.armed = false;
1318 }
1319 }
1320}
1321
1322impl Drop for RegisteredPeerCleanupGuard {
1323 fn drop(&mut self) {
1324 if !self.armed {
1325 return;
1326 }
1327
1328 self.disconnect_token.cancel();
1329 self.remove_admitted_services();
1330
1331 let supervisor = self.supervisor.clone();
1332 let peer_id = self.peer_id.clone();
1333 let conn_id = self.conn_id;
1334 if let Ok(handle) = tokio::runtime::Handle::try_current() {
1335 let _task = handle.spawn(async move {
1336 supervisor.deregister(&peer_id, conn_id).await;
1337 });
1338 }
1339 }
1340}
1341
1342#[derive(Clone, Debug)]
1343struct ZakuraConnTrace {
1344 id: u64,
1345 peer_label: Option<Arc<str>>,
1346}
1347
1348impl ZakuraConnTrace {
1349 fn new(trace: &ZakuraTrace, id: u64, peer_id: &ZakuraPeerId) -> Self {
1350 Self {
1351 id,
1352 peer_label: trace
1353 .is_enabled()
1354 .then(|| Arc::<str>::from(trace_peer_label(peer_id))),
1355 }
1356 }
1357
1358 fn without_peer(id: u64) -> Self {
1359 Self {
1360 id,
1361 peer_label: None,
1362 }
1363 }
1364
1365 #[cfg(any(test, feature = "zakura-testkit"))]
1366 fn placeholder() -> Self {
1367 Self::without_peer(0)
1368 }
1369
1370 fn peer(&self) -> Option<&str> {
1371 self.peer_label.as_deref()
1372 }
1373
1374 fn event<'a>(&'a self, event: &'static str) -> ZakuraTraceEvent<'a> {
1375 ZakuraTraceEvent::new(event)
1376 .conn(self.id)
1377 .maybe_peer(self.peer())
1378 }
1379}
1380
1381struct StreamAdmission<'a> {
1382 trace: ZakuraTrace,
1383 conn: ZakuraConnTrace,
1384 peer_id: &'a ZakuraPeerId,
1385 stream_sem: &'a Arc<Semaphore>,
1386 open_limiter: &'a mut TokenBucket,
1387 message_buckets: &'a mut MessageRateBuckets,
1388 workers: &'a mut JoinSet<()>,
1389 limits: ZakuraConnectionLimits,
1390 accepted_capabilities: u64,
1391 connection_token: CancellationToken,
1392 close_cause: CloseCause,
1393 freshness_tx: watch::Sender<Instant>,
1394}
1395
1396impl StreamAdmission<'_> {
1397 fn event(&self, event: &'static str, stream_id: u64) -> ZakuraTraceEvent<'_> {
1398 self.conn.event(event).stream(stream_id)
1399 }
1400}
1401
1402struct ConnectionServeContext {
1403 limits: ZakuraConnectionLimits,
1404 accepted_capabilities: u64,
1405 role: &'static str,
1406 direction: ServicePeerDirection,
1407 transcript_hash: [u8; TRANSCRIPT_HASH_BYTES],
1408 i_open_collision_winner: bool,
1414 conn: ZakuraConnTrace,
1415}
1416
1417struct RegisteredConnectionServeContext {
1418 limits: ZakuraConnectionLimits,
1419 conn: ZakuraConnTrace,
1420 conn_id: ZakuraConnId,
1421 connection_token: CancellationToken,
1422 close_cause: CloseCause,
1423 accepted_capabilities: u64,
1424 is_initiator: bool,
1429 i_open_collision_winner: bool,
1431 direction: ServicePeerDirection,
1432}
1433
1434struct StreamWorkerContext {
1435 trace: ZakuraTrace,
1436 conn: ZakuraConnTrace,
1437 peer_id: ZakuraPeerId,
1438 stream_id: u64,
1439 _permit: OwnedSemaphorePermit,
1440 limits: ZakuraConnectionLimits,
1441 message_bucket: SharedMessageBucket,
1442 connection_token: CancellationToken,
1443 stream_token: CancellationToken,
1444 close_cause: CloseCause,
1445 freshness_tx: watch::Sender<Instant>,
1446}
1447
1448impl StreamWorkerContext {
1449 fn event(&self, event: &'static str) -> ZakuraTraceEvent<'_> {
1450 self.conn.event(event).stream(self.stream_id)
1451 }
1452}
1453
1454struct AdmittedOrderedStream {
1455 kind: u16,
1456 session_id: u64,
1457 recv: FramedRecv,
1458 send: FramedSend,
1459 cancel_token: CancellationToken,
1460}
1461
1462#[derive(Copy, Clone, Debug)]
1463struct OrderedStreamExit {
1464 stream: Stream,
1465 session_id: u64,
1466 opened_locally: bool,
1467}
1468
1469fn should_reopen_ordered_stream(
1478 exited: OrderedStreamExit,
1479 is_initiator: bool,
1480 connection_cancelled: bool,
1481) -> bool {
1482 exited.opened_locally
1483 && is_initiator
1484 && exited.stream.kind == ZAKURA_STREAM_HEADER_SYNC
1485 && !connection_cancelled
1486}
1487
1488fn ordered_stream_reopen_backoff(attempts: u32) -> Duration {
1491 ORDERED_STREAM_REOPEN_BACKOFF
1493 .saturating_mul(1u32 << attempts.min(8))
1494 .min(ORDERED_STREAM_REOPEN_BACKOFF_CAP)
1495}
1496
1497#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1498enum InboundMessageAdmission {
1499 Admit,
1500 Oversize,
1501 Throttled,
1502}
1503
1504fn admit_inbound_message(
1505 payload_len: usize,
1506 context: &StreamWorkerContext,
1507 stream_kind: u16,
1508) -> InboundMessageAdmission {
1509 let stream_kind = stream_kind_label(stream_kind);
1510 let max_message_bytes = usize::try_from(context.limits.max_message_bytes)
1511 .expect("u32 message byte limit fits in usize");
1512 if payload_len > max_message_bytes {
1513 metrics::counter!(
1514 "zakura.p2p.ratelimit.message.oversize",
1515 "stream_kind" => stream_kind,
1516 )
1517 .increment(1);
1518 context.trace.emit(
1519 RATELIMIT_TABLE,
1520 context.event("message.oversize").stream_kind(stream_kind),
1521 );
1522 return InboundMessageAdmission::Oversize;
1523 }
1524
1525 let admitted = {
1526 let mut bucket = context
1527 .message_bucket
1528 .lock()
1529 .expect("Zakura message-rate bucket mutex is never poisoned");
1530 bucket.try_take()
1531 };
1532 if !admitted {
1533 metrics::counter!(
1534 "zakura.p2p.ratelimit.message.throttled",
1535 "stream_kind" => stream_kind,
1536 )
1537 .increment(1);
1538 context.trace.emit(
1539 RATELIMIT_TABLE,
1540 context.event("message.throttled").stream_kind(stream_kind),
1541 );
1542 return InboundMessageAdmission::Throttled;
1543 }
1544
1545 InboundMessageAdmission::Admit
1546}
1547
1548#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1549pub(crate) struct NativeHandshakeNegotiated {
1550 pub(crate) limits: ZakuraAcceptedLimits,
1551 pub(crate) accepted_capabilities: u64,
1552}
1553
1554pub(crate) fn service_registry(
1555 _supervisor: &ZakuraSupervisorHandle,
1556 header_sync: Option<super::HeaderSyncHandle>,
1557 block_sync: Option<BlockSyncHandle>,
1558 block_sync_config: ZakuraBlockSyncConfig,
1559 legacy_service: Arc<dyn Service>,
1560 discovery_service: Arc<dyn Service>,
1561) -> Result<Arc<ServiceRegistry>, BoxError> {
1562 let mut services = vec![legacy_service.clone(), discovery_service];
1563 if let Some(header_sync) = &header_sync {
1564 services.push(Arc::new(HeaderSyncService::new(header_sync.clone())) as Arc<dyn Service>);
1565 } else {
1566 services
1567 .push(Arc::new(HeaderSyncPassthroughService::new(legacy_service)) as Arc<dyn Service>);
1568 }
1569 let block_sync = match block_sync {
1570 Some(block_sync) => BlockSyncService::new_with_handle(block_sync_config, block_sync),
1571 None => match header_sync.as_ref() {
1572 Some(header_sync) => BlockSyncService::new_with_header_tip(
1573 block_sync_config,
1574 header_sync.subscribe_tip(),
1575 ),
1576 None => BlockSyncService::new(block_sync_config),
1577 },
1578 };
1579 services.push(Arc::new(block_sync) as Arc<dyn Service>);
1580
1581 Ok(Arc::new(
1582 ServiceRegistry::new(services).map_err(|error| -> BoxError { Box::new(error) })?,
1583 ))
1584}
1585
1586#[derive(Debug, Clone)]
1588pub struct ZakuraProtocolHandler {
1589 supervisor: ZakuraSupervisorHandle,
1590 handshake_config: ZakuraHandshakeConfig,
1591 limits: ZakuraLocalLimits,
1592 registry: Arc<ServiceRegistry>,
1593 trace: ZakuraTrace,
1594 next_conn_id: Arc<AtomicU64>,
1595 next_stream_id: Arc<AtomicU64>,
1596 admission: Arc<Semaphore>,
1597 pending_handshakes: Arc<Semaphore>,
1598 shutdown: CancellationToken,
1599 endpoint: Option<Endpoint>,
1607}
1608
1609fn random_stream_session_seed() -> u64 {
1610 let mut rng = OsRng;
1611 loop {
1612 let seed = rng.next_u64();
1613 if seed != 0 {
1614 return seed;
1615 }
1616 }
1617}
1618
1619fn inbound_remote_ip(endpoint: &Endpoint, node_id: NodeId) -> Option<IpAddr> {
1631 if let Some(mut conn_type) = endpoint.conn_type(node_id) {
1632 match conn_type.get() {
1633 ConnectionType::Direct(addr) | ConnectionType::Mixed(addr, _) => {
1634 return Some(addr.ip())
1635 }
1636 ConnectionType::Relay(_) | ConnectionType::None => {}
1637 }
1638 }
1639 endpoint.remote_info(node_id).and_then(|info| {
1640 info.addrs
1641 .into_iter()
1642 .filter(|addr| addr.last_payload.is_some())
1643 .min_by_key(|addr| addr.last_payload)
1646 .map(|addr| addr.addr.ip())
1647 })
1648}
1649
1650fn native_connection_transcript_hash(
1651 direction: ServicePeerDirection,
1652 local_node_id: &NodeId,
1653 remote_node_id: &NodeId,
1654) -> [u8; TRANSCRIPT_HASH_BYTES] {
1655 let initiator = match direction {
1656 ServicePeerDirection::Inbound => remote_node_id,
1657 ServicePeerDirection::Outbound => local_node_id,
1658 };
1659 *initiator.as_bytes()
1660}
1661
1662fn i_open_collision_winner(local_node_id: &NodeId, remote_node_id: &NodeId) -> bool {
1669 local_node_id.as_bytes() < remote_node_id.as_bytes()
1670}
1671
1672impl ZakuraProtocolHandler {
1673 pub fn new(
1675 supervisor: ZakuraSupervisorHandle,
1676 network: Network,
1677 handshake_config: ZakuraHandshakeConfig,
1678 limits: ZakuraLocalLimits,
1679 ) -> Self {
1680 Self::new_with_registry(
1681 supervisor,
1682 network,
1683 handshake_config,
1684 limits,
1685 Arc::new(ServiceRegistry::default()),
1686 )
1687 }
1688
1689 pub fn new_with_registry(
1691 supervisor: ZakuraSupervisorHandle,
1692 network: Network,
1693 handshake_config: ZakuraHandshakeConfig,
1694 limits: ZakuraLocalLimits,
1695 registry: Arc<ServiceRegistry>,
1696 ) -> Self {
1697 Self::new_with_registry_and_trace(
1698 supervisor,
1699 network,
1700 handshake_config,
1701 limits,
1702 registry,
1703 ZakuraTrace::noop(),
1704 )
1705 }
1706
1707 pub fn new_with_registry_and_trace(
1709 supervisor: ZakuraSupervisorHandle,
1710 _network: Network,
1711 handshake_config: ZakuraHandshakeConfig,
1712 limits: ZakuraLocalLimits,
1713 registry: Arc<ServiceRegistry>,
1714 trace: ZakuraTrace,
1715 ) -> Self {
1716 let mut handshake_config = handshake_config;
1717 handshake_config.supported_capabilities = registry.supported_capabilities();
1718 Self {
1719 supervisor,
1720 handshake_config,
1721 registry,
1722 trace,
1723 next_conn_id: Arc::new(AtomicU64::new(1)),
1724 next_stream_id: Arc::new(AtomicU64::new(random_stream_session_seed())),
1728 admission: Arc::new(Semaphore::new(limits.max_connections)),
1729 pending_handshakes: Arc::new(Semaphore::new(limits.max_pending_handshakes)),
1730 shutdown: CancellationToken::new(),
1731 limits,
1732 endpoint: None,
1733 }
1734 }
1735
1736 #[cfg(any(test, feature = "zakura-testkit"))]
1741 pub(crate) fn with_supported_capabilities(mut self, supported_capabilities: u64) -> Self {
1742 self.handshake_config.supported_capabilities &= supported_capabilities;
1743 self
1744 }
1745
1746 pub fn with_endpoint(mut self, endpoint: Endpoint) -> Self {
1749 self.endpoint = Some(endpoint);
1750 self
1751 }
1752
1753 async fn accept_connection(&self, connection: Connection) -> Result<(), AcceptError> {
1754 let conn_id = self.next_conn_id.fetch_add(1, Ordering::Relaxed);
1755 let Ok(_admission) = self.admission.clone().try_acquire_owned() else {
1756 metrics::counter!("zakura.p2p.conn.rejected.admission").increment(1);
1757 let conn = ZakuraConnTrace::without_peer(conn_id);
1758 self.trace.emit(
1759 CONN_TABLE,
1760 conn.event("rejected.admission")
1761 .direction("inbound")
1762 .reason("admission"),
1763 );
1764 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"admission");
1765 return Ok(());
1766 };
1767
1768 let remote_node_id = connection.remote_node_id()?;
1769 let remote_peer_id =
1770 ZakuraPeerId::new(remote_node_id.as_bytes().to_vec()).map_err(AcceptError::from_err)?;
1771 let conn = ZakuraConnTrace::new(&self.trace, conn_id, &remote_peer_id);
1772
1773 let negotiated = match self
1774 .run_native_responder_handshake_with_permit(&connection, &remote_peer_id, &conn)
1775 .await
1776 {
1777 Ok(negotiated) => negotiated,
1778 Err(ZakuraHandlerError::ResourceLimit("pending handshake")) => return Ok(()),
1779 Err(error) => {
1780 debug!(?error, "Zakura control handshake failed");
1781 connection.close(VarInt::from_u32(ZAKURA_CLOSE_NEUTRAL), b"control handshake");
1782 return Ok(());
1783 }
1784 };
1785
1786 let conn_limits = self.limits.clamp(&negotiated.limits);
1787 let remote_ip = self
1796 .endpoint
1797 .as_ref()
1798 .and_then(|endpoint| inbound_remote_ip(endpoint, remote_node_id));
1799 let local_node_id = self
1800 .endpoint
1801 .as_ref()
1802 .map(|endpoint| endpoint.node_id())
1803 .unwrap_or(remote_node_id);
1804 let direction = ServicePeerDirection::Inbound;
1805 self.register_and_serve(
1806 connection,
1807 remote_peer_id,
1808 remote_ip,
1809 ConnectionServeContext {
1810 limits: conn_limits,
1811 accepted_capabilities: negotiated.accepted_capabilities,
1812 role: "responder",
1813 direction,
1814 transcript_hash: native_connection_transcript_hash(
1815 direction,
1816 &local_node_id,
1817 &remote_node_id,
1818 ),
1819 i_open_collision_winner: i_open_collision_winner(&local_node_id, &remote_node_id),
1820 conn,
1821 },
1822 )
1823 .await
1824 .map_err(AcceptError::from_err)
1825 }
1826
1827 async fn run_native_responder_handshake_with_permit(
1828 &self,
1829 connection: &Connection,
1830 remote_peer_id: &ZakuraPeerId,
1831 conn: &ZakuraConnTrace,
1832 ) -> Result<NativeHandshakeNegotiated, ZakuraHandlerError> {
1833 let Ok(_handshake) = self.pending_handshakes.clone().try_acquire_owned() else {
1834 metrics::counter!("zakura.p2p.conn.rejected.pending_handshake").increment(1);
1835 self.trace.emit(
1836 CONN_TABLE,
1837 conn.event("rejected.admission")
1838 .direction("inbound")
1839 .reason("pending_handshake"),
1840 );
1841 connection.close(
1842 VarInt::from_u32(ZAKURA_CLOSE_RESOURCE),
1843 b"pending handshake",
1844 );
1845 return Err(ZakuraHandlerError::ResourceLimit("pending handshake"));
1846 };
1847
1848 self.run_native_responder_handshake(connection, remote_peer_id, conn)
1849 .await
1850 }
1851
1852 async fn run_native_responder_handshake(
1853 &self,
1854 connection: &Connection,
1855 remote_peer_id: &ZakuraPeerId,
1856 conn: &ZakuraConnTrace,
1857 ) -> Result<NativeHandshakeNegotiated, ZakuraHandlerError> {
1858 self.trace.emit(
1859 HANDSHAKE_TABLE,
1860 conn.event("control.started")
1861 .role("responder")
1862 .phase("control")
1863 .network(self.handshake_config.network_label()),
1864 );
1865 let (mut send, mut recv) = timeout(self.limits.control_timeout, connection.accept_bi())
1866 .await
1867 .map_err(|_| ZakuraHandlerError::Timeout("accept control stream"))??;
1868
1869 let hello_bytes = read_control_payload(
1870 &mut recv,
1871 self.handshake_config.max_control_frame_bytes,
1872 self.limits.control_timeout,
1873 )
1874 .await?;
1875 let hello = ZakuraControlHello::decode(&hello_bytes)?;
1876 let expected = ZakuraControlValidation {
1877 local: &self.handshake_config,
1878 authenticated_remote_id: remote_peer_id.as_bytes(),
1879 selected_zakura_protocol: ZAKURA_PROTOCOL_VERSION_1,
1880 handshake_path: ZakuraHandshakePath::Native,
1881 remote_role: ZakuraControlRole::Initiator,
1882 initiator_upgrade_nonce: [0; 32],
1883 responder_upgrade_nonce: [0; 32],
1884 legacy_upgrade_transcript: [0; 32],
1885 };
1886 hello.validate(&expected)?;
1887
1888 let mut local_nonce = [0; 32];
1889 OsRng.fill_bytes(&mut local_nonce);
1890
1891 let accepted_limits = self.accepted_limits_for(&hello.initial_limits);
1892 let ack = ZakuraControlAck {
1893 magic: CONTROL_ACK_MAGIC,
1894 control_version: CONTROL_VERSION,
1895 selected_zakura_protocol: hello.selected_zakura_protocol,
1896 peer_nonce: local_nonce,
1897 remote_peer_nonce: hello.peer_nonce,
1898 accepted_capabilities: hello.capabilities
1899 & self.handshake_config.supported_capabilities,
1900 accepted_channels: hello.required_channels & self.handshake_config.supported_channels,
1901 accepted_limits,
1902 };
1903 write_control_payload(&mut send, &ack.encode()?, self.limits.control_timeout).await?;
1904 self.trace.emit(
1905 HANDSHAKE_TABLE,
1906 conn.event("control.succeeded")
1907 .role("responder")
1908 .phase("control")
1909 .selected_protocol(ack.selected_zakura_protocol)
1910 .network(self.handshake_config.network_label()),
1911 );
1912 Ok(NativeHandshakeNegotiated {
1913 limits: accepted_limits,
1914 accepted_capabilities: ack.accepted_capabilities,
1915 })
1916 }
1917
1918 fn accepted_limits_for(&self, remote_limits: &ZakuraInitialLimits) -> ZakuraAcceptedLimits {
1919 ZakuraLimits {
1920 max_frame_bytes: remote_limits
1921 .max_frame_bytes
1922 .min(self.limits.max_frame_bytes),
1923 max_message_bytes: remote_limits
1924 .max_message_bytes
1925 .min(self.limits.max_message_bytes),
1926 max_open_streams: remote_limits
1927 .max_open_streams
1928 .min(self.limits.max_open_streams),
1929 max_inbound_queue_depth: remote_limits
1930 .max_inbound_queue_depth
1931 .min(self.limits.max_inbound_queue_depth),
1932 idle_timeout_millis: remote_limits
1933 .idle_timeout_millis
1934 .min(self.limits.initial_limits().idle_timeout_millis),
1935 }
1936 }
1937
1938 async fn serve_connection(
1939 &self,
1940 connection: Connection,
1941 peer_id: ZakuraPeerId,
1942 remote_ip: Option<IpAddr>,
1943 mut outbound_rx: mpsc::Receiver<ZakuraOutboundFrame>,
1944 context: RegisteredConnectionServeContext,
1945 mut cleanup_guard: RegisteredPeerCleanupGuard,
1946 ) -> Result<(), ZakuraHandlerError> {
1947 let limits = context.limits;
1948 let conn = context.conn;
1949 let conn_id = context.conn_id;
1950 let connection_token = context.connection_token;
1951 let close_cause = context.close_cause;
1952 let accepted_capabilities = context.accepted_capabilities;
1953 let stream_sem = Arc::new(Semaphore::new(usize::from(limits.max_open_streams)));
1954 let mut workers = JoinSet::new();
1955 let (ordered_stream_exit_tx, mut ordered_stream_exit_rx) = mpsc::unbounded_channel();
1956 let (ordered_stream_reopen_tx, mut ordered_stream_reopen_rx) = mpsc::unbounded_channel();
1957 let mut ordered_stream_reopen_attempts: HashMap<u16, u32> = HashMap::new();
1960 let mut open_limiter = TokenBucket::new(limits.stream_open_rate_per_second);
1961 let mut message_buckets = MessageRateBuckets::new();
1962 let (freshness_tx, freshness_rx) = watch::channel(Instant::now());
1963 let negotiated_ordered_streams = self
1964 .registry
1965 .ordered_streams_for_negotiated(accepted_capabilities);
1966 let ordered_streams: Vec<Stream> = self
1975 .registry
1976 .ordered_streams_for_escalation(accepted_capabilities, &peer_id, context.direction)
1977 .into_iter()
1978 .filter(|stream| {
1979 context.is_initiator || (!cfg!(test) && stream.kind == ZAKURA_STREAM_BLOCK_SYNC)
1982 })
1983 .collect();
1984 let request_response_stream_count = self
1985 .registry
1986 .request_response_streams_for_negotiated(accepted_capabilities)
1987 .len();
1988 if ordered_streams.len() > usize::from(limits.max_open_streams) {
1989 debug!(
1990 max_open_streams = limits.max_open_streams,
1991 ordered_stream_count = ordered_streams.len(),
1992 "closing Zakura peer because negotiated ordered streams exceed max-open-streams"
1993 );
1994 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"ordered streams");
1995 close_cause.record("resource_ordered_streams");
1996 connection_token.cancel();
1997 } else if !ordered_streams.is_empty()
1998 && usize::from(limits.max_inbound_queue_depth) < ordered_streams.len()
1999 {
2000 debug!(
2001 max_inbound_queue_depth = limits.max_inbound_queue_depth,
2002 ordered_stream_count = ordered_streams.len(),
2003 "closing Zakura peer because inbound queue depth cannot be split across ordered streams"
2004 );
2005 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"queue split");
2006 close_cause.record("resource_queue_split");
2007 connection_token.cancel();
2008 }
2009 let selected_ordered_streams: HashMap<u16, Stream> = negotiated_ordered_streams
2010 .iter()
2011 .map(|stream| (stream.kind, *stream))
2012 .collect();
2013 let ordered_kinds: HashSet<u16> = selected_ordered_streams.keys().copied().collect();
2014 let queue_split_stream_count = negotiated_ordered_streams.len().max(ordered_streams.len());
2015 let per_stream_queue_depth = per_stream_inbound_queue_depth(
2016 limits.max_inbound_queue_depth,
2017 queue_split_stream_count,
2018 );
2019 let mut service_streams = HashMap::new();
2020 let mut opened_kinds: HashSet<u16> = HashSet::new();
2023 let mut accepted_ordered_kinds = HashSet::new();
2024 let run_freshness_reaper =
2025 should_run_freshness_reaper(queue_split_stream_count, request_response_stream_count);
2026
2027 if negotiated_ordered_streams.is_empty() {
2028 self.registry
2029 .add_peer(Peer::new_with_conn_id_and_direction_and_close_cause(
2030 conn_id,
2031 peer_id.clone(),
2032 remote_ip,
2033 accepted_capabilities,
2034 context.direction,
2035 HashMap::new(),
2036 connection_token.clone(),
2037 close_cause.clone(),
2038 ));
2039 cleanup_guard.add_admitted_capabilities(accepted_capabilities);
2040 } else if !connection_token.is_cancelled() {
2041 let mut opened_capabilities = 0;
2042 for stream in ordered_streams {
2043 opened_capabilities |= stream.capability;
2044 let admitted = match self
2045 .open_ordered_service_stream(
2046 &connection,
2047 stream,
2048 &mut workers,
2049 &stream_sem,
2050 &mut message_buckets,
2051 limits,
2052 per_stream_queue_depth,
2053 connection_token.clone(),
2054 close_cause.clone(),
2055 freshness_tx.clone(),
2056 conn.clone(),
2057 peer_id.clone(),
2058 ordered_stream_exit_tx.clone(),
2059 )
2060 .await
2061 {
2062 Ok(admitted) => admitted,
2063 Err(error) => {
2064 debug!(
2065 ?error,
2066 stream_kind = stream.kind,
2067 "closing Zakura peer after ordered stream setup failed"
2068 );
2069 connection.close(
2070 VarInt::from_u32(ZAKURA_CLOSE_RESOURCE),
2071 b"ordered stream setup",
2072 );
2073 close_cause.record("stream_setup_failed");
2074 connection_token.cancel();
2075 break;
2076 }
2077 };
2078 opened_kinds.insert(admitted.kind);
2079 service_streams.insert(
2080 admitted.kind,
2081 ServiceStream::new(
2082 admitted.session_id,
2083 admitted.recv,
2084 admitted.send,
2085 admitted.cancel_token,
2086 ),
2087 );
2088 }
2089 if !connection_token.is_cancelled() {
2090 let admitted_capabilities =
2094 self.registry
2095 .add_escalated_peer(Peer::new_with_service_streams(
2096 conn_id,
2097 peer_id.clone(),
2098 remote_ip,
2099 opened_capabilities,
2100 context.direction,
2101 std::mem::take(&mut service_streams),
2102 connection_token.clone(),
2103 close_cause.clone(),
2104 ));
2105 cleanup_guard.add_admitted_capabilities(admitted_capabilities);
2106 }
2107 }
2108
2109 loop {
2110 tokio::select! {
2111 biased;
2112 _ = connection_token.cancelled() => break,
2113 _ = freshness_reaper(freshness_rx.clone(), limits.idle_timeout), if run_freshness_reaper => {
2114 connection.close(VarInt::from_u32(ZAKURA_CLOSE_NEUTRAL), b"idle");
2115 close_cause.record("idle_timeout");
2116 break;
2117 }
2118 Some(exited) = ordered_stream_exit_rx.recv() => {
2119 if exited.opened_locally {
2123 opened_kinds.remove(&exited.stream.kind);
2124 } else {
2125 accepted_ordered_kinds.remove(&exited.stream.kind);
2126 }
2127 if should_reopen_ordered_stream(
2128 exited,
2129 context.is_initiator,
2130 connection_token.is_cancelled(),
2131 ) {
2132 let attempts = ordered_stream_reopen_attempts
2133 .entry(exited.stream.kind)
2134 .or_insert(0);
2135 let delay = ordered_stream_reopen_backoff(*attempts);
2136 *attempts = attempts.saturating_add(1);
2137 debug!(
2138 stream_kind = exited.stream.kind,
2139 stream_version = exited.stream.version,
2140 session_id = exited.session_id,
2141 ?delay,
2142 "scheduling Zakura ordered stream reopen"
2143 );
2144 let reopen = ordered_stream_reopen_tx.clone();
2145 tokio::spawn(async move {
2146 tokio::time::sleep(delay).await;
2147 let _ = reopen.send(exited.stream);
2148 });
2149 }
2150 }
2151 Some(stream) = ordered_stream_reopen_rx.recv() => {
2152 if connection_token.is_cancelled()
2153 || !context.is_initiator
2154 || opened_kinds.contains(&stream.kind)
2155 || selected_ordered_streams.get(&stream.kind) != Some(&stream)
2156 || !self.registry.wants_ordered_stream(
2157 stream.kind,
2158 accepted_capabilities,
2159 &peer_id,
2160 context.direction,
2161 )
2162 {
2163 continue;
2164 }
2165
2166 match self
2167 .open_ordered_service_stream(
2168 &connection,
2169 stream,
2170 &mut workers,
2171 &stream_sem,
2172 &mut message_buckets,
2173 limits,
2174 per_stream_queue_depth,
2175 connection_token.clone(),
2176 close_cause.clone(),
2177 freshness_tx.clone(),
2178 conn.clone(),
2179 peer_id.clone(),
2180 ordered_stream_exit_tx.clone(),
2181 )
2182 .await
2183 {
2184 Ok(admitted) => {
2185 ordered_stream_reopen_attempts.remove(&admitted.kind);
2186 opened_kinds.insert(admitted.kind);
2187 let service_streams = HashMap::from([(
2188 admitted.kind,
2189 ServiceStream::new(
2190 admitted.session_id,
2191 admitted.recv,
2192 admitted.send,
2193 admitted.cancel_token,
2194 ),
2195 )]);
2196 let admitted_capabilities = self.registry.add_escalated_peer(
2197 Peer::new_with_service_streams(
2198 conn_id,
2199 peer_id.clone(),
2200 remote_ip,
2201 accepted_capabilities,
2202 context.direction,
2203 service_streams,
2204 connection_token.clone(),
2205 close_cause.clone(),
2206 ),
2207 );
2208 cleanup_guard.add_admitted_capabilities(admitted_capabilities);
2209 }
2210 Err(error) => {
2211 let attempts = ordered_stream_reopen_attempts
2212 .entry(stream.kind)
2213 .or_insert(0);
2214 let delay = ordered_stream_reopen_backoff(*attempts);
2215 *attempts = attempts.saturating_add(1);
2216 debug!(
2217 ?error,
2218 stream_kind = stream.kind,
2219 ?delay,
2220 "retrying Zakura ordered stream reopen after local failure"
2221 );
2222 let reopen = ordered_stream_reopen_tx.clone();
2223 tokio::spawn(async move {
2224 tokio::time::sleep(delay).await;
2225 let _ = reopen.send(stream);
2226 });
2227 }
2228 }
2229 }
2230 Some(joined) = workers.join_next() => {
2231 if let Err(error) = joined {
2232 debug!(?error, "Zakura stream worker exited unexpectedly");
2233 }
2234 }
2235 accepted = connection.accept_bi() => {
2236 match accepted {
2237 Ok((send, recv)) => {
2238 let mut admission = StreamAdmission {
2239 trace: self.trace.clone(),
2240 conn: conn.clone(),
2241 peer_id: &peer_id,
2242 stream_sem: &stream_sem,
2243 open_limiter: &mut open_limiter,
2244 message_buckets: &mut message_buckets,
2245 workers: &mut workers,
2246 limits,
2247 accepted_capabilities,
2248 connection_token: connection_token.clone(),
2249 close_cause: close_cause.clone(),
2250 freshness_tx: freshness_tx.clone(),
2251 };
2252 if let Some(admitted) = self
2253 .admit_bi_stream(
2254 send,
2255 recv,
2256 &mut admission,
2257 per_stream_queue_depth,
2258 ordered_stream_exit_tx.clone(),
2259 )
2260 .await
2261 {
2262 let kind = admitted.kind;
2263
2264 if !ordered_kinds.contains(&kind) {
2268 debug!(
2269 stream_kind = kind,
2270 "closing peer after unexpected ordered stream"
2271 );
2272 close_cause.record("unexpected_stream");
2273 connection_token.cancel();
2274 continue;
2275 }
2276 let is_block_sync = kind == ZAKURA_STREAM_BLOCK_SYNC;
2277 let is_block_sync_collision =
2280 is_block_sync && opened_kinds.contains(&kind);
2281
2282 if is_block_sync_collision && context.i_open_collision_winner {
2287 debug!(
2288 stream_kind = kind,
2289 "winning block-sync collision; parking peer's duplicate"
2290 );
2291 admitted.cancel_token.cancel();
2292 continue;
2293 }
2294
2295 if !is_block_sync && context.is_initiator {
2303 debug!(
2304 stream_kind = kind,
2305 "closing peer after unexpected ordered stream"
2306 );
2307 close_cause.record("unexpected_stream");
2308 connection_token.cancel();
2309 continue;
2310 }
2311
2312 if !accepted_ordered_kinds.insert(kind) {
2316 debug!(
2317 stream_kind = kind,
2318 "closing peer after duplicate ordered stream"
2319 );
2320 close_cause.record("duplicate_stream");
2321 connection_token.cancel();
2322 continue;
2323 }
2324
2325 if !is_block_sync_collision
2331 && !self.registry.wants_ordered_stream(
2332 kind,
2333 accepted_capabilities,
2334 &peer_id,
2335 context.direction,
2336 )
2337 {
2338 metrics::counter!(
2339 "zakura.p2p.stream.parked.no_demand",
2340 "stream_kind" => stream_kind_label(kind),
2341 )
2342 .increment(1);
2343 info!(
2344 ?peer_id,
2345 stream_kind = kind,
2346 "locally parking ordered service stream because the service has no demand"
2347 );
2348 admitted.cancel_token.cancel();
2349 continue;
2350 }
2351
2352 let service_streams = HashMap::from([(
2353 kind,
2354 ServiceStream::new(
2355 admitted.session_id,
2356 admitted.recv,
2357 admitted.send,
2358 admitted.cancel_token.clone(),
2359 ),
2360 )]);
2361 let admitted_capabilities = self.registry.add_escalated_peer(
2374 Peer::new_with_service_streams(
2375 conn_id,
2376 peer_id.clone(),
2377 remote_ip,
2378 accepted_capabilities,
2379 context.direction,
2380 service_streams,
2381 connection_token.clone(),
2382 close_cause.clone(),
2383 ),
2384 );
2385 cleanup_guard.add_admitted_capabilities(admitted_capabilities);
2386 }
2387 }
2388 Err(error) => {
2389 debug!(?error, "Zakura connection stopped accepting streams");
2390 close_cause.record("accept_failed");
2391 break;
2392 }
2393 }
2394 }
2395 outbound = outbound_rx.recv() => {
2396 let Some(outbound) = outbound else {
2397 close_cause.record("outbound_closed");
2398 break;
2399 };
2400 match outbound {
2401 ZakuraOutboundFrame::Request {
2402 stream_kind,
2403 request_id,
2404 message_type,
2405 flags,
2406 payload,
2407 mut completion,
2408 } => {
2409 let result = tokio::select! {
2410 biased;
2411 _ = completion.closed() => {
2412 Err(OutboundRequestError::Local(
2413 "Zakura outbound request receiver dropped".into(),
2414 ))
2415 }
2416 result = write_outbound_request_frame(
2417 &connection,
2418 limits,
2419 stream_kind,
2420 request_id,
2421 message_type,
2422 flags,
2423 payload,
2424 ) => result,
2425 };
2426 match result {
2427 Ok(frames) => {
2428 let _ = completion.send(Ok(frames));
2429 }
2430 Err(OutboundRequestError::Local(error)) => {
2431 let _ = completion.send(Err(error));
2432 }
2433 Err(OutboundRequestError::Fatal(error)) => {
2434 connection.close(
2435 VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE),
2436 b"malformed response",
2437 );
2438 close_cause.record("bad_response");
2439 connection_token.cancel();
2440 let _ = completion.send(Err(error));
2441 }
2442 }
2443 }
2444 }
2445 }
2446 }
2447 }
2448
2449 connection_token.cancel();
2450 while let Some(joined) = timeout(STREAM_WORKER_DRAIN_TIMEOUT, workers.join_next())
2451 .await
2452 .ok()
2453 .flatten()
2454 {
2455 if let Err(error) = joined {
2456 debug!(?error, "Zakura stream worker failed during shutdown");
2457 }
2458 }
2459 workers.abort_all();
2460 cleanup_guard.cleanup_registered_peer().await;
2461 metrics::counter!("zakura.p2p.conn.closed.neutral").increment(1);
2462 self.trace.emit(
2463 CONN_TABLE,
2464 conn.event("closed.neutral")
2465 .reason(close_cause.get_or("cancelled")),
2466 );
2467 Ok(())
2468 }
2469
2470 #[allow(clippy::too_many_arguments)]
2471 async fn open_ordered_service_stream(
2472 &self,
2473 connection: &Connection,
2474 stream: Stream,
2475 workers: &mut JoinSet<()>,
2476 stream_sem: &Arc<Semaphore>,
2477 message_buckets: &mut MessageRateBuckets,
2478 limits: ZakuraConnectionLimits,
2479 per_stream_queue_depth: usize,
2480 connection_token: CancellationToken,
2481 close_cause: CloseCause,
2482 freshness_tx: watch::Sender<Instant>,
2483 conn: ZakuraConnTrace,
2484 peer_id: ZakuraPeerId,
2485 ordered_stream_exit_tx: mpsc::UnboundedSender<OrderedStreamExit>,
2486 ) -> Result<AdmittedOrderedStream, ZakuraHandlerError> {
2487 let stream_id = self.next_stream_id.fetch_add(1, Ordering::Relaxed);
2488 let permit = stream_sem
2489 .clone()
2490 .try_acquire_owned()
2491 .map_err(|_| ZakuraHandlerError::ResourceLimit("ordered stream permit"))?;
2492 let (mut send, recv) = timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, connection.open_bi())
2493 .await
2494 .map_err(|_| ZakuraHandlerError::Timeout("open ordered service stream"))??;
2495 let prelude = StreamPrelude {
2496 magic: STREAM_PRELUDE_MAGIC,
2497 stream_kind: stream.kind,
2498 stream_version: stream.version,
2499 request_id: None,
2500 max_frame_bytes: app_frame_cap_for_stream_kind(&limits, stream.kind),
2501 };
2502 let prelude_bytes = prelude.encode()?;
2503 timeout(
2504 OUTBOUND_STREAM_WRITE_TIMEOUT,
2505 send.write_all(&prelude_bytes),
2506 )
2507 .await
2508 .map_err(|_| ZakuraHandlerError::Timeout("ordered stream prelude write"))??;
2509
2510 let message_bucket = message_bucket_for(
2511 message_buckets,
2512 stream.kind,
2513 limits.message_rate_per_second,
2514 RealClock,
2515 );
2516 let stream_token = connection_token.child_token();
2517 let context = StreamWorkerContext {
2518 trace: self.trace.clone(),
2519 conn: conn.clone(),
2520 peer_id,
2521 stream_id,
2522 _permit: permit,
2523 limits,
2524 message_bucket,
2525 connection_token,
2526 stream_token,
2527 close_cause,
2528 freshness_tx,
2529 };
2530
2531 metrics::counter!(
2532 "zakura.p2p.stream.accepted",
2533 "stream_kind" => stream_kind_label(stream.kind),
2534 )
2535 .increment(1);
2536 self.trace.emit(
2537 STREAM_TABLE,
2538 conn.event("accepted")
2539 .stream(stream_id)
2540 .stream_kind(stream_kind_label(stream.kind)),
2541 );
2542
2543 Ok(spawn_persistent_stream_worker(
2544 workers,
2545 send,
2546 recv,
2547 stream,
2548 prelude,
2549 context,
2550 per_stream_queue_depth,
2551 true,
2552 ordered_stream_exit_tx,
2553 ))
2554 }
2555
2556 async fn admit_bi_stream(
2557 &self,
2558 mut send: SendStream,
2559 mut recv: RecvStream,
2560 admission: &mut StreamAdmission<'_>,
2561 per_stream_queue_depth: usize,
2562 ordered_stream_exit_tx: mpsc::UnboundedSender<OrderedStreamExit>,
2563 ) -> Option<AdmittedOrderedStream> {
2564 let stream_id = self.next_stream_id.fetch_add(1, Ordering::Relaxed);
2565 let Ok(permit) = admission.stream_sem.clone().try_acquire_owned() else {
2566 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE));
2567 metrics::counter!("zakura.p2p.stream.rejected.semaphore").increment(1);
2568 admission.trace.emit(
2569 STREAM_TABLE,
2570 admission.event("rejected.semaphore", stream_id),
2571 );
2572 return None;
2573 };
2574
2575 if !admission.open_limiter.try_take() {
2583 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RATE_LIMIT));
2584 metrics::counter!("zakura.p2p.stream.rejected.open_rate").increment(1);
2585 admission.trace.emit(
2586 STREAM_TABLE,
2587 admission.event("rejected.open_rate", stream_id),
2588 );
2589 return None;
2590 }
2591
2592 let prelude = match read_stream_prelude(&mut recv, admission.limits.prelude_timeout).await {
2593 Ok(prelude) => prelude,
2594 Err(error) => {
2595 debug!(?error, "rejecting Zakura stream with bad prelude");
2596 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
2597 metrics::counter!("zakura.p2p.stream.rejected.prelude").increment(1);
2598 admission
2599 .trace
2600 .emit(STREAM_TABLE, admission.event("rejected.prelude", stream_id));
2601 return None;
2602 }
2603 };
2604 let stream_kind = stream_kind_label(prelude.stream_kind);
2605
2606 let Some(stream) = self
2607 .registry
2608 .stream(prelude.stream_kind, prelude.stream_version)
2609 else {
2610 debug!(
2611 stream_kind = prelude.stream_kind,
2612 stream_version = prelude.stream_version,
2613 "rejecting Zakura stream with unknown kind or unsupported version"
2614 );
2615 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_UNKNOWN_STREAM));
2616 metrics::counter!(
2617 "zakura.p2p.stream.rejected.unknown_kind",
2618 "stream_kind" => stream_kind,
2619 )
2620 .increment(1);
2621 admission.trace.emit(
2622 STREAM_TABLE,
2623 admission
2624 .event("rejected.unknown_kind", stream_id)
2625 .stream_kind(stream_kind),
2626 );
2627 return None;
2628 };
2629
2630 if admission.accepted_capabilities & stream.capability != stream.capability {
2631 debug!(
2632 stream_kind = prelude.stream_kind,
2633 stream_version = prelude.stream_version,
2634 accepted_capabilities = admission.accepted_capabilities,
2635 "rejecting Zakura stream that was not negotiated for this peer"
2636 );
2637 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_UNKNOWN_STREAM));
2638 metrics::counter!(
2639 "zakura.p2p.stream.rejected.unnegotiated_capability",
2640 "stream_kind" => stream_kind,
2641 )
2642 .increment(1);
2643 admission.trace.emit(
2644 STREAM_TABLE,
2645 admission
2646 .event("rejected.unnegotiated_capability", stream_id)
2647 .stream_kind(stream_kind),
2648 );
2649 return None;
2650 }
2651
2652 if stream.mode != StreamMode::RequestResponse && prelude.request_id.is_some() {
2653 debug!("rejecting non-request Zakura stream with request id");
2654 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
2655 admission.close_cause.record("unexpected_request_id");
2656 admission.connection_token.cancel();
2657 metrics::counter!("zakura.p2p.stream.rejected.unexpected_request_id").increment(1);
2658 admission.trace.emit(
2659 STREAM_TABLE,
2660 admission
2661 .event("rejected.unexpected_request_id", stream_id)
2662 .stream_kind(stream_kind),
2663 );
2664 return None;
2665 }
2666
2667 if stream.mode == StreamMode::RequestResponse && prelude.request_id.is_none() {
2668 debug!("rejecting Zakura request stream without request id");
2669 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
2670 admission.close_cause.record("request_without_id");
2671 admission.connection_token.cancel();
2672 metrics::counter!("zakura.p2p.stream.rejected.request_without_id").increment(1);
2673 admission.trace.emit(
2674 STREAM_TABLE,
2675 admission
2676 .event("rejected.request_without_id", stream_id)
2677 .stream_kind(stream_kind),
2678 );
2679 return None;
2680 }
2681
2682 metrics::counter!(
2683 "zakura.p2p.stream.accepted",
2684 "stream_kind" => stream_kind,
2685 )
2686 .increment(1);
2687 admission.trace.emit(
2688 STREAM_TABLE,
2689 admission
2690 .event("accepted", stream_id)
2691 .stream_kind(stream_kind),
2692 );
2693
2694 let message_bucket = message_bucket_for(
2695 admission.message_buckets,
2696 prelude.stream_kind,
2697 admission.limits.message_rate_per_second,
2698 RealClock,
2699 );
2700 let stream_token = admission.connection_token.child_token();
2701
2702 let context = StreamWorkerContext {
2703 trace: admission.trace.clone(),
2704 conn: admission.conn.clone(),
2705 peer_id: admission.peer_id.clone(),
2706 stream_id,
2707 _permit: permit,
2708 limits: admission.limits,
2709 message_bucket,
2710 connection_token: admission.connection_token.clone(),
2711 stream_token,
2712 close_cause: admission.close_cause.clone(),
2713 freshness_tx: admission.freshness_tx.clone(),
2714 };
2715
2716 if stream.mode == StreamMode::RequestResponse {
2717 admission.workers.spawn(request_stream_worker(
2718 send,
2719 recv,
2720 prelude,
2721 context,
2722 self.registry.clone(),
2723 ));
2724 None
2725 } else {
2726 Some(spawn_persistent_stream_worker(
2727 admission.workers,
2728 send,
2729 recv,
2730 stream,
2731 prelude,
2732 context,
2733 per_stream_queue_depth,
2734 false,
2735 ordered_stream_exit_tx,
2736 ))
2737 }
2738 }
2739
2740 async fn register_and_serve(
2741 &self,
2742 connection: Connection,
2743 peer_id: ZakuraPeerId,
2744 remote_ip: Option<IpAddr>,
2745 context: ConnectionServeContext,
2746 ) -> Result<(), ZakuraHandlerError> {
2747 let ordered_stream_count = self
2751 .registry
2752 .ordered_streams_for_escalation(
2753 context.accepted_capabilities,
2754 &peer_id,
2755 context.direction,
2756 )
2757 .len();
2758 if ordered_stream_count > usize::from(context.limits.max_open_streams) {
2759 debug!(
2760 max_open_streams = context.limits.max_open_streams,
2761 ordered_stream_count,
2762 "rejecting Zakura peer before registration because negotiated ordered streams exceed max-open-streams"
2763 );
2764 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"ordered streams");
2765 return Ok(());
2766 }
2767 if ordered_stream_count > 0
2768 && usize::from(context.limits.max_inbound_queue_depth) < ordered_stream_count
2769 {
2770 debug!(
2771 max_inbound_queue_depth = context.limits.max_inbound_queue_depth,
2772 ordered_stream_count,
2773 "rejecting Zakura peer before registration because inbound queue depth cannot be split"
2774 );
2775 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"queue split");
2776 return Ok(());
2777 }
2778
2779 let (outbound_tx, outbound_rx) =
2780 mpsc::channel(usize::from(context.limits.max_inbound_queue_depth));
2781 let close_cause = CloseCause::new();
2782 let outbound_handle = ZakuraPeerHandle {
2783 peer_id: peer_id.clone(),
2784 sender: outbound_tx,
2785 };
2786 let connection_token = self.shutdown.child_token();
2787 let registration = self
2788 .supervisor
2789 .register(
2790 context.conn.id,
2791 peer_id,
2792 remote_ip,
2793 context.transcript_hash,
2794 outbound_handle,
2795 connection_token.clone(),
2796 context.accepted_capabilities,
2797 )
2798 .await;
2799
2800 match registration {
2801 ZakuraRegistration::Registered {
2802 conn_id,
2803 peer_id,
2804 remote_ip,
2805 disconnect_token,
2806 } => {
2807 metrics::counter!("zakura.p2p.conn.accepted", "role" => context.role).increment(1);
2808 self.trace.emit(
2809 CONN_TABLE,
2810 context
2811 .conn
2812 .event("accepted")
2813 .role(context.role)
2814 .direction(context.direction.trace_label()),
2815 );
2816 let cleanup_guard = RegisteredPeerCleanupGuard::new(
2817 self.supervisor.clone(),
2818 self.registry.clone(),
2819 peer_id.clone(),
2820 conn_id,
2821 disconnect_token.clone(),
2822 );
2823 self.serve_connection(
2824 connection,
2825 peer_id,
2826 remote_ip,
2827 outbound_rx,
2828 RegisteredConnectionServeContext {
2829 limits: context.limits,
2830 conn: context.conn,
2831 conn_id,
2832 connection_token: disconnect_token,
2833 close_cause,
2834 accepted_capabilities: context.accepted_capabilities,
2835 is_initiator: context.role == "initiator",
2836 i_open_collision_winner: context.i_open_collision_winner,
2837 direction: context.direction,
2838 },
2839 cleanup_guard,
2840 )
2841 .await
2842 }
2843 ZakuraRegistration::Duplicate { peer_id } => {
2844 debug!(?peer_id, "closing duplicate Zakura peer neutrally");
2845 metrics::counter!("zakura.p2p.conn.duplicate").increment(1);
2846 self.trace.emit(
2847 CONN_TABLE,
2848 context
2849 .conn
2850 .event("duplicate")
2851 .role(context.role)
2852 .direction(context.direction.trace_label()),
2853 );
2854 connection.close(VarInt::from_u32(ZAKURA_CLOSE_NEUTRAL), b"duplicate");
2855 Ok(())
2856 }
2857 ZakuraRegistration::Rejected(reason) => {
2858 debug!(
2859 ?reason,
2860 "closing Zakura peer neutrally after registration rejection"
2861 );
2862 self.trace.emit(
2863 CONN_TABLE,
2864 context
2865 .conn
2866 .event("rejected.admission")
2867 .role(context.role)
2868 .direction(context.direction.trace_label())
2869 .reason(reject_reason_label(reason)),
2870 );
2871 connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"registration");
2872 Ok(())
2873 }
2874 }
2875 }
2876}
2877
2878impl ProtocolHandler for ZakuraProtocolHandler {
2879 async fn accept(&self, connection: Connection) -> Result<(), AcceptError> {
2880 self.accept_connection(connection).await
2881 }
2882
2883 async fn shutdown(&self) {
2884 self.shutdown.cancel();
2885 }
2886}
2887
2888const ZAKURA_LOOPBACK_BIND_V4: SocketAddrV4 = SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0);
2892const ZAKURA_LOOPBACK_BIND_V6: SocketAddrV6 = SocketAddrV6::new(Ipv6Addr::LOCALHOST, 0, 0, 0);
2894
2895fn bind_native_endpoint(
2905 builder: iroh::endpoint::Builder,
2906 listen_addr: Option<SocketAddr>,
2907) -> iroh::endpoint::Builder {
2908 match listen_addr {
2909 Some(SocketAddr::V4(addr)) => builder.bind_addr_v4(addr),
2910 Some(SocketAddr::V6(addr)) => builder.bind_addr_v6(addr),
2911 None => builder
2912 .bind_addr_v4(ZAKURA_LOOPBACK_BIND_V4)
2913 .bind_addr_v6(ZAKURA_LOOPBACK_BIND_V6),
2914 }
2915}
2916
2917pub async fn spawn_zakura_endpoint(
2919 config: &Config,
2920 sink_factory: impl FnOnce(ZakuraSupervisorHandle, ZakuraTrace) -> Arc<dyn Service>,
2921) -> Result<Option<ZakuraEndpoint>, BoxError> {
2922 spawn_zakura_endpoint_with_header_sync_driver(config, sink_factory, None).await
2923}
2924
2925pub async fn spawn_zakura_endpoint_with_header_sync_driver(
2927 config: &Config,
2928 sink_factory: impl FnOnce(ZakuraSupervisorHandle, ZakuraTrace) -> Arc<dyn Service>,
2929 header_sync_driver_startup: Option<ZakuraHeaderSyncDriverStartup>,
2930) -> Result<Option<ZakuraEndpoint>, BoxError> {
2931 if !config.v2_p2p() {
2932 return Ok(None);
2933 }
2934
2935 let limits = ZakuraLocalLimits::from_config(config);
2936 validate_idle_invariant(&limits)?;
2937 let secret_key = zakura_secret_key(config)?;
2938 let discovery_secret_key = secret_key.clone();
2939 let builder = direct_endpoint_builder(secret_key).transport_config(limits.transport_config());
2940 let builder = bind_native_endpoint(builder, config.zakura.listen_addr);
2944 let endpoint = builder.bind().await?;
2945 let supervisor = ZakuraSupervisorHandle::new(config.zakura.max_connections_per_ip());
2946 let tracer = config
2947 .zakura
2948 .trace_dir
2949 .clone()
2950 .map(zakura_jsonl_trace::JsonlTracer::spawn)
2951 .unwrap_or_else(zakura_jsonl_trace::JsonlTracer::noop);
2952 let trace = ZakuraTrace::new(tracer, zakura_jsonl_trace::node_id());
2953 let handshake_config = ZakuraHandshakeConfig::for_network_with_dev_cohort(
2954 &config.network,
2955 config.zakura.dev_network.as_deref(),
2956 );
2957 let discovery = super::discovery::build_discovery_handle(
2958 discovery_secret_key,
2959 config.zakura.listen_addr.into_iter().collect(),
2960 super::discovery::default_advertised_services(),
2961 &handshake_config,
2962 config.zakura.max_connections,
2963 config.zakura.bootstrap_peers.len(),
2964 supervisor.subscribe(),
2965 )?;
2966 let anchor = config.zakura.header_sync.anchor(&config.network)?;
2967 let frontiers = header_sync_driver_startup.as_ref().map_or(
2968 HeaderSyncFrontiers {
2969 finalized_height: anchor.0,
2970 verified_block_tip: anchor.0,
2971 verified_block_hash: anchor.1,
2972 },
2973 |startup| startup.frontiers,
2974 );
2975 let best_header_tip = header_sync_driver_startup
2976 .as_ref()
2977 .map_or(Some(anchor), |startup| startup.best_header_tip);
2978 let sync_frontier = header_sync_driver_startup.as_ref().map(|driver_startup| {
2979 let best_header_tip = driver_startup.best_header_tip.unwrap_or(anchor);
2980 let initial = FrontierUpdate {
2981 frontier: crate::zakura::chain_frontier_from_parts(
2982 driver_startup.frontiers.finalized_height,
2983 Frontier::new(
2984 driver_startup.frontiers.verified_block_tip,
2985 driver_startup.verified_block_tip_hash,
2986 ),
2987 Frontier::new(best_header_tip.0, best_header_tip.1),
2988 ),
2989 change: FrontierChange::Snapshot,
2990 };
2991 ZakuraSyncExchange::new(initial, trace.clone())
2992 });
2993 let mut startup = HeaderSyncStartup::new(
2994 config.network.clone(),
2995 anchor,
2996 frontiers,
2997 best_header_tip,
2998 config.zakura.header_sync.clone(),
2999 limits.max_frame_bytes,
3000 );
3001 startup.trace = trace.clone();
3002 startup.frontier_updates = sync_frontier
3003 .as_ref()
3004 .map(ZakuraSyncExchange::subscribe_frontier);
3005 let header_sync_shutdown = CancellationToken::new();
3006 startup.shutdown = header_sync_shutdown.clone();
3007 if header_sync_driver_startup.is_some() {
3008 startup.range_state_actions_enabled = true;
3009 startup.inbound_new_block_acceptance_enabled = config.zakura.header_sync.accept_new_blocks;
3010 }
3011 let (header_sync, header_sync_actions, header_sync_task) = spawn_header_sync_reactor(startup)?;
3012 let block_sync_driver_enabled = header_sync_driver_startup.is_some();
3013 let (block_sync, block_sync_actions, block_sync_task) =
3014 if let Some(driver_startup) = header_sync_driver_startup.as_ref() {
3015 let best_header_tip = driver_startup.best_header_tip.unwrap_or(anchor);
3016 let frontier_updates = sync_frontier
3017 .as_ref()
3018 .expect("sync frontier is initialized when block sync driver is enabled")
3019 .subscribe_frontier();
3020 let mut startup = BlockSyncStartup::new_with_exchange(
3021 BlockSyncFrontiers {
3022 finalized_height: driver_startup.frontiers.finalized_height,
3023 verified_block_tip: driver_startup.frontiers.verified_block_tip,
3024 verified_block_hash: driver_startup.verified_block_tip_hash,
3025 },
3026 best_header_tip,
3027 frontier_updates,
3028 config.zakura.block_sync.clone(),
3029 );
3030 startup.shutdown = header_sync_shutdown.clone();
3031 startup.trace = trace.clone();
3032 let (handle, actions, task) = spawn_block_sync_reactor(startup);
3033 (Some(handle), Some(actions), Some(task))
3034 } else {
3035 (None, None, None)
3036 };
3037 let discovery_service = Arc::new(super::DiscoveryService::with_sync_services(
3038 discovery.clone(),
3039 header_sync.clone(),
3040 block_sync.clone(),
3041 )) as Arc<dyn Service>;
3042 let legacy_service = sink_factory(supervisor.clone(), trace.clone());
3043 let registry = service_registry(
3044 &supervisor,
3045 Some(header_sync.clone()),
3046 block_sync.clone(),
3047 config.zakura.block_sync.clone(),
3048 legacy_service,
3049 discovery_service,
3050 )?;
3051 let mut tasks = vec![header_sync_task];
3052 if let Some(task) = block_sync_task {
3053 tasks.push(task);
3054 }
3055 let header_sync_actions = if block_sync_driver_enabled {
3056 Some(Arc::new(Mutex::new(Some(header_sync_actions))))
3057 } else {
3058 let action_driver_task = tokio::spawn(drive_header_sync_actions(
3059 header_sync_actions,
3060 header_sync.clone(),
3061 supervisor.clone(),
3062 header_sync_shutdown.clone(),
3063 ));
3064 tasks.push(action_driver_task);
3065 None
3066 };
3067 let block_sync_actions = block_sync_actions
3068 .filter(|_| block_sync_driver_enabled)
3069 .map(|actions| Arc::new(Mutex::new(Some(actions))));
3070 let header_sync_tasks = Arc::new(HeaderSyncBackgroundTasks {
3071 shutdown: header_sync_shutdown,
3072 tasks: Mutex::new(tasks),
3073 });
3074 let handler = ZakuraProtocolHandler::new_with_registry_and_trace(
3075 supervisor.clone(),
3076 config.network.clone(),
3077 handshake_config,
3078 limits.clone(),
3079 registry,
3080 trace,
3081 )
3082 .with_endpoint(endpoint.clone());
3085 let router = Router::builder(endpoint)
3086 .accept(P2P_V2_ALPN, handler.clone())
3087 .spawn();
3088 let endpoint = ZakuraEndpoint {
3089 router,
3090 supervisor,
3091 handler,
3092 header_sync: Some(header_sync),
3093 block_sync,
3094 sync_frontier,
3095 header_sync_tasks: Some(header_sync_tasks),
3096 header_sync_actions,
3097 block_sync_actions,
3098 upgrade_dials: Arc::new(StdMutex::new(HashMap::new())),
3099 };
3100
3101 {
3106 let shutdown = endpoint.background_shutdown_token();
3107 let log_endpoint = endpoint.clone();
3108 let task = tokio::spawn(async move {
3109 tokio::select! {
3110 biased;
3111 _ = shutdown.cancelled() => {}
3112 node_addr = log_endpoint.node_addr() => {
3113 let direct_addresses: Vec<String> = node_addr
3114 .direct_addresses()
3115 .map(|addr| addr.to_string())
3116 .collect();
3117 info!(
3118 node_id = %node_addr.node_id,
3119 ?direct_addresses,
3120 "Zakura P2P endpoint ready; advertise <node_id>@<direct_addr> as a bootstrap peer",
3121 );
3122 }
3123 }
3124 });
3125 endpoint.push_header_sync_task(task).await;
3126 }
3127
3128 super::discovery::insert_static_bootstrap_candidates(
3129 &discovery,
3130 &config.zakura.bootstrap_peers,
3131 )
3132 .await;
3133 for task in spawn_native_bootstrap_dialer(
3139 endpoint.clone(),
3140 config.zakura.bootstrap_peers.clone(),
3141 limits.clone(),
3142 ) {
3143 endpoint.push_header_sync_task(task).await;
3144 }
3145 let discovery_dialer =
3146 super::discovery::spawn_native_discovery_dialer(endpoint.clone(), discovery, limits);
3147 endpoint.push_header_sync_task(discovery_dialer).await;
3148 Ok(Some(endpoint))
3149}
3150
3151pub(crate) async fn serve_native_dial_connection(
3152 endpoint: &ZakuraEndpoint,
3153 node_addr: NodeAddr,
3154 limits: &ZakuraLocalLimits,
3155) -> Result<(), ZakuraHandlerError> {
3156 let conn_id = endpoint
3157 .handler
3158 .next_conn_id
3159 .fetch_add(1, Ordering::Relaxed);
3160 let _admission = endpoint
3161 .handler
3162 .admission
3163 .clone()
3164 .try_acquire_owned()
3165 .map_err(|_| ZakuraHandlerError::ResourceLimit("admission"))?;
3166 let remote_ip = node_addr.direct_addresses().next().map(|addr| addr.ip());
3167 let connection = timeout(
3168 limits.control_timeout,
3169 endpoint.router.endpoint().connect(node_addr, P2P_V2_ALPN),
3170 )
3171 .await
3172 .map_err(|_| ZakuraHandlerError::Timeout("native dial"))??;
3173 let remote_node_id = connection.remote_node_id()?;
3174 let peer_id = ZakuraPeerId::new(remote_node_id.as_bytes().to_vec())?;
3175 let conn = ZakuraConnTrace::new(&endpoint.handler.trace, conn_id, &peer_id);
3176 let local_node_id = endpoint.router.endpoint().node_id();
3177 let local_peer_id = ZakuraPeerId::new(local_node_id.as_bytes().to_vec())?;
3178 let negotiated = {
3179 let _handshake = endpoint
3180 .handler
3181 .pending_handshakes
3182 .clone()
3183 .try_acquire_owned()
3184 .map_err(|_| ZakuraHandlerError::ResourceLimit("pending handshake"))?;
3185 run_native_initiator_handshake(
3186 &connection,
3187 limits,
3188 &endpoint.handler.handshake_config,
3189 &local_peer_id,
3190 &endpoint.handler.trace,
3191 &conn,
3192 )
3193 .await?
3194 };
3195 let conn_limits = limits.clamp(&negotiated.limits);
3196 let direction = ServicePeerDirection::Outbound;
3197 endpoint
3198 .handler
3199 .register_and_serve(
3200 connection,
3201 peer_id,
3202 remote_ip,
3203 ConnectionServeContext {
3204 limits: conn_limits,
3205 accepted_capabilities: negotiated.accepted_capabilities,
3206 role: "initiator",
3207 direction,
3208 transcript_hash: native_connection_transcript_hash(
3209 direction,
3210 &local_node_id,
3211 &remote_node_id,
3212 ),
3213 i_open_collision_winner: i_open_collision_winner(&local_node_id, &remote_node_id),
3214 conn,
3215 },
3216 )
3217 .await
3218}
3219
3220#[cfg(any(test, feature = "zakura-testkit"))]
3221pub(crate) async fn run_native_initiator_handshake_without_trace(
3222 connection: &Connection,
3223 limits: &ZakuraLocalLimits,
3224 handshake_config: &ZakuraHandshakeConfig,
3225 local_peer_id: &ZakuraPeerId,
3226) -> Result<NativeHandshakeNegotiated, ZakuraHandlerError> {
3227 run_native_initiator_handshake(
3228 connection,
3229 limits,
3230 handshake_config,
3231 local_peer_id,
3232 &ZakuraTrace::noop(),
3233 &ZakuraConnTrace::placeholder(),
3234 )
3235 .await
3236}
3237
3238async fn run_native_initiator_handshake(
3239 connection: &Connection,
3240 limits: &ZakuraLocalLimits,
3241 handshake_config: &ZakuraHandshakeConfig,
3242 local_peer_id: &ZakuraPeerId,
3243 trace: &ZakuraTrace,
3244 conn: &ZakuraConnTrace,
3245) -> Result<NativeHandshakeNegotiated, ZakuraHandlerError> {
3246 trace.emit(
3247 HANDSHAKE_TABLE,
3248 conn.event("control.started")
3249 .role("initiator")
3250 .phase("control")
3251 .network(handshake_config.network_label()),
3252 );
3253 let (mut send, mut recv) = timeout(limits.control_timeout, connection.open_bi())
3254 .await
3255 .map_err(|_| ZakuraHandlerError::Timeout("open control stream"))??;
3256 let mut local_nonce = [0; 32];
3257 OsRng.fill_bytes(&mut local_nonce);
3258
3259 let hello = ZakuraControlHello {
3260 magic: CONTROL_HELLO_MAGIC,
3261 control_version: CONTROL_VERSION,
3262 selected_zakura_protocol: ZAKURA_PROTOCOL_VERSION_1,
3263 handshake_path: ZakuraHandshakePath::Native,
3264 role: ZakuraControlRole::Initiator,
3265 network_id: handshake_config.network_id,
3266 chain_id: handshake_config.chain_id,
3267 iroh_node_id: local_peer_id.as_bytes().to_vec(),
3268 peer_nonce: local_nonce,
3269 initiator_upgrade_nonce: [0; 32],
3270 responder_upgrade_nonce: [0; 32],
3271 legacy_upgrade_transcript: [0; 32],
3272 capabilities: handshake_config.supported_capabilities,
3273 required_channels: 0,
3274 initial_limits: limits.initial_limits(),
3275 };
3276
3277 write_control_payload(&mut send, &hello.encode()?, limits.control_timeout).await?;
3278 let ack_bytes = read_control_payload(
3279 &mut recv,
3280 handshake_config.max_control_frame_bytes,
3281 limits.control_timeout,
3282 )
3283 .await?;
3284 let ack = ZakuraControlAck::decode(&ack_bytes)?;
3285 ack.validate(
3286 ZAKURA_PROTOCOL_VERSION_1,
3287 local_nonce,
3288 ack.peer_nonce,
3289 &limits.initial_limits(),
3290 handshake_config,
3291 )?;
3292 trace.emit(
3293 HANDSHAKE_TABLE,
3294 conn.event("control.succeeded")
3295 .role("initiator")
3296 .phase("control")
3297 .selected_protocol(ack.selected_zakura_protocol)
3298 .network(handshake_config.network_label()),
3299 );
3300 Ok(NativeHandshakeNegotiated {
3301 limits: ack.accepted_limits,
3302 accepted_capabilities: ack.accepted_capabilities,
3303 })
3304}
3305
3306#[allow(clippy::too_many_arguments)]
3307fn spawn_persistent_stream_worker(
3308 workers: &mut JoinSet<()>,
3309 send: SendStream,
3310 recv: RecvStream,
3311 stream: Stream,
3312 prelude: StreamPrelude,
3313 context: StreamWorkerContext,
3314 queue_depth: usize,
3315 opened_locally: bool,
3316 ordered_stream_exit_tx: mpsc::UnboundedSender<OrderedStreamExit>,
3317) -> AdmittedOrderedStream {
3318 let (to_service_tx, to_service_rx) = mpsc::channel(queue_depth);
3319 let (from_service_tx, from_service_rx) = mpsc::channel(queue_depth);
3320 let admitted = AdmittedOrderedStream {
3321 kind: prelude.stream_kind,
3322 session_id: context.stream_id,
3323 recv: FramedRecv::new(to_service_rx),
3324 send: FramedSend::new(from_service_tx),
3325 cancel_token: context.stream_token.clone(),
3326 };
3327
3328 let exit = OrderedStreamExit {
3329 stream,
3330 session_id: admitted.session_id,
3331 opened_locally,
3332 };
3333 workers.spawn(async move {
3334 persistent_stream_worker(
3335 send,
3336 recv,
3337 prelude,
3338 context,
3339 to_service_tx,
3340 from_service_rx,
3341 queue_depth,
3342 )
3343 .await;
3344 let _ = ordered_stream_exit_tx.send(exit);
3345 });
3346
3347 admitted
3348}
3349
3350async fn persistent_stream_worker(
3351 mut send: SendStream,
3352 recv: RecvStream,
3353 prelude: StreamPrelude,
3354 context: StreamWorkerContext,
3355 inbound_tx: mpsc::Sender<Frame>,
3356 outbound_rx: mpsc::Receiver<Frame>,
3357 queue_depth_limit: usize,
3358) {
3359 let context = Arc::new(context);
3360 let stream_kind = prelude.stream_kind;
3361
3362 let (frame_tx, mut frame_rx) = mpsc::channel::<Result<Frame, ZakuraHandlerError>>(1);
3374 let reader_context = Arc::clone(&context);
3375 let reader = tokio::spawn(async move {
3376 let mut recv = recv;
3377 loop {
3378 let frame = tokio::select! {
3379 biased;
3380 _ = reader_context.connection_token.cancelled() => break,
3381 _ = reader_context.stream_token.cancelled() => break,
3382 frame = read_frame(
3383 &mut recv,
3384 inbound_frame_cap_for_stream_kind(&reader_context.limits, stream_kind),
3385 reader_context.limits.idle_timeout,
3386 None,
3392 ) => frame,
3393 };
3394 let message = match frame {
3402 Ok(frame) => {
3403 let _ = reader_context.freshness_tx.send(Instant::now());
3404 match admit_inbound_message(frame.payload.len(), &reader_context, stream_kind) {
3405 InboundMessageAdmission::Admit => Ok(frame),
3406 InboundMessageAdmission::Oversize => Err(ZakuraHandlerError::Oversize),
3407 InboundMessageAdmission::Throttled => Err(ZakuraHandlerError::RateLimited),
3408 }
3409 }
3410 Err(error) => {
3411 if let Some((payload_len, frame_len, max_frame_bytes)) =
3415 error.oversize_frame_details()
3416 {
3417 reader_context.trace.emit(
3418 RATELIMIT_TABLE,
3419 reader_context
3420 .event("frame.oversize")
3421 .stream_kind(stream_kind_label(stream_kind))
3422 .payload_len(payload_len)
3423 .frame_len(frame_len)
3424 .max_frame_bytes(max_frame_bytes),
3425 );
3426 }
3427 Err(error)
3428 }
3429 };
3430 let is_terminal = message.is_err();
3437 let must_disconnect =
3438 matches!(&message, Err(error) if !matches!(error, ZakuraHandlerError::Closed));
3439 let forward_failed = frame_tx.send(message).await.is_err();
3440 if must_disconnect {
3441 reader_context.close_cause.record("ordered_read_error");
3442 reader_context.connection_token.cancel();
3443 }
3444 if forward_failed || is_terminal {
3445 break;
3446 }
3447 }
3448 });
3449
3450 let mut outbound_rx = Some(outbound_rx);
3451 loop {
3452 tokio::select! {
3453 biased;
3454 _ = context.connection_token.cancelled() => break,
3455 _ = context.stream_token.cancelled() => break,
3456 outbound = async {
3457 match outbound_rx.as_mut() {
3458 Some(outbound_rx) => outbound_rx.recv().await,
3459 None => future::pending().await,
3460 }
3461 } => {
3462 match outbound {
3463 Some(frame) => {
3464 if let Err(error) = write_ordered_frame(&mut send, frame, context.limits, stream_kind).await {
3465 if ordered_stream_write_was_stopped(&error) {
3466 debug!(?error, "closing Zakura ordered stream after peer stopped receiving");
3467 break;
3468 }
3469 debug!(?error, "closing Zakura ordered stream writer");
3470 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
3471 context.close_cause.record("ordered_write_error");
3472 context.connection_token.cancel();
3473 break;
3474 }
3475 }
3476 None => {
3477 outbound_rx = None;
3478 }
3479 }
3480 }
3481 inbound = frame_rx.recv() => {
3482 match inbound {
3483 Some(Ok(frame)) => {
3485 if inbound_tx.send(frame).await.is_err() {
3486 debug!(
3487 stream_kind,
3488 "closing Zakura ordered stream after local service receiver dropped"
3489 );
3490 break;
3491 }
3492 metrics::gauge!(
3493 "zakura.p2p.queue.depth",
3494 "stream_kind" => stream_kind_label(stream_kind),
3495 )
3496 .set(queue_depth_limit.saturating_sub(inbound_tx.capacity()) as f64);
3497 }
3498 Some(Err(ZakuraHandlerError::Oversize)) => {
3500 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_OVERSIZE));
3501 context.close_cause.record("ordered_oversize");
3502 context.connection_token.cancel();
3503 break;
3504 }
3505 Some(Err(ZakuraHandlerError::RateLimited)) => {
3506 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RATE_LIMIT));
3507 context.close_cause.record("ordered_rate_limited");
3508 context.connection_token.cancel();
3509 break;
3510 }
3511 Some(Err(ZakuraHandlerError::Closed)) | None => {
3512 break;
3513 }
3514 Some(Err(error)) => {
3516 debug!(?error, "closing Zakura stream worker");
3517 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
3518 context.close_cause.record("ordered_read_error");
3519 context.connection_token.cancel();
3520 break;
3521 }
3522 }
3523 }
3524 }
3525 }
3526
3527 reader.abort();
3531}
3532
3533fn ordered_stream_write_was_stopped(error: &BoxError) -> bool {
3534 matches!(
3535 error.downcast_ref::<iroh::endpoint::WriteError>(),
3536 Some(iroh::endpoint::WriteError::Stopped(_))
3537 )
3538}
3539
3540async fn request_stream_worker(
3541 mut send: SendStream,
3542 mut recv: RecvStream,
3543 prelude: StreamPrelude,
3544 context: StreamWorkerContext,
3545 registry: Arc<ServiceRegistry>,
3546) {
3547 let Some(request_id) = prelude.request_id else {
3548 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
3549 context.close_cause.record("request_without_id");
3550 context.connection_token.cancel();
3551 return;
3552 };
3553
3554 let frame = tokio::select! {
3555 biased;
3556 _ = context.connection_token.cancelled() => return,
3557 frame = read_frame(
3558 &mut recv,
3559 inbound_frame_cap_for_stream_kind(&context.limits, prelude.stream_kind),
3560 context.limits.idle_timeout,
3561 Some(context.limits.idle_timeout),
3565 ) => frame,
3566 };
3567
3568 let frame = match frame {
3569 Ok(frame) => frame,
3570 Err(error) => {
3571 debug!(
3572 ?error,
3573 "closing Zakura request stream with invalid request frame"
3574 );
3575 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
3576 context.close_cause.record("request_read_error");
3577 context.connection_token.cancel();
3578 return;
3579 }
3580 };
3581
3582 let _ = context.freshness_tx.send(Instant::now());
3583 match admit_inbound_message(frame.payload.len(), &context, prelude.stream_kind) {
3584 InboundMessageAdmission::Admit => {}
3585 InboundMessageAdmission::Oversize => {
3586 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_OVERSIZE));
3587 context.close_cause.record("request_oversize");
3588 context.connection_token.cancel();
3589 return;
3590 }
3591 InboundMessageAdmission::Throttled => {
3592 }
3597 }
3598
3599 let response_frames = match registry
3600 .request(
3601 context.peer_id.clone(),
3602 prelude.stream_kind,
3603 request_id,
3604 app_frame_cap_for_stream_kind(&context.limits, prelude.stream_kind),
3605 context.limits.max_message_bytes,
3606 frame,
3607 )
3608 .await
3609 {
3610 Ok(frames) => frames,
3611 Err(SinkReject::Protocol(error)) => {
3612 debug!(
3613 ?error,
3614 "Zakura inbound sink rejected protocol-invalid request"
3615 );
3616 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
3617 context.close_cause.record("request_protocol_reject");
3618 context.connection_token.cancel();
3619 return;
3620 }
3621 Err(SinkReject::Local(error)) => {
3622 debug!(
3623 ?error,
3624 "Zakura inbound sink could not answer request locally"
3625 );
3626 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE));
3627 return;
3628 }
3629 };
3630
3631 for frame in response_frames {
3632 if let Err(error) = write_response_frame(&mut send, frame, context.limits).await {
3633 debug!(?error, "failed to write Zakura request response frame");
3634 let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE));
3635 return;
3636 }
3637 }
3638
3639 let _ = send.finish();
3640}
3641
3642async fn freshness_reaper(mut freshness_rx: watch::Receiver<Instant>, idle_timeout: Duration) {
3643 loop {
3644 let last = *freshness_rx.borrow_and_update();
3645 let elapsed = last.elapsed();
3646 if elapsed >= idle_timeout {
3647 return;
3648 }
3649 tokio::select! {
3650 _ = tokio::time::sleep(idle_timeout - elapsed) => return,
3651 changed = freshness_rx.changed() => {
3652 if changed.is_err() {
3653 return;
3654 }
3655 }
3656 }
3657 }
3658}
3659
3660async fn read_stream_prelude(
3661 recv: &mut RecvStream,
3662 prelude_timeout: Duration,
3663) -> Result<StreamPrelude, ZakuraHandlerError> {
3664 let mut fixed = [0; STREAM_PRELUDE_FIXED_BYTES];
3665 timeout(prelude_timeout, recv.read_exact(&mut fixed))
3666 .await
3667 .map_err(|_| ZakuraHandlerError::Timeout("stream prelude"))??;
3668
3669 let mut fixed_reader = &fixed[..];
3670 let mut magic = [0; 4];
3671 fixed_reader.read_exact(&mut magic)?;
3672 if magic != STREAM_PRELUDE_MAGIC {
3673 return Err(ZakuraProtocolError::InvalidMagic.into());
3674 }
3675 let stream_kind = fixed_reader.read_u16::<LittleEndian>()?;
3676 let stream_version = fixed_reader.read_u16::<LittleEndian>()?;
3677 let request_id = match fixed[STREAM_PRELUDE_REQUEST_ID_FLAG_OFFSET] {
3678 0 => None,
3679 1 => {
3680 let mut bytes = [0; STREAM_PRELUDE_REQUEST_ID_BYTES];
3681 timeout(prelude_timeout, recv.read_exact(&mut bytes))
3682 .await
3683 .map_err(|_| ZakuraHandlerError::Timeout("stream request id"))??;
3684 let mut request_id_reader = &bytes[..];
3685 Some(request_id_reader.read_u64::<LittleEndian>()?)
3686 }
3687 flag => return Err(ZakuraProtocolError::InvalidFlag(flag).into()),
3688 };
3689
3690 let mut cap = [0; STREAM_PRELUDE_CAP_BYTES];
3691 timeout(prelude_timeout, recv.read_exact(&mut cap))
3692 .await
3693 .map_err(|_| ZakuraHandlerError::Timeout("stream prelude cap"))??;
3694 let mut cap_reader = &cap[..];
3695 let max_frame_bytes = cap_reader.read_u32::<LittleEndian>()?;
3696
3697 Ok(StreamPrelude {
3698 magic,
3699 stream_kind,
3700 stream_version,
3701 request_id,
3702 max_frame_bytes,
3703 })
3704}
3705
3706async fn read_frame(
3726 recv: &mut RecvStream,
3727 max_frame_bytes: u32,
3728 read_timeout: Duration,
3729 first_byte_timeout: Option<Duration>,
3730) -> Result<Frame, ZakuraHandlerError> {
3731 let mut header = [0; FRAME_HEADER_BYTES];
3732 let first_byte = recv.read_exact(&mut header[..1]);
3737 match first_byte_timeout {
3738 Some(first_byte_timeout) => match timeout(first_byte_timeout, first_byte).await {
3739 Ok(Ok(())) => {}
3740 Ok(Err(_)) => return Err(ZakuraHandlerError::Closed),
3741 Err(_) => return Err(ZakuraHandlerError::Timeout("frame header")),
3742 },
3743 None => match first_byte.await {
3744 Ok(()) => {}
3745 Err(_) => return Err(ZakuraHandlerError::Closed),
3746 },
3747 }
3748 match timeout(read_timeout, recv.read_exact(&mut header[1..])).await {
3751 Ok(Ok(())) => {}
3752 Ok(Err(_)) => return Err(ZakuraHandlerError::Closed),
3753 Err(_) => return Err(ZakuraHandlerError::Timeout("frame header")),
3754 }
3755 let mut reader = &header[..];
3756 let message_type = reader.read_u16::<LittleEndian>()?;
3757 let flags = reader.read_u16::<LittleEndian>()?;
3758 let payload_len = usize::try_from(reader.read_u32::<LittleEndian>()?)
3759 .expect("u32 payload lengths fit usize on supported targets");
3760 let max_frame_bytes =
3761 usize::try_from(max_frame_bytes).expect("u32 frame cap fits usize on supported targets");
3762 let frame_len = FRAME_HEADER_BYTES.saturating_add(payload_len);
3763 if frame_len > max_frame_bytes {
3764 metrics::counter!("zakura.p2p.ratelimit.frame.oversize").increment(1);
3765 return Err(ZakuraHandlerError::OversizeFrame {
3766 payload_len,
3767 frame_len,
3768 max_frame_bytes,
3769 });
3770 }
3771 let mut payload = vec![0; payload_len];
3772 timeout(read_timeout, recv.read_exact(&mut payload))
3773 .await
3774 .map_err(|_| ZakuraHandlerError::Timeout("frame payload"))??;
3775 Ok(Frame {
3776 message_type,
3777 flags,
3778 payload,
3779 })
3780}
3781
3782async fn read_control_payload(
3783 recv: &mut RecvStream,
3784 max_bytes: u32,
3785 read_timeout: Duration,
3786) -> Result<Vec<u8>, ZakuraHandlerError> {
3787 let max_bytes = max_bytes.min(MAX_CONTROL_PAYLOAD_BYTES as u32);
3796 let mut len_bytes = [0; CONTROL_LENGTH_BYTES];
3797 timeout(read_timeout, recv.read_exact(&mut len_bytes))
3798 .await
3799 .map_err(|_| ZakuraHandlerError::Timeout("control length"))??;
3800 let len = (&len_bytes[..]).read_u32::<LittleEndian>()?;
3801 if len == 0 || len > max_bytes {
3802 return Err(ZakuraHandlerError::Oversize);
3803 }
3804 let mut bytes = vec![0; len as usize];
3805 timeout(read_timeout, recv.read_exact(&mut bytes))
3806 .await
3807 .map_err(|_| ZakuraHandlerError::Timeout("control payload"))??;
3808 Ok(bytes)
3809}
3810
3811async fn write_control_payload(
3812 send: &mut SendStream,
3813 bytes: &[u8],
3814 write_timeout: Duration,
3815) -> Result<(), ZakuraHandlerError> {
3816 let mut len = Vec::with_capacity(CONTROL_LENGTH_BYTES);
3817 len.write_u32::<LittleEndian>(bytes.len() as u32)?;
3818 timeout(write_timeout, send.write_all(&len))
3819 .await
3820 .map_err(|_| ZakuraHandlerError::Timeout("control length write"))??;
3821 timeout(write_timeout, send.write_all(bytes))
3822 .await
3823 .map_err(|_| ZakuraHandlerError::Timeout("control payload write"))??;
3824 let _ = send.finish();
3825 Ok(())
3826}
3827
3828async fn write_ordered_frame(
3829 send: &mut SendStream,
3830 frame: Frame,
3831 limits: ZakuraConnectionLimits,
3832 stream_kind: u16,
3833) -> Result<(), BoxError> {
3834 if frame.payload.len() > limits.max_message_bytes as usize {
3841 return Err(format!(
3842 "Zakura outbound ordered frame payload {} exceeds negotiated max_message_bytes {}",
3843 frame.payload.len(),
3844 limits.max_message_bytes,
3845 )
3846 .into());
3847 }
3848 let frame = frame.encode(app_frame_cap_for_stream_kind(&limits, stream_kind))?;
3849 timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, send.write_all(&frame))
3850 .await
3851 .map_err(|_| -> BoxError { "Zakura outbound frame write timed out".into() })??;
3852 Ok(())
3853}
3854
3855async fn write_outbound_request_frame(
3856 connection: &Connection,
3857 limits: ZakuraConnectionLimits,
3858 stream_kind: u16,
3859 request_id: u64,
3860 message_type: u16,
3861 flags: u16,
3862 payload: Vec<u8>,
3863) -> Result<Vec<Frame>, OutboundRequestError> {
3864 timeout(
3865 OUTBOUND_REQUEST_RESPONSE_TIMEOUT,
3866 write_outbound_request_frame_inner(
3867 connection,
3868 limits,
3869 stream_kind,
3870 request_id,
3871 message_type,
3872 flags,
3873 payload,
3874 ),
3875 )
3876 .await
3877 .map_err(|_| OutboundRequestError::Local("Zakura outbound request/response timed out".into()))?
3878}
3879
3880async fn write_outbound_request_frame_inner(
3881 connection: &Connection,
3882 limits: ZakuraConnectionLimits,
3883 stream_kind: u16,
3884 request_id: u64,
3885 message_type: u16,
3886 flags: u16,
3887 payload: Vec<u8>,
3888) -> Result<Vec<Frame>, OutboundRequestError> {
3889 let mut legacy_state = LegacyResponseReadState::new(LegacyResponseBudget::from_request(
3891 message_type,
3892 &payload,
3893 limits,
3894 )?);
3895 let (mut send, mut recv) = timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, connection.open_bi())
3896 .await
3897 .map_err(|_| -> BoxError { "Zakura outbound request stream open timed out".into() })
3898 .map_err(OutboundRequestError::Local)?
3899 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
3900 let prelude = StreamPrelude {
3901 magic: STREAM_PRELUDE_MAGIC,
3902 stream_kind,
3903 stream_version: ZAKURA_STREAM_VERSION_1,
3904 request_id: Some(request_id),
3905 max_frame_bytes: app_frame_cap_for_stream_kind(&limits, stream_kind),
3906 };
3907 let frame = Frame {
3908 message_type,
3909 flags,
3910 payload,
3911 };
3912 let prelude = prelude.encode().map_err(|error| {
3913 OutboundRequestError::Local(BoxError::from(format!("failed to encode prelude: {error}")))
3914 })?;
3915 let frame = frame
3916 .encode(app_frame_cap_for_stream_kind(&limits, stream_kind))
3917 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
3918 timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, send.write_all(&prelude))
3919 .await
3920 .map_err(|_| -> BoxError { "Zakura outbound request prelude write timed out".into() })
3921 .map_err(OutboundRequestError::Local)?
3922 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
3923 timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, send.write_all(&frame))
3924 .await
3925 .map_err(|_| -> BoxError { "Zakura outbound request frame write timed out".into() })
3926 .map_err(OutboundRequestError::Local)?
3927 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
3928 let _ = send.finish();
3929
3930 let mut frames = Vec::new();
3931 loop {
3932 match read_frame(
3933 &mut recv,
3934 inbound_frame_cap_for_stream_kind(&limits, stream_kind),
3935 limits.idle_timeout,
3936 Some(limits.idle_timeout),
3940 )
3941 .await
3942 {
3943 Ok(frame) => {
3944 legacy_state.validate_frame(request_id, &frame)?;
3945 frames.push(frame);
3946 }
3947 Err(ZakuraHandlerError::Closed) => {
3948 legacy_state.finish()?;
3949 return Ok(frames);
3950 }
3951 Err(ZakuraHandlerError::Timeout(_)) => {
3952 return Err(OutboundRequestError::Local(Box::new(
3953 ZakuraHandlerError::Timeout("outbound response"),
3954 )));
3955 }
3956 Err(error @ ZakuraHandlerError::OversizeFrame { .. })
3957 | Err(error @ ZakuraHandlerError::Oversize) => {
3958 return Err(OutboundRequestError::Fatal(Box::new(error)));
3959 }
3960 Err(error) => return Err(OutboundRequestError::Fatal(Box::new(error))),
3961 }
3962 }
3963}
3964
3965#[derive(Debug)]
3966enum OutboundRequestError {
3967 Local(BoxError),
3968 Fatal(BoxError),
3969}
3970
3971#[derive(Copy, Clone, Debug, Eq, PartialEq)]
3972enum LegacyResponseKind {
3973 Blocks,
3974 Transactions,
3975 BlockHashes,
3976 BlockHeaders,
3977 TransactionIds,
3978 Pong,
3979 Nil,
3980}
3981
3982#[derive(Copy, Clone, Debug)]
3983struct LegacyResponseBudget {
3984 kind: LegacyResponseKind,
3985 max_items: usize,
3986 max_frames: usize,
3987 max_bytes: usize,
3988 max_message_bytes: usize,
3989}
3990
3991impl LegacyResponseBudget {
3992 fn from_request(
3993 message_type: u16,
3994 payload: &[u8],
3995 limits: ZakuraConnectionLimits,
3996 ) -> Result<Self, OutboundRequestError> {
3997 let kind = match message_type {
3998 LEGACY_REQUEST_BLOCKS_BY_HASH => LegacyResponseKind::Blocks,
3999 LEGACY_REQUEST_TRANSACTIONS_BY_ID => LegacyResponseKind::Transactions,
4000 LEGACY_REQUEST_FIND_BLOCKS => LegacyResponseKind::BlockHashes,
4001 LEGACY_REQUEST_FIND_HEADERS => LegacyResponseKind::BlockHeaders,
4002 LEGACY_REQUEST_MEMPOOL_TRANSACTION_IDS => LegacyResponseKind::TransactionIds,
4003 LEGACY_REQUEST_PING => LegacyResponseKind::Pong,
4004 LEGACY_REQUEST_PUSH_TRANSACTION => LegacyResponseKind::Nil,
4005 _ => {
4006 return Err(OutboundRequestError::Local(
4007 format!("unsupported legacy request message type: {message_type}").into(),
4008 ));
4009 }
4010 };
4011 let max_message_bytes = usize::try_from(limits.max_message_bytes)
4012 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
4013 let max_inventory_items = usize::try_from(MAX_TX_INV_IN_SENT_MESSAGE)
4014 .map_err(|error| OutboundRequestError::Local(Box::new(error)))?;
4015 let (max_items, max_frames, max_bytes) = match kind {
4016 LegacyResponseKind::Blocks | LegacyResponseKind::Transactions => {
4017 let item_count = legacy_inventory_count(payload, kind)?;
4018 let max_frames = item_count
4019 .saturating_mul(LEGACY_RESPONSE_MAX_FRAMES_PER_ITEM)
4020 .saturating_add(1);
4021 let item_bytes = match kind {
4022 LegacyResponseKind::Blocks => LEGACY_BLOCK_HASH_BYTES,
4023 LegacyResponseKind::Transactions => LEGACY_INVENTORY_HASH_BYTES,
4024 _ => unreachable!("matched inventory response kind"),
4025 };
4026 let missing_bytes = LEGACY_RESPONSE_REQUEST_ID_BYTES
4027 .saturating_add(LEGACY_COMPACT_SIZE_PREFIX_BYTES)
4028 .saturating_add(item_count.saturating_mul(item_bytes));
4029 let max_bytes = item_count
4030 .saturating_mul(max_message_bytes)
4031 .saturating_add(missing_bytes)
4032 .max(LEGACY_RESPONSE_CHUNK_HEADER_BYTES);
4033 (item_count, max_frames, max_bytes)
4034 }
4035 LegacyResponseKind::BlockHashes => {
4036 let max_bytes = LEGACY_RESPONSE_REQUEST_ID_BYTES
4037 .saturating_add(LEGACY_COMPACT_SIZE_PREFIX_BYTES)
4038 .saturating_add(max_inventory_items.saturating_mul(LEGACY_BLOCK_HASH_BYTES));
4039 (max_inventory_items, 1, max_bytes)
4040 }
4041 LegacyResponseKind::BlockHeaders => (
4042 MAX_HEADERS_PER_MESSAGE,
4043 1,
4044 LEGACY_RESPONSE_REQUEST_ID_BYTES.saturating_add(max_message_bytes),
4045 ),
4046 LegacyResponseKind::TransactionIds => {
4047 let max_bytes = LEGACY_RESPONSE_REQUEST_ID_BYTES
4048 .saturating_add(LEGACY_COMPACT_SIZE_PREFIX_BYTES)
4049 .saturating_add(
4050 max_inventory_items.saturating_mul(LEGACY_INVENTORY_HASH_BYTES),
4051 );
4052 (max_inventory_items, 1, max_bytes)
4053 }
4054 LegacyResponseKind::Pong | LegacyResponseKind::Nil => {
4055 (1, 1, LEGACY_RESPONSE_REQUEST_ID_BYTES)
4056 }
4057 };
4058
4059 let max_bytes = max_bytes.min(LEGACY_RESPONSE_MAX_AGGREGATE_BYTES);
4068
4069 Ok(Self {
4070 kind,
4071 max_items,
4072 max_frames,
4073 max_bytes,
4074 max_message_bytes,
4075 })
4076 }
4077}
4078
4079#[derive(Debug)]
4080struct LegacyResponseReadState {
4081 budget: LegacyResponseBudget,
4082 frames: usize,
4083 bytes: usize,
4084 items: usize,
4085 active_chunk_type: Option<u16>,
4086 active_chunk: Vec<u8>,
4087}
4088
4089impl LegacyResponseReadState {
4090 fn new(budget: LegacyResponseBudget) -> Self {
4091 Self {
4092 budget,
4093 frames: 0,
4094 bytes: 0,
4095 items: 0,
4096 active_chunk_type: None,
4097 active_chunk: Vec::new(),
4098 }
4099 }
4100
4101 fn validate_frame(
4102 &mut self,
4103 request_id: u64,
4104 frame: &Frame,
4105 ) -> Result<(), OutboundRequestError> {
4106 if frame.flags != 0 {
4107 return Err(OutboundRequestError::Fatal(
4108 format!("unsupported legacy response flags: {}", frame.flags).into(),
4109 ));
4110 }
4111
4112 self.frames = self.frames.saturating_add(1);
4113 if self.frames > self.budget.max_frames {
4114 return Err(OutboundRequestError::Fatal(
4115 "too many legacy response frames".into(),
4116 ));
4117 }
4118
4119 self.bytes = self.bytes.saturating_add(frame.payload.len());
4120 if self.bytes > self.budget.max_bytes {
4121 return Err(OutboundRequestError::Fatal(
4122 "legacy response exceeded cumulative byte budget".into(),
4123 ));
4124 }
4125
4126 match frame.message_type {
4127 LEGACY_RESPONSE_BLOCK => {
4128 self.validate_available_chunk(request_id, frame, LegacyResponseKind::Blocks)
4129 }
4130 LEGACY_RESPONSE_TRANSACTION => {
4131 self.validate_available_chunk(request_id, frame, LegacyResponseKind::Transactions)
4132 }
4133 LEGACY_RESPONSE_MISSING_BLOCKS => {
4134 self.validate_missing(request_id, &frame.payload, LegacyResponseKind::Blocks)
4135 }
4136 LEGACY_RESPONSE_MISSING_TRANSACTIONS => {
4137 self.validate_missing(request_id, &frame.payload, LegacyResponseKind::Transactions)
4138 }
4139 LEGACY_RESPONSE_BLOCK_HASHES => self.validate_id_prefixed_hashes(
4140 request_id,
4141 &frame.payload,
4142 LegacyResponseKind::BlockHashes,
4143 ),
4144 LEGACY_RESPONSE_BLOCK_HEADERS => self.validate_headers(request_id, &frame.payload),
4145 LEGACY_RESPONSE_TRANSACTION_IDS => self.validate_id_prefixed_hashes(
4146 request_id,
4147 &frame.payload,
4148 LegacyResponseKind::TransactionIds,
4149 ),
4150 LEGACY_RESPONSE_PONG => {
4151 self.validate_id_only(request_id, &frame.payload, LegacyResponseKind::Pong)
4152 }
4153 LEGACY_RESPONSE_NIL => self.validate_nil(request_id, &frame.payload),
4154 message_type => Err(OutboundRequestError::Fatal(
4155 format!("unknown legacy response message type: {message_type}").into(),
4156 )),
4157 }
4158 }
4159
4160 fn finish(self) -> Result<(), OutboundRequestError> {
4161 if self.active_chunk_type.is_some() {
4162 return Err(OutboundRequestError::Fatal(
4163 "incomplete legacy response chunk".into(),
4164 ));
4165 }
4166 Ok(())
4167 }
4168
4169 fn validate_available_chunk(
4170 &mut self,
4171 request_id: u64,
4172 frame: &Frame,
4173 kind: LegacyResponseKind,
4174 ) -> Result<(), OutboundRequestError> {
4175 if self.budget.kind != kind {
4176 return Err(OutboundRequestError::Fatal(
4177 "legacy response kind does not match request".into(),
4178 ));
4179 }
4180
4181 let (response_id, is_last, bytes) = legacy_response_chunk_header(&frame.payload)?;
4182 if response_id != request_id {
4183 return Err(OutboundRequestError::Fatal(
4184 format!(
4185 "wrong legacy response request id: expected {request_id}, got {response_id}"
4186 )
4187 .into(),
4188 ));
4189 }
4190
4191 match self.active_chunk_type {
4192 Some(active) if active != frame.message_type => {
4193 return Err(OutboundRequestError::Fatal(
4194 "interleaved legacy response chunks".into(),
4195 ));
4196 }
4197 None => self.active_chunk_type = Some(frame.message_type),
4198 _ => {}
4199 }
4200
4201 let new_len = self.active_chunk.len().saturating_add(bytes.len());
4202 if new_len > self.budget.max_message_bytes {
4203 return Err(OutboundRequestError::Fatal(
4204 "legacy response item exceeded message byte cap".into(),
4205 ));
4206 }
4207 self.active_chunk.extend_from_slice(bytes);
4208
4209 if is_last {
4210 self.validate_completed_item(kind)?;
4211 self.active_chunk_type = None;
4212 self.active_chunk.clear();
4213 self.add_items(1)?;
4214 }
4215
4216 Ok(())
4217 }
4218
4219 fn validate_completed_item(
4220 &self,
4221 kind: LegacyResponseKind,
4222 ) -> Result<(), OutboundRequestError> {
4223 match kind {
4224 LegacyResponseKind::Blocks => {
4225 Block::zcash_deserialize(&mut Cursor::new(self.active_chunk.as_slice()))
4226 .map(|_| ())
4227 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))
4228 }
4229 LegacyResponseKind::Transactions => {
4230 Transaction::zcash_deserialize(&mut Cursor::new(self.active_chunk.as_slice()))
4231 .map(|_| ())
4232 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))
4233 }
4234 LegacyResponseKind::BlockHashes
4235 | LegacyResponseKind::BlockHeaders
4236 | LegacyResponseKind::TransactionIds
4237 | LegacyResponseKind::Pong
4238 | LegacyResponseKind::Nil => unreachable!("non-chunk legacy response kind"),
4239 }
4240 }
4241
4242 fn validate_missing(
4243 &mut self,
4244 request_id: u64,
4245 payload: &[u8],
4246 kind: LegacyResponseKind,
4247 ) -> Result<(), OutboundRequestError> {
4248 let payload = self.begin_id_prefixed_response(kind, request_id, payload, "missing")?;
4249 let count = legacy_inventory_count(payload, kind)?;
4250 self.add_items(count)
4251 }
4252
4253 fn validate_id_prefixed_hashes(
4254 &mut self,
4255 request_id: u64,
4256 payload: &[u8],
4257 kind: LegacyResponseKind,
4258 ) -> Result<(), OutboundRequestError> {
4259 let payload = self.begin_id_prefixed_response(kind, request_id, payload, "list")?;
4260 let count = legacy_inventory_count(payload, kind)?;
4261 self.add_items(count)
4262 }
4263
4264 fn validate_headers(
4265 &mut self,
4266 request_id: u64,
4267 payload: &[u8],
4268 ) -> Result<(), OutboundRequestError> {
4269 let payload = self.begin_id_prefixed_response(
4270 LegacyResponseKind::BlockHeaders,
4271 request_id,
4272 payload,
4273 "headers",
4274 )?;
4275 let count = legacy_header_count(payload)?;
4276 self.add_items(count)
4277 }
4278
4279 fn validate_id_only(
4280 &mut self,
4281 request_id: u64,
4282 payload: &[u8],
4283 kind: LegacyResponseKind,
4284 ) -> Result<(), OutboundRequestError> {
4285 let payload =
4286 self.begin_id_prefixed_response(kind, request_id, payload, "acknowledgement")?;
4287 if !payload.is_empty() {
4288 return Err(OutboundRequestError::Fatal(
4289 "legacy acknowledgement response has trailing bytes".into(),
4290 ));
4291 }
4292 self.add_items(1)
4293 }
4294
4295 fn validate_nil(
4307 &mut self,
4308 request_id: u64,
4309 payload: &[u8],
4310 ) -> Result<(), OutboundRequestError> {
4311 match self.budget.kind {
4312 LegacyResponseKind::BlockHashes
4313 | LegacyResponseKind::BlockHeaders
4314 | LegacyResponseKind::TransactionIds
4315 | LegacyResponseKind::Nil => {}
4316 LegacyResponseKind::Blocks
4317 | LegacyResponseKind::Transactions
4318 | LegacyResponseKind::Pong => {
4319 return Err(OutboundRequestError::Fatal(
4320 "unexpected legacy nil response for inventory or ping request".into(),
4321 ));
4322 }
4323 }
4324 if self.active_chunk_type.is_some() {
4325 return Err(OutboundRequestError::Fatal(
4326 "legacy nil response interleaved with response chunk".into(),
4327 ));
4328 }
4329 let (response_id, payload) = legacy_response_id(payload)?;
4330 if response_id != request_id {
4331 return Err(OutboundRequestError::Fatal(
4332 format!(
4333 "wrong legacy nil response request id: expected {request_id}, got {response_id}"
4334 )
4335 .into(),
4336 ));
4337 }
4338 if !payload.is_empty() {
4339 return Err(OutboundRequestError::Fatal(
4340 "legacy nil response has trailing bytes".into(),
4341 ));
4342 }
4343 self.add_items(1)
4344 }
4345
4346 fn begin_id_prefixed_response<'a>(
4347 &self,
4348 kind: LegacyResponseKind,
4349 request_id: u64,
4350 payload: &'a [u8],
4351 label: &'static str,
4352 ) -> Result<&'a [u8], OutboundRequestError> {
4353 if self.budget.kind != kind {
4354 return Err(OutboundRequestError::Fatal(
4355 format!("legacy {label} response kind does not match request").into(),
4356 ));
4357 }
4358 if self.active_chunk_type.is_some() {
4359 return Err(OutboundRequestError::Fatal(
4360 format!("legacy {label} response interleaved with response chunk").into(),
4361 ));
4362 }
4363 let (response_id, payload) = legacy_response_id(payload)?;
4364 if response_id != request_id {
4365 return Err(OutboundRequestError::Fatal(
4366 format!(
4367 "wrong legacy {label} response request id: expected {request_id}, got {response_id}"
4368 )
4369 .into(),
4370 ));
4371 }
4372 Ok(payload)
4373 }
4374
4375 fn add_items(&mut self, count: usize) -> Result<(), OutboundRequestError> {
4376 self.items = self.items.saturating_add(count);
4377 if self.items > self.budget.max_items {
4378 return Err(OutboundRequestError::Fatal(
4379 "legacy response contained more items than requested".into(),
4380 ));
4381 }
4382 Ok(())
4383 }
4384}
4385
4386fn legacy_inventory_count(
4387 payload: &[u8],
4388 kind: LegacyResponseKind,
4389) -> Result<usize, OutboundRequestError> {
4390 let mut reader = Cursor::new(payload);
4391 let count = usize::from(
4392 CompactSizeMessage::zcash_deserialize(&mut reader)
4393 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?,
4394 );
4395
4396 match kind {
4397 LegacyResponseKind::Blocks | LegacyResponseKind::BlockHashes => {
4398 let consumed = usize::try_from(reader.position())
4399 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?;
4400 let expected_len =
4401 consumed.saturating_add(count.saturating_mul(LEGACY_BLOCK_HASH_BYTES));
4402 if expected_len != payload.len() {
4403 return Err(OutboundRequestError::Fatal(
4404 "legacy block inventory has trailing or truncated bytes".into(),
4405 ));
4406 }
4407 }
4408 LegacyResponseKind::Transactions | LegacyResponseKind::TransactionIds => {
4409 for _ in 0..count {
4410 let inventory = InventoryHash::zcash_deserialize(&mut reader)
4411 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?;
4412 if inventory.unmined_tx_id().is_none() {
4413 return Err(OutboundRequestError::Fatal(
4414 "legacy transaction inventory contained a non-transaction item".into(),
4415 ));
4416 }
4417 }
4418 reject_legacy_inventory_trailing(payload, &reader)?;
4419 }
4420 LegacyResponseKind::BlockHeaders | LegacyResponseKind::Pong | LegacyResponseKind::Nil => {
4421 return Err(OutboundRequestError::Local(
4422 "legacy response kind does not use inventory counts".into(),
4423 ));
4424 }
4425 }
4426
4427 Ok(count)
4428}
4429
4430fn legacy_header_count(payload: &[u8]) -> Result<usize, OutboundRequestError> {
4431 let mut reader = Cursor::new(payload);
4432 let count = usize::from(
4433 CompactSizeMessage::zcash_deserialize(&mut reader)
4434 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?,
4435 );
4436 if count > MAX_HEADERS_PER_MESSAGE {
4437 return Err(OutboundRequestError::Fatal(
4438 "legacy headers response exceeded header count cap".into(),
4439 ));
4440 }
4441 for _ in 0..count {
4442 CountedHeader::zcash_deserialize(&mut reader)
4443 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?;
4444 }
4445 reject_legacy_inventory_trailing(payload, &reader)?;
4446 Ok(count)
4447}
4448
4449fn reject_legacy_inventory_trailing(
4450 payload: &[u8],
4451 reader: &Cursor<&[u8]>,
4452) -> Result<(), OutboundRequestError> {
4453 if usize::try_from(reader.position())
4454 .map_err(|error| OutboundRequestError::Fatal(Box::new(error)))?
4455 != payload.len()
4456 {
4457 return Err(OutboundRequestError::Fatal(
4458 "legacy transaction inventory has trailing bytes".into(),
4459 ));
4460 }
4461 Ok(())
4462}
4463
4464fn legacy_response_id(payload: &[u8]) -> Result<(u64, &[u8]), OutboundRequestError> {
4465 if payload.len() < LEGACY_RESPONSE_REQUEST_ID_BYTES {
4466 return Err(OutboundRequestError::Fatal(
4467 "truncated legacy response id".into(),
4468 ));
4469 }
4470 let mut id = [0; LEGACY_RESPONSE_REQUEST_ID_BYTES];
4471 id.copy_from_slice(&payload[..LEGACY_RESPONSE_REQUEST_ID_BYTES]);
4472 Ok((
4473 u64::from_le_bytes(id),
4474 &payload[LEGACY_RESPONSE_REQUEST_ID_BYTES..],
4475 ))
4476}
4477
4478fn legacy_response_chunk_header(
4479 payload: &[u8],
4480) -> Result<(u64, bool, &[u8]), OutboundRequestError> {
4481 if payload.len() < LEGACY_RESPONSE_CHUNK_HEADER_BYTES {
4482 return Err(OutboundRequestError::Fatal(
4483 "truncated legacy response chunk".into(),
4484 ));
4485 }
4486 let (request_id, payload) = legacy_response_id(payload)?;
4487 Ok((request_id, payload[0] != 0, &payload[1..]))
4488}
4489
4490async fn write_response_frame(
4491 send: &mut SendStream,
4492 frame: Frame,
4493 limits: ZakuraConnectionLimits,
4494) -> Result<(), ZakuraHandlerError> {
4495 if frame.payload.len() > limits.max_message_bytes as usize {
4496 return Err(ZakuraHandlerError::Oversize);
4497 }
4498 let frame = frame.encode(limits.max_frame_bytes)?;
4499 timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, send.write_all(&frame))
4500 .await
4501 .map_err(|_| ZakuraHandlerError::Timeout("frame write"))??;
4502 Ok(())
4503}
4504
4505fn validate_idle_invariant(limits: &ZakuraLocalLimits) -> Result<(), ZakuraHandlerError> {
4506 if limits.keep_alive_interval >= limits.quic_idle_timeout {
4507 return Err(ZakuraHandlerError::InvalidLocalLimits);
4508 }
4509 if limits.initial_limits().idle_timeout_millis as u128 >= limits.quic_idle_timeout.as_millis() {
4510 return Err(ZakuraHandlerError::InvalidLocalLimits);
4511 }
4512 Ok(())
4513}
4514
4515fn zakura_secret_key(config: &Config) -> Result<SecretKey, ZakuraHandlerError> {
4516 config
4519 .zakura_secret_key()
4520 .map_err(|_| ZakuraHandlerError::InvalidSecretKey)
4521}
4522
4523fn stream_kind_label(stream_kind: u16) -> &'static str {
4524 match stream_kind {
4525 0 => "control",
4526 1 => "request",
4527 LEGACY_GOSSIP_STREAM_KIND => "gossip",
4528 LEGACY_REQUEST_STREAM_KIND => "legacy_request",
4529 DISCOVERY_STREAM_KIND => "discovery",
4530 HEADER_SYNC_STREAM_KIND => "header_sync",
4531 ZAKURA_STREAM_BLOCK_SYNC => "block_sync",
4532 _ => "unknown",
4533 }
4534}
4535
4536fn app_frame_cap_for_stream_kind(limits: &ZakuraConnectionLimits, stream_kind: u16) -> u32 {
4537 match stream_kind {
4538 HEADER_SYNC_STREAM_KIND => {
4539 let header_sync_cap =
4540 u32::try_from(MAX_HS_MESSAGE_BYTES.saturating_add(FRAME_HEADER_BYTES))
4541 .expect("header-sync frame cap fits in u32");
4542 limits.max_frame_bytes.min(header_sync_cap)
4543 }
4544 ZAKURA_STREAM_BLOCK_SYNC => limits.max_frame_bytes.min(MAX_BS_FRAME_BYTES),
4545 _ => limits.max_frame_bytes.min(LOCAL_MAX_CONTROL_FRAME_BYTES),
4546 }
4547 .max(1)
4548}
4549
4550fn inbound_frame_cap_for_stream_kind(limits: &ZakuraConnectionLimits, stream_kind: u16) -> u32 {
4560 let frame_header_bytes =
4561 u32::try_from(FRAME_HEADER_BYTES).expect("frame header byte count fits in u32");
4562 app_frame_cap_for_stream_kind(limits, stream_kind)
4563 .min(limits.max_message_bytes.saturating_add(frame_header_bytes))
4564}
4565
4566fn per_stream_inbound_queue_depth(
4567 max_inbound_queue_depth: u16,
4568 ordered_stream_count: usize,
4569) -> usize {
4570 let total = usize::from(max_inbound_queue_depth).max(1);
4571 if ordered_stream_count == 0 {
4572 return total;
4573 }
4574
4575 total.saturating_div(ordered_stream_count).max(1)
4576}
4577
4578fn should_run_freshness_reaper(
4579 ordered_stream_count: usize,
4580 request_response_stream_count: usize,
4581) -> bool {
4582 ordered_stream_count > 0 || request_response_stream_count == 0
4583}
4584
4585const ZAKURA_STREAM_VERSION_1: u16 = 1;
4589const ZAKURA_STREAM_VERSION_7: u16 = 7;
4590
4591#[cfg(test)]
4599fn is_supported_stream(registry: &ServiceRegistry, stream_kind: u16, stream_version: u16) -> bool {
4600 registry
4601 .capability_for_stream(stream_kind, stream_version)
4602 .is_some()
4603}
4604
4605type SharedMessageBucket<C = RealClock> = Arc<std::sync::Mutex<TokenBucket<C>>>;
4612
4613type MessageRateBuckets<C = RealClock> = HashMap<u16, SharedMessageBucket<C>>;
4617
4618fn message_bucket_for<C: Clock>(
4625 buckets: &mut MessageRateBuckets<C>,
4626 stream_kind: u16,
4627 message_rate_per_second: u32,
4628 clock: C,
4629) -> SharedMessageBucket<C> {
4630 buckets
4631 .entry(stream_kind)
4632 .or_insert_with(|| {
4633 Arc::new(std::sync::Mutex::new(TokenBucket::with_clock(
4634 message_rate_per_second,
4635 clock,
4636 )))
4637 })
4638 .clone()
4639}
4640
4641#[derive(Clone, Debug)]
4642struct TokenBucket<C = RealClock> {
4643 capacity: u32,
4644 tokens: u32,
4645 refill_per_second: u32,
4646 last_refill: Instant,
4647 clock: C,
4648}
4649
4650impl TokenBucket<RealClock> {
4651 fn new(refill_per_second: u32) -> Self {
4652 Self::with_clock(refill_per_second, RealClock)
4653 }
4654}
4655
4656impl<C: Clock> TokenBucket<C> {
4657 fn with_clock(refill_per_second: u32, clock: C) -> Self {
4658 let capacity = refill_per_second.max(1);
4659 Self {
4660 capacity,
4661 tokens: capacity,
4662 refill_per_second: capacity,
4663 last_refill: clock.now(),
4664 clock,
4665 }
4666 }
4667
4668 fn try_take(&mut self) -> bool {
4669 self.refill(self.clock.now());
4670 if self.tokens == 0 {
4671 return false;
4672 }
4673 self.tokens -= 1;
4674 true
4675 }
4676
4677 fn refill(&mut self, now: Instant) {
4678 let elapsed = now.saturating_duration_since(self.last_refill);
4679 let elapsed_nanos = elapsed.as_nanos();
4680 let new_tokens =
4681 elapsed_nanos.saturating_mul(u128::from(self.refill_per_second)) / 1_000_000_000;
4682 if new_tokens == 0 {
4683 return;
4684 }
4685 let new_tokens = new_tokens.min(u128::from(u32::MAX)) as u32;
4686 self.tokens = self.capacity.min(self.tokens.saturating_add(new_tokens));
4687 self.last_refill = now;
4688 }
4689}
4690
4691#[derive(Debug, Error)]
4693pub enum ZakuraHandlerError {
4694 #[error("Zakura {0} timed out")]
4696 Timeout(&'static str),
4697 #[error("Zakura payload exceeded its cap")]
4699 Oversize,
4700 #[error(
4702 "Zakura frame length {frame_len} exceeded cap {max_frame_bytes} \
4703 (payload length {payload_len})"
4704 )]
4705 OversizeFrame {
4706 payload_len: usize,
4708 frame_len: usize,
4710 max_frame_bytes: usize,
4712 },
4713 #[error("Zakura stream closed")]
4715 Closed,
4716 #[error("invalid Zakura bootstrap peer")]
4718 InvalidBootstrapPeer,
4719 #[error("invalid Zakura iroh secret key")]
4721 InvalidSecretKey,
4722 #[error("invalid Zakura local limits")]
4724 InvalidLocalLimits,
4725 #[error("Zakura resource limit exceeded: {0}")]
4727 ResourceLimit(&'static str),
4728 #[error("Zakura message rate exceeded")]
4730 RateLimited,
4731 #[error(transparent)]
4733 IrohConnection(#[from] iroh::endpoint::ConnectionError),
4734 #[error(transparent)]
4736 IrohConnect(#[from] iroh::endpoint::ConnectError),
4737 #[error(transparent)]
4739 IrohRemoteId(#[from] iroh::endpoint::RemoteNodeIdError),
4740 #[error(transparent)]
4742 IrohWrite(#[from] iroh::endpoint::WriteError),
4743 #[error(transparent)]
4745 IrohRead(#[from] iroh::endpoint::ReadExactError),
4746 #[error(transparent)]
4748 IrohClosedStream(#[from] iroh::endpoint::ClosedStream),
4749 #[error(transparent)]
4751 Protocol(#[from] ZakuraProtocolError),
4752 #[error(transparent)]
4754 Validation(#[from] super::ZakuraValidationError),
4755 #[error(transparent)]
4757 Io(#[from] std::io::Error),
4758}
4759
4760impl ZakuraHandlerError {
4761 fn oversize_frame_details(&self) -> Option<(u64, u64, u64)> {
4762 let Self::OversizeFrame {
4763 payload_len,
4764 frame_len,
4765 max_frame_bytes,
4766 } = self
4767 else {
4768 return None;
4769 };
4770
4771 Some((
4772 u64::try_from(*payload_len).unwrap_or(u64::MAX),
4773 u64::try_from(*frame_len).unwrap_or(u64::MAX),
4774 u64::try_from(*max_frame_bytes).unwrap_or(u64::MAX),
4775 ))
4776 }
4777}
4778
4779#[cfg(test)]
4780mod tests {
4781 use super::*;
4782 use crate::{
4783 protocol::internal::{InventoryResponse, Response},
4784 zakura::{
4785 legacy_gossip::{LegacyRequestFrame, LegacyRequestKind, LegacyResponseCodec},
4786 testkit::LocalEndpointFactory,
4787 HeaderSyncEvent, HeaderSyncMessage, HeaderSyncMisbehavior, HeaderSyncPeerSession,
4788 HeaderSyncRequestId, HeaderSyncStatus, ServicePeerLimits, ZakuraDiscoveryConfig,
4789 ZakuraDiscoveryHandle, ZakuraDiscoveryLocalConfig, LOCAL_MAX_MESSAGE_BYTES,
4790 MAX_HS_MESSAGE_BYTES, MSG_HS_STATUS, ZAKURA_CAP_DISCOVERY, ZAKURA_CAP_HEADER_SYNC,
4791 ZAKURA_CAP_LEGACY_GOSSIP,
4792 },
4793 P2pStack,
4794 };
4795 use iroh::{
4796 endpoint::Connection,
4797 protocol::{AcceptError, ProtocolHandler},
4798 };
4799 use zakura_chain::{
4800 block::{self, Block},
4801 serialization::{ZcashDeserialize, MAX_PROTOCOL_MESSAGE_LEN},
4802 transaction::{self, UnminedTxId},
4803 };
4804 use zakura_test::vectors::BLOCK_TESTNET_141042_BYTES;
4805
4806 #[tokio::test]
4812 async fn unset_listen_addr_binds_loopback_not_unspecified() {
4813 let builder = direct_endpoint_builder(SecretKey::generate(OsRng));
4814 let builder = bind_native_endpoint(builder, None);
4815 let endpoint = builder.bind().await.expect("loopback bind should succeed");
4816
4817 let sockets = endpoint.bound_sockets();
4818 assert!(
4819 !sockets.is_empty(),
4820 "endpoint should bind at least one socket"
4821 );
4822 for socket in &sockets {
4823 assert!(
4824 socket.ip().is_loopback(),
4825 "unset listen_addr must bind loopback only, but bound {socket} \
4826 (exposes P2P_V2_ALPN on all interfaces)"
4827 );
4828 }
4829
4830 endpoint.close().await;
4831 }
4832
4833 #[derive(Debug, Clone)]
4834 struct CaptureConnection {
4835 connection_tx: mpsc::Sender<Connection>,
4836 stream_tx: mpsc::Sender<(SendStream, RecvStream)>,
4837 }
4838
4839 impl ProtocolHandler for CaptureConnection {
4840 async fn accept(&self, connection: Connection) -> Result<(), AcceptError> {
4841 let _ = self.connection_tx.send(connection.clone()).await;
4842 for _ in 0..2 {
4843 let Ok(streams) = connection.accept_bi().await else {
4844 break;
4845 };
4846 let _ = self.stream_tx.send(streams).await;
4847 }
4848 Ok(())
4849 }
4850 }
4851
4852 #[derive(Debug, Default)]
4853 struct RecordingService {
4854 deliveries: std::sync::Mutex<Vec<(ZakuraPeerId, u16, u16)>>,
4855 }
4856
4857 impl RecordingService {
4858 fn deliveries(&self) -> Vec<(ZakuraPeerId, u16, u16)> {
4859 self.deliveries
4860 .lock()
4861 .expect("recording sink mutex is never poisoned")
4862 .clone()
4863 }
4864 }
4865
4866 impl Service for RecordingService {
4867 fn name(&self) -> &'static str {
4868 "recording"
4869 }
4870
4871 fn streams(&self) -> &[Stream] {
4872 crate::zakura::legacy_gossip::legacy_gossip_streams()
4873 }
4874
4875 fn add_peer(&self, _peer: Peer) {}
4876
4877 fn remove_peer(&self, _peer: &ZakuraPeerId, _conn_id: ZakuraConnId) {}
4878
4879 fn deliver_frame(
4880 &self,
4881 peer_id: ZakuraPeerId,
4882 stream_kind: u16,
4883 frame: Frame,
4884 ) -> Result<(), SinkReject> {
4885 self.deliveries
4886 .lock()
4887 .map_err(|error| SinkReject::local(format!("recording sink poisoned: {error}")))?
4888 .push((peer_id, stream_kind, frame.message_type));
4889 Ok(())
4890 }
4891 }
4892
4893 #[derive(Debug)]
4894 struct NoopService;
4895
4896 impl Service for NoopService {
4897 fn name(&self) -> &'static str {
4898 "noop"
4899 }
4900
4901 fn streams(&self) -> &[Stream] {
4902 &[]
4903 }
4904
4905 fn add_peer(&self, _peer: Peer) {}
4906
4907 fn remove_peer(&self, _peer: &ZakuraPeerId, _conn_id: ZakuraConnId) {}
4908 }
4909
4910 #[derive(Debug)]
4911 struct DeclaredStreamService {
4912 streams: Vec<Stream>,
4913 }
4914
4915 impl Service for DeclaredStreamService {
4916 fn name(&self) -> &'static str {
4917 "declared-stream"
4918 }
4919
4920 fn streams(&self) -> &[Stream] {
4921 &self.streams
4922 }
4923
4924 fn add_peer(&self, _peer: Peer) {}
4925
4926 fn remove_peer(&self, _peer: &ZakuraPeerId, _conn_id: ZakuraConnId) {}
4927 }
4928
4929 #[derive(Debug)]
4930 struct GenerationGuardedRecordingService {
4931 streams: Vec<Stream>,
4932 active: std::sync::Mutex<HashMap<ZakuraPeerId, (ZakuraConnId, CancellationToken)>>,
4933 disconnected: std::sync::Mutex<Vec<(ZakuraPeerId, ZakuraConnId)>>,
4934 }
4935
4936 impl GenerationGuardedRecordingService {
4937 fn new(streams: Vec<Stream>) -> Arc<Self> {
4938 Arc::new(Self {
4939 streams,
4940 active: std::sync::Mutex::new(HashMap::new()),
4941 disconnected: std::sync::Mutex::new(Vec::new()),
4942 })
4943 }
4944
4945 fn active_conn(&self, peer: &ZakuraPeerId) -> Option<ZakuraConnId> {
4946 self.active
4947 .lock()
4948 .expect("recording service active map is never poisoned")
4949 .get(peer)
4950 .map(|(conn_id, _token)| *conn_id)
4951 }
4952
4953 fn active_token_cancelled(&self, peer: &ZakuraPeerId) -> Option<bool> {
4954 self.active
4955 .lock()
4956 .expect("recording service active map is never poisoned")
4957 .get(peer)
4958 .map(|(_conn_id, token)| token.is_cancelled())
4959 }
4960
4961 fn disconnected(&self) -> Vec<(ZakuraPeerId, ZakuraConnId)> {
4962 self.disconnected
4963 .lock()
4964 .expect("recording service disconnected list is never poisoned")
4965 .clone()
4966 }
4967 }
4968
4969 impl Service for GenerationGuardedRecordingService {
4970 fn name(&self) -> &'static str {
4971 "generation-recording"
4972 }
4973
4974 fn streams(&self) -> &[Stream] {
4975 &self.streams
4976 }
4977
4978 fn add_peer(&self, peer: Peer) {
4979 let mut active = self
4980 .active
4981 .lock()
4982 .expect("recording service active map is never poisoned");
4983 if active
4984 .get(&peer.id)
4985 .is_some_and(|(active_conn_id, _token)| *active_conn_id > peer.conn_id)
4986 {
4987 peer.service_cancel_token().cancel();
4988 return;
4989 }
4990 if let Some((_old_conn_id, old_token)) =
4991 active.insert(peer.id.clone(), (peer.conn_id, peer.service_cancel_token()))
4992 {
4993 old_token.cancel();
4994 }
4995 }
4996
4997 fn remove_peer(&self, peer: &ZakuraPeerId, conn_id: ZakuraConnId) {
4998 let removed = {
4999 let mut active = self
5000 .active
5001 .lock()
5002 .expect("recording service active map is never poisoned");
5003 if active
5004 .get(peer)
5005 .is_some_and(|(active_conn_id, _token)| *active_conn_id == conn_id)
5006 {
5007 active.remove(peer)
5008 } else {
5009 None
5010 }
5011 };
5012
5013 if let Some((_active_conn_id, token)) = removed {
5014 token.cancel();
5015 self.disconnected
5016 .lock()
5017 .expect("recording service disconnected list is never poisoned")
5018 .push((peer.clone(), conn_id));
5019 }
5020 }
5021 }
5022
5023 fn test_peer(byte: u8) -> ZakuraPeerId {
5024 ZakuraPeerId::new(vec![byte; 32]).expect("32-byte node id is valid")
5025 }
5026
5027 fn test_conn_id() -> ZakuraConnId {
5028 static NEXT_TEST_CONN_ID: AtomicU64 = AtomicU64::new(1);
5029 NEXT_TEST_CONN_ID.fetch_add(1, Ordering::Relaxed)
5030 }
5031
5032 fn registered_conn_id(registration: ZakuraRegistration) -> ZakuraConnId {
5033 match registration {
5034 ZakuraRegistration::Registered { conn_id, .. } => conn_id,
5035 other => panic!("expected peer registration to succeed, got {other:?}"),
5036 }
5037 }
5038
5039 #[test]
5040 fn native_duplicate_tie_breaker_converges_for_simultaneous_open() {
5041 let node_a = LocalEndpointFactory::secret_key(1).public();
5042 let node_b = LocalEndpointFactory::secret_key(2).public();
5043 let a_outbound =
5044 native_connection_transcript_hash(ServicePeerDirection::Outbound, &node_a, &node_b);
5045 let a_inbound =
5046 native_connection_transcript_hash(ServicePeerDirection::Inbound, &node_a, &node_b);
5047 let b_outbound =
5048 native_connection_transcript_hash(ServicePeerDirection::Outbound, &node_b, &node_a);
5049 let b_inbound =
5050 native_connection_transcript_hash(ServicePeerDirection::Inbound, &node_b, &node_a);
5051
5052 assert_eq!(a_outbound, b_inbound);
5053 assert_eq!(a_inbound, b_outbound);
5054 assert_ne!(a_outbound, a_inbound);
5055
5056 let winning_key = a_outbound.min(a_inbound);
5057 let losing_key = a_outbound.max(a_inbound);
5058 let peer = test_peer(7);
5059 let mut supervisor_a = ZakuraPeerSupervisor::default();
5060 let mut supervisor_b = ZakuraPeerSupervisor::default();
5061 assert!(matches!(
5062 supervisor_a.register_authenticated(peer.clone(), a_outbound),
5063 ZakuraUpgradeOutcome::Upgraded { .. }
5064 ));
5065 let _ = supervisor_a.register_authenticated(peer.clone(), a_inbound);
5066 assert!(matches!(
5067 supervisor_b.register_authenticated(peer.clone(), b_inbound),
5068 ZakuraUpgradeOutcome::Upgraded { .. }
5069 ));
5070 let _ = supervisor_b.register_authenticated(peer.clone(), b_outbound);
5071
5072 assert!(matches!(
5073 supervisor_a.register_authenticated(peer.clone(), winning_key),
5074 ZakuraUpgradeOutcome::Duplicate { .. }
5075 ));
5076 assert!(matches!(
5077 supervisor_b.register_authenticated(peer.clone(), winning_key),
5078 ZakuraUpgradeOutcome::Duplicate { .. }
5079 ));
5080 assert!(matches!(
5081 supervisor_a.register_authenticated(peer.clone(), losing_key),
5082 ZakuraUpgradeOutcome::Duplicate { .. }
5083 ));
5084 assert!(matches!(
5085 supervisor_b.register_authenticated(peer, losing_key),
5086 ZakuraUpgradeOutcome::Duplicate { .. }
5087 ));
5088 }
5089
5090 #[tokio::test]
5091 async fn upgraded_winner_registers_second_loser_cleanup_is_generation_guarded() {
5092 let supervisor = ZakuraSupervisorHandle::new(1);
5093 let peer = test_peer(11);
5094 let remote_ip: IpAddr = "203.0.113.11".parse().expect("test ip parses");
5095 let losing_trace_conn = 100;
5096 let winning_trace_conn = 1;
5097 let losing_hash = [0x80; TRANSCRIPT_HASH_BYTES];
5098 let winning_hash = [0x10; TRANSCRIPT_HASH_BYTES];
5099 let losing_token = CancellationToken::new();
5100 let winning_token = CancellationToken::new();
5101
5102 async fn register_with_hash(
5103 supervisor: &ZakuraSupervisorHandle,
5104 conn_id: ZakuraConnId,
5105 peer: &ZakuraPeerId,
5106 remote_ip: IpAddr,
5107 transcript_hash: [u8; TRANSCRIPT_HASH_BYTES],
5108 token: CancellationToken,
5109 ) -> ZakuraRegistration {
5110 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
5111 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
5112 supervisor
5113 .register(
5114 conn_id,
5115 peer.clone(),
5116 Some(remote_ip),
5117 transcript_hash,
5118 outbound_handle,
5119 token,
5120 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
5121 )
5122 .await
5123 }
5124
5125 let losing_conn = registered_conn_id(
5126 register_with_hash(
5127 &supervisor,
5128 losing_trace_conn,
5129 &peer,
5130 remote_ip,
5131 losing_hash,
5132 losing_token.clone(),
5133 )
5134 .await,
5135 );
5136 let winning_conn = registered_conn_id(
5137 register_with_hash(
5138 &supervisor,
5139 winning_trace_conn,
5140 &peer,
5141 remote_ip,
5142 winning_hash,
5143 winning_token.clone(),
5144 )
5145 .await,
5146 );
5147 assert!(
5148 winning_trace_conn < losing_trace_conn,
5149 "the test deliberately inverts raw trace-id order",
5150 );
5151 assert!(
5152 winning_conn > losing_conn,
5153 "service generations follow successful registration order",
5154 );
5155
5156 assert!(
5157 losing_token.is_cancelled(),
5158 "upgraded winner explicitly cancels the incumbent loser",
5159 );
5160 assert!(
5161 !winning_token.is_cancelled(),
5162 "the winning replacement remains live",
5163 );
5164
5165 supervisor.deregister(&peer, losing_conn).await;
5166
5167 {
5168 let state = supervisor.inner.lock().await;
5169 let entry = state
5170 .active_by_peer
5171 .get(&peer)
5172 .expect("winner remains registered after loser cleanup");
5173 assert_eq!(entry.conn_id, winning_conn);
5174 assert_eq!(entry.remote_ip, Some(remote_ip));
5175 assert_eq!(state.active_by_ip.get(&remote_ip), Some(&1));
5176 state.debug_assert_accounting();
5177 }
5178 assert_eq!(supervisor.registered_ids().await, vec![peer.clone()]);
5179 assert_eq!(supervisor.outbound_peer_handles().await.len(), 1);
5180
5181 supervisor.deregister(&peer, winning_conn).await;
5182 let state = supervisor.inner.lock().await;
5183 assert!(state.active_by_peer.is_empty());
5184 assert!(state.active_by_ip.is_empty());
5185 state.debug_assert_accounting();
5186 }
5187
5188 #[tokio::test]
5189 async fn upgraded_replacement_from_full_different_ip_is_rejected() {
5190 let supervisor = ZakuraSupervisorHandle::new(1);
5191 let peer_a = test_peer(21);
5192 let peer_b = test_peer(22);
5193 let ip1: IpAddr = "203.0.113.21".parse().expect("test ip parses");
5194 let ip2: IpAddr = "203.0.113.22".parse().expect("test ip parses");
5195 let peer_a_losing_conn = test_conn_id();
5196 let peer_b_conn = test_conn_id();
5197 let peer_a_winning_conn = test_conn_id();
5198 let peer_a_losing_hash = [0x80; TRANSCRIPT_HASH_BYTES];
5199 let peer_a_winning_hash = [0x10; TRANSCRIPT_HASH_BYTES];
5200 let peer_b_hash = [0x40; TRANSCRIPT_HASH_BYTES];
5201 let peer_a_losing_token = CancellationToken::new();
5202 let peer_b_token = CancellationToken::new();
5203 let peer_a_winning_token = CancellationToken::new();
5204
5205 let peer_a_losing_generation = registered_conn_id(
5206 register_test_peer_with_hash_and_ip(
5207 &supervisor,
5208 peer_a_losing_conn,
5209 &peer_a,
5210 Some(ip1),
5211 peer_a_losing_hash,
5212 peer_a_losing_token.clone(),
5213 )
5214 .await,
5215 );
5216 let peer_b_generation = registered_conn_id(
5217 register_test_peer_with_hash_and_ip(
5218 &supervisor,
5219 peer_b_conn,
5220 &peer_b,
5221 Some(ip2),
5222 peer_b_hash,
5223 peer_b_token.clone(),
5224 )
5225 .await,
5226 );
5227
5228 let registration = register_test_peer_with_hash_and_ip(
5229 &supervisor,
5230 peer_a_winning_conn,
5231 &peer_a,
5232 Some(ip2),
5233 peer_a_winning_hash,
5234 peer_a_winning_token.clone(),
5235 )
5236 .await;
5237 assert!(
5238 matches!(
5239 registration,
5240 ZakuraRegistration::Rejected(ZakuraRejectReason::ResourceLimit)
5241 ),
5242 "an upgraded same-peer replacement from a different full IP bucket must be rejected",
5243 );
5244 assert!(
5245 !peer_a_losing_token.is_cancelled(),
5246 "the incumbent stays live because the target IP bucket has no room",
5247 );
5248 assert!(
5249 !peer_b_token.is_cancelled(),
5250 "the unrelated peer that fills the target IP bucket stays live",
5251 );
5252 assert!(
5253 !peer_a_winning_token.is_cancelled(),
5254 "the rejected replacement was never registered",
5255 );
5256
5257 let state = supervisor.inner.lock().await;
5258 assert_eq!(
5259 state
5260 .active_by_peer
5261 .get(&peer_a)
5262 .expect("peer A incumbent remains registered")
5263 .conn_id,
5264 peer_a_losing_generation
5265 );
5266 assert_eq!(
5267 state
5268 .active_by_peer
5269 .get(&peer_b)
5270 .expect("peer B remains registered")
5271 .conn_id,
5272 peer_b_generation
5273 );
5274 assert_eq!(state.active_by_ip.get(&ip1), Some(&1));
5275 assert_eq!(state.active_by_ip.get(&ip2), Some(&1));
5276 state.debug_assert_accounting();
5277 }
5278
5279 #[tokio::test]
5280 async fn loser_cleanup_does_not_disconnect_upgraded_winner_service_session() {
5281 let supervisor = ZakuraSupervisorHandle::new(2);
5282 let peer = test_peer(23);
5283 let remote_ip: IpAddr = "203.0.113.23".parse().expect("test ip parses");
5284 let losing_conn = test_conn_id();
5285 let winning_conn = test_conn_id();
5286 let losing_hash = [0x80; TRANSCRIPT_HASH_BYTES];
5287 let winning_hash = [0x10; TRANSCRIPT_HASH_BYTES];
5288 let losing_token = CancellationToken::new();
5289 let winning_token = CancellationToken::new();
5290
5291 let losing_conn = registered_conn_id(
5292 register_test_peer_with_hash_and_ip(
5293 &supervisor,
5294 losing_conn,
5295 &peer,
5296 Some(remote_ip),
5297 losing_hash,
5298 losing_token.clone(),
5299 )
5300 .await,
5301 );
5302 let winning_conn = registered_conn_id(
5303 register_test_peer_with_hash_and_ip(
5304 &supervisor,
5305 winning_conn,
5306 &peer,
5307 Some(remote_ip),
5308 winning_hash,
5309 winning_token.clone(),
5310 )
5311 .await,
5312 );
5313 assert!(
5314 winning_conn > losing_conn,
5315 "winner must have a later service generation",
5316 );
5317
5318 let stream = Stream {
5319 kind: 61,
5320 version: 1,
5321 frame_cap: 1024,
5322 capability: ZAKURA_CAP_LEGACY_GOSSIP,
5323 mode: StreamMode::Ordered,
5324 };
5325 let service = GenerationGuardedRecordingService::new(vec![stream]);
5326 let registry = Arc::new(
5327 ServiceRegistry::new(vec![service.clone()])
5328 .expect("test service declares a valid stream"),
5329 );
5330 let (_inbound_tx, inbound_rx) = crate::zakura::framed_channel(1);
5331 let (outbound_tx, _outbound_rx) = crate::zakura::framed_channel(1);
5332 let streams = HashMap::from([(stream.kind, (inbound_rx, outbound_tx))]);
5333 let winning_service_token = CancellationToken::new();
5334
5335 registry.add_peer(Peer::new_with_conn_id_and_direction(
5336 winning_conn,
5337 peer.clone(),
5338 Some(remote_ip),
5339 ZAKURA_CAP_LEGACY_GOSSIP,
5340 ServicePeerDirection::Inbound,
5341 streams,
5342 winning_service_token.clone(),
5343 ));
5344 assert_eq!(service.active_conn(&peer), Some(winning_conn));
5345
5346 let (_stale_inbound_tx, stale_inbound_rx) = crate::zakura::framed_channel(1);
5347 let (stale_outbound_tx, _stale_outbound_rx) = crate::zakura::framed_channel(1);
5348 let stale_streams = HashMap::from([(stream.kind, (stale_inbound_rx, stale_outbound_tx))]);
5349 let stale_service_token = CancellationToken::new();
5350 registry.add_peer(Peer::new_with_conn_id_and_direction(
5351 losing_conn,
5352 peer.clone(),
5353 Some(remote_ip),
5354 ZAKURA_CAP_LEGACY_GOSSIP,
5355 ServicePeerDirection::Inbound,
5356 stale_streams,
5357 stale_service_token.clone(),
5358 ));
5359 assert_eq!(
5360 service.active_conn(&peer),
5361 Some(winning_conn),
5362 "stale loser add must not overwrite the winner's service session",
5363 );
5364 assert!(
5365 matches!(service.active_token_cancelled(&peer), Some(false)),
5366 "stale loser add must not cancel the winner's service token",
5367 );
5368
5369 registry.remove_peer(&peer, losing_conn, ZAKURA_CAP_LEGACY_GOSSIP);
5370 supervisor.deregister(&peer, losing_conn).await;
5371
5372 assert_eq!(
5373 service.active_conn(&peer),
5374 Some(winning_conn),
5375 "stale loser cleanup must not remove the winner's service session",
5376 );
5377 assert!(
5378 service.disconnected().is_empty(),
5379 "no service disconnect should be emitted for the winner generation",
5380 );
5381 assert!(
5382 matches!(service.active_token_cancelled(&peer), Some(false)),
5383 "the winner's service session token must remain live",
5384 );
5385
5386 {
5387 let state = supervisor.inner.lock().await;
5388 let entry = state
5389 .active_by_peer
5390 .get(&peer)
5391 .expect("winner remains registered after loser cleanup");
5392 assert_eq!(entry.conn_id, winning_conn);
5393 assert_eq!(state.active_by_ip.get(&remote_ip), Some(&1));
5394 state.debug_assert_accounting();
5395 }
5396
5397 registry.remove_peer(&peer, winning_conn, ZAKURA_CAP_LEGACY_GOSSIP);
5398 supervisor.deregister(&peer, winning_conn).await;
5399 }
5400
5401 #[tokio::test]
5402 async fn registered_peer_cleanup_guard_drop_deregisters_and_allows_redial() {
5403 let supervisor = ZakuraSupervisorHandle::new(2);
5404 let peer = test_peer(24);
5405 let transcript_hash = [0x24; TRANSCRIPT_HASH_BYTES];
5406 let disconnect_token = CancellationToken::new();
5407 let conn_id = registered_conn_id(
5408 register_test_peer_with_hash_and_ip(
5409 &supervisor,
5410 test_conn_id(),
5411 &peer,
5412 None,
5413 transcript_hash,
5414 disconnect_token.clone(),
5415 )
5416 .await,
5417 );
5418
5419 let duplicate = register_test_peer_with_hash_and_ip(
5420 &supervisor,
5421 test_conn_id(),
5422 &peer,
5423 None,
5424 transcript_hash,
5425 CancellationToken::new(),
5426 )
5427 .await;
5428 assert!(
5429 matches!(duplicate, ZakuraRegistration::Duplicate { .. }),
5430 "an exact same-direction redial is a duplicate while the incumbent is registered",
5431 );
5432
5433 let stream = Stream {
5434 kind: 62,
5435 version: 1,
5436 frame_cap: 1024,
5437 capability: ZAKURA_CAP_LEGACY_GOSSIP,
5438 mode: StreamMode::Ordered,
5439 };
5440 let service = GenerationGuardedRecordingService::new(vec![stream]);
5441 let registry = Arc::new(
5442 ServiceRegistry::new(vec![service.clone()])
5443 .expect("test service declares a valid stream"),
5444 );
5445 let (_inbound_tx, inbound_rx) = crate::zakura::framed_channel(1);
5446 let (outbound_tx, _outbound_rx) = crate::zakura::framed_channel(1);
5447 registry.add_peer(Peer::new_with_conn_id_and_direction(
5448 conn_id,
5449 peer.clone(),
5450 None,
5451 ZAKURA_CAP_LEGACY_GOSSIP,
5452 ServicePeerDirection::Inbound,
5453 HashMap::from([(stream.kind, (inbound_rx, outbound_tx))]),
5454 disconnect_token.clone(),
5455 ));
5456 assert_eq!(service.active_conn(&peer), Some(conn_id));
5457
5458 let mut cleanup_guard = RegisteredPeerCleanupGuard::new(
5459 supervisor.clone(),
5460 registry,
5461 peer.clone(),
5462 conn_id,
5463 disconnect_token.clone(),
5464 );
5465 cleanup_guard.add_admitted_capabilities(ZAKURA_CAP_LEGACY_GOSSIP);
5466 drop(cleanup_guard);
5467
5468 assert!(
5469 disconnect_token.is_cancelled(),
5470 "drop cleanup cancels the registered connection token",
5471 );
5472 assert_eq!(
5473 service.active_conn(&peer),
5474 None,
5475 "drop cleanup removes admitted service state synchronously",
5476 );
5477 assert_eq!(service.disconnected(), vec![(peer.clone(), conn_id)]);
5478
5479 tokio::time::timeout(Duration::from_secs(1), async {
5480 loop {
5481 if supervisor.registered_ids().await.is_empty() {
5482 break;
5483 }
5484 tokio::task::yield_now().await;
5485 }
5486 })
5487 .await
5488 .expect("drop cleanup deregisters the supervisor entry promptly");
5489
5490 let redial = register_test_peer_with_hash_and_ip(
5491 &supervisor,
5492 test_conn_id(),
5493 &peer,
5494 None,
5495 transcript_hash,
5496 CancellationToken::new(),
5497 )
5498 .await;
5499 assert!(
5500 matches!(redial, ZakuraRegistration::Registered { .. }),
5501 "after drop cleanup, the deterministic same-direction redial can register again",
5502 );
5503 }
5504
5505 #[tokio::test]
5506 async fn registration_generation_exhaustion_rejects_without_mutating_state() {
5507 let supervisor = ZakuraSupervisorHandle::new(2);
5508 let peer = test_peer(25);
5509
5510 {
5511 let mut state = supervisor.inner.lock().await;
5512 state.next_registration_id = u64::MAX;
5513 }
5514
5515 let registration = register_test_peer_with_hash_and_ip(
5516 &supervisor,
5517 test_conn_id(),
5518 &peer,
5519 None,
5520 [0x25; TRANSCRIPT_HASH_BYTES],
5521 CancellationToken::new(),
5522 )
5523 .await;
5524
5525 assert!(
5526 matches!(
5527 registration,
5528 ZakuraRegistration::Rejected(ZakuraRejectReason::TemporaryUnavailable)
5529 ),
5530 "exhausted registration generations reject the new connection instead of panicking",
5531 );
5532 assert!(
5533 supervisor.registered_ids().await.is_empty(),
5534 "failed registration must not mutate active supervisor state",
5535 );
5536 }
5537
5538 #[test]
5539 fn ordered_stream_collision_winner_is_mirror_stable() {
5540 let node_a = LocalEndpointFactory::secret_key(1).public();
5544 let node_b = LocalEndpointFactory::secret_key(2).public();
5545 assert_ne!(node_a, node_b);
5546 assert_ne!(
5547 i_open_collision_winner(&node_a, &node_b),
5548 i_open_collision_winner(&node_b, &node_a),
5549 "exactly one side must win a same-kind ordered-stream collision",
5550 );
5551 let a_wins = node_a.as_bytes() < node_b.as_bytes();
5553 assert_eq!(i_open_collision_winner(&node_a, &node_b), a_wins);
5554 }
5555
5556 fn header_sync_test_session(
5557 peer: ZakuraPeerId,
5558 ) -> (HeaderSyncPeerSession, crate::zakura::FramedRecv) {
5559 let (send, recv) = crate::zakura::framed_channel(32);
5560 (
5561 HeaderSyncPeerSession::from_parts(peer, send, CancellationToken::new()),
5562 recv,
5563 )
5564 }
5565
5566 fn test_discovery_service(supervisor: &ZakuraSupervisorHandle) -> Arc<dyn Service> {
5567 let (_handle, service) =
5568 test_discovery_service_with_peer_limits(supervisor, ServicePeerLimits::default());
5569 service
5570 }
5571
5572 fn test_discovery_service_with_peer_limits(
5573 supervisor: &ZakuraSupervisorHandle,
5574 peer_limits: ServicePeerLimits,
5575 ) -> (ZakuraDiscoveryHandle, Arc<dyn Service>) {
5576 let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
5577 let handle = ZakuraDiscoveryHandle::new(
5578 ZakuraDiscoveryLocalConfig {
5579 secret_key: SecretKey::from_bytes(&[7u8; 32]),
5580 direct_addrs: Vec::new(),
5581 services: crate::zakura::discovery::default_advertised_services(),
5582 zakura_protocol_min: handshake.zakura_protocol_min,
5583 zakura_protocol_max: handshake.zakura_protocol_max,
5584 network_id: handshake.network_id,
5585 chain_id: handshake.chain_id,
5586 last_authored_sequence: None,
5587 },
5588 ZakuraDiscoveryConfig {
5589 peer_limits,
5590 ..ZakuraDiscoveryConfig::default()
5591 },
5592 supervisor.subscribe(),
5593 )
5594 .expect("test discovery handle builds");
5595 let service =
5596 Arc::new(crate::zakura::DiscoveryService::new(handle.clone())) as Arc<dyn Service>;
5597 (handle, service)
5598 }
5599
5600 fn header_sync_startup(shutdown: CancellationToken) -> HeaderSyncStartup {
5601 let network = Network::Mainnet;
5602 let anchor = (block::Height(0), network.genesis_hash());
5603 let mut startup = HeaderSyncStartup::new(
5604 network,
5605 anchor,
5606 HeaderSyncFrontiers {
5607 finalized_height: anchor.0,
5608 verified_block_tip: anchor.0,
5609 verified_block_hash: anchor.1,
5610 },
5611 Some(anchor),
5612 ZakuraHeaderSyncConfig::default(),
5613 LOCAL_MAX_MESSAGE_BYTES,
5614 );
5615 startup.shutdown = shutdown;
5616 startup
5617 }
5618
5619 async fn next_header_sync_action(
5620 actions: &mut mpsc::Receiver<HeaderSyncAction>,
5621 ) -> HeaderSyncAction {
5622 tokio::time::timeout(Duration::from_secs(2), actions.recv())
5623 .await
5624 .expect("header-sync action arrives before timeout")
5625 .expect("header-sync action channel stays open")
5626 }
5627
5628 fn status_at_genesis(network: &Network) -> HeaderSyncStatus {
5629 HeaderSyncStatus {
5630 tip_height: block::Height(0),
5631 tip_hash: network.genesis_hash(),
5632 anchor_height: block::Height(0),
5633 ..HeaderSyncStatus::default()
5634 }
5635 }
5636
5637 async fn register_test_peer(
5638 supervisor: &ZakuraSupervisorHandle,
5639 peer: ZakuraPeerId,
5640 disconnect_token: CancellationToken,
5641 ) {
5642 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
5643 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
5644 let registration = supervisor
5645 .register(
5646 test_conn_id(),
5647 peer.clone(),
5648 None,
5649 [peer.as_bytes()[0]; TRANSCRIPT_HASH_BYTES],
5650 outbound_handle,
5651 disconnect_token,
5652 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
5653 )
5654 .await;
5655
5656 assert!(
5657 matches!(registration, ZakuraRegistration::Registered { .. }),
5658 "test peer should register once"
5659 );
5660 }
5661
5662 async fn register_test_peer_with_hash_and_ip(
5663 supervisor: &ZakuraSupervisorHandle,
5664 conn_id: ZakuraConnId,
5665 peer: &ZakuraPeerId,
5666 remote_ip: Option<IpAddr>,
5667 transcript_hash: [u8; TRANSCRIPT_HASH_BYTES],
5668 disconnect_token: CancellationToken,
5669 ) -> ZakuraRegistration {
5670 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
5671 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
5672 supervisor
5673 .register(
5674 conn_id,
5675 peer.clone(),
5676 remote_ip,
5677 transcript_hash,
5678 outbound_handle,
5679 disconnect_token,
5680 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
5681 )
5682 .await
5683 }
5684
5685 #[test]
5686 fn local_limits_clamp_negotiated_values_down() {
5687 let config = Config::default();
5688 let limits = ZakuraLocalLimits::from_config(&config);
5689 let negotiated = ZakuraAcceptedLimits {
5690 max_frame_bytes: u32::MAX,
5691 max_message_bytes: u32::MAX,
5692 max_open_streams: u16::MAX,
5693 max_inbound_queue_depth: u16::MAX,
5694 idle_timeout_millis: u32::MAX,
5695 };
5696
5697 let clamped = limits.clamp(&negotiated);
5698
5699 assert_eq!(clamped.max_frame_bytes, limits.max_frame_bytes);
5700 assert_eq!(clamped.max_message_bytes, limits.max_message_bytes);
5701 assert_eq!(clamped.max_open_streams, limits.max_open_streams);
5702 assert_eq!(
5703 clamped.max_inbound_queue_depth,
5704 limits.max_inbound_queue_depth
5705 );
5706 assert!(clamped.idle_timeout < limits.quic_idle_timeout);
5707 }
5708
5709 #[test]
5710 fn inbound_queue_depth_is_split_across_ordered_streams() {
5711 assert_eq!(per_stream_inbound_queue_depth(64, 2), 32);
5712 assert_eq!(per_stream_inbound_queue_depth(63, 2), 31);
5713 assert_eq!(per_stream_inbound_queue_depth(64, 0), 64);
5714 assert!(per_stream_inbound_queue_depth(63, 2) * 2 <= 63);
5715 }
5716
5717 #[test]
5718 fn request_response_only_peers_do_not_use_ordered_stream_freshness_reaper() {
5719 assert!(!should_run_freshness_reaper(0, 1));
5720 assert!(should_run_freshness_reaper(1, 1));
5721 assert!(should_run_freshness_reaper(1, 0));
5722 assert!(should_run_freshness_reaper(0, 0));
5723 }
5724
5725 #[tokio::test]
5726 async fn v2_p2p_false_leaves_header_sync_disabled() -> Result<(), BoxError> {
5727 let _guard = zakura_test::init();
5728 let config = Config::for_test(P2pStack::Legacy);
5729
5730 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
5731 Arc::new(NoopService) as Arc<dyn Service>
5732 })
5733 .await?;
5734
5735 assert!(endpoint.is_none());
5736 Ok(())
5737 }
5738
5739 #[tokio::test]
5740 async fn v2_p2p_true_starts_header_sync_handle() -> Result<(), BoxError> {
5741 let _guard = zakura_test::init();
5742 let config = Config::for_test(P2pStack::Dual);
5743
5744 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
5745 Arc::new(NoopService) as Arc<dyn Service>
5746 })
5747 .await?
5748 .expect("v2_p2p is enabled in test config");
5749
5750 assert!(endpoint.header_sync().is_some());
5751 endpoint.shutdown().await;
5752 Ok(())
5753 }
5754
5755 #[tokio::test]
5756 async fn endpoint_shutdown_stops_header_sync_task() -> Result<(), BoxError> {
5757 let _guard = zakura_test::init();
5758 let config = Config::for_test(P2pStack::Dual);
5759 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
5760 Arc::new(NoopService) as Arc<dyn Service>
5761 })
5762 .await?
5763 .expect("v2_p2p is enabled in test config");
5764 let header_sync = endpoint
5765 .header_sync()
5766 .expect("header sync handle exists for a v2 endpoint");
5767
5768 endpoint.shutdown().await;
5769
5770 let (session, _recv) = header_sync_test_session(
5771 ZakuraPeerId::new(vec![5u8; 32]).expect("32-byte node id is valid"),
5772 );
5773 let send_result = tokio::time::timeout(
5774 Duration::from_secs(1),
5775 header_sync.send(HeaderSyncEvent::PeerConnected(session)),
5776 )
5777 .await
5778 .expect("send returns promptly after header-sync shutdown");
5779 assert!(send_result.is_err());
5780
5781 Ok(())
5782 }
5783
5784 #[tokio::test]
5792 async fn endpoint_shutdown_stops_maintained_native_dial_loop() -> Result<(), BoxError> {
5793 let _guard = zakura_test::init();
5794 let config = Config::for_test(P2pStack::Dual);
5795 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
5796 Arc::new(NoopService) as Arc<dyn Service>
5797 })
5798 .await?
5799 .expect("v2_p2p is enabled in test config");
5800
5801 let unreachable_addr: SocketAddr = "192.0.2.1:65535".parse().expect("valid test address");
5804 let unreachable = NodeAddr::new(LocalEndpointFactory::secret_key(987_654).public())
5805 .with_direct_addresses([unreachable_addr]);
5806 let dial = endpoint.spawn_native_dial(unreachable);
5807
5808 tokio::time::sleep(Duration::from_millis(50)).await;
5810 assert!(
5811 !dial.is_finished(),
5812 "maintained dial loop to an unreachable peer should still be running"
5813 );
5814
5815 endpoint.shutdown().await;
5816
5817 tokio::time::timeout(Duration::from_secs(5), dial)
5818 .await
5819 .expect("maintained native dial loop must exit after endpoint shutdown")
5820 .expect("native dial task must not panic");
5821 Ok(())
5822 }
5823
5824 #[tokio::test]
5830 async fn endpoint_shutdown_stops_discovery_candidate_dialer() -> Result<(), BoxError> {
5831 let _guard = zakura_test::init();
5832 let config = Config::for_test(P2pStack::Dual);
5833 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
5834 Arc::new(NoopService) as Arc<dyn Service>
5835 })
5836 .await?
5837 .expect("v2_p2p is enabled in test config");
5838
5839 let limits = ZakuraLocalLimits::from_config(&config);
5840 let handshake = ZakuraHandshakeConfig::for_network(&config.network);
5841 let discovery = crate::zakura::discovery::build_discovery_handle(
5842 SecretKey::generate(OsRng),
5843 Vec::new(),
5844 crate::zakura::discovery::default_advertised_services(),
5845 &handshake,
5846 limits.max_connections,
5847 0,
5848 endpoint.supervisor().subscribe(),
5849 )?;
5850 let dialer = crate::zakura::discovery::spawn_native_discovery_dialer(
5851 endpoint.clone(),
5852 discovery,
5853 limits,
5854 );
5855
5856 tokio::time::sleep(Duration::from_millis(50)).await;
5857 assert!(
5858 !dialer.is_finished(),
5859 "discovery candidate dialer should still be running"
5860 );
5861
5862 endpoint.shutdown().await;
5863
5864 tokio::time::timeout(Duration::from_secs(5), dialer)
5865 .await
5866 .expect("discovery candidate dialer must exit after endpoint shutdown")
5867 .expect("discovery dialer task must not panic");
5868 Ok(())
5869 }
5870
5871 #[tokio::test(start_paused = true)]
5883 async fn failed_legacy_upgrade_does_not_leak_maintained_dial() -> Result<(), BoxError> {
5884 let _guard = zakura_test::init();
5885 let config = Config::for_test(P2pStack::Dual);
5886 let endpoint = spawn_zakura_endpoint(&config, |_supervisor, _trace| {
5887 Arc::new(NoopService) as Arc<dyn Service>
5888 })
5889 .await?
5890 .expect("v2_p2p is enabled in test config");
5891
5892 let node_id = LocalEndpointFactory::secret_key(0x0BAD_C0DE)
5896 .public()
5897 .as_bytes()
5898 .to_vec();
5899 let peer_id = ZakuraPeerId::new(node_id.clone()).expect("32-byte node id is valid");
5900 let direct_addresses = vec![b"192.0.2.1:1".to_vec()];
5901
5902 let connector =
5903 crate::zakura::ZakuraHandshakeConnector::new_with_endpoint(endpoint.clone());
5904 let upgraded = connector
5905 .spawn_zakura_dial_to_hints_and_wait(&peer_id, &node_id, &direct_addresses)
5906 .await;
5907
5908 assert!(
5909 !upgraded,
5910 "an unreachable upgrade peer must not report a completed hand-off",
5911 );
5912 assert!(
5913 endpoint
5914 .upgrade_dials
5915 .lock()
5916 .expect("Zakura upgrade dial registry mutex is never poisoned")
5917 .is_empty(),
5918 "a failed legacy upgrade leaked a maintained native dial / upgrade_dials entry",
5919 );
5920
5921 endpoint.shutdown().await;
5922 Ok(())
5923 }
5924
5925 #[tokio::test]
5926 async fn registry_routes_legacy_and_header_sync_and_drops_unknown() -> Result<(), BoxError> {
5927 let _guard = zakura_test::init();
5928 let shutdown = CancellationToken::new();
5929 let startup = header_sync_startup(shutdown.clone());
5930 let (header_sync, mut actions, task) = spawn_header_sync_reactor(startup)?;
5931 let recorder = Arc::new(RecordingService::default());
5932 let supervisor = ZakuraSupervisorHandle::new(1);
5933 let registry = service_registry(
5934 &supervisor,
5935 Some(header_sync.clone()),
5936 None,
5937 ZakuraBlockSyncConfig::default(),
5938 recorder.clone(),
5939 test_discovery_service(&supervisor),
5940 )?;
5941 let peer = test_peer(6);
5942
5943 let (session, _recv) = header_sync_test_session(peer.clone());
5944 header_sync
5945 .send(HeaderSyncEvent::PeerConnected(session))
5946 .await?;
5947 assert!(matches!(
5948 next_header_sync_action(&mut actions).await,
5949 HeaderSyncAction::SendMessage {
5950 msg: HeaderSyncMessage::Status(_),
5951 ..
5952 }
5953 ));
5954
5955 let gossip_frame = Frame {
5956 message_type: 11,
5957 flags: 0,
5958 payload: Vec::new(),
5959 };
5960 let request_frame = Frame {
5961 message_type: 12,
5962 flags: 0,
5963 payload: Vec::new(),
5964 };
5965 let get_headers_frame = HeaderSyncMessage::GetHeaders {
5966 start_height: block::Height(1),
5967 count: 1,
5968 want_tree_aux_roots: false,
5969 }
5970 .encode_frame(Some(
5971 HeaderSyncRequestId::new(1).expect("non-zero request id"),
5972 ))?;
5973
5974 let (gossip_result, request_result, header_sync_result) = tokio::join!(
5975 async { registry.deliver(peer.clone(), LEGACY_GOSSIP_STREAM_KIND, gossip_frame) },
5976 async { registry.deliver(peer.clone(), LEGACY_REQUEST_STREAM_KIND, request_frame) },
5977 async { registry.deliver(peer.clone(), HEADER_SYNC_STREAM_KIND, get_headers_frame) },
5978 );
5979
5980 gossip_result?;
5981 request_result?;
5982 header_sync_result?;
5983 registry.deliver(
5984 peer.clone(),
5985 99,
5986 Frame {
5987 message_type: 13,
5988 flags: 0,
5989 payload: Vec::new(),
5990 },
5991 )?;
5992
5993 let action = next_header_sync_action(&mut actions).await;
5994 assert!(
5995 matches!(
5996 action,
5997 HeaderSyncAction::Misbehavior {
5998 reason: HeaderSyncMisbehavior::GetHeadersSpam,
5999 ..
6000 }
6001 ),
6002 "kind-5 GetHeaders must reach the header-sync reactor, got {action:?}"
6003 );
6004
6005 let deliveries = recorder.deliveries();
6006 assert_eq!(deliveries.len(), 2);
6007 assert!(deliveries.iter().any(|(_, kind, message_type)| *kind
6008 == LEGACY_GOSSIP_STREAM_KIND
6009 && *message_type == 11));
6010 assert!(deliveries.iter().any(|(_, kind, message_type)| *kind
6011 == LEGACY_REQUEST_STREAM_KIND
6012 && *message_type == 12));
6013 assert!(!deliveries
6014 .iter()
6015 .any(|(_, kind, _)| *kind == HEADER_SYNC_STREAM_KIND));
6016
6017 let rejected = registry
6018 .request(
6019 peer,
6020 HEADER_SYNC_STREAM_KIND,
6021 99,
6022 LOCAL_MAX_CONTROL_FRAME_BYTES,
6023 LOCAL_MAX_CONTROL_FRAME_BYTES,
6024 Frame {
6025 message_type: 1,
6026 flags: 0,
6027 payload: Vec::new(),
6028 },
6029 )
6030 .await;
6031 assert!(matches!(rejected, Err(SinkReject::Protocol(_))));
6032
6033 shutdown.cancel();
6034 task.await?;
6035 Ok(())
6036 }
6037
6038 #[tokio::test]
6039 async fn malformed_header_sync_frame_is_protocol_reject_when_reactor_queue_closed(
6040 ) -> Result<(), BoxError> {
6041 let shutdown = CancellationToken::new();
6042 let startup = header_sync_startup(shutdown.clone());
6043 let (header_sync, _actions, task) = spawn_header_sync_reactor(startup)?;
6044 let service = HeaderSyncService::new(header_sync);
6045 let peer = test_peer(11);
6046
6047 shutdown.cancel();
6048 task.await?;
6049
6050 let valid_status_frame =
6051 HeaderSyncMessage::Status(status_at_genesis(&Network::Mainnet)).encode_frame(None)?;
6052 let valid_result =
6053 service.deliver_frame(peer.clone(), HEADER_SYNC_STREAM_KIND, valid_status_frame);
6054 assert!(
6055 matches!(valid_result, Err(SinkReject::Local(_))),
6056 "valid header-sync frames depend on local reactor queue availability"
6057 );
6058
6059 let malformed_frame = Frame {
6060 message_type: u16::from(MSG_HS_STATUS),
6061 flags: 0,
6062 payload: Vec::new(),
6063 };
6064 let malformed_result =
6065 service.deliver_frame(peer, HEADER_SYNC_STREAM_KIND, malformed_frame);
6066 assert!(
6067 matches!(malformed_result, Err(SinkReject::Protocol(_))),
6068 "malformed header-sync frames must disconnect independently of reactor queue availability"
6069 );
6070
6071 Ok(())
6072 }
6073
6074 #[tokio::test]
6075 async fn header_sync_peer_connected_has_ready_outbound_source() -> Result<(), BoxError> {
6076 let _guard = zakura_test::init();
6077 let shutdown = CancellationToken::new();
6078 let startup = header_sync_startup(shutdown.clone());
6079 let (header_sync, mut actions, reactor_task) = spawn_header_sync_reactor(startup)?;
6080 let service = HeaderSyncService::new(header_sync);
6081 let peer = test_peer(12);
6082 let cancel_token = CancellationToken::new();
6083 let (_inbound_tx, inbound_rx) = crate::zakura::framed_channel(1);
6084 let (outbound_tx, mut outbound_rx) = crate::zakura::framed_channel(1);
6085
6086 let mut streams = HashMap::new();
6087 streams.insert(HEADER_SYNC_STREAM_KIND, (inbound_rx, outbound_tx));
6088 service.add_peer(Peer::new(
6089 peer.clone(),
6090 None,
6091 ZAKURA_CAP_HEADER_SYNC,
6092 streams,
6093 cancel_token.clone(),
6094 ));
6095
6096 let received = tokio::time::timeout(Duration::from_secs(1), outbound_rx.recv())
6097 .await
6098 .expect("header-sync outbound source is immediately ready")
6099 .expect("header-sync outbound receiver stays open");
6100 assert_eq!(received.message_type, u16::from(MSG_HS_STATUS));
6101
6102 assert!(matches!(
6103 next_header_sync_action(&mut actions).await,
6104 HeaderSyncAction::SendMessage {
6105 msg: HeaderSyncMessage::Status(_),
6106 ..
6107 }
6108 ));
6109
6110 cancel_token.cancel();
6111 service.remove_peer(&peer, 0);
6112 shutdown.cancel();
6113 reactor_task.await?;
6114 Ok(())
6115 }
6116
6117 #[tokio::test]
6118 async fn lazy_escalation_service_demand_respects_directional_caps() -> Result<(), BoxError> {
6119 let _guard = zakura_test::init();
6120 let shutdown = CancellationToken::new();
6121 let mut startup = header_sync_startup(shutdown.clone());
6122 startup.config.peer_limits = ServicePeerLimits {
6123 max_inbound_peers: 0,
6124 max_outbound_peers: 1,
6125 ..ServicePeerLimits::default()
6126 };
6127 let (header_sync, _actions, reactor_task) = spawn_header_sync_reactor(startup)?;
6128 let service = HeaderSyncService::new(header_sync);
6129 let peer = test_peer(17);
6130
6131 assert!(
6132 !service.wants_peer(&peer, ZAKURA_CAP_HEADER_SYNC, ServicePeerDirection::Inbound,),
6133 "a full inbound cap must not open a new header-sync stream"
6134 );
6135 assert!(
6136 service.wants_peer(
6137 &peer,
6138 ZAKURA_CAP_HEADER_SYNC,
6139 ServicePeerDirection::Outbound,
6140 ),
6141 "an outbound slot should still allow lazy header-sync escalation"
6142 );
6143
6144 shutdown.cancel();
6145 reactor_task.await?;
6146 Ok(())
6147 }
6148
6149 #[tokio::test]
6150 async fn registry_add_remove_updates_header_sync_peers() -> Result<(), BoxError> {
6151 let _guard = zakura_test::init();
6152 let shutdown = CancellationToken::new();
6153 let startup = header_sync_startup(shutdown.clone());
6154 let (header_sync, mut actions, reactor_task) = spawn_header_sync_reactor(startup)?;
6155 let peer = test_peer(7);
6156 let supervisor = ZakuraSupervisorHandle::new(1);
6157 let registry = service_registry(
6158 &supervisor,
6159 Some(header_sync.clone()),
6160 None,
6161 ZakuraBlockSyncConfig::default(),
6162 Arc::new(RecordingService::default()),
6163 test_discovery_service(&supervisor),
6164 )?;
6165 let (_inbound_tx, inbound_rx) = crate::zakura::framed_channel(1);
6166 let (outbound_tx, _outbound_rx) = crate::zakura::framed_channel(1);
6167 let mut streams = HashMap::new();
6168 streams.insert(HEADER_SYNC_STREAM_KIND, (inbound_rx, outbound_tx));
6169 registry.add_peer(Peer::new(
6170 peer.clone(),
6171 None,
6172 ZAKURA_CAP_HEADER_SYNC,
6173 streams,
6174 CancellationToken::new(),
6175 ));
6176 assert!(matches!(
6177 next_header_sync_action(&mut actions).await,
6178 HeaderSyncAction::SendMessage {
6179 msg: HeaderSyncMessage::Status(_),
6180 ..
6181 }
6182 ));
6183
6184 header_sync
6185 .send(HeaderSyncEvent::WireMessage {
6186 peer: peer.clone(),
6187 msg: HeaderSyncMessage::Status(status_at_genesis(&Network::Mainnet)),
6188 })
6189 .await?;
6190 assert_eq!(header_sync.peer_snapshot().inbound_peers, 1);
6191
6192 registry.remove_peer(&peer, 0, ZAKURA_CAP_HEADER_SYNC);
6193
6194 timeout(Duration::from_secs(2), async {
6198 while header_sync.peer_snapshot().inbound_peers != 0 {
6199 tokio::time::sleep(Duration::from_millis(10)).await;
6200 }
6201 })
6202 .await
6203 .expect("deregistered peer must be removed from header-sync state");
6204
6205 shutdown.cancel();
6206 reactor_task.await?;
6207 Ok(())
6208 }
6209
6210 #[tokio::test]
6211 async fn discovery_admission_is_independent_when_header_sync_is_full() -> Result<(), BoxError> {
6212 let _guard = zakura_test::init();
6213 let shutdown = CancellationToken::new();
6214 let mut startup = header_sync_startup(shutdown.clone());
6215 startup.config.peer_limits = ServicePeerLimits {
6216 max_inbound_peers: 0,
6217 ..ServicePeerLimits::default()
6218 };
6219 let (header_sync, mut actions, reactor_task) = spawn_header_sync_reactor(startup)?;
6220 let supervisor = ZakuraSupervisorHandle::new(1);
6221 let (discovery_handle, discovery_service) = test_discovery_service_with_peer_limits(
6222 &supervisor,
6223 ServicePeerLimits {
6224 max_inbound_peers: 1,
6225 ..ServicePeerLimits::default()
6226 },
6227 );
6228 let registry = service_registry(
6229 &supervisor,
6230 Some(header_sync.clone()),
6231 None,
6232 ZakuraBlockSyncConfig::default(),
6233 Arc::new(RecordingService::default()),
6234 discovery_service,
6235 )?;
6236 let peer_node_id = SecretKey::from_bytes(&[13u8; 32]).public();
6237 let peer = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
6238 let (_hs_inbound_tx, hs_inbound_rx) = crate::zakura::framed_channel(1);
6239 let (hs_outbound_tx, mut hs_outbound_rx) = crate::zakura::framed_channel(1);
6240 let (_discovery_inbound_tx, discovery_inbound_rx) = crate::zakura::framed_channel(1);
6241 let (discovery_outbound_tx, mut discovery_outbound_rx) = crate::zakura::framed_channel(4);
6242 let streams = HashMap::from([
6243 (HEADER_SYNC_STREAM_KIND, (hs_inbound_rx, hs_outbound_tx)),
6244 (
6245 DISCOVERY_STREAM_KIND,
6246 (discovery_inbound_rx, discovery_outbound_tx),
6247 ),
6248 ]);
6249
6250 registry.add_peer(Peer::new(
6251 peer.clone(),
6252 None,
6253 ZAKURA_CAP_HEADER_SYNC | ZAKURA_CAP_DISCOVERY,
6254 streams,
6255 CancellationToken::new(),
6256 ));
6257
6258 tokio::time::timeout(Duration::from_secs(1), async {
6259 let mut snapshots = discovery_handle.subscribe_peer_snapshot();
6260 while snapshots.borrow().inbound_peers != 1 {
6261 snapshots
6262 .changed()
6263 .await
6264 .expect("discovery snapshot channel stays open");
6265 }
6266 })
6267 .await
6268 .expect("discovery admits while header-sync is full");
6269
6270 assert_eq!(header_sync.peer_snapshot().inbound_peers, 0);
6271 assert_eq!(header_sync.peer_snapshot().inbound_slots_free, 0);
6272 assert!(
6273 !matches!(
6274 tokio::time::timeout(Duration::from_millis(100), hs_outbound_rx.recv()).await,
6275 Ok(Some(_))
6276 ),
6277 "header-sync rejected peer must not receive Status"
6278 );
6279 let discovery_frame =
6280 tokio::time::timeout(Duration::from_secs(1), discovery_outbound_rx.recv())
6281 .await
6282 .expect("discovery source sends its normal hello")
6283 .expect("discovery outbound stream stays open");
6284 assert_eq!(discovery_frame.message_type, 1);
6285
6286 while let Ok(Some(action)) =
6287 tokio::time::timeout(Duration::from_millis(50), actions.recv()).await
6288 {
6289 assert!(
6290 !matches!(action, HeaderSyncAction::Misbehavior { peer: action_peer, .. } if action_peer == peer),
6291 "header-sync cap rejection must not score peer misbehavior"
6292 );
6293 }
6294
6295 registry.remove_peer(&peer, 0, ZAKURA_CAP_HEADER_SYNC | ZAKURA_CAP_DISCOVERY);
6296 tokio::time::timeout(Duration::from_secs(1), async {
6297 let mut snapshots = discovery_handle.subscribe_peer_snapshot();
6298 while snapshots.borrow().inbound_peers != 0 {
6299 snapshots
6300 .changed()
6301 .await
6302 .expect("discovery snapshot channel stays open");
6303 }
6304 })
6305 .await
6306 .expect("discovery peer state is removed");
6307
6308 shutdown.cancel();
6309 reactor_task.await?;
6310 Ok(())
6311 }
6312
6313 #[tokio::test]
6314 async fn header_sync_admission_is_independent_when_discovery_is_full() -> Result<(), BoxError> {
6315 let _guard = zakura_test::init();
6316 let shutdown = CancellationToken::new();
6317 let mut startup = header_sync_startup(shutdown.clone());
6318 startup.config.peer_limits = ServicePeerLimits {
6319 max_inbound_peers: 1,
6320 ..ServicePeerLimits::default()
6321 };
6322 let (header_sync, mut actions, reactor_task) = spawn_header_sync_reactor(startup)?;
6323 let supervisor = ZakuraSupervisorHandle::new(1);
6324 let (discovery_handle, discovery_service) = test_discovery_service_with_peer_limits(
6325 &supervisor,
6326 ServicePeerLimits {
6327 max_inbound_peers: 0,
6328 ..ServicePeerLimits::default()
6329 },
6330 );
6331 let registry = service_registry(
6332 &supervisor,
6333 Some(header_sync.clone()),
6334 None,
6335 ZakuraBlockSyncConfig::default(),
6336 Arc::new(RecordingService::default()),
6337 discovery_service,
6338 )?;
6339 let peer_node_id = SecretKey::from_bytes(&[14u8; 32]).public();
6340 let peer = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
6341 let (_hs_inbound_tx, hs_inbound_rx) = crate::zakura::framed_channel(1);
6342 let (hs_outbound_tx, mut hs_outbound_rx) = crate::zakura::framed_channel(1);
6343 let (_discovery_inbound_tx, discovery_inbound_rx) = crate::zakura::framed_channel(1);
6344 let (discovery_outbound_tx, mut discovery_outbound_rx) = crate::zakura::framed_channel(1);
6345 let streams = HashMap::from([
6346 (HEADER_SYNC_STREAM_KIND, (hs_inbound_rx, hs_outbound_tx)),
6347 (
6348 DISCOVERY_STREAM_KIND,
6349 (discovery_inbound_rx, discovery_outbound_tx),
6350 ),
6351 ]);
6352
6353 registry.add_peer(Peer::new(
6354 peer.clone(),
6355 None,
6356 ZAKURA_CAP_HEADER_SYNC | ZAKURA_CAP_DISCOVERY,
6357 streams,
6358 CancellationToken::new(),
6359 ));
6360
6361 let header_frame = tokio::time::timeout(Duration::from_secs(1), hs_outbound_rx.recv())
6362 .await
6363 .expect("header-sync sends immediate Status")
6364 .expect("header-sync outbound stream stays open");
6365 assert_eq!(header_frame.message_type, u16::from(MSG_HS_STATUS));
6366 assert!(matches!(
6367 next_header_sync_action(&mut actions).await,
6368 HeaderSyncAction::SendMessage {
6369 peer: action_peer,
6370 msg: HeaderSyncMessage::Status(_),
6371 ..
6372 } if action_peer == peer
6373 ));
6374 assert_eq!(header_sync.peer_snapshot().inbound_peers, 1);
6375 assert_eq!(discovery_handle.peer_snapshot().inbound_peers, 0);
6376 assert!(
6377 !matches!(
6378 tokio::time::timeout(Duration::from_millis(100), discovery_outbound_rx.recv())
6379 .await,
6380 Ok(Some(_))
6381 ),
6382 "discovery rejected peer must not receive discovery source messages"
6383 );
6384
6385 registry.remove_peer(&peer, 0, ZAKURA_CAP_HEADER_SYNC | ZAKURA_CAP_DISCOVERY);
6386 tokio::time::sleep(Duration::from_millis(50)).await;
6387 assert_eq!(header_sync.peer_snapshot().inbound_peers, 0);
6388
6389 shutdown.cancel();
6390 reactor_task.await?;
6391 Ok(())
6392 }
6393
6394 #[tokio::test]
6395 async fn supervisor_disconnect_peer_cancels_registered_token() -> Result<(), BoxError> {
6396 let supervisor = ZakuraSupervisorHandle::new(4);
6397 let peer = test_peer(8);
6398 let disconnect_token = CancellationToken::new();
6399 register_test_peer(&supervisor, peer.clone(), disconnect_token.clone()).await;
6400
6401 assert!(supervisor.disconnect_peer(&peer).await);
6402 tokio::time::timeout(Duration::from_secs(1), disconnect_token.cancelled())
6403 .await
6404 .expect("disconnect token is cancelled promptly");
6405 assert!(!supervisor.disconnect_peer(&test_peer(9)).await);
6406
6407 Ok(())
6408 }
6409
6410 #[tokio::test(start_paused = true)]
6411 async fn duplicate_evicts_stale_incumbent_but_keeps_fresh_one() -> Result<(), BoxError> {
6412 let supervisor = ZakuraSupervisorHandle::new(4);
6418
6419 async fn register_duplicate(
6420 supervisor: &ZakuraSupervisorHandle,
6421 peer: &ZakuraPeerId,
6422 token: CancellationToken,
6423 ) -> ZakuraRegistration {
6424 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
6425 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
6426 supervisor
6427 .register(
6428 test_conn_id(),
6429 peer.clone(),
6430 None,
6431 [peer.as_bytes()[0]; TRANSCRIPT_HASH_BYTES],
6432 outbound_handle,
6433 token,
6434 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
6435 )
6436 .await
6437 }
6438
6439 let fresh_peer = test_peer(8);
6442 let fresh_incumbent = CancellationToken::new();
6443 register_test_peer(&supervisor, fresh_peer.clone(), fresh_incumbent.clone()).await;
6444 let registration =
6445 register_duplicate(&supervisor, &fresh_peer, CancellationToken::new()).await;
6446 assert!(matches!(registration, ZakuraRegistration::Duplicate { .. }));
6447 assert!(
6448 !fresh_incumbent.is_cancelled(),
6449 "a young incumbent is kept so simultaneous-open races do not flap",
6450 );
6451
6452 let stale_peer = test_peer(9);
6455 let stale_incumbent = CancellationToken::new();
6456 register_test_peer(&supervisor, stale_peer.clone(), stale_incumbent.clone()).await;
6457 tokio::time::advance(ZAKURA_DUPLICATE_EVICT_MIN_AGE + Duration::from_secs(1)).await;
6458 let newcomer = CancellationToken::new();
6459 let registration = register_duplicate(&supervisor, &stale_peer, newcomer.clone()).await;
6460 assert!(matches!(registration, ZakuraRegistration::Duplicate { .. }));
6461 assert!(
6462 stale_incumbent.is_cancelled(),
6463 "a stale incumbent is evicted so the restarted peer's redial reclaims the slot",
6464 );
6465 assert!(
6466 !newcomer.is_cancelled(),
6467 "the rejected newcomer's token is never registered, so it is left to redial",
6468 );
6469
6470 Ok(())
6471 }
6472
6473 #[tokio::test(start_paused = true)]
6486 async fn same_peer_duplicate_at_per_ip_cap_still_evicts_stale_incumbent() -> Result<(), BoxError>
6487 {
6488 async fn register_from_ip(
6489 supervisor: &ZakuraSupervisorHandle,
6490 peer: &ZakuraPeerId,
6491 ip: IpAddr,
6492 token: CancellationToken,
6493 ) -> ZakuraRegistration {
6494 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
6495 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
6496 supervisor
6497 .register(
6498 test_conn_id(),
6499 peer.clone(),
6500 Some(ip),
6501 [peer.as_bytes()[0]; TRANSCRIPT_HASH_BYTES],
6502 outbound_handle,
6503 token,
6504 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
6505 )
6506 .await
6507 }
6508
6509 let supervisor = ZakuraSupervisorHandle::new(1);
6510 let ip: IpAddr = "203.0.113.9".parse().expect("test ip parses");
6511 let peer = test_peer(42);
6512
6513 let incumbent = CancellationToken::new();
6516 let registration = register_from_ip(&supervisor, &peer, ip, incumbent.clone()).await;
6517 assert!(
6518 matches!(registration, ZakuraRegistration::Registered { .. }),
6519 "the first connection from the IP registers",
6520 );
6521
6522 tokio::time::advance(ZAKURA_DUPLICATE_EVICT_MIN_AGE + Duration::from_secs(1)).await;
6525
6526 let newcomer = CancellationToken::new();
6530 let registration = register_from_ip(&supervisor, &peer, ip, newcomer.clone()).await;
6531 assert!(
6532 matches!(registration, ZakuraRegistration::Duplicate { .. }),
6533 "a same-peer redial from a capped IP must reach duplicate handling, not be \
6534 rejected as a resource limit before stale eviction can run",
6535 );
6536 assert!(
6537 incumbent.is_cancelled(),
6538 "the stale incumbent must be evicted so the restarted peer's redial reclaims \
6539 the slot in milliseconds instead of waiting for the QUIC idle timeout",
6540 );
6541 assert!(
6542 !newcomer.is_cancelled(),
6543 "the rejected newcomer's token is never registered, so it is left to redial",
6544 );
6545
6546 Ok(())
6547 }
6548
6549 #[tokio::test(start_paused = true)]
6558 async fn same_ip_duplicate_uses_short_eviction_gate() -> Result<(), BoxError> {
6559 async fn register_from_ip(
6560 supervisor: &ZakuraSupervisorHandle,
6561 peer: &ZakuraPeerId,
6562 ip: Option<IpAddr>,
6563 token: CancellationToken,
6564 ) -> ZakuraRegistration {
6565 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
6566 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
6567 supervisor
6568 .register(
6569 test_conn_id(),
6570 peer.clone(),
6571 ip,
6572 [peer.as_bytes()[0]; TRANSCRIPT_HASH_BYTES],
6573 outbound_handle,
6574 token,
6575 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
6576 )
6577 .await
6578 }
6579
6580 let supervisor = ZakuraSupervisorHandle::new(4);
6581 let ip_a: IpAddr = "203.0.113.10".parse().expect("test ip parses");
6582 let peer = test_peer(50);
6583
6584 let incumbent = CancellationToken::new();
6586 let registration =
6587 register_from_ip(&supervisor, &peer, Some(ip_a), incumbent.clone()).await;
6588 assert!(matches!(
6589 registration,
6590 ZakuraRegistration::Registered { .. }
6591 ));
6592
6593 tokio::time::advance(ZAKURA_SAME_IP_DUPLICATE_EVICT_MIN_AGE - Duration::from_secs(1)).await;
6596 let early = CancellationToken::new();
6597 let registration = register_from_ip(&supervisor, &peer, Some(ip_a), early.clone()).await;
6598 assert!(matches!(registration, ZakuraRegistration::Duplicate { .. }));
6599 assert!(
6600 !incumbent.is_cancelled(),
6601 "a same-IP duplicate below the short gate keeps the incumbent",
6602 );
6603
6604 tokio::time::advance(Duration::from_secs(2)).await;
6608 let redial = CancellationToken::new();
6609 let registration = register_from_ip(&supervisor, &peer, Some(ip_a), redial.clone()).await;
6610 assert!(matches!(registration, ZakuraRegistration::Duplicate { .. }));
6611 assert!(
6612 incumbent.is_cancelled(),
6613 "a same-IP duplicate past the short gate evicts the stale incumbent",
6614 );
6615 assert!(
6616 !redial.is_cancelled(),
6617 "the rejected newcomer's token is never registered, so it is left to redial",
6618 );
6619
6620 let peer2 = test_peer(51);
6623 let ip_b: IpAddr = "203.0.113.11".parse().expect("test ip parses");
6624 let ip_c: IpAddr = "203.0.113.12".parse().expect("test ip parses");
6625 let incumbent2 = CancellationToken::new();
6626 let registration =
6627 register_from_ip(&supervisor, &peer2, Some(ip_b), incumbent2.clone()).await;
6628 assert!(matches!(
6629 registration,
6630 ZakuraRegistration::Registered { .. }
6631 ));
6632 tokio::time::advance(ZAKURA_SAME_IP_DUPLICATE_EVICT_MIN_AGE + Duration::from_secs(1)).await;
6633 let diff_ip = CancellationToken::new();
6634 let registration = register_from_ip(&supervisor, &peer2, Some(ip_c), diff_ip.clone()).await;
6635 assert!(matches!(registration, ZakuraRegistration::Duplicate { .. }));
6636 assert!(
6637 !incumbent2.is_cancelled(),
6638 "a different-IP duplicate keeps the long gate and does not evict a young incumbent",
6639 );
6640
6641 Ok(())
6642 }
6643
6644 #[tokio::test]
6645 async fn header_sync_misbehavior_action_does_not_disconnect_peer() -> Result<(), BoxError> {
6646 let _guard = zakura_test::init();
6647 let reactor_shutdown = CancellationToken::new();
6648 let startup = header_sync_startup(reactor_shutdown.clone());
6649 let (header_sync, _reactor_actions, reactor_task) = spawn_header_sync_reactor(startup)?;
6650 let supervisor = ZakuraSupervisorHandle::new(4);
6651 let peer = test_peer(10);
6652 let disconnect_token = CancellationToken::new();
6653 register_test_peer(&supervisor, peer.clone(), disconnect_token.clone()).await;
6654
6655 let (actions_tx, actions_rx) = mpsc::channel(4);
6656 let driver_shutdown = CancellationToken::new();
6657 let driver_task = tokio::spawn(drive_header_sync_actions(
6658 actions_rx,
6659 header_sync,
6660 supervisor,
6661 driver_shutdown.clone(),
6662 ));
6663 actions_tx
6664 .send(HeaderSyncAction::Misbehavior {
6665 peer,
6666 reason: HeaderSyncMisbehavior::MalformedMessage,
6667 })
6668 .await?;
6669
6670 assert!(
6671 tokio::time::timeout(Duration::from_millis(100), disconnect_token.cancelled())
6672 .await
6673 .is_err(),
6674 "misbehavior is record-only and must not cancel the registered connection",
6675 );
6676
6677 driver_shutdown.cancel();
6678 driver_task.await?;
6679 reactor_shutdown.cancel();
6680 reactor_task.await?;
6681 Ok(())
6682 }
6683
6684 #[test]
6688 fn ordered_stream_reopen_is_limited_to_local_header_initiator() {
6689 let header = Stream {
6690 kind: ZAKURA_STREAM_HEADER_SYNC,
6691 version: ZAKURA_HEADER_SYNC_STREAM_VERSION,
6692 frame_cap: 1,
6693 capability: ZAKURA_CAP_HEADER_SYNC,
6694 mode: StreamMode::Ordered,
6695 };
6696 let local_header = OrderedStreamExit {
6697 stream: header,
6698 session_id: 1,
6699 opened_locally: true,
6700 };
6701 assert!(should_reopen_ordered_stream(local_header, true, false));
6702 assert!(!should_reopen_ordered_stream(local_header, false, false));
6703 assert!(!should_reopen_ordered_stream(local_header, true, true));
6704
6705 let accepted_header = OrderedStreamExit {
6706 opened_locally: false,
6707 ..local_header
6708 };
6709 assert!(!should_reopen_ordered_stream(accepted_header, true, false));
6710
6711 let gossip = OrderedStreamExit {
6712 stream: Stream {
6713 kind: LEGACY_GOSSIP_STREAM_KIND,
6714 ..header
6715 },
6716 ..local_header
6717 };
6718 assert!(!should_reopen_ordered_stream(gossip, true, false));
6719 }
6720
6721 #[test]
6722 fn ordered_stream_reopen_backoff_grows_and_caps() {
6723 assert_eq!(
6724 ordered_stream_reopen_backoff(0),
6725 ORDERED_STREAM_REOPEN_BACKOFF
6726 );
6727 assert_eq!(
6728 ordered_stream_reopen_backoff(1),
6729 ORDERED_STREAM_REOPEN_BACKOFF * 2
6730 );
6731 assert_eq!(
6732 ordered_stream_reopen_backoff(5),
6733 ORDERED_STREAM_REOPEN_BACKOFF_CAP
6734 );
6735 assert_eq!(
6736 ordered_stream_reopen_backoff(u32::MAX),
6737 ORDERED_STREAM_REOPEN_BACKOFF_CAP
6738 );
6739 }
6740
6741 #[tokio::test]
6742 async fn stream_cancel_closes_ordered_worker_without_connection_cancel() -> Result<(), BoxError>
6743 {
6744 const ALPN: &[u8] = b"/zakura/testkit/stream-cancel/0";
6745
6746 let _guard = zakura_test::init();
6747 let server = LocalEndpointFactory::new().endpoint(50).await?;
6748 let (conn_tx, mut conn_rx) = mpsc::channel(1);
6749 let (stream_tx, mut stream_rx) = mpsc::channel(2);
6750 let router = Router::builder(server)
6751 .accept(
6752 ALPN,
6753 CaptureConnection {
6754 connection_tx: conn_tx,
6755 stream_tx,
6756 },
6757 )
6758 .spawn();
6759 let client = LocalEndpointFactory::new().endpoint(51).await?;
6760 let server_addr = router.endpoint().node_addr().initialized().await;
6761 client.add_node_addr(server_addr.clone())?;
6762
6763 let client_conn = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
6764 .await
6765 .expect("client connects to loopback capture endpoint")?;
6766 let server_conn = timeout(Duration::from_secs(10), conn_rx.recv())
6767 .await
6768 .expect("server connection is captured")
6769 .expect("capture handler sends the accepted connection");
6770 drop(server_conn);
6771 let (mut client_send, _client_recv) =
6772 timeout(Duration::from_secs(1), client_conn.open_bi())
6773 .await
6774 .expect("client opens the worker stream")?;
6775 let test_frame = Frame {
6776 message_type: 1,
6777 flags: 0,
6778 payload: Vec::new(),
6779 }
6780 .encode(LOCAL_MAX_CONTROL_FRAME_BYTES)?;
6781 timeout(Duration::from_secs(1), client_send.write_all(&test_frame))
6782 .await
6783 .expect("client writes the worker stream frame")?;
6784 let (server_send, server_recv) = timeout(Duration::from_secs(1), stream_rx.recv())
6785 .await
6786 .expect("server accepts the worker stream")
6787 .expect("capture handler sends the worker stream");
6788
6789 let mut limits = test_connection_limits();
6790 limits.idle_timeout = Duration::from_millis(50);
6791 let stream_kind = DISCOVERY_STREAM_KIND;
6792 let connection_token = CancellationToken::new();
6793 let stream_token = connection_token.child_token();
6794 let (freshness_tx, _freshness_rx) = watch::channel(Instant::now());
6795 let permit = Arc::new(Semaphore::new(1))
6796 .try_acquire_owned()
6797 .expect("test semaphore starts with one permit");
6798 let context = StreamWorkerContext {
6799 trace: ZakuraTrace::noop(),
6800 conn: ZakuraConnTrace::without_peer(1),
6801 peer_id: test_peer(55),
6802 stream_id: 1,
6803 _permit: permit,
6804 limits,
6805 message_bucket: Arc::new(std::sync::Mutex::new(TokenBucket::new(128))),
6806 connection_token: connection_token.clone(),
6807 stream_token: stream_token.clone(),
6808 close_cause: CloseCause::new(),
6809 freshness_tx,
6810 };
6811 let prelude = StreamPrelude {
6812 magic: STREAM_PRELUDE_MAGIC,
6813 stream_kind,
6814 stream_version: ZAKURA_STREAM_VERSION_1,
6815 request_id: None,
6816 max_frame_bytes: app_frame_cap_for_stream_kind(&limits, stream_kind),
6817 };
6818
6819 let stream = Stream {
6820 kind: stream_kind,
6821 version: ZAKURA_STREAM_VERSION_1,
6822 frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES,
6823 capability: ZAKURA_CAP_DISCOVERY,
6824 mode: StreamMode::Ordered,
6825 };
6826 let mut workers = JoinSet::new();
6827 let (ordered_stream_exit_tx, mut ordered_stream_exit_rx) = mpsc::unbounded_channel();
6828 let admitted = spawn_persistent_stream_worker(
6829 &mut workers,
6830 server_send,
6831 server_recv,
6832 stream,
6833 prelude,
6834 context,
6835 1,
6836 true,
6837 ordered_stream_exit_tx,
6838 );
6839
6840 admitted.cancel_token.cancel();
6841 let exited = timeout(Duration::from_secs(1), ordered_stream_exit_rx.recv())
6844 .await
6845 .expect("stream cancellation reports worker exit")
6846 .expect("exit channel stays open");
6847 assert_eq!(exited.stream, stream);
6848 assert_eq!(exited.session_id, 1);
6849 assert!(exited.opened_locally);
6850 timeout(Duration::from_secs(1), workers.join_next())
6851 .await
6852 .expect("stream cancellation stops the worker promptly")
6853 .expect("worker exit is observed")?;
6854 tokio::time::sleep(Duration::from_millis(100)).await;
6855 assert!(
6856 !connection_token.is_cancelled(),
6857 "stream-local cancellation must not cancel the shared connection"
6858 );
6859
6860 let (mut sibling_send, _sibling_recv) =
6861 timeout(Duration::from_secs(1), client_conn.open_bi())
6862 .await
6863 .expect("connection survives to open a sibling stream")?;
6864 timeout(Duration::from_secs(1), sibling_send.write_all(&test_frame))
6865 .await
6866 .expect("client writes the sibling stream frame")?;
6867 timeout(Duration::from_secs(1), stream_rx.recv())
6868 .await
6869 .expect("server accepts the sibling stream")
6870 .expect("capture handler sends the sibling stream");
6871
6872 client_conn.close(0u32.into(), b"done");
6873 client.close().await;
6874 router.shutdown().await?;
6875 Ok(())
6876 }
6877
6878 #[tokio::test]
6903 async fn quiet_ordered_stream_does_not_cancel_connection_with_active_sibling(
6904 ) -> Result<(), BoxError> {
6905 const ALPN: &[u8] = b"/zakura/testkit/quiet-ordered-stream/0";
6906
6907 let _guard = zakura_test::init();
6908 let server = LocalEndpointFactory::new().endpoint(80).await?;
6909 let (conn_tx, _conn_rx) = mpsc::channel(1);
6910 let (stream_tx, mut stream_rx) = mpsc::channel(2);
6911 let router = Router::builder(server)
6912 .accept(
6913 ALPN,
6914 CaptureConnection {
6915 connection_tx: conn_tx,
6916 stream_tx,
6917 },
6918 )
6919 .spawn();
6920 let client = LocalEndpointFactory::new().endpoint(81).await?;
6921 let server_addr = router.endpoint().node_addr().initialized().await;
6922 client.add_node_addr(server_addr.clone())?;
6923
6924 let client_conn = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
6925 .await
6926 .expect("client connects to loopback capture endpoint")?;
6927
6928 let frame = Frame {
6929 message_type: 1,
6930 flags: 0,
6931 payload: Vec::new(),
6932 }
6933 .encode(LOCAL_MAX_CONTROL_FRAME_BYTES)?;
6934
6935 let (mut active_send, _active_recv) =
6941 timeout(Duration::from_secs(1), client_conn.open_bi())
6942 .await
6943 .expect("client opens the active stream")?;
6944 timeout(Duration::from_secs(1), active_send.write_all(&frame))
6945 .await
6946 .expect("client writes the active stream's first frame")?;
6947 let (active_server_send, active_server_recv) =
6948 timeout(Duration::from_secs(1), stream_rx.recv())
6949 .await
6950 .expect("server accepts the active stream")
6951 .expect("capture handler forwards the active stream");
6952
6953 let (mut quiet_send, _quiet_recv) = timeout(Duration::from_secs(1), client_conn.open_bi())
6954 .await
6955 .expect("client opens the quiet stream")?;
6956 timeout(Duration::from_secs(1), quiet_send.write_all(&frame))
6957 .await
6958 .expect("client writes the quiet stream's only frame")?;
6959 let (quiet_server_send, quiet_server_recv) =
6960 timeout(Duration::from_secs(1), stream_rx.recv())
6961 .await
6962 .expect("server accepts the quiet stream")
6963 .expect("capture handler forwards the quiet stream");
6964
6965 let idle_timeout = Duration::from_millis(250);
6969 let mut limits = test_connection_limits();
6970 limits.idle_timeout = idle_timeout;
6971 let stream_kind = LEGACY_GOSSIP_STREAM_KIND;
6972 let connection_token = CancellationToken::new();
6973 let (freshness_tx, freshness_rx) = watch::channel(Instant::now());
6974
6975 let spawn_ordered_worker = |send: SendStream, recv: RecvStream, stream_id: u64| {
6976 let permit = Arc::new(Semaphore::new(1))
6977 .try_acquire_owned()
6978 .expect("test semaphore starts with one permit");
6979 let context = StreamWorkerContext {
6980 trace: ZakuraTrace::noop(),
6981 conn: ZakuraConnTrace::without_peer(1),
6982 peer_id: test_peer(80),
6983 stream_id,
6984 _permit: permit,
6985 limits,
6986 message_bucket: Arc::new(std::sync::Mutex::new(TokenBucket::new(
6987 limits.message_rate_per_second,
6988 ))),
6989 connection_token: connection_token.clone(),
6990 stream_token: connection_token.child_token(),
6991 close_cause: CloseCause::new(),
6992 freshness_tx: freshness_tx.clone(),
6993 };
6994 let prelude = StreamPrelude {
6995 magic: STREAM_PRELUDE_MAGIC,
6996 stream_kind,
6997 stream_version: ZAKURA_STREAM_VERSION_1,
6998 request_id: None,
6999 max_frame_bytes: app_frame_cap_for_stream_kind(&limits, stream_kind),
7000 };
7001 let (inbound_tx, mut inbound_rx) = mpsc::channel(16);
7005 let (_outbound_tx, outbound_rx) = mpsc::channel(1);
7006 tokio::spawn(async move { while inbound_rx.recv().await.is_some() {} });
7007 tokio::spawn(persistent_stream_worker(
7008 send,
7009 recv,
7010 prelude,
7011 context,
7012 inbound_tx,
7013 outbound_rx,
7014 16,
7015 ))
7016 };
7017
7018 let _active_worker = spawn_ordered_worker(active_server_send, active_server_recv, 1);
7019 let _quiet_worker = spawn_ordered_worker(quiet_server_send, quiet_server_recv, 2);
7020
7021 let serve = tokio::spawn({
7026 let connection_token = connection_token.clone();
7027 async move {
7028 tokio::select! {
7029 biased;
7030 _ = connection_token.cancelled() => {}
7031 _ = freshness_reaper(freshness_rx, idle_timeout) => connection_token.cancel(),
7032 }
7033 }
7034 });
7035
7036 let keepalive_frame = frame.clone();
7039 let keepalive = tokio::spawn(async move {
7040 for _ in 0..16 {
7041 tokio::time::sleep(Duration::from_millis(80)).await;
7042 if active_send.write_all(&keepalive_frame).await.is_err() {
7043 break;
7044 }
7045 }
7046 });
7047
7048 tokio::time::sleep(idle_timeout * 4).await;
7051 assert!(
7052 !connection_token.is_cancelled(),
7053 "a quiet ordered stream must not cancel a connection whose sibling stream is \
7054 actively transferring; the shared connection_token was cancelled within {:?}",
7055 idle_timeout * 4,
7056 );
7057
7058 connection_token.cancel();
7059 keepalive.abort();
7060 let _ = serve.await;
7061 client_conn.close(0u32.into(), b"done");
7062 client.close().await;
7063 router.shutdown().await?;
7064 Ok(())
7065 }
7066
7067 #[tokio::test]
7080 async fn write_ordered_frame_rejects_payload_over_message_cap() -> Result<(), BoxError> {
7081 const ALPN: &[u8] = b"/zakura/testkit/ordered-message-cap/0";
7082
7083 let _guard = zakura_test::init();
7084 let server = LocalEndpointFactory::new().endpoint(74).await?;
7085 let (conn_tx, _conn_rx) = mpsc::channel(1);
7086 let (stream_tx, _stream_rx) = mpsc::channel(2);
7087 let router = Router::builder(server)
7088 .accept(
7089 ALPN,
7090 CaptureConnection {
7091 connection_tx: conn_tx,
7092 stream_tx,
7093 },
7094 )
7095 .spawn();
7096 let client = LocalEndpointFactory::new().endpoint(75).await?;
7097 let server_addr = router.endpoint().node_addr().initialized().await;
7098 client.add_node_addr(server_addr.clone())?;
7099
7100 let client_conn = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
7101 .await
7102 .expect("client connects to the ordered-message-cap endpoint")?;
7103 let (mut send, _recv) = timeout(Duration::from_secs(1), client_conn.open_bi())
7104 .await
7105 .expect("client opens an ordered stream")?;
7106
7107 let mut limits = test_connection_limits();
7110 limits.max_message_bytes = 256;
7111 let stream_kind = DISCOVERY_STREAM_KIND;
7112
7113 let oversized = Frame {
7116 message_type: 1,
7117 flags: 0,
7118 payload: vec![0xab; 4096],
7119 };
7120 assert!(
7121 oversized.payload.len() <= app_frame_cap_for_stream_kind(&limits, stream_kind) as usize,
7122 "test payload must fit the frame cap so only the message cap can reject it"
7123 );
7124
7125 let result = write_ordered_frame(&mut send, oversized, limits, stream_kind).await;
7126 assert!(
7127 result.is_err(),
7128 "write_ordered_frame must reject a payload over the negotiated max_message_bytes \
7129 before encoding/writing it, mirroring write_response_frame; got {result:?}"
7130 );
7131
7132 let within_cap = Frame {
7134 message_type: 1,
7135 flags: 0,
7136 payload: vec![0xcd; 128],
7137 };
7138 write_ordered_frame(&mut send, within_cap, limits, stream_kind)
7139 .await
7140 .expect("a frame within the negotiated message cap must still be written");
7141
7142 client_conn.close(0u32.into(), b"done");
7143 client.close().await;
7144 router.shutdown().await?;
7145 Ok(())
7146 }
7147
7148 #[tokio::test]
7156 async fn read_control_payload_enforces_hard_cap_before_reading_body() -> Result<(), BoxError> {
7157 const ALPN: &[u8] = b"/zakura/testkit/control-hard-cap/0";
7158
7159 let _guard = zakura_test::init();
7160 let server = LocalEndpointFactory::new().endpoint(70).await?;
7161 let (conn_tx, _conn_rx) = mpsc::channel(1);
7162 let (stream_tx, mut stream_rx) = mpsc::channel(2);
7163 let router = Router::builder(server)
7164 .accept(
7165 ALPN,
7166 CaptureConnection {
7167 connection_tx: conn_tx,
7168 stream_tx,
7169 },
7170 )
7171 .spawn();
7172 let client = LocalEndpointFactory::new().endpoint(71).await?;
7173 let server_addr = router.endpoint().node_addr().initialized().await;
7174 client.add_node_addr(server_addr.clone())?;
7175
7176 let client_conn = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
7177 .await
7178 .expect("client connects to the control-hard-cap endpoint")?;
7179
7180 let oversized_len =
7184 u32::try_from(MAX_CONTROL_PAYLOAD_BYTES + 1).expect("control hard cap + 1 fits in u32");
7185 let (mut over_send, _over_recv) = timeout(Duration::from_secs(1), client_conn.open_bi())
7186 .await
7187 .expect("client opens the oversized control stream")?;
7188 timeout(
7189 Duration::from_secs(1),
7190 over_send.write_all(&oversized_len.to_le_bytes()),
7191 )
7192 .await
7193 .expect("client writes the oversized control length")?;
7194 let _ = over_send.finish();
7195
7196 let (_over_server_send, mut over_server_recv) =
7197 timeout(Duration::from_secs(1), stream_rx.recv())
7198 .await
7199 .expect("server accepts the oversized control stream")
7200 .expect("capture handler forwards the oversized control stream");
7201
7202 let oversized = read_control_payload(
7205 &mut over_server_recv,
7206 LOCAL_MAX_CONTROL_FRAME_BYTES,
7207 Duration::from_secs(2),
7208 )
7209 .await;
7210 assert!(
7211 matches!(oversized, Err(ZakuraHandlerError::Oversize)),
7212 "a control length over the 16 KiB hard cap must be rejected as Oversize \
7213 on the length prefix, before the body is read; got {oversized:?}"
7214 );
7215
7216 let valid_body = vec![0xa5u8; 128];
7219 let valid_len = u32::try_from(valid_body.len()).expect("128 fits in u32");
7220 let (mut ok_send, _ok_recv) = timeout(Duration::from_secs(1), client_conn.open_bi())
7221 .await
7222 .expect("client opens the in-cap control stream")?;
7223 timeout(
7224 Duration::from_secs(1),
7225 ok_send.write_all(&valid_len.to_le_bytes()),
7226 )
7227 .await
7228 .expect("client writes the in-cap control length")?;
7229 timeout(Duration::from_secs(1), ok_send.write_all(&valid_body))
7230 .await
7231 .expect("client writes the in-cap control body")?;
7232 let _ = ok_send.finish();
7233
7234 let (_ok_server_send, mut ok_server_recv) =
7235 timeout(Duration::from_secs(1), stream_rx.recv())
7236 .await
7237 .expect("server accepts the in-cap control stream")
7238 .expect("capture handler forwards the in-cap control stream");
7239 let read_back = read_control_payload(
7240 &mut ok_server_recv,
7241 LOCAL_MAX_CONTROL_FRAME_BYTES,
7242 Duration::from_secs(2),
7243 )
7244 .await
7245 .expect("a control payload within the hard cap is read");
7246 assert_eq!(
7247 read_back, valid_body,
7248 "an in-cap control payload must round-trip unchanged"
7249 );
7250
7251 client_conn.close(0u32.into(), b"done");
7252 client.close().await;
7253 router.shutdown().await?;
7254 Ok(())
7255 }
7256
7257 #[tokio::test]
7267 async fn inbound_frame_cap_rejects_over_message_frame_before_reading_payload(
7268 ) -> Result<(), BoxError> {
7269 const ALPN: &[u8] = b"/zakura/testkit/late-message-cap/0";
7270 const MAX_FRAME_BYTES: u32 = 64 * 1024;
7273 const MAX_MESSAGE_BYTES: u32 = 1024;
7274 const OVER_MESSAGE_PAYLOAD_LEN: u32 = 2048;
7277 let stream_kind = LEGACY_GOSSIP_STREAM_KIND;
7278
7279 let limits = ZakuraConnectionLimits {
7280 max_frame_bytes: MAX_FRAME_BYTES,
7281 max_message_bytes: MAX_MESSAGE_BYTES,
7282 ..test_connection_limits()
7283 };
7284 let inbound_cap = inbound_frame_cap_for_stream_kind(&limits, stream_kind);
7287 let raw_cap = app_frame_cap_for_stream_kind(&limits, stream_kind);
7288 assert!(
7289 inbound_cap < raw_cap,
7290 "the inbound cap must be tighter than the raw frame cap when the caps diverge \
7291 (inbound_cap={inbound_cap}, raw_cap={raw_cap})"
7292 );
7293
7294 let _guard = zakura_test::init();
7295 let server = LocalEndpointFactory::new().endpoint(72).await?;
7296 let (conn_tx, _conn_rx) = mpsc::channel(2);
7297 let (stream_tx, mut stream_rx) = mpsc::channel(4);
7298 let router = Router::builder(server)
7299 .accept(
7300 ALPN,
7301 CaptureConnection {
7302 connection_tx: conn_tx,
7303 stream_tx,
7304 },
7305 )
7306 .spawn();
7307 let client = LocalEndpointFactory::new().endpoint(73).await?;
7308 let server_addr = router.endpoint().node_addr().initialized().await;
7309 client.add_node_addr(server_addr.clone())?;
7310
7311 let frame_header = |payload_len: u32| -> Vec<u8> {
7313 let mut header = Vec::with_capacity(FRAME_HEADER_BYTES);
7314 header.extend_from_slice(&0u16.to_le_bytes());
7315 header.extend_from_slice(&0u16.to_le_bytes());
7316 header.extend_from_slice(&payload_len.to_le_bytes());
7317 header
7318 };
7319
7320 let conn_a = timeout(
7324 Duration::from_secs(10),
7325 client.connect(server_addr.clone(), ALPN),
7326 )
7327 .await
7328 .expect("client connects for the oversized streams")?;
7329
7330 let (mut over_send, _over_recv) =
7333 timeout(Duration::from_secs(1), conn_a.open_bi())
7334 .await
7335 .expect("client opens the oversized inbound-cap stream")?;
7336 timeout(
7337 Duration::from_secs(1),
7338 over_send.write_all(&frame_header(OVER_MESSAGE_PAYLOAD_LEN)),
7339 )
7340 .await
7341 .expect("client writes the oversized frame header")?;
7342 let _ = over_send.finish();
7343 let (_s1_send, mut s1_recv) = timeout(Duration::from_secs(1), stream_rx.recv())
7344 .await
7345 .expect("server accepts the oversized inbound-cap stream")
7346 .expect("capture handler forwards the oversized inbound-cap stream");
7347 let rejected = read_frame(
7348 &mut s1_recv,
7349 inbound_cap,
7350 Duration::from_secs(2),
7351 Some(Duration::from_secs(2)),
7352 )
7353 .await;
7354 assert!(
7355 matches!(rejected, Err(ZakuraHandlerError::OversizeFrame { .. })),
7356 "a frame whose payload exceeds max_message_bytes must be rejected as Oversize \
7357 on the header alone with the inbound (message-limited) cap, before the payload \
7358 is allocated and read; got {rejected:?}"
7359 );
7360
7361 let (mut raw_send, _raw_recv) = timeout(Duration::from_secs(1), conn_a.open_bi())
7367 .await
7368 .expect("client opens the oversized raw-cap stream")?;
7369 timeout(
7370 Duration::from_secs(1),
7371 raw_send.write_all(&frame_header(OVER_MESSAGE_PAYLOAD_LEN)),
7372 )
7373 .await
7374 .expect("client writes the oversized frame header again")?;
7375 let _ = raw_send.finish();
7376 let (_s2_send, mut s2_recv) = timeout(Duration::from_secs(1), stream_rx.recv())
7377 .await
7378 .expect("server accepts the oversized raw-cap stream")
7379 .expect("capture handler forwards the oversized raw-cap stream");
7380 let allocated = read_frame(
7381 &mut s2_recv,
7382 raw_cap,
7383 Duration::from_secs(2),
7384 Some(Duration::from_secs(2)),
7385 )
7386 .await;
7387 assert!(
7388 allocated.is_err()
7389 && !matches!(
7390 allocated,
7391 Err(ZakuraHandlerError::Oversize | ZakuraHandlerError::OversizeFrame { .. })
7392 ),
7393 "with the raw frame cap the same oversized frame passes the size check and \
7394 read_frame proceeds to allocate/read the payload (it is not rejected as \
7395 Oversize), proving the message cap is enforced too late; got {allocated:?}"
7396 );
7397
7398 let conn_b = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
7401 .await
7402 .expect("client connects for the in-cap stream")?;
7403 let valid_payload = vec![0x5au8; (MAX_MESSAGE_BYTES / 2) as usize];
7404 let mut valid_bytes =
7405 frame_header(u32::try_from(valid_payload.len()).expect("in-cap payload len fits u32"));
7406 valid_bytes.extend_from_slice(&valid_payload);
7407 let (mut ok_send, _ok_recv) = timeout(Duration::from_secs(1), conn_b.open_bi())
7408 .await
7409 .expect("client opens the in-cap stream")?;
7410 timeout(Duration::from_secs(1), ok_send.write_all(&valid_bytes))
7411 .await
7412 .expect("client writes the in-cap frame")?;
7413 let _ = ok_send.finish();
7414 let (_s3_send, mut s3_recv) = timeout(Duration::from_secs(1), stream_rx.recv())
7415 .await
7416 .expect("server accepts the in-cap stream")
7417 .expect("capture handler forwards the in-cap stream");
7418 let frame = read_frame(
7419 &mut s3_recv,
7420 inbound_cap,
7421 Duration::from_secs(2),
7422 Some(Duration::from_secs(2)),
7423 )
7424 .await
7425 .expect("a frame within the message cap is read with the inbound cap");
7426 assert_eq!(
7427 frame.payload, valid_payload,
7428 "an in-cap frame payload must round-trip unchanged"
7429 );
7430
7431 conn_a.close(0u32.into(), b"done");
7432 conn_b.close(0u32.into(), b"done");
7433 client.close().await;
7434 router.shutdown().await?;
7435 Ok(())
7436 }
7437
7438 #[tokio::test]
7439 async fn invalid_prelude_stream_churn_charges_open_rate_token() -> Result<(), BoxError> {
7440 const ALPN: &[u8] = b"/zakura/testkit/open-rate-churn/0";
7449
7450 let _guard = zakura_test::init();
7451 let server = LocalEndpointFactory::new().endpoint(90).await?;
7452 let (conn_tx, _conn_rx) = mpsc::channel(1);
7453 let (stream_tx, mut stream_rx) = mpsc::channel(2);
7454 let router = Router::builder(server)
7455 .accept(
7456 ALPN,
7457 CaptureConnection {
7458 connection_tx: conn_tx,
7459 stream_tx,
7460 },
7461 )
7462 .spawn();
7463 let client = LocalEndpointFactory::new().endpoint(91).await?;
7464 let server_addr = router.endpoint().node_addr().initialized().await;
7465 client.add_node_addr(server_addr.clone())?;
7466 let client_conn = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN))
7467 .await
7468 .expect("client connects to the open-rate-churn endpoint")?;
7469
7470 let prelude = StreamPrelude {
7474 magic: STREAM_PRELUDE_MAGIC,
7475 stream_kind: 9,
7476 stream_version: 1,
7477 request_id: None,
7478 max_frame_bytes: 1024,
7479 };
7480 let (mut client_send, _client_recv) =
7481 timeout(Duration::from_secs(1), client_conn.open_bi())
7482 .await
7483 .expect("client opens the unknown-kind stream")?;
7484 timeout(
7485 Duration::from_secs(1),
7486 client_send.write_all(&prelude.encode()?),
7487 )
7488 .await
7489 .expect("client writes the unknown-kind prelude")?;
7490 let _ = client_send.finish();
7491
7492 let (server_send, server_recv) = timeout(Duration::from_secs(1), stream_rx.recv())
7493 .await
7494 .expect("server accepts the unknown-kind stream")
7495 .expect("capture handler forwards the unknown-kind stream");
7496
7497 let supervisor = ZakuraSupervisorHandle::new(16);
7498 let handler = ZakuraProtocolHandler::new(
7499 supervisor,
7500 Network::Mainnet,
7501 ZakuraHandshakeConfig::for_network(&Network::Mainnet),
7502 ZakuraLocalLimits::from_config(&Config::default()),
7503 );
7504
7505 let peer_id = test_peer(9);
7506 let stream_sem = Arc::new(Semaphore::new(16));
7507 let mut open_limiter = TokenBucket::new(4);
7510 let mut message_buckets = MessageRateBuckets::new();
7511 let mut workers = JoinSet::new();
7512 let connection_token = CancellationToken::new();
7513 let (freshness_tx, _freshness_rx) = watch::channel(Instant::now());
7514
7515 let mut admission = StreamAdmission {
7516 trace: handler.trace.clone(),
7517 conn: ZakuraConnTrace::placeholder(),
7518 peer_id: &peer_id,
7519 stream_sem: &stream_sem,
7520 open_limiter: &mut open_limiter,
7521 message_buckets: &mut message_buckets,
7522 workers: &mut workers,
7523 limits: test_connection_limits(),
7524 accepted_capabilities: 0,
7525 connection_token: connection_token.clone(),
7526 close_cause: CloseCause::new(),
7527 freshness_tx,
7528 };
7529 let (ordered_stream_exit_tx, _ordered_stream_exit_rx) = mpsc::unbounded_channel();
7530 let admitted = handler
7531 .admit_bi_stream(
7532 server_send,
7533 server_recv,
7534 &mut admission,
7535 16,
7536 ordered_stream_exit_tx,
7537 )
7538 .await;
7539
7540 assert!(
7541 admitted.is_none(),
7542 "an unknown-kind stream must be rejected, not admitted"
7543 );
7544 assert!(
7545 !connection_token.is_cancelled(),
7546 "an unknown-kind stream is reset stream-only and must keep the connection alive"
7547 );
7548 assert_eq!(
7549 admission.open_limiter.tokens, 3,
7550 "the protocol-invalid stream open must spend exactly one open-rate token \
7551 (capacity 4 -> 3); before the fix the unknown-kind reject returned before \
7552 reaching the limiter, leaving the bucket full at 4"
7553 );
7554
7555 client.close().await;
7556 router.shutdown().await?;
7557 Ok(())
7558 }
7559
7560 #[test]
7561 fn local_application_frame_cap_admits_default_header_sync_response() {
7562 let limits = ZakuraLocalLimits::from_config(&Config::default());
7563 let default_header_sync_frame_bytes =
7564 u32::try_from(MAX_HS_MESSAGE_BYTES.saturating_add(FRAME_HEADER_BYTES))
7565 .expect("header-sync frame cap fits in u32");
7566
7567 assert!(limits.max_frame_bytes >= default_header_sync_frame_bytes);
7568 assert!(limits.initial_limits().max_frame_bytes >= default_header_sync_frame_bytes);
7569 }
7570
7571 #[test]
7572 fn ordered_stream_stopped_write_is_stream_local() {
7573 let error: BoxError = iroh::endpoint::WriteError::Stopped(VarInt::from_u32(0)).into();
7574
7575 assert!(ordered_stream_write_was_stopped(&error));
7576 }
7577
7578 #[test]
7579 fn stream_specific_application_frame_caps_keep_gossip_and_discovery_tight() {
7580 let limits = ZakuraLocalLimits::from_config(&Config::default());
7581 let negotiated = limits.clamp(&ZakuraAcceptedLimits {
7582 max_frame_bytes: u32::MAX,
7583 max_message_bytes: u32::MAX,
7584 max_open_streams: u16::MAX,
7585 max_inbound_queue_depth: u16::MAX,
7586 idle_timeout_millis: u32::MAX,
7587 });
7588 let header_sync_frame_bytes =
7589 u32::try_from(MAX_HS_MESSAGE_BYTES.saturating_add(FRAME_HEADER_BYTES))
7590 .expect("header-sync frame cap fits in u32");
7591
7592 assert_eq!(
7593 app_frame_cap_for_stream_kind(&negotiated, LEGACY_GOSSIP_STREAM_KIND),
7594 LOCAL_MAX_CONTROL_FRAME_BYTES
7595 );
7596 assert_eq!(
7597 app_frame_cap_for_stream_kind(&negotiated, LEGACY_REQUEST_STREAM_KIND),
7598 LOCAL_MAX_CONTROL_FRAME_BYTES
7599 );
7600 assert_eq!(
7601 app_frame_cap_for_stream_kind(&negotiated, DISCOVERY_STREAM_KIND),
7602 LOCAL_MAX_CONTROL_FRAME_BYTES
7603 );
7604 assert_eq!(
7605 app_frame_cap_for_stream_kind(&negotiated, HEADER_SYNC_STREAM_KIND),
7606 header_sync_frame_bytes
7607 );
7608
7609 let over_tight_cap = usize::try_from(LOCAL_MAX_CONTROL_FRAME_BYTES).unwrap() + 1;
7610 let header_sync_cap = usize::try_from(header_sync_frame_bytes).unwrap();
7611 let gossip_frame = Frame {
7612 message_type: 1,
7613 flags: 0,
7614 payload: vec![0; over_tight_cap.saturating_sub(FRAME_HEADER_BYTES)],
7615 };
7616 let header_sync_frame = Frame {
7617 message_type: 1,
7618 flags: 0,
7619 payload: vec![0; header_sync_cap.saturating_sub(FRAME_HEADER_BYTES)],
7620 };
7621
7622 assert!(
7623 gossip_frame
7624 .encode(app_frame_cap_for_stream_kind(
7625 &negotiated,
7626 LEGACY_GOSSIP_STREAM_KIND
7627 ))
7628 .is_err(),
7629 "gossip frames over the tight stream cap must be rejected"
7630 );
7631 assert!(
7632 gossip_frame
7633 .encode(app_frame_cap_for_stream_kind(
7634 &negotiated,
7635 DISCOVERY_STREAM_KIND
7636 ))
7637 .is_err(),
7638 "discovery frames over the tight stream cap must be rejected"
7639 );
7640 assert!(
7641 gossip_frame
7642 .encode(app_frame_cap_for_stream_kind(
7643 &negotiated,
7644 LEGACY_REQUEST_STREAM_KIND
7645 ))
7646 .is_err(),
7647 "legacy request frames over the tight stream cap must be rejected"
7648 );
7649 assert!(
7650 header_sync_frame
7651 .encode(app_frame_cap_for_stream_kind(
7652 &negotiated,
7653 HEADER_SYNC_STREAM_KIND
7654 ))
7655 .is_ok(),
7656 "header-sync frames up to MAX_HS_MESSAGE_BYTES must still be accepted"
7657 );
7658 }
7659
7660 #[test]
7661 fn token_bucket_rejects_churn_until_refill() {
7662 let clock = crate::zakura::testkit::TestClock::new();
7663 let mut bucket = TokenBucket::with_clock(2, clock.clone());
7664
7665 assert!(bucket.try_take());
7666 assert!(bucket.try_take());
7667 assert!(!bucket.try_take());
7668 clock.advance(Duration::from_millis(500));
7669 assert!(bucket.try_take());
7670 assert!(!bucket.try_take());
7671 clock.advance(Duration::from_millis(500));
7672 assert!(bucket.try_take());
7673 }
7674
7675 #[test]
7676 fn supported_stream_accepts_registered_kinds_at_declared_version_only() {
7677 let registry = ServiceRegistry::new(vec![Arc::new(DeclaredStreamService {
7678 streams: vec![
7679 Stream {
7680 kind: LEGACY_GOSSIP_STREAM_KIND,
7681 version: ZAKURA_STREAM_VERSION_1,
7682 frame_cap: 1024,
7683 capability: ZAKURA_CAP_LEGACY_GOSSIP,
7684 mode: StreamMode::Ordered,
7685 },
7686 Stream {
7687 kind: LEGACY_REQUEST_STREAM_KIND,
7688 version: ZAKURA_STREAM_VERSION_1,
7689 frame_cap: 1024,
7690 capability: ZAKURA_CAP_LEGACY_GOSSIP,
7691 mode: StreamMode::RequestResponse,
7692 },
7693 Stream {
7694 kind: DISCOVERY_STREAM_KIND,
7695 version: ZAKURA_STREAM_VERSION_1,
7696 frame_cap: 1024,
7697 capability: ZAKURA_CAP_DISCOVERY,
7698 mode: StreamMode::Ordered,
7699 },
7700 Stream {
7701 kind: HEADER_SYNC_STREAM_KIND,
7702 version: ZAKURA_STREAM_VERSION_7,
7703 frame_cap: 1024,
7704 capability: ZAKURA_CAP_HEADER_SYNC,
7705 mode: StreamMode::Ordered,
7706 },
7707 Stream {
7708 kind: ZAKURA_STREAM_BLOCK_SYNC,
7709 version: ZAKURA_STREAM_VERSION_1,
7710 frame_cap: MAX_BS_FRAME_BYTES,
7711 capability: crate::zakura::ZAKURA_CAP_BLOCK_SYNC,
7712 mode: StreamMode::Ordered,
7713 },
7714 ],
7715 }) as Arc<dyn Service>])
7716 .expect("test registry declares unique stream kinds");
7717
7718 for (kind, version) in [
7719 (LEGACY_GOSSIP_STREAM_KIND, ZAKURA_STREAM_VERSION_1),
7720 (LEGACY_REQUEST_STREAM_KIND, ZAKURA_STREAM_VERSION_1),
7721 (DISCOVERY_STREAM_KIND, ZAKURA_STREAM_VERSION_1),
7722 (HEADER_SYNC_STREAM_KIND, ZAKURA_STREAM_VERSION_7),
7723 (ZAKURA_STREAM_BLOCK_SYNC, ZAKURA_STREAM_VERSION_1),
7724 ] {
7725 assert!(
7726 is_supported_stream(®istry, kind, version),
7727 "registered kind {kind} at declared version {version} must be supported"
7728 );
7729 assert!(
7730 !is_supported_stream(®istry, kind, 0),
7731 "registered kind {kind} at version 0 must be rejected"
7732 );
7733 assert!(
7734 !is_supported_stream(®istry, kind, version.saturating_add(1)),
7735 "registered kind {kind} at an unsupported version must be rejected"
7736 );
7737 }
7738
7739 assert!(
7740 !is_supported_stream(®istry, HEADER_SYNC_STREAM_KIND, ZAKURA_STREAM_VERSION_1),
7741 "header-sync v1 is rejected because native header sync uses exact version matching"
7742 );
7743 assert!(
7744 !is_supported_stream(®istry, HEADER_SYNC_STREAM_KIND, 5),
7745 "header-sync v5 is rejected after the expanded Ironwood root-record wire break"
7746 );
7747
7748 assert_eq!(stream_kind_label(2), "gossip");
7749 assert_eq!(stream_kind_label(3), "legacy_request");
7750 assert_eq!(stream_kind_label(4), "discovery");
7751 assert_eq!(stream_kind_label(5), "header_sync");
7752 assert_eq!(stream_kind_label(6), "block_sync");
7753
7754 for kind in [0u16, 1, 7, 255, u16::MAX] {
7755 assert!(
7756 !is_supported_stream(®istry, kind, ZAKURA_STREAM_VERSION_1),
7757 "unknown kind {kind} must be rejected even at version 1"
7758 );
7759 }
7760
7761 for kind in [7u16, 255, u16::MAX] {
7762 assert_eq!(stream_kind_label(kind), "unknown");
7763 }
7764 }
7765
7766 #[test]
7767 fn legacy_transport_response_validator_accepts_codec_frames() -> Result<(), BoxError> {
7768 let block = Arc::new(Block::zcash_deserialize(
7769 BLOCK_TESTNET_141042_BYTES.as_slice(),
7770 )?);
7771 let header = block::CountedHeader {
7772 header: block.header.clone(),
7773 };
7774 let tx_id = legacy_tx_id(7);
7775
7776 assert_codec_frames_validate_at_transport(
7777 LegacyRequestFrame::BlocksByHash(vec![block.hash()]),
7778 LegacyRequestKind::Blocks,
7779 Response::Blocks(vec![InventoryResponse::Available((block, None))]),
7780 )?;
7781 assert_codec_frames_validate_at_transport(
7782 LegacyRequestFrame::TransactionsById(vec![tx_id]),
7783 LegacyRequestKind::Transactions,
7784 Response::Transactions(vec![InventoryResponse::Missing(tx_id)]),
7785 )?;
7786 assert_codec_frames_validate_at_transport(
7787 LegacyRequestFrame::FindBlocks {
7788 known_blocks: vec![block_hash(1)],
7789 stop: None,
7790 },
7791 LegacyRequestKind::FindBlocks,
7792 Response::BlockHashes(vec![block_hash(2)]),
7793 )?;
7794 assert_codec_frames_validate_at_transport(
7795 LegacyRequestFrame::FindHeaders {
7796 known_blocks: vec![block_hash(1)],
7797 stop: None,
7798 },
7799 LegacyRequestKind::FindHeaders,
7800 Response::BlockHeaders(vec![header]),
7801 )?;
7802 assert_codec_frames_validate_at_transport(
7803 LegacyRequestFrame::MempoolTransactionIds,
7804 LegacyRequestKind::MempoolTransactionIds,
7805 Response::TransactionIds(vec![tx_id]),
7806 )?;
7807 assert_codec_frames_validate_at_transport(
7808 LegacyRequestFrame::Ping,
7809 LegacyRequestKind::Ping,
7810 Response::Pong(Duration::ZERO),
7811 )?;
7812 assert_codec_frames_validate_at_transport(
7813 LegacyRequestFrame::PushTransaction(empty_v5_transaction(8).into()),
7814 LegacyRequestKind::PushTransaction,
7815 Response::Nil,
7816 )?;
7817
7818 assert_codec_frames_validate_at_transport(
7823 LegacyRequestFrame::FindBlocks {
7824 known_blocks: vec![block_hash(1)],
7825 stop: None,
7826 },
7827 LegacyRequestKind::FindBlocks,
7828 Response::Nil,
7829 )?;
7830 assert_codec_frames_validate_at_transport(
7831 LegacyRequestFrame::FindHeaders {
7832 known_blocks: vec![block_hash(1)],
7833 stop: None,
7834 },
7835 LegacyRequestKind::FindHeaders,
7836 Response::Nil,
7837 )?;
7838 assert_codec_frames_validate_at_transport(
7839 LegacyRequestFrame::MempoolTransactionIds,
7840 LegacyRequestKind::MempoolTransactionIds,
7841 Response::Nil,
7842 )?;
7843
7844 Ok(())
7845 }
7846
7847 fn assert_codec_frames_validate_at_transport(
7848 request: LegacyRequestFrame,
7849 request_kind: LegacyRequestKind,
7850 response: Response,
7851 ) -> Result<(), BoxError> {
7852 let limits = test_connection_limits();
7853 let request_id = 99;
7854 let request_frame = request.encode_frame()?;
7855 let budget = LegacyResponseBudget::from_request(
7856 request_frame.message_type,
7857 &request_frame.payload,
7858 limits,
7859 )
7860 .map_err(|error| -> BoxError { format!("{error:?}").into() })?;
7861 let frames = LegacyResponseCodec::encode_response(
7862 request_id,
7863 response,
7864 limits.max_frame_bytes,
7865 limits.max_message_bytes,
7866 )?;
7867 LegacyResponseCodec::decode_response(request_id, request_kind, frames.clone(), None)?;
7868
7869 let mut state = LegacyResponseReadState::new(budget);
7870 for frame in &frames {
7871 state
7872 .validate_frame(request_id, frame)
7873 .map_err(|error| -> BoxError { format!("{error:?}").into() })?;
7874 }
7875 state
7876 .finish()
7877 .map_err(|error| -> BoxError { format!("{error:?}").into() })?;
7878
7879 Ok(())
7880 }
7881
7882 fn test_connection_limits() -> ZakuraConnectionLimits {
7883 let max_protocol_message_len =
7884 u32::try_from(MAX_PROTOCOL_MESSAGE_LEN).expect("protocol message length fits in u32");
7885
7886 ZakuraConnectionLimits {
7887 max_frame_bytes: max_protocol_message_len,
7888 max_message_bytes: max_protocol_message_len,
7889 max_open_streams: 16,
7890 max_inbound_queue_depth: 16,
7891 idle_timeout: Duration::from_secs(1),
7892 prelude_timeout: Duration::from_secs(1),
7893 control_timeout: Duration::from_secs(1),
7894 stream_open_rate_per_second: 16,
7895 message_rate_per_second: 128,
7896 }
7897 }
7898
7899 fn block_hash(byte: u8) -> block::Hash {
7900 block::Hash([byte; 32])
7901 }
7902
7903 fn legacy_tx_id(byte: u8) -> UnminedTxId {
7904 UnminedTxId::from_legacy_id(transaction::Hash([byte; 32]))
7905 }
7906
7907 fn empty_v5_transaction(byte: u8) -> transaction::Transaction {
7908 transaction::Transaction::V5 {
7909 network_upgrade: zakura_chain::parameters::NetworkUpgrade::Nu5,
7910 lock_time: transaction::LockTime::min_lock_time_timestamp(),
7911 expiry_height: block::Height(u32::from(byte)),
7912 inputs: Vec::new(),
7913 outputs: Vec::new(),
7914 sapling_shielded_data: None,
7915 orchard_shielded_data: None,
7916 }
7917 }
7918
7919 #[test]
7920 fn same_kind_streams_share_one_connection_message_budget() {
7921 let clock = crate::zakura::testkit::TestClock::new();
7926 let mut buckets: MessageRateBuckets<crate::zakura::testkit::TestClock> =
7927 MessageRateBuckets::new();
7928 let rate = 4;
7929
7930 let stream_a = message_bucket_for(&mut buckets, 1, rate, clock.clone());
7932 let stream_b = message_bucket_for(&mut buckets, 1, rate, clock.clone());
7933 assert!(
7934 Arc::ptr_eq(&stream_a, &stream_b),
7935 "same kind must reuse one shared bucket"
7936 );
7937 assert_eq!(
7938 buckets.len(),
7939 1,
7940 "no extra bucket created for the same kind"
7941 );
7942
7943 let take = |bucket: &SharedMessageBucket<crate::zakura::testkit::TestClock>| {
7945 bucket
7946 .lock()
7947 .expect("test bucket mutex is never poisoned")
7948 .try_take()
7949 };
7950 let mut accepted = 0;
7951 for _ in 0..rate as usize * 4 {
7952 if take(&stream_a) {
7953 accepted += 1;
7954 }
7955 if take(&stream_b) {
7956 accepted += 1;
7957 }
7958 }
7959 assert_eq!(
7960 accepted, rate as usize,
7961 "the second same-kind stream must NOT get a fresh full budget"
7962 );
7963
7964 clock.advance(Duration::from_secs(1));
7967 let mut refilled = 0;
7968 for _ in 0..rate as usize * 4 {
7969 if take(&stream_a) {
7970 refilled += 1;
7971 }
7972 if take(&stream_b) {
7973 refilled += 1;
7974 }
7975 }
7976 assert_eq!(
7977 refilled, rate as usize,
7978 "refill restores one shared budget, not one-per-stream"
7979 );
7980 }
7981
7982 #[test]
7983 fn different_kinds_get_independent_message_budgets() {
7984 let clock = crate::zakura::testkit::TestClock::new();
7987 let mut buckets: MessageRateBuckets<crate::zakura::testkit::TestClock> =
7988 MessageRateBuckets::new();
7989
7990 let request = message_bucket_for(&mut buckets, 1, 2, clock.clone());
7991 let gossip = message_bucket_for(&mut buckets, 2, 2, clock.clone());
7992 assert!(
7993 !Arc::ptr_eq(&request, &gossip),
7994 "distinct kinds must not share a bucket"
7995 );
7996 assert_eq!(buckets.len(), 2);
7997
7998 let take = |bucket: &SharedMessageBucket<crate::zakura::testkit::TestClock>| {
7999 bucket
8000 .lock()
8001 .expect("test bucket mutex is never poisoned")
8002 .try_take()
8003 };
8004 assert!(take(&request));
8006 assert!(take(&request));
8007 assert!(!take(&request));
8008 assert!(take(&gossip));
8009 assert!(take(&gossip));
8010 assert!(!take(&gossip));
8011 }
8012
8013 #[test]
8014 fn idle_invariant_keeps_app_timeout_below_quic_timeout() {
8015 let limits = ZakuraLocalLimits::from_config(&Config::default());
8016
8017 validate_idle_invariant(&limits).expect("default Zakura limits satisfy idle invariant");
8018 assert!(
8019 (limits.initial_limits().idle_timeout_millis as u128)
8020 < limits.quic_idle_timeout.as_millis()
8021 );
8022 }
8023
8024 #[test]
8043 fn requester_response_budget_is_capped_for_large_inventory_request() {
8044 let limits = test_connection_limits();
8045 assert_eq!(
8046 limits.max_message_bytes as usize, MAX_PROTOCOL_MESSAGE_LEN,
8047 "fixture should negotiate the protocol-max message cap so the unclamped \
8048 budget is maximal",
8049 );
8050
8051 let max_items =
8052 usize::try_from(MAX_TX_INV_IN_SENT_MESSAGE).expect("inventory cap fits in usize");
8053
8054 let blocks = LegacyRequestFrame::BlocksByHash(vec![block_hash(7); max_items])
8057 .encode_frame()
8058 .expect("max-inventory blocks request encodes");
8059 let blocks_budget =
8060 LegacyResponseBudget::from_request(blocks.message_type, &blocks.payload, limits)
8061 .expect("budget derives from a max-inventory blocks request");
8062 assert!(
8063 blocks_budget.max_bytes <= LEGACY_RESPONSE_MAX_AGGREGATE_BYTES,
8064 "BlocksByHash retained-frame budget {} must be clamped to the aggregate cap {}",
8065 blocks_budget.max_bytes,
8066 LEGACY_RESPONSE_MAX_AGGREGATE_BYTES,
8067 );
8068
8069 let txs = LegacyRequestFrame::TransactionsById(vec![legacy_tx_id(7); max_items])
8071 .encode_frame()
8072 .expect("max-inventory transactions request encodes");
8073 let txs_budget = LegacyResponseBudget::from_request(txs.message_type, &txs.payload, limits)
8074 .expect("budget derives from a max-inventory transactions request");
8075 assert!(
8076 txs_budget.max_bytes <= LEGACY_RESPONSE_MAX_AGGREGATE_BYTES,
8077 "TransactionsById retained-frame budget {} must be clamped to the aggregate cap {}",
8078 txs_budget.max_bytes,
8079 LEGACY_RESPONSE_MAX_AGGREGATE_BYTES,
8080 );
8081
8082 let small = LegacyRequestFrame::BlocksByHash(vec![block_hash(1)])
8086 .encode_frame()
8087 .expect("single-item blocks request encodes");
8088 let small_budget =
8089 LegacyResponseBudget::from_request(small.message_type, &small.payload, limits)
8090 .expect("budget derives from a single-item blocks request");
8091 assert!(
8092 small_budget.max_bytes >= limits.max_message_bytes as usize,
8093 "a single-item request must still permit one full response message; budget was {}",
8094 small_budget.max_bytes,
8095 );
8096 assert!(
8097 small_budget.max_bytes <= LEGACY_RESPONSE_MAX_AGGREGATE_BYTES,
8098 "even a single-item budget stays within the aggregate cap",
8099 );
8100 }
8101
8102 #[test]
8121 fn nil_response_to_inventory_or_ping_request_is_not_fail_closed() {
8122 let limits = test_connection_limits();
8123 let request_id = 99;
8124
8125 let cases: [(LegacyRequestFrame, LegacyRequestKind); 3] = [
8126 (
8127 LegacyRequestFrame::BlocksByHash(vec![block_hash(1)]),
8128 LegacyRequestKind::Blocks,
8129 ),
8130 (
8131 LegacyRequestFrame::TransactionsById(vec![legacy_tx_id(2)]),
8132 LegacyRequestKind::Transactions,
8133 ),
8134 (LegacyRequestFrame::Ping, LegacyRequestKind::Ping),
8135 ];
8136
8137 for (request, request_kind) in cases {
8138 let request_frame = request.encode_frame().expect("request frame encodes");
8139 let budget = LegacyResponseBudget::from_request(
8140 request_frame.message_type,
8141 &request_frame.payload,
8142 limits,
8143 )
8144 .expect("budget derives from request");
8145
8146 let nil_frames = LegacyResponseCodec::encode_response(
8149 request_id,
8150 Response::Nil,
8151 limits.max_frame_bytes,
8152 limits.max_message_bytes,
8153 )
8154 .expect("nil response encodes");
8155
8156 let decoded = LegacyResponseCodec::decode_response(
8158 request_id,
8159 request_kind,
8160 nil_frames.clone(),
8161 None,
8162 );
8163 assert!(
8164 decoded.is_err(),
8165 "decode_response must reject a bare NIL for {request_kind:?}",
8166 );
8167
8168 let mut state = LegacyResponseReadState::new(budget);
8172 let mut validate = Ok(());
8173 for frame in &nil_frames {
8174 validate = state.validate_frame(request_id, frame);
8175 if validate.is_err() {
8176 break;
8177 }
8178 }
8179 let validate = validate.and_then(|()| state.finish());
8180 assert!(
8181 matches!(validate, Err(OutboundRequestError::Fatal(_))),
8182 "transport must fail closed (Fatal) on a NIL answer to {request_kind:?} so the \
8183 request stream worker disconnects the peer; got {validate:?}",
8184 );
8185 }
8186 }
8187
8188 #[tokio::test]
8209 async fn inbound_accept_enforces_per_ip_cap() -> Result<(), BoxError> {
8210 let _guard = zakura_test::init();
8211
8212 async fn try_register(
8216 supervisor: &ZakuraSupervisorHandle,
8217 peer: &ZakuraPeerId,
8218 remote_ip: Option<IpAddr>,
8219 ) -> ZakuraRegistration {
8220 let (outbound_tx, _outbound_rx) = mpsc::channel(1);
8221 let outbound_handle = ZakuraPeerHandle::new_for_tests(peer.clone(), outbound_tx);
8222 supervisor
8223 .register(
8224 test_conn_id(),
8225 peer.clone(),
8226 remote_ip,
8227 [peer.as_bytes()[0]; TRANSCRIPT_HASH_BYTES],
8228 outbound_handle,
8229 CancellationToken::new(),
8230 ZAKURA_CAP_LEGACY_GOSSIP | ZAKURA_CAP_HEADER_SYNC,
8231 )
8232 .await
8233 }
8234 let ip: IpAddr = "203.0.113.7".parse().expect("test ip parses");
8235 let supervisor = ZakuraSupervisorHandle::new(1);
8236 assert!(
8237 matches!(
8238 try_register(&supervisor, &test_peer(1), Some(ip)).await,
8239 ZakuraRegistration::Registered { .. }
8240 ),
8241 "first identity from the IP registers",
8242 );
8243 assert!(
8244 matches!(
8245 try_register(&supervisor, &test_peer(2), Some(ip)).await,
8246 ZakuraRegistration::Rejected(ZakuraRejectReason::ResourceLimit)
8247 ),
8248 "a second distinct identity from the same IP must be Rejected(ResourceLimit) at cap 1",
8249 );
8250
8251 let limits = ZakuraLocalLimits::from_config(&Config::default());
8256 assert!(
8257 limits.max_connections > 1,
8258 "global admission cap must exceed the per-IP cap so the second same-IP identity is \
8259 turned away by the per-IP cap rather than the global gate",
8260 );
8261 let server_ep = LocalEndpointFactory::with_transport_config(limits.transport_config())
8262 .endpoint(880)
8263 .await?;
8264 let supervisor = ZakuraSupervisorHandle::new(1);
8265 let handler = ZakuraProtocolHandler::new(
8266 supervisor.clone(),
8267 Network::Mainnet,
8268 ZakuraHandshakeConfig::for_network(&Network::Mainnet),
8269 limits.clone(),
8270 )
8271 .with_endpoint(server_ep.clone());
8272 let router = Router::builder(server_ep)
8273 .accept(P2P_V2_ALPN, handler)
8274 .spawn();
8275 let full_addr = router.endpoint().node_addr().initialized().await;
8281 let loopback_addr = NodeAddr::new(full_addr.node_id).with_direct_addresses(
8282 full_addr
8283 .direct_addresses()
8284 .copied()
8285 .filter(|addr| addr.is_ipv4() && addr.ip().is_loopback()),
8286 );
8287 assert!(
8288 loopback_addr.direct_addresses().next().is_some(),
8289 "server must advertise an IPv4 loopback direct address",
8290 );
8291 let server_addr = loopback_addr;
8292
8293 async fn connect_native(
8297 server_addr: &NodeAddr,
8298 seed: u64,
8299 limits: &ZakuraLocalLimits,
8300 ) -> Result<(Endpoint, Connection), BoxError> {
8301 let endpoint = LocalEndpointFactory::with_transport_config(limits.transport_config())
8302 .endpoint(seed)
8303 .await?;
8304 endpoint.add_node_addr(server_addr.clone())?;
8305 let connection = endpoint.connect(server_addr.clone(), P2P_V2_ALPN).await?;
8306 Ok((endpoint, connection))
8307 }
8308 async fn run_handshake(
8309 endpoint: &Endpoint,
8310 connection: &Connection,
8311 limits: &ZakuraLocalLimits,
8312 ) -> Result<(), BoxError> {
8313 let config = ZakuraHandshakeConfig::for_network(&Config::default().network);
8314 let local_peer_id = ZakuraPeerId::new(endpoint.node_id().as_bytes().to_vec())?;
8315 run_native_initiator_handshake_without_trace(
8316 connection,
8317 limits,
8318 &config,
8319 &local_peer_id,
8320 )
8321 .await?;
8322 Ok(())
8323 }
8324
8325 let (_ep1, _conn1) = connect_native(&server_addr, 881, &limits).await?;
8328 run_handshake(&_ep1, &_conn1, &limits).await?;
8329 let mut first_registered = 0;
8330 for _ in 0..200 {
8331 first_registered = supervisor.registered_ids().await.len();
8332 if first_registered >= 1 {
8333 break;
8334 }
8335 tokio::time::sleep(Duration::from_millis(25)).await;
8336 }
8337 assert_eq!(
8338 first_registered, 1,
8339 "the first inbound identity from the source IP must register (a real source IP was \
8340 resolved from the endpoint and counted against the per-IP cap)",
8341 );
8342
8343 let (_ep2, conn2) = connect_native(&server_addr, 882, &limits).await?;
8350 let _ = run_handshake(&_ep2, &conn2, &limits).await;
8351 let mut rejected_close = false;
8352 for _ in 0..400 {
8353 if supervisor.registered_ids().await.len() >= 2 {
8354 break;
8355 }
8356 if matches!(
8357 conn2.close_reason(),
8358 Some(iroh::endpoint::ConnectionError::ApplicationClosed(ref close))
8359 if close.error_code == VarInt::from_u32(ZAKURA_CLOSE_RESOURCE)
8360 ) {
8361 rejected_close = true;
8362 break;
8363 }
8364 tokio::time::sleep(Duration::from_millis(25)).await;
8365 }
8366 let registered = supervisor.registered_ids().await.len();
8367 assert!(
8368 rejected_close && registered == 1,
8369 "a second distinct identity from the same source IP must be rejected by the per-IP \
8370 cap with a resource-limit close (resource_close={rejected_close}, \
8371 registered={registered}); before the fix the inbound accept passed remote_ip = None, \
8372 so both identities registered and one source IP could exhaust the connection budget",
8373 );
8374
8375 router.shutdown().await?;
8376 Ok(())
8377 }
8378}