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