dig_nat/mux.rs
1//! Stream multiplexing + byte-range streams over a single established peer connection.
2//!
3//! Every traversal tier (direct / UPnP / NAT-PMP / PCP / hole-punch / relayed) yields the SAME
4//! thing: one mTLS byte stream to the peer. On top of that single stream this module layers
5//! **[`yamux`]** multiplexing so the content/download layer can open **many cheap concurrent logical
6//! streams** to the peer with no head-of-line blocking — the transport is **streaming-first**, never
7//! "send request, buffer the whole response in memory".
8//!
9//! Two capabilities:
10//!
11//! 1. **Multiplexing** — [`PeerSession::open_stream`] opens an independent bidirectional
12//! [`PeerStream`] (a tokio [`AsyncRead`] + [`AsyncWrite`]); open N of them concurrently and read
13//! each incrementally with natural backpressure (yamux windows).
14//! 2. **Byte-range streams** — [`PeerSession::open_range_stream`] opens a stream scoped to a
15//! `[offset, offset+len)` range of a named resource by writing a small [`RangeRequest`] preamble,
16//! then hands back the stream so the caller reads exactly those bytes as they arrive. A downloader
17//! opens range streams to DIFFERENT peers in parallel and reassembles — multi-source parallel
18//! download falls out of "streams are cheap + multiplexed + range-scoped".
19//!
20//! The uniform abstraction holds regardless of how the connection was established, and regardless of
21//! whether the underlying byte stream is direct or (tier-6) relay-proxied.
22//!
23//! ## Wire alignment (normative)
24//!
25//! The control + range types here conform to the published **L7 peer-network spec** (docs.dig.net
26//! "L7 · DIG Node peer network", §8 streaming, §9 byte-range fetch + availability). The shapes are
27//! the `dig.getAvailability` / `dig.fetchRange` request/response and the streamed `RangeFrame`
28//! (`{offset, length, bytes, complete}`, plus the fixed-size identity set `root` + `total_length` +
29//! `chunk_count` on every frame and the resource-scaling `chunk_lens` + `inclusion_proof` once per
30//! stream, paged by `chunk_lens_offset`). Per-chunk integrity (split by `chunk_lens`, verify
31//! the whole-resource inclusion proof vs the chain-anchored `root`, AES-256-GCM-SIV-open) is done by
32//! the CONTENT layer above dig-nat; dig-nat carries these frames faithfully over the mux transport.
33
34use std::io;
35use std::sync::atomic::{AtomicBool, Ordering};
36use std::sync::Arc;
37
38use serde::{Deserialize, Serialize};
39use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
40use tokio::sync::Notify;
41use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
42
43use crate::safe_text::SafeText;
44
45/// One logical, bidirectional stream to the peer — a tokio [`AsyncRead`] + [`AsyncWrite`]. Reads
46/// deliver bytes incrementally as they arrive (streaming, with yamux-window backpressure); many
47/// [`PeerStream`]s coexist on one [`PeerSession`] without head-of-line blocking.
48///
49/// yamux streams are `futures` streams; this is the tokio-trait view via `tokio-util` compat.
50pub type PeerStream = Compat<yamux::Stream>;
51
52/// One item in a [`dig.getAvailability`](AvailabilityRequest) batch — a resource key at store, root,
53/// or capsule/resource granularity (inferred from which fields are present, per the L7 spec §9):
54/// `store_id` only → *has_store*; `+ root` → *has_root* (the capsule `store_id:root`); `+
55/// retrieval_key` → *has_resource*. Hashes are 64-hex.
56#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
57#[non_exhaustive]
58pub struct AvailabilityItem {
59 /// The store id (64-hex). Always present.
60 pub store_id: String,
61 /// The generation root (64-hex). Present for root/resource granularity.
62 #[serde(default, skip_serializing_if = "Option::is_none")]
63 pub root: Option<String>,
64 /// The resource retrieval key (64-hex). Present for resource granularity.
65 #[serde(default, skip_serializing_if = "Option::is_none")]
66 pub retrieval_key: Option<String>,
67}
68
69/// The **availability pre-check** (`dig.getAvailability`, L7 spec §9) — asked BEFORE any range fetch.
70/// A multi-source download batches candidate peers × items in one call each and only fans byte-range
71/// requests at peers that answer *available* — never opening range streams to peers that may not hold
72/// the content. A message-style control call over the mux'd mTLS connection.
73#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
74#[non_exhaustive]
75pub struct AvailabilityRequest {
76 /// The items to check, batched.
77 pub items: Vec<AvailabilityItem>,
78}
79
80/// One answer in an [`AvailabilityResponse`], positionally aligned with the request `items`.
81#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
82#[non_exhaustive]
83pub struct AvailabilityAnswer {
84 /// Whether the peer holds the queried item. Always present.
85 pub available: bool,
86 /// (store granularity) generation roots the peer holds for the store, newest-first.
87 #[serde(default, skip_serializing_if = "Option::is_none")]
88 pub roots: Option<Vec<String>>,
89 /// (root/resource granularity) the ciphertext length — lets the caller plan its ranges.
90 #[serde(default, skip_serializing_if = "Option::is_none")]
91 pub total_length: Option<u64>,
92 /// (root/resource granularity) the chunk count.
93 #[serde(default, skip_serializing_if = "Option::is_none")]
94 pub chunk_count: Option<u64>,
95 /// Whether the peer holds the FULL resource/capsule (`true`) or only part (`false`).
96 #[serde(default, skip_serializing_if = "Option::is_none")]
97 pub complete: Option<bool>,
98}
99
100/// The peer's answer to an [`AvailabilityRequest`]: one [`AvailabilityAnswer`] per queried item,
101/// positionally aligned with the request's `items`.
102#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
103#[non_exhaustive]
104pub struct AvailabilityResponse {
105 /// One answer per queried item, in request order.
106 pub items: Vec<AvailabilityAnswer>,
107}
108
109/// A byte-range request (`dig.fetchRange`, L7 spec §9) written at the start of a range-scoped stream.
110/// Identifies a resource (`store_id` + `retrieval_key` [+ `root`]) or a whole capsule
111/// (`capsule: true`, identified by `store_id` [+ `root`]) and the `[offset, offset+length)` range.
112#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
113#[non_exhaustive]
114pub struct RangeRequest {
115 /// The store id (64-hex).
116 pub store_id: String,
117 /// The resource retrieval key (64-hex). Omitted when `capsule` is true.
118 #[serde(default, skip_serializing_if = "Option::is_none")]
119 pub retrieval_key: Option<String>,
120 /// The generation root (64-hex). Optional — defaults to the chain-anchored tip.
121 #[serde(default, skip_serializing_if = "Option::is_none")]
122 pub root: Option<String>,
123 /// Fetch a whole capsule / `.dig` (identified by `store_id` [+ `root`]) rather than one resource.
124 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
125 pub capsule: bool,
126 /// Start offset (bytes) into the resource ciphertext. Default `0`.
127 #[serde(default)]
128 pub offset: u64,
129 /// Length (bytes) to return (widened to whole-chunk boundaries; clamped to the node window).
130 pub length: u64,
131 /// Suppress the resource-scaling layout metadata (`chunk_lens` + `inclusion_proof`) on this
132 /// stream's frames, because the client already holds the commitment for this `root`.
133 ///
134 /// A client that has already read the layout once — a resumed download, a second range of the same
135 /// resource, a parallel fetch from another holder — does not need it again, and for a large
136 /// resource re-sending it costs a paged prologue per stream. Absent or `false` preserves the
137 /// pre-0.13.0 behaviour, so an older holder that ignores this field is never broken by it: it
138 /// simply sends metadata the client discards. The fixed-size identity fields (`root`,
139 /// `total_length`, `chunk_count`, `chunk_index`) are NOT suppressed — they are what detects a
140 /// wrong-generation holder on arrival.
141 #[serde(default, skip_serializing_if = "Option::is_none")]
142 pub skip_layout: Option<bool>,
143}
144
145/// One streamed `dig.fetchRange` frame (L7 spec §8 framing). Frames arrive in ascending `offset`
146/// order and tile the requested range exactly; the caller reassembles by `offset` and stops on
147/// `complete`.
148///
149/// ## The metadata splits into two sets, and which frames carry them differs
150///
151/// Read this before implementing either side — the two halves have different rules, and conforming to
152/// the wrong half produces a verification miss on a multi-frame read rather than a clean failure.
153///
154/// - **Identity — fixed-size, on EVERY frame:** `root`, `total_length`, `chunk_count`, plus
155/// `chunk_index` where the window is chunk-aligned. These are what let a reader reject a
156/// wrong-generation or wrong-layout holder the moment a frame arrives, which is a property the
157/// once-per-stream set can never have. They cost a bounded number of bytes, so carrying them
158/// everywhere is cheap. Set them with [`with_identity`](Self::with_identity) +
159/// [`with_chunk_index`](Self::with_chunk_index).
160/// - **Prologue — resource-scaling, ONCE per range stream:** `chunk_lens` (paged, located by
161/// `chunk_lens_offset`) and `inclusion_proof`. Their size is a function of the RESOURCE rather than of
162/// the frame, so they ride the first frame or a paged prologue and **MUST NOT** be repeated on later
163/// frames. Set them with [`with_chunk_lens_page`](Self::with_chunk_lens_page) +
164/// [`with_inclusion_proof`](Self::with_inclusion_proof), or omit them entirely when the request set
165/// [`RangeRequest::skip_layout`].
166///
167/// Before 0.13.0 every one of these was "first frame only", because the whole layout had to fit one
168/// frame or the range was unservable (#1640). `SPEC.md` §5.1.1 is normative.
169#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
170#[non_exhaustive]
171pub struct RangeFrame {
172 /// This frame's start offset within the requested range.
173 pub offset: u64,
174 /// This frame's byte length.
175 pub length: u64,
176 /// The raw ciphertext bytes. On the wire they are **base64** (`base64_bytes`) — the canonical
177 /// `dig.fetchRange` frame encoding every producer emits.
178 #[serde(with = "base64_bytes")]
179 pub bytes: Vec<u8>,
180 /// Whether this is the final frame of the range.
181 pub complete: bool,
182 /// **(identity — every frame)** the full resource ciphertext length.
183 #[serde(default, skip_serializing_if = "Option::is_none")]
184 pub total_length: Option<u64>,
185 /// **(prologue — once per stream)** per-chunk ciphertext lengths of the whole resource, in order,
186 /// or ONE PAGE of them beginning at `chunk_lens_offset`.
187 ///
188 /// Resource-scaling, so it MUST NOT be repeated on later frames: on a later frame it was only a
189 /// redundant equality check on an array the client already held. Omitted entirely when the request
190 /// set [`RangeRequest::skip_layout`].
191 #[serde(default, skip_serializing_if = "Option::is_none")]
192 pub chunk_lens: Option<Vec<u64>>,
193 /// **(identity — every frame, where the window is chunk-aligned)** index into the resource's
194 /// `chunk_lens` array of the first chunk in THIS frame.
195 ///
196 /// Fixed-size, and about this frame rather than about the resource, so a chunk-aligned continuation
197 /// frame states it too — see [`with_chunk_index`](Self::with_chunk_index), which is deliberately
198 /// separate from the proof setter so stating alignment never costs a repeated proof.
199 #[serde(default, skip_serializing_if = "Option::is_none")]
200 pub chunk_index: Option<u64>,
201 /// **(prologue — once per stream)** merkle inclusion proof of the whole resource vs the generation
202 /// `root` (base64, relayed verbatim); `null`/absent for `capsule: true` (self-verifying on install).
203 ///
204 /// Resource-scaling and capped at [`MAX_INCLUSION_PROOF_B64`]. At 4,096 B it would consume the
205 /// entire remaining frame budget if repeated, which is the concrete reason "once per stream" is a
206 /// MUST NOT rather than a preference.
207 #[serde(default, skip_serializing_if = "Option::is_none")]
208 pub inclusion_proof: Option<String>,
209 /// **(identity — every frame)** the generation root (64-hex) this range is served from, and which
210 /// the inclusion proof is against.
211 #[serde(default, skip_serializing_if = "Option::is_none")]
212 pub root: Option<String>,
213 /// **(identity — every frame)** the resource's TOTAL chunk count — how many entries the
214 /// reassembled `chunk_lens` array has.
215 ///
216 /// Together with `root` and `total_length` it is what lets a reader detect a wrong-generation or
217 /// wrong-layout holder the moment a frame arrives. It is also how a reader sizes the array it is
218 /// paging in, and how it knows the prologue is complete.
219 #[serde(default, skip_serializing_if = "Option::is_none")]
220 pub chunk_count: Option<u64>,
221 /// **(prologue — once per page)** the index into the resource's `chunk_lens` array at which THIS
222 /// frame's `chunk_lens` page begins — how a **paged prologue** is located and reassembled.
223 ///
224 /// A resource whose layout exceeds [`MAX_CHUNK_LENS_PER_FRAME`] cannot state it on one frame, so
225 /// the sender pages it: successive frames each carry up to that many entries, stamped with the
226 /// offset they start at. A reader places each page at its offset and has the whole array once it
227 /// holds `chunk_count` entries. Absent means "this frame's `chunk_lens`, if any, begins at 0" — the
228 /// single-frame layout, which is the pre-0.13.0 shape.
229 #[serde(default, skip_serializing_if = "Option::is_none")]
230 pub chunk_lens_offset: Option<u64>,
231}
232
233/// Serde for [`RangeFrame::bytes`]: **base64** on the wire, raw `Vec<u8>` in Rust.
234///
235/// The `dig.fetchRange` frame is JSON, and the canonical wire type
236/// (`dig_rpc_protocol::types::RangeFrame`, "this window's ciphertext, base64") — and every real
237/// producer, including the dig-node peer serve path — encodes the window as a base64 STRING. Reading
238/// it with `serde_bytes` instead yielded the string's literal characters, so a served window arrived
239/// as its own base64 text and the reassembler rejected the frame (#1586, the read-leg blocker).
240///
241/// Deserialization is tolerant: a base64 string (canonical) OR a byte array (what an older dig-nat
242/// emitted) both decode, so a mixed-version peer is never dropped.
243mod base64_bytes {
244 use base64::Engine as _;
245 use serde::de::{SeqAccess, Visitor};
246 use serde::{Deserializer, Serializer};
247 use std::fmt;
248
249 pub fn serialize<S: Serializer>(bytes: &[u8], s: S) -> Result<S::Ok, S::Error> {
250 s.serialize_str(&base64::engine::general_purpose::STANDARD.encode(bytes))
251 }
252
253 pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<u8>, D::Error> {
254 d.deserialize_any(Base64OrArray)
255 }
256
257 struct Base64OrArray;
258
259 impl<'de> Visitor<'de> for Base64OrArray {
260 type Value = Vec<u8>;
261
262 fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
263 f.write_str("base64-encoded ciphertext (or a legacy byte array)")
264 }
265
266 fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<Self::Value, E> {
267 base64::engine::general_purpose::STANDARD
268 .decode(v)
269 .map_err(|e| E::custom(format!("range frame bytes are not valid base64: {e}")))
270 }
271
272 fn visit_bytes<E: serde::de::Error>(self, v: &[u8]) -> Result<Self::Value, E> {
273 Ok(v.to_vec())
274 }
275
276 fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value, A::Error> {
277 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or_default());
278 while let Some(b) = seq.next_element::<u8>()? {
279 out.push(b);
280 }
281 Ok(out)
282 }
283 }
284}
285
286impl AvailabilityItem {
287 /// A *has_store* query — does the peer hold anything for this store?
288 pub fn store(store_id: impl Into<String>) -> Self {
289 AvailabilityItem {
290 store_id: store_id.into(),
291 root: None,
292 retrieval_key: None,
293 }
294 }
295
296 /// Narrow the query to one generation `root` — a *has_root* query for the capsule `store_id:root`.
297 pub fn with_root(mut self, root: impl Into<String>) -> Self {
298 self.root = Some(root.into());
299 self
300 }
301
302 /// Narrow the query to one resource — a *has_resource* query.
303 pub fn with_retrieval_key(mut self, retrieval_key: impl Into<String>) -> Self {
304 self.retrieval_key = Some(retrieval_key.into());
305 self
306 }
307}
308
309impl AvailabilityRequest {
310 /// A batched availability pre-check over `items`.
311 pub fn new(items: Vec<AvailabilityItem>) -> Self {
312 AvailabilityRequest { items }
313 }
314
315 /// Serialize as a `u32` big-endian length prefix + JSON body (the uniform control framing).
316 pub fn encode(&self) -> io::Result<Vec<u8>> {
317 encode_framed(self)
318 }
319 /// Read + decode an [`AvailabilityRequest`] from `r` (the peer/serving side).
320 pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
321 decode_framed(r).await
322 }
323}
324
325impl AvailabilityAnswer {
326 /// The peer holds the queried item. Describe WHAT it holds with the `with_*` setters.
327 pub fn available() -> Self {
328 AvailabilityAnswer::with_availability(true)
329 }
330
331 /// The peer does not hold the queried item — the whole answer, since no other field is meaningful.
332 pub fn unavailable() -> Self {
333 AvailabilityAnswer::with_availability(false)
334 }
335
336 fn with_availability(available: bool) -> Self {
337 AvailabilityAnswer {
338 available,
339 roots: None,
340 total_length: None,
341 chunk_count: None,
342 complete: None,
343 }
344 }
345
346 /// (store granularity) the generation roots held for the store, newest-first.
347 pub fn with_roots(mut self, roots: Vec<String>) -> Self {
348 self.roots = Some(roots);
349 self
350 }
351
352 /// (root/resource granularity) the resource's ciphertext length, so the caller can plan its ranges.
353 pub fn with_total_length(mut self, total_length: u64) -> Self {
354 self.total_length = Some(total_length);
355 self
356 }
357
358 /// (root/resource granularity) the resource's chunk count.
359 pub fn with_chunk_count(mut self, chunk_count: u64) -> Self {
360 self.chunk_count = Some(chunk_count);
361 self
362 }
363
364 /// Whether the peer holds the FULL resource/capsule rather than only part of it.
365 pub fn with_complete(mut self, complete: bool) -> Self {
366 self.complete = Some(complete);
367 self
368 }
369}
370
371impl AvailabilityResponse {
372 /// The answers to an [`AvailabilityRequest`], positionally aligned with its `items`.
373 pub fn new(items: Vec<AvailabilityAnswer>) -> Self {
374 AvailabilityResponse { items }
375 }
376
377 /// Serialize as a `u32` big-endian length prefix + JSON body.
378 pub fn encode(&self) -> io::Result<Vec<u8>> {
379 encode_framed(self)
380 }
381 /// Read + decode an [`AvailabilityResponse`] from `r` (the requesting side).
382 pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
383 decode_framed(r).await
384 }
385}
386
387impl RangeRequest {
388 /// A range request for a content resource (`store_id` + `retrieval_key`).
389 pub fn resource(
390 store_id: impl Into<String>,
391 retrieval_key: impl Into<String>,
392 offset: u64,
393 length: u64,
394 ) -> Self {
395 RangeRequest::resource_or_capsule(
396 store_id,
397 Some(retrieval_key.into()),
398 false,
399 offset,
400 length,
401 )
402 }
403
404 /// A range request for a whole capsule / `.dig`, identified by `store_id` (plus a `root` if the
405 /// caller pins a generation).
406 pub fn capsule(store_id: impl Into<String>, offset: u64, length: u64) -> Self {
407 RangeRequest::resource_or_capsule(store_id, None, true, offset, length)
408 }
409
410 fn resource_or_capsule(
411 store_id: impl Into<String>,
412 retrieval_key: Option<String>,
413 capsule: bool,
414 offset: u64,
415 length: u64,
416 ) -> Self {
417 RangeRequest {
418 store_id: store_id.into(),
419 retrieval_key,
420 root: None,
421 capsule,
422 offset,
423 length,
424 skip_layout: None,
425 }
426 }
427
428 /// Pin the request to one generation `root` (64-hex) rather than the chain-anchored tip.
429 pub fn with_root(mut self, root: impl Into<String>) -> Self {
430 self.root = Some(root.into());
431 self
432 }
433
434 /// Ask the holder to omit the resource-scaling layout metadata, because this client already holds
435 /// the commitment for this `root`. See [`RangeRequest::skip_layout`].
436 pub fn with_skip_layout(mut self, skip_layout: bool) -> Self {
437 self.skip_layout = Some(skip_layout);
438 self
439 }
440
441 /// Serialize as a `u32` big-endian length prefix + JSON body — the preamble a peer reads to learn
442 /// the resource + range before streaming the frames.
443 pub fn encode(&self) -> io::Result<Vec<u8>> {
444 encode_framed(self)
445 }
446 /// Read + decode a [`RangeRequest`] preamble from `r` (the serving side of a range stream).
447 pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
448 decode_framed(r).await
449 }
450}
451
452impl RangeFrame {
453 /// A **data frame**: `bytes` at `offset`, with no metadata — the shape of every continuation frame.
454 ///
455 /// `length` is derived from `bytes`, since a frame that misdeclares its own payload length is never
456 /// what a serve path wants. The metadata a FIRST frame or a prologue page carries is layered on with
457 /// the `with_*` setters, so a continuation frame cannot accidentally claim a layout it is not
458 /// stating — the two shapes are different call chains, not one call with a pile of `None`s.
459 pub fn data(offset: u64, bytes: Vec<u8>) -> Self {
460 RangeFrame {
461 offset,
462 length: bytes.len() as u64,
463 bytes,
464 complete: false,
465 total_length: None,
466 chunk_lens: None,
467 chunk_index: None,
468 inclusion_proof: None,
469 root: None,
470 chunk_count: None,
471 chunk_lens_offset: None,
472 }
473 }
474
475 /// Mark this as the final frame of the range.
476 pub fn with_complete(mut self, complete: bool) -> Self {
477 self.complete = complete;
478 self
479 }
480
481 /// The fixed-size identity metadata EVERY frame of a range should carry: the generation `root`
482 /// (64-hex) the range is served from, the resource's ciphertext `total_length`, and its
483 /// `chunk_count`.
484 ///
485 /// These three are what let a reader reject a wrong-generation or wrong-layout holder the moment a
486 /// frame arrives — which the resource-scaling metadata never could, since it arrives once. They are
487 /// fixed-size, so carrying them everywhere costs a bounded number of bytes.
488 pub fn with_identity(
489 mut self,
490 root: impl Into<String>,
491 total_length: u64,
492 chunk_count: u64,
493 ) -> Self {
494 self.root = Some(root.into());
495 self.total_length = Some(total_length);
496 self.chunk_count = Some(chunk_count);
497 self
498 }
499
500 /// A page of the resource's `chunk_lens` array, beginning at entry `chunk_lens_offset`.
501 ///
502 /// Call it once with offset `0` for a layout that fits one frame (at most
503 /// [`MAX_CHUNK_LENS_PER_FRAME`] entries), or once per page of a **paged prologue**. `chunk_lens` is
504 /// a DECRYPT input — per-chunk AES-GCM-SIV needs the WHOLE array, and a reader rejects an array
505 /// whose sum differs from `total_length` — so a page is only ever useful as part of a complete set.
506 pub fn with_chunk_lens_page(mut self, chunk_lens_offset: u64, chunk_lens: Vec<u64>) -> Self {
507 self.chunk_lens_offset = Some(chunk_lens_offset);
508 self.chunk_lens = Some(chunk_lens);
509 self
510 }
511
512 /// Split a whole resource's `chunk_lens` into the pages of a **paged prologue**: `(offset, page)`
513 /// pairs, in order, each ready for [`with_chunk_lens_page`](Self::with_chunk_lens_page).
514 ///
515 /// This is the encoder half of the paging rule, and it exists so that no serve path re-derives that
516 /// rule by hand. #1640 was an encode/decode asymmetry, so the split and its mirror
517 /// [`ChunkLensAssembler`] are published together from one module: the pages this returns are exactly
518 /// the pages the assembler requires — aligned offsets, [`MAX_CHUNK_LENS_PER_FRAME`] entries each
519 /// except a shorter tail, tiling the array with no gap and no overlap.
520 ///
521 /// A layout at or under the threshold yields a single page at offset 0, which is the single-frame
522 /// pre-0.13.0 shape; an empty layout yields no pages at all.
523 pub fn split_chunk_lens_pages(chunk_lens: &[u64]) -> Vec<(u64, Vec<u64>)> {
524 chunk_lens
525 .chunks(MAX_CHUNK_LENS_PER_FRAME)
526 .enumerate()
527 .map(|(page, entries)| ((page * MAX_CHUNK_LENS_PER_FRAME) as u64, entries.to_vec()))
528 .collect()
529 }
530
531 /// Where this frame starts in the resource's chunk sequence — the index into `chunk_lens` of its
532 /// first chunk, for a chunk-aligned window.
533 ///
534 /// Fixed-size **identity**, so a chunk-aligned CONTINUATION frame states it too, not only the first
535 /// frame. It has its own setter for exactly that reason: it used to be a parameter of
536 /// [`with_inclusion_proof`](Self::with_inclusion_proof), which meant the one field every reader wants
537 /// on every frame could not be stated without repeating a once-per-stream proof — a `SPEC.md` §5.1.1
538 /// MUST NOT, and 4,096 B per frame against a budget with zero slack.
539 pub fn with_chunk_index(mut self, chunk_index: u64) -> Self {
540 self.chunk_index = Some(chunk_index);
541 self
542 }
543
544 /// The merkle inclusion proof of the whole resource against the generation `root` (base64, relayed
545 /// verbatim).
546 ///
547 /// Resource-scaling, so it rides the first frame or the prologue — once per range stream, never
548 /// repeated. Base64 longer than [`MAX_INCLUSION_PROOF_B64`] is refused by
549 /// [`encode`](Self::encode): such a resource has no conforming range stream at all, and the holder
550 /// must say so rather than stream frames a reader cannot verify.
551 pub fn with_inclusion_proof(mut self, inclusion_proof: impl Into<String>) -> Self {
552 self.inclusion_proof = Some(inclusion_proof.into());
553 self
554 }
555
556 /// Override the declared `length`, which [`data`](Self::data) otherwise derives from `bytes`.
557 ///
558 /// For budget fixtures that hold every scalar at its widest legal value. A serve path has no reason
559 /// to call this: a frame whose `length` disagrees with its payload is a frame the reader distrusts.
560 pub fn with_declared_length(mut self, length: u64) -> Self {
561 self.length = length;
562 self
563 }
564
565 /// Serialize as a `u32` big-endian length prefix + JSON body (one framed frame on the stream).
566 ///
567 /// Fails with [`io::ErrorKind::InvalidData`] if [`bytes`](Self::bytes) exceeds
568 /// [`MAX_RANGE_FRAME_PAYLOAD`], if [`inclusion_proof`](Self::inclusion_proof) exceeds
569 /// [`MAX_INCLUSION_PROOF_B64`], or if the serialized body exceeds [`MAX_FRAMED_BODY`] for any
570 /// other reason (an unusually large `chunk_lens` page). A serving peer therefore cannot
571 /// emit a frame [`decode`](Self::decode) is required to reject: it splits its resource on
572 /// [`MAX_RANGE_FRAME_PAYLOAD`] or it learns about it here, at the send site, with the ceiling
573 /// named in the error.
574 ///
575 /// The payload is checked separately from the body because the body check alone is too weak: a
576 /// payload well over the ceiling still fits in [`MAX_FRAMED_BODY`] once base64'd when the frame
577 /// carries no metadata, so it would encode here and then overflow the moment the same span rode a
578 /// FIRST frame with a chunk table attached — a size-dependent, intermittent failure. One explicit
579 /// ceiling on `bytes` makes the limit the same for every frame. The proof is checked for the same
580 /// reason: it is **premise 2** of [`MAX_FIRST_FRAME_CHUNK_LENS`], and a premise only a test believes
581 /// is not a bound (#1655).
582 pub fn encode(&self) -> io::Result<Vec<u8>> {
583 if self.bytes.len() > MAX_RANGE_FRAME_PAYLOAD {
584 return Err(io::Error::new(
585 io::ErrorKind::InvalidData,
586 format!(
587 "RangeFrame payload {} exceeds MAX_RANGE_FRAME_PAYLOAD {MAX_RANGE_FRAME_PAYLOAD}; \
588 split the range into ceiling-sized frames",
589 self.bytes.len()
590 ),
591 ));
592 }
593 if let Some(proof) = &self.inclusion_proof {
594 if proof.len() > MAX_INCLUSION_PROOF_B64 {
595 return Err(io::Error::new(
596 io::ErrorKind::InvalidData,
597 format!(
598 "RangeFrame inclusion_proof {} exceeds MAX_INCLUSION_PROOF_B64 \
599 {MAX_INCLUSION_PROOF_B64}; the resource has no conforming range stream, so \
600 answer RANGE_METADATA_UNREPRESENTABLE rather than streaming frames",
601 proof.len()
602 ),
603 ));
604 }
605 }
606 encode_framed(self)
607 }
608 /// Read + decode one [`RangeFrame`] from `r`. Returns `Ok(None)` at clean end-of-stream (the
609 /// reader hit EOF on a frame boundary), so a consumer loops until `None` or `complete`.
610 pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Option<Self>> {
611 decode_framed_opt(r).await
612 }
613}
614
615/// Maximum length-prefixed frame BODY, in bytes — the one number both sides of the framing contract
616/// obey. A decoder rejects a longer body (guarding against a malicious length prefix forcing a huge
617/// allocation) and, since #1640, an encoder refuses to produce one, so a sender can never emit a
618/// frame a conforming receiver is required to reject.
619///
620/// This is a **shared byte-identical wire constant**: any second implementation of DIG peer framing
621/// — including `dig-node`'s own `write_framed`/`read_framed` — MUST use this exact value.
622pub const MAX_FRAMED_BODY: usize = 64 * 1024;
623
624/// Maximum raw [`RangeFrame::bytes`] length, in bytes, that a single frame may carry — the number a
625/// serving peer splits a resource on. **32 KiB.**
626///
627/// ## Why it is not ~48 KiB
628///
629/// `bytes` travels base64 (4 output bytes per 3 input bytes), so [`MAX_FRAMED_BODY`] of body would
630/// hold at most ~48 KiB of raw payload *if the payload were the only thing in the frame*. It is not.
631/// The frame is JSON and, per #1577, the FIRST frame of every range additionally carries `root`,
632/// `total_length`, `chunk_index`, a base64 `inclusion_proof`, and **the entire `chunk_lens` array of
633/// the whole resource** — whose size is driven by the RESOURCE, not by the frame.
634///
635/// So the ceiling is deliberately CONSERVATIVE rather than exact-fit:
636///
637/// | component | budget |
638/// |---|---|
639/// | base64 of 32 KiB of `bytes` | 43,692 B |
640/// | remaining allowance for JSON keys + `chunk_lens` + proof + root | 21,844 B |
641/// | **total** | **65,536 B** = [`MAX_FRAMED_BODY`] |
642///
643/// **Resist tightening it.** An exact-fit constant satisfies a naive round-trip test and then
644/// overflows in production on the first resource with a large chunk table — which is precisely the
645/// class of defect #1640 was.
646///
647/// ## The allowance does NOT cover every permitted resource
648///
649/// It is bounded, and the bound is [`MAX_FIRST_FRAME_CHUNK_LENS`] — read that before assuming this
650/// ceiling makes any legal range answerable. It does not.
651pub const MAX_RANGE_FRAME_PAYLOAD: usize = 32 * 1024;
652
653/// Maximum length of a base64 [`RangeFrame::inclusion_proof`], in bytes: **4,096** — enough for a
654/// 96-level merkle path, and **premise 2 of [`MAX_FIRST_FRAME_CHUNK_LENS`]**.
655///
656/// Published and enforced from 0.13.0. Before that it lived only in a test-local `const` and in prose,
657/// so the word GUARANTEED on the entry bound rested on a cap nothing checked: an 8 KiB proof made a
658/// *sub*-bound resource unsendable, which the published bound says cannot happen (#1655).
659/// [`RangeFrame::encode`] now refuses an over-cap proof at the SENDER, naming this constant.
660///
661/// Widening it is not a local decision: every byte added here comes out of
662/// [`MAX_FIRST_FRAME_CHUNK_LENS`], which must then be re-derived and re-published. A proof that does
663/// not fit is a resource with NO conforming range stream — the holder answers a structured
664/// `RANGE_METADATA_UNREPRESENTABLE` error rather than streaming an unverifiable frame.
665///
666/// This is a **shared byte-identical wire constant**: every implementation of DIG peer framing MUST
667/// use this exact value.
668pub const MAX_INCLUSION_PROOF_B64: usize = 4096;
669
670/// The `chunk_lens` entries a serving peer puts on ONE frame before it starts **paging the prologue**:
671/// **2,048**.
672///
673/// This is the SENDER's threshold, deliberately below the hard arithmetic ceiling
674/// [`MAX_FIRST_FRAME_CHUNK_LENS`]; the gap between the two is margin, and collapsing one into the other
675/// removes the only slack the budget has. A resource whose layout exceeds this is described by a
676/// **paged prologue**: successive frames each carry up to this many entries, and every page states the
677/// [`RangeFrame::chunk_lens_offset`] it begins at, so the reader reassembles one array of
678/// [`RangeFrame::chunk_count`] entries before it decrypts.
679///
680/// This is a **shared byte-identical wire constant**: every implementation of DIG peer framing MUST
681/// use this exact value.
682pub const MAX_CHUNK_LENS_PER_FRAME: usize = 2048;
683
684/// The largest `chunk_lens` array, in entries, that is GUARANTEED to fit on a first frame: **2,486**.
685///
686/// ## The four maxima this is derived against — all of them, simultaneously
687///
688/// Read these before using or changing the number. The same budget has been derived wrong three times,
689/// each time because ONE of these stayed implicit, and each time in the UNSAFE direction — too
690/// generous, which lets a conforming sender emit a frame the receiver must reject, i.e. exactly the
691/// defect #1640 exists to close.
692///
693/// 1. **Payload at its cap** — [`MAX_RANGE_FRAME_PAYLOAD`] (32,768 B raw → 43,692 B base64).
694/// 2. **Inclusion proof at its cap** — `MAX_INCLUSION_PROOF_B64` = 4,096 B. *(The 2,891 figure this
695/// constant replaces held the proof at ~1,400 B, so it only fit a small-proof frame: at 2,891
696/// entries the body is 64,199 B with no proof, 65,223 B at a 1,024 B proof, and 68,295 B at the
697/// real cap.)*
698/// 3. **Entries at their maximum decimal WIDTH** — a chunk may legally be 256 KiB, and `chunk_lens` is
699/// JSON decimal, so a worst-case entry is SIX digits. *(The 3,373 figure assumed five, from a wrong
700/// premise that the chunker targets 256 KiB; it is FastCDC target **64 KiB**, min 16 KiB, max
701/// 256 KiB — `digs crates/digstore-chunker/src/config.rs`.)*
702/// 4. **Every `u64` scalar at max width, including the 0.13.0 ones** — `offset`, `length`,
703/// `total_length`, `chunk_index`,
704/// `chunk_count`, `chunk_lens_offset`, each at 20 digits. Deliberately stricter than the
705/// protocol-tight widths (`length` ≤ 32,768 is 5 digits, `chunk_index` and `chunk_count`
706/// ≤ 1,048,576 are 7). That slack is given up on purpose: a bound that depends on a scalar
707/// happening to be small is a bound with a fifth implicit premise. Holding `chunk_count` at 7
708/// digits rather than 20 is worth 13 B — very nearly two entries — which is exactly the size of
709/// mistake this rule exists to prevent.
710///
711/// Measured on the **real 0.13.0 struct**, `chunk_count` and `chunk_lens_offset` included. It was
712/// pre-derived against those fields before they existed — via a mirror struct, which reserved a measured
713/// 76 B — and the reservation proved exact: the number did NOT move when the fields landed. That is the
714/// technique worth reusing, because this bound has **zero** slack (65,536 B, the cap exactly), so any
715/// further field WILL move it, and only the both-sides pin in `tests/framing_ceiling.rs` will say so.
716///
717/// ## Measured, never argued
718///
719/// Every figure here comes from serializing the real struct. Bodies at the four maxima, by entry
720/// width, on the 0.13.0 field set:
721///
722/// | `chunk_lens` entries | 5-digit entries | 6-digit entries |
723/// |---|---|---|
724/// | 2,048 | 60,422 B | 62,470 B |
725/// | **2,486** | 63,050 B | **65,536 B — the bound, exactly at the cap** |
726/// | 2,487 | 63,056 B | 65,543 B (over by 7) |
727/// | 2,495 | 63,104 B | 65,599 B (over by 63) |
728/// | 2,891 | 65,480 B | 68,371 B (over by 2,835) |
729/// | 4,096 | 72,710 B | 76,806 B |
730///
731/// At five-digit entries 2,900 would fit, but that is a property of the DATA, not of the protocol —
732/// never rely on it.
733///
734/// In resource terms 2,486 entries is about **155 MiB** at the 64 KiB target chunk size and about
735/// **38 MiB** at the 16 KiB minimum.
736///
737/// ## Not the same number as the sender's paging threshold
738///
739/// `MAX_CHUNK_LENS_PER_FRAME` = 2,048 is the 0.13.0 SENDER threshold — where a serve path starts
740/// paging the prologue (62,470 B at these same maxima). This constant is the hard ARITHMETIC ceiling.
741/// Two numbers doing two different jobs, and **the gap between them is deliberate margin**; do not
742/// collapse one into the other.
743///
744/// ## This is a CEILING, not the largest servable resource (#1640, resolved in 0.13.0)
745///
746/// The ecosystem permits resources far past this bound — `digstore-host` sets
747/// `MAX_MODULE_BYTES = 256 MiB` (~4,096 chunks at the 64 KiB target, the last row above) and
748/// dig-download accepts `MAX_MODULE_CHUNK_COUNT = 1,048,576` — and they are all servable, because the
749/// resource-scaling metadata no longer has to fit one frame: it rides a **paged prologue**
750/// ([`MAX_CHUNK_LENS_PER_FRAME`] entries per page, each located by
751/// [`RangeFrame::chunk_lens_offset`]). A 1,048,576-chunk layout costs roughly 7.3 MB across about 512
752/// pages, once per range stream — and nothing at all when the client sets
753/// [`RangeRequest::skip_layout`] because it already holds the layout.
754///
755/// So this constant governs how much layout ONE frame may state, which is why it is still the number a
756/// sender must respect and never the number a resource must fit under. Before 0.13.0 the two were the
757/// same thing, and a resource past the bound simply had no conforming first frame: [`RangeFrame::encode`]
758/// refused it, hard-failing LOUDLY at the sender rather than corrupting a read at the receiver — the
759/// correct half of the trade, but only half. Surrendering the payload entirely bought only so much room:
760/// past **8,727 entries** the metadata ALONE fills the body, which is why no value of
761/// [`MAX_RANGE_FRAME_PAYLOAD`] was ever the fix and the wire shape had to change.
762///
763/// Two workarounds remain forbidden, for reasons that did not go away: raising [`MAX_FRAMED_BODY`] is a
764/// RECEIVER bound (no sender may exceed it until every receiver is deployed, and no finite cap holds a
765/// million entries without abandoning the bounded-allocation property the cap exists for), and
766/// truncating `chunk_lens` is not an option either — it is a DECRYPT input, since per-chunk
767/// AES-256-GCM-SIV needs the whole array and a reader rejects any array whose sum differs from
768/// `total_length`. The ONE remaining unrepresentable input is an `inclusion_proof` over
769/// [`MAX_INCLUSION_PROOF_B64`]; such a resource has no conforming range stream at all, and the holder
770/// says so with a structured `RANGE_METADATA_UNREPRESENTABLE` error instead of streaming frames.
771pub const MAX_FIRST_FRAME_CHUNK_LENS: usize = 2_486;
772
773/// The largest `chunk_count` a reader will reassemble a layout for: **1,048,576 entries**.
774///
775/// It bounds the ONE allocation a paged prologue makes from a peer-declared number. At the ceiling the
776/// array is 1,048,576 × 8 B = **8 MB** of `u64`, which is the whole derivation: a bounded, affordable
777/// worst case per in-flight stream. Above it a reader refuses rather than sizes itself to a stranger's
778/// claim — a ~64-byte frame declaring a vast count is otherwise a memory-exhaustion primitive, and one
779/// such frame has already aborted a node.
780///
781/// It mirrors dig-download's `MAX_MODULE_CHUNK_COUNT`, deliberately: the two are the same bound seen
782/// from the transport and from the content layer, so they must not drift apart. At the ~64 KiB FastCDC
783/// target this ceiling is about 64 GiB of resource, far past any size the ecosystem permits
784/// (`digstore-host`'s `MAX_MODULE_BYTES` is 256 MiB), so it constrains a liar and never an honest holder.
785///
786/// This is a **shared byte-identical wire constant**: every implementation of DIG paged-prologue
787/// reassembly MUST use this exact value.
788pub const MAX_RESOURCE_CHUNK_COUNT: usize = 1_048_576;
789
790/// Why a [`ChunkLensAssembler`] refused a page, or refused to yield an array.
791///
792/// Each variant is ONE rejection with its own name. That is deliberate: a guard justified by a single
793/// attacker behaviour ("a liar sends a short page") is bypassed by the next variant of it — a *middle*
794/// page rather than a last one, a repeated page rather than a missing one, an offset off by one rather
795/// than wildly wrong. Naming each member of the class separately is what lets a consumer report, and a
796/// test pin, the specific rule that was broken instead of a generic "bad prologue".
797#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
798pub enum ChunkLensError {
799 /// The declared `chunk_count` exceeds [`MAX_RESOURCE_CHUNK_COUNT`]. Refused BEFORE allocating.
800 #[error(
801 "declared chunk_count {chunk_count} exceeds MAX_RESOURCE_CHUNK_COUNT {MAX_RESOURCE_CHUNK_COUNT}"
802 )]
803 ChunkCountTooLarge {
804 /// The count the sender declared.
805 chunk_count: usize,
806 },
807
808 /// The array for a legal `chunk_count` could not be allocated on this host. A resource ceiling is
809 /// not a memory guarantee, so the reservation is fallible and its failure is an error rather than an
810 /// abort.
811 #[error("cannot reserve a {chunk_count}-entry chunk_lens array ({bytes} bytes)")]
812 AllocationFailed {
813 /// The count that could not be reserved.
814 chunk_count: usize,
815 /// The reservation size in bytes.
816 bytes: usize,
817 },
818
819 /// A page carrying no entries. It fills nothing, so accepting it would let a sender stream frames
820 /// indefinitely without ever completing the prologue.
821 #[error("chunk_lens page at offset {offset} is empty")]
822 EmptyPage {
823 /// The offset the empty page claimed.
824 offset: u64,
825 },
826
827 /// A page carrying more than [`MAX_CHUNK_LENS_PER_FRAME`] entries — a wire-constant violation,
828 /// independent of where the page claims to sit.
829 #[error(
830 "chunk_lens page of {entries} entries exceeds MAX_CHUNK_LENS_PER_FRAME {MAX_CHUNK_LENS_PER_FRAME}"
831 )]
832 PageTooLarge {
833 /// The page's entry count.
834 entries: usize,
835 },
836
837 /// A `chunk_lens_offset` that is not a multiple of [`MAX_CHUNK_LENS_PER_FRAME`]. Pages are placed by
838 /// page index, so a misaligned page would straddle two slots and make occupancy unanswerable.
839 #[error(
840 "chunk_lens_offset {offset} is not a multiple of MAX_CHUNK_LENS_PER_FRAME {MAX_CHUNK_LENS_PER_FRAME}"
841 )]
842 MisalignedOffset {
843 /// The offset the page claimed.
844 offset: u64,
845 },
846
847 /// A `chunk_lens_offset` at or beyond the declared `chunk_count` — there is no such entry to fill.
848 #[error("chunk_lens_offset {offset} is beyond the declared chunk_count {chunk_count}")]
849 OffsetOutOfRange {
850 /// The offset the page claimed.
851 offset: u64,
852 /// The declared total entry count.
853 chunk_count: usize,
854 },
855
856 /// A page whose extent runs past the end of the array — an aligned offset is not sufficient, since a
857 /// full page at the LAST page's offset covers entries that do not exist.
858 #[error(
859 "chunk_lens page of {entries} entries at offset {offset} extends past the declared chunk_count \
860 {chunk_count}"
861 )]
862 PageExtendsPastEnd {
863 /// The offset the page claimed.
864 offset: u64,
865 /// The page's entry count.
866 entries: usize,
867 /// The declared total entry count.
868 chunk_count: usize,
869 },
870
871 /// A page that is neither full nor the tail. Pages must tile the array exactly, so a short
872 /// non-final page leaves a gap that no page-aligned page can ever fill; it is refused on arrival
873 /// rather than surfacing later as an unexplained incompleteness.
874 #[error(
875 "chunk_lens page at offset {offset} has {entries} entries; this page must have exactly {expected}"
876 )]
877 UnexpectedPageLength {
878 /// The offset the page claimed.
879 offset: u64,
880 /// The page's entry count.
881 entries: usize,
882 /// The entry count this page slot requires.
883 expected: usize,
884 },
885
886 /// A page for a slot that is already filled. A duplicate or overlapping page is a REJECT, never an
887 /// overwrite — otherwise the LAST sender of a page decides its contents.
888 #[error("a chunk_lens page at offset {offset} was already accepted")]
889 DuplicatePage {
890 /// The offset of the already-filled slot.
891 offset: u64,
892 },
893
894 /// The prologue ended short of `chunk_count`, so there is no array to yield.
895 #[error("incomplete chunk_lens prologue: {have} of {want} entries")]
896 Incomplete {
897 /// Entries actually accumulated.
898 have: usize,
899 /// Entries the declared `chunk_count` requires.
900 want: usize,
901 },
902}
903
904/// Reassembles one resource's `chunk_lens` array from the pages of a **paged prologue** — the
905/// decode-side mirror of [`RangeFrame::split_chunk_lens_pages`] (`SPEC.md` §5.1.1 is normative).
906///
907/// ## Why it lives here, beside the encoder
908///
909/// #1640 was an encode/decode asymmetry across a crate boundary: dig-nat capped DECODE and capped
910/// ENCODE at nothing, so every read past ~48 KiB failed. Both halves of the paging rule therefore live
911/// in this one module, next to the constants they are derived from. A second implementation of these
912/// placement, duplicate and completeness rules would be that same defect class waiting to recur.
913///
914/// ## Fail-closed, with no partial results
915///
916/// `chunk_lens` is a **DECRYPT** input: per-chunk AES-256-GCM-SIV needs the WHOLE array, and its
917/// entries must sum to the resource's `total_length`. A partial layout is therefore unusable rather
918/// than partially useful, which is why [`into_chunk_lens`](Self::into_chunk_lens) yields NOTHING until
919/// every page has landed, and why every irregular page is refused on arrival instead of adopted and
920/// checked later. The array is reserved fallibly and bounded by [`MAX_RESOURCE_CHUNK_COUNT`], so a
921/// declared count is never allowed to become an allocation this host cannot survive.
922///
923/// ```no_run
924/// use dig_nat::mux::{ChunkLensAssembler, RangeFrame};
925///
926/// # fn f(frames: Vec<RangeFrame>) -> Result<(), Box<dyn std::error::Error>> {
927/// let mut assembler = ChunkLensAssembler::new(5_000)?;
928/// for frame in frames {
929/// if let Some(page) = &frame.chunk_lens {
930/// assembler.accept_page(frame.chunk_lens_offset.unwrap_or(0), page)?;
931/// }
932/// }
933/// let chunk_lens = assembler.into_chunk_lens()?; // errors unless the prologue is complete
934/// # Ok(())
935/// # }
936/// ```
937#[derive(Debug, Clone)]
938pub struct ChunkLensAssembler {
939 /// The array under construction, pre-sized to `chunk_count`. Unfilled entries are zero, which is
940 /// never mistaken for data because occupancy is tracked separately in `filled`.
941 lens: Vec<u64>,
942 /// One flag per page slot — the authoritative occupancy record, so a repeated page is detectable
943 /// even when it carries the same bytes as the page already accepted.
944 filled: Vec<bool>,
945 /// How many page slots are filled, so completeness is answered without rescanning `filled`.
946 filled_pages: usize,
947}
948
949impl ChunkLensAssembler {
950 /// Start reassembling the layout of a resource whose frames declare `chunk_count` entries.
951 ///
952 /// Refuses a `chunk_count` above [`MAX_RESOURCE_CHUNK_COUNT`] BEFORE it allocates anything, and
953 /// reserves the array with `try_reserve_exact` so even a legal count cannot abort the process on a
954 /// host that cannot spare it.
955 pub fn new(chunk_count: usize) -> Result<Self, ChunkLensError> {
956 if chunk_count > MAX_RESOURCE_CHUNK_COUNT {
957 return Err(ChunkLensError::ChunkCountTooLarge { chunk_count });
958 }
959
960 let mut lens = Vec::new();
961 lens.try_reserve_exact(chunk_count)
962 .map_err(|_| ChunkLensError::AllocationFailed {
963 chunk_count,
964 bytes: chunk_count * std::mem::size_of::<u64>(),
965 })?;
966 lens.resize(chunk_count, 0);
967
968 let page_count = chunk_count.div_ceil(MAX_CHUNK_LENS_PER_FRAME);
969 let mut filled = Vec::new();
970 filled
971 .try_reserve_exact(page_count)
972 .map_err(|_| ChunkLensError::AllocationFailed {
973 chunk_count,
974 bytes: page_count,
975 })?;
976 filled.resize(page_count, false);
977
978 Ok(ChunkLensAssembler {
979 lens,
980 filled,
981 filled_pages: 0,
982 })
983 }
984
985 /// Place one `chunk_lens` page, which begins at entry `offset` of the resource's array.
986 ///
987 /// A page is accepted only if it satisfies EVERY rule of the paged prologue; each violation has its
988 /// own [`ChunkLensError`] variant, and a rejected page leaves the assembler exactly as it was.
989 ///
990 /// - non-empty, and at most [`MAX_CHUNK_LENS_PER_FRAME`] entries
991 /// - `offset` a multiple of [`MAX_CHUNK_LENS_PER_FRAME`], and inside the declared count
992 /// - an extent that stays within the array, and a length that exactly fills its page slot
993 /// - a slot not already filled — a duplicate or overlapping page is REJECTED, never an overwrite
994 pub fn accept_page(&mut self, offset: u64, page: &[u64]) -> Result<(), ChunkLensError> {
995 let chunk_count = self.lens.len();
996 let page_size = MAX_CHUNK_LENS_PER_FRAME as u64;
997
998 if page.is_empty() {
999 return Err(ChunkLensError::EmptyPage { offset });
1000 }
1001 if page.len() > MAX_CHUNK_LENS_PER_FRAME {
1002 return Err(ChunkLensError::PageTooLarge {
1003 entries: page.len(),
1004 });
1005 }
1006 // `%` rather than `u64::is_multiple_of`, which is not stable at this crate's MSRV (1.75).
1007 if offset % page_size != 0 {
1008 return Err(ChunkLensError::MisalignedOffset { offset });
1009 }
1010 // Compared as u64 so an offset near `u64::MAX` cannot wrap into a valid index on a 64-bit host.
1011 if offset >= chunk_count as u64 {
1012 return Err(ChunkLensError::OffsetOutOfRange {
1013 offset,
1014 chunk_count,
1015 });
1016 }
1017
1018 let start = offset as usize;
1019 let end = start + page.len();
1020 if end > chunk_count {
1021 return Err(ChunkLensError::PageExtendsPastEnd {
1022 offset,
1023 entries: page.len(),
1024 chunk_count,
1025 });
1026 }
1027 // Every page but the tail must be exactly full, or it leaves a gap no aligned page can fill.
1028 let expected = MAX_CHUNK_LENS_PER_FRAME.min(chunk_count - start);
1029 if page.len() != expected {
1030 return Err(ChunkLensError::UnexpectedPageLength {
1031 offset,
1032 entries: page.len(),
1033 expected,
1034 });
1035 }
1036
1037 let slot = start / MAX_CHUNK_LENS_PER_FRAME;
1038 if self.filled[slot] {
1039 return Err(ChunkLensError::DuplicatePage { offset });
1040 }
1041
1042 self.lens[start..end].copy_from_slice(page);
1043 self.filled[slot] = true;
1044 self.filled_pages += 1;
1045 Ok(())
1046 }
1047
1048 /// Whether every page slot has been filled, i.e. the layout is whole and decryptable.
1049 ///
1050 /// A zero-chunk resource is complete immediately: there is no layout to reassemble, and reporting it
1051 /// forever-incomplete would stall a stream over a resource already fully described.
1052 pub fn is_complete(&self) -> bool {
1053 self.filled_pages == self.filled.len()
1054 }
1055
1056 /// The reassembled `chunk_lens` array, or [`ChunkLensError::Incomplete`] if any page is missing.
1057 ///
1058 /// There is no partial variant of this call, by design: a layout short even one entry cannot decrypt
1059 /// the resource, and handing back what arrived so far would invite a caller to treat a hostile or
1060 /// truncated prologue as usable.
1061 pub fn into_chunk_lens(self) -> Result<Vec<u64>, ChunkLensError> {
1062 if !self.is_complete() {
1063 let have = self
1064 .filled
1065 .iter()
1066 .enumerate()
1067 .filter(|(_, filled)| **filled)
1068 .map(|(slot, _)| {
1069 let start = slot * MAX_CHUNK_LENS_PER_FRAME;
1070 MAX_CHUNK_LENS_PER_FRAME.min(self.lens.len() - start)
1071 })
1072 .sum();
1073 return Err(ChunkLensError::Incomplete {
1074 have,
1075 want: self.lens.len(),
1076 });
1077 }
1078 Ok(self.lens)
1079 }
1080}
1081
1082/// Serialize `value` as a `u32` big-endian length prefix + JSON body — the uniform framing for every
1083/// control message on a stream (availability + range preambles, and the range frames themselves).
1084///
1085/// Fails with [`io::ErrorKind::InvalidData`] if the body would exceed [`MAX_FRAMED_BODY`]. That check
1086/// is the whole point: the decoders below MUST reject such a body, so producing one is a bug at the
1087/// SENDER and belongs at the sender's call site — not as an opaque `InvalidData` surfacing on some
1088/// remote peer's read.
1089fn encode_framed<T: Serialize>(value: &T) -> io::Result<Vec<u8>> {
1090 let body =
1091 serde_json::to_vec(value).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
1092 if body.len() > MAX_FRAMED_BODY {
1093 return Err(io::Error::new(
1094 io::ErrorKind::InvalidData,
1095 format!(
1096 "framed body {} exceeds MAX_FRAMED_BODY {MAX_FRAMED_BODY}; split the payload on \
1097 MAX_RANGE_FRAME_PAYLOAD ({MAX_RANGE_FRAME_PAYLOAD})",
1098 body.len()
1099 ),
1100 ));
1101 }
1102 let mut out = Vec::with_capacity(4 + body.len());
1103 out.extend_from_slice(&(body.len() as u32).to_be_bytes());
1104 out.extend_from_slice(&body);
1105 Ok(out)
1106}
1107
1108/// Read + decode a length-prefixed JSON control message from `r`, bounded by [`MAX_FRAMED_BODY`].
1109async fn decode_framed<T: for<'de> Deserialize<'de>, R: AsyncRead + Unpin>(
1110 r: &mut R,
1111) -> io::Result<T> {
1112 let mut len_buf = [0u8; 4];
1113 r.read_exact(&mut len_buf).await?;
1114 let len = u32::from_be_bytes(len_buf) as usize;
1115 if len > MAX_FRAMED_BODY {
1116 return Err(io::Error::new(
1117 io::ErrorKind::InvalidData,
1118 "control message too large",
1119 ));
1120 }
1121 let mut body = vec![0u8; len];
1122 r.read_exact(&mut body).await?;
1123 serde_json::from_slice(&body).map_err(malformed_frame)
1124}
1125
1126/// Like [`decode_framed`] but returns `Ok(None)` on a CLEAN end-of-stream at a frame boundary (the
1127/// length prefix read hits immediate EOF), so a streaming consumer can loop until the stream ends.
1128async fn decode_framed_opt<T: for<'de> Deserialize<'de>, R: AsyncRead + Unpin>(
1129 r: &mut R,
1130) -> io::Result<Option<T>> {
1131 let mut len_buf = [0u8; 4];
1132 match r.read_exact(&mut len_buf).await {
1133 Ok(_) => {}
1134 Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
1135 Err(e) => return Err(e),
1136 }
1137 let len = u32::from_be_bytes(len_buf) as usize;
1138 if len > MAX_FRAMED_BODY {
1139 return Err(io::Error::new(
1140 io::ErrorKind::InvalidData,
1141 "control message too large",
1142 ));
1143 }
1144 let mut body = vec![0u8; len];
1145 r.read_exact(&mut body).await?;
1146 serde_json::from_slice(&body)
1147 .map(Some)
1148 .map_err(malformed_frame)
1149}
1150
1151/// Turn a `serde_json` decode failure on a PEER-SUPPLIED frame into an `io::Error` that says what
1152/// went wrong without repeating what the peer sent (#1674).
1153///
1154/// The naive `io::Error::new(InvalidData, e)` keeps serde's own message, and serde's message quotes
1155/// the offending input back — so every node that logged a malformed frame logged text of the
1156/// sender's choosing. Routing through [`SafeText::describing_json_error`] keeps the diagnosis
1157/// (category, line, column) and drops the quotation.
1158fn malformed_frame(error: serde_json::Error) -> io::Error {
1159 io::Error::new(
1160 io::ErrorKind::InvalidData,
1161 format!(
1162 "malformed frame: {}",
1163 SafeText::describing_json_error(&error)
1164 ),
1165 )
1166}
1167
1168/// One command to the yamux driver task. yamux 0.13 has no `Control` handle, so we drive the
1169/// [`yamux::Connection`] in a task and talk to it over this channel.
1170enum MuxCommand {
1171 /// Open a new outbound stream; the resulting [`yamux::Stream`] (or error) comes back on the sender.
1172 OpenOutbound(tokio::sync::oneshot::Sender<Result<yamux::Stream, String>>),
1173}
1174
1175/// A multiplexed session over one peer connection: open many concurrent logical [`PeerStream`]s.
1176///
1177/// yamux 0.13 exposes a poll-based [`yamux::Connection`] (no `Control` handle), so a background
1178/// driver task owns the connection and serves open-stream requests over a command channel; inbound
1179/// streams are surfaced on [`Self::inbound_rx`] for a serving node. Dropping the session closes the
1180/// command channel, which ends the driver and tears down the underlying byte stream.
1181pub struct PeerSession {
1182 cmd_tx: tokio::sync::mpsc::Sender<MuxCommand>,
1183 /// Inbound streams opened BY the peer (server role / bidirectional use). A pure client can
1184 /// ignore this; a serving node reads accepted range-request streams from here.
1185 inbound_rx: tokio::sync::mpsc::Receiver<PeerStream>,
1186 /// Set + notified when the driver task ends (the underlying byte stream closed — a clean close or
1187 /// a transport error). Observed via [`Self::closed_handle`] so fast-connect can detect a transport
1188 /// dying and fall back without holding the session lock.
1189 closed_flag: Arc<AtomicBool>,
1190 closed_notify: Arc<Notify>,
1191}
1192
1193/// A cheap, cloneable observer of a [`PeerSession`]'s underlying byte stream closing (the mux driver
1194/// task ending — a clean close OR a transport error). Fast-connect's promotion guard holds one to
1195/// detect the active transport dying and fall back to another tier, without locking the session.
1196#[derive(Clone)]
1197pub struct ClosedHandle {
1198 flag: Arc<AtomicBool>,
1199 notify: Arc<Notify>,
1200}
1201
1202impl ClosedHandle {
1203 /// Whether the session's transport has already closed.
1204 pub fn is_closed(&self) -> bool {
1205 self.flag.load(Ordering::Acquire)
1206 }
1207
1208 /// Resolve once the session's transport has closed (returns immediately if already closed).
1209 pub async fn closed(&self) {
1210 loop {
1211 if self.flag.load(Ordering::Acquire) {
1212 return;
1213 }
1214 // Arm the wait, then re-check the flag: the driver sets the flag BEFORE
1215 // `notify_waiters`, so a close that races this arming is caught by the recheck (no lost
1216 // wakeup).
1217 let notified = self.notify.notified();
1218 if self.flag.load(Ordering::Acquire) {
1219 return;
1220 }
1221 notified.await;
1222 }
1223 }
1224}
1225
1226impl std::fmt::Debug for PeerSession {
1227 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1228 f.debug_struct("PeerSession").finish_non_exhaustive()
1229 }
1230}
1231
1232impl PeerSession {
1233 /// Wrap an established mTLS byte stream in yamux as the **client** (outbound-stream opener) and
1234 /// spawn the driver. `io` is any tokio duplex stream (the mTLS [`tokio_rustls::client::TlsStream`]
1235 /// or, in tests, a loopback stream). Returns the session; open streams with
1236 /// [`Self::open_stream`] / [`Self::open_range_stream`].
1237 pub fn client<S>(io: S) -> Self
1238 where
1239 S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
1240 {
1241 Self::new(io, yamux::Mode::Client)
1242 }
1243
1244 /// Wrap an established byte stream in yamux as the **server** (accepts inbound streams). Inbound
1245 /// streams the peer opens are delivered via [`Self::accept_stream`]. Provided for symmetry + the
1246 /// serving side of tests.
1247 pub fn server<S>(io: S) -> Self
1248 where
1249 S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
1250 {
1251 Self::new(io, yamux::Mode::Server)
1252 }
1253
1254 fn new<S>(io: S, mode: yamux::Mode) -> Self
1255 where
1256 S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
1257 {
1258 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::channel::<MuxCommand>(64);
1259 let (inbound_tx, inbound_rx) = tokio::sync::mpsc::channel::<PeerStream>(64);
1260 let conn = yamux::Connection::new(io.compat(), yamux::Config::default(), mode);
1261 let closed_flag = Arc::new(AtomicBool::new(false));
1262 let closed_notify = Arc::new(Notify::new());
1263 tokio::spawn(drive_connection(
1264 conn,
1265 cmd_rx,
1266 inbound_tx,
1267 Arc::clone(&closed_flag),
1268 Arc::clone(&closed_notify),
1269 ));
1270 PeerSession {
1271 cmd_tx,
1272 inbound_rx,
1273 closed_flag,
1274 closed_notify,
1275 }
1276 }
1277
1278 /// A cloneable observer of this session's transport closing — see [`ClosedHandle`].
1279 pub fn closed_handle(&self) -> ClosedHandle {
1280 ClosedHandle {
1281 flag: Arc::clone(&self.closed_flag),
1282 notify: Arc::clone(&self.closed_notify),
1283 }
1284 }
1285
1286 /// Open a new outbound logical stream to the peer. Cheap — open as many as you need to run
1287 /// concurrent transfers without head-of-line blocking.
1288 pub async fn open_stream(&mut self) -> io::Result<PeerStream> {
1289 let (tx, rx) = tokio::sync::oneshot::channel();
1290 self.cmd_tx
1291 .send(MuxCommand::OpenOutbound(tx))
1292 .await
1293 .map_err(|_| io::Error::other("mux driver closed"))?;
1294 let stream = rx
1295 .await
1296 .map_err(|_| io::Error::other("mux driver dropped request"))?
1297 .map_err(io::Error::other)?;
1298 Ok(stream.compat())
1299 }
1300
1301 /// Accept the next inbound logical stream the peer opened (server side). Returns `None` when the
1302 /// connection has closed. A pure client never calls this.
1303 pub async fn accept_stream(&mut self) -> Option<PeerStream> {
1304 self.inbound_rx.recv().await
1305 }
1306
1307 /// Open a `dig.fetchRange` stream for `req`: opens a fresh logical stream, writes the
1308 /// [`RangeRequest`] preamble, and returns the stream for the caller to read [`RangeFrame`]s from
1309 /// (via [`RangeFrame::decode`]) as they arrive. The building block for multi-source parallel
1310 /// range downloads — open one of these per (peer, range) and read them concurrently.
1311 pub async fn open_range_stream(&mut self, req: &RangeRequest) -> io::Result<PeerStream> {
1312 let mut stream = self.open_stream().await?;
1313 stream.write_all(&req.encode()?).await?;
1314 stream.flush().await?;
1315 Ok(stream)
1316 }
1317
1318 /// **Availability pre-check** (`dig.getAvailability`) — ask the peer which of `items` it holds,
1319 /// BEFORE opening any range streams. Opens a short-lived control stream, writes the batched
1320 /// [`AvailabilityRequest`], reads the [`AvailabilityResponse`]. A multi-source downloader runs
1321 /// this against candidate peers and only range-fetches from holders — the normative flow is:
1322 /// discover peers → `query_availability` (batch) → fan byte-ranges across holders → verify each
1323 /// vs the chain-anchored root → retry a bad range from another holder → reassemble.
1324 pub async fn query_availability(
1325 &mut self,
1326 items: Vec<AvailabilityItem>,
1327 ) -> io::Result<AvailabilityResponse> {
1328 let req = AvailabilityRequest { items };
1329 let mut stream = self.open_stream().await?;
1330 stream.write_all(&req.encode()?).await?;
1331 stream.flush().await?;
1332 AvailabilityResponse::decode(&mut stream).await
1333 }
1334}
1335
1336/// Drive one yamux [`Connection`](yamux::Connection): concurrently service open-outbound commands
1337/// and surface inbound streams, until the command channel closes (session dropped) or the connection
1338/// errors. This is the task that replaces yamux 0.12's `Control`.
1339///
1340/// `T` is the futures-io view of the byte stream (a `tokio-util` [`Compat`] of the tokio mTLS
1341/// stream), since yamux operates on `futures::AsyncRead + AsyncWrite`.
1342async fn drive_connection<T>(
1343 mut conn: yamux::Connection<T>,
1344 mut cmd_rx: tokio::sync::mpsc::Receiver<MuxCommand>,
1345 inbound_tx: tokio::sync::mpsc::Sender<PeerStream>,
1346 closed_flag: Arc<AtomicBool>,
1347 closed_notify: Arc<Notify>,
1348) where
1349 T: futures::AsyncRead + futures::AsyncWrite + Send + Unpin + 'static,
1350{
1351 use std::future::poll_fn;
1352
1353 loop {
1354 tokio::select! {
1355 // An open-outbound request from the session.
1356 cmd = cmd_rx.recv() => {
1357 match cmd {
1358 Some(MuxCommand::OpenOutbound(reply)) => {
1359 let res = poll_fn(|cx| conn.poll_new_outbound(cx)).await;
1360 let _ = reply.send(res.map_err(|e| e.to_string()));
1361 }
1362 None => {
1363 // Session dropped — close the connection and end the driver.
1364 let _ = poll_fn(|cx| conn.poll_close(cx)).await;
1365 break;
1366 }
1367 }
1368 }
1369 // An inbound stream opened by the peer.
1370 inbound = poll_fn(|cx| conn.poll_next_inbound(cx)) => {
1371 match inbound {
1372 Some(Ok(stream)) => {
1373 // Deliver to a serving node; if no one is accepting, the stream is dropped.
1374 let _ = inbound_tx.try_send(stream.compat());
1375 }
1376 Some(Err(_)) | None => {
1377 // Connection closed / errored — end the driver.
1378 break;
1379 }
1380 }
1381 }
1382 }
1383 }
1384
1385 // The transport is gone: publish closure (flag BEFORE notify, so `ClosedHandle::closed`'s
1386 // arm-then-recheck can never miss the wakeup).
1387 closed_flag.store(true, Ordering::Release);
1388 closed_notify.notify_waiters();
1389}