Skip to main content

zakura_network/zakura/
handler.rs

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