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