Skip to main content

zakura_network/zakura/
handler.rs

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