Skip to main content

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
43/// One logical, bidirectional stream to the peer — a tokio [`AsyncRead`] + [`AsyncWrite`]. Reads
44/// deliver bytes incrementally as they arrive (streaming, with yamux-window backpressure); many
45/// [`PeerStream`]s coexist on one [`PeerSession`] without head-of-line blocking.
46///
47/// yamux streams are `futures` streams; this is the tokio-trait view via `tokio-util` compat.
48pub type PeerStream = Compat<yamux::Stream>;
49
50/// One item in a [`dig.getAvailability`](AvailabilityRequest) batch — a resource key at store, root,
51/// or capsule/resource granularity (inferred from which fields are present, per the L7 spec §9):
52/// `store_id` only → *has_store*; `+ root` → *has_root* (the capsule `store_id:root`); `+
53/// retrieval_key` → *has_resource*. Hashes are 64-hex.
54#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
55#[non_exhaustive]
56pub struct AvailabilityItem {
57    /// The store id (64-hex). Always present.
58    pub store_id: String,
59    /// The generation root (64-hex). Present for root/resource granularity.
60    #[serde(default, skip_serializing_if = "Option::is_none")]
61    pub root: Option<String>,
62    /// The resource retrieval key (64-hex). Present for resource granularity.
63    #[serde(default, skip_serializing_if = "Option::is_none")]
64    pub retrieval_key: Option<String>,
65}
66
67/// The **availability pre-check** (`dig.getAvailability`, L7 spec §9) — asked BEFORE any range fetch.
68/// A multi-source download batches candidate peers × items in one call each and only fans byte-range
69/// requests at peers that answer *available* — never opening range streams to peers that may not hold
70/// the content. A message-style control call over the mux'd mTLS connection.
71#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
72#[non_exhaustive]
73pub struct AvailabilityRequest {
74    /// The items to check, batched.
75    pub items: Vec<AvailabilityItem>,
76}
77
78/// One answer in an [`AvailabilityResponse`], positionally aligned with the request `items`.
79#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
80#[non_exhaustive]
81pub struct AvailabilityAnswer {
82    /// Whether the peer holds the queried item. Always present.
83    pub available: bool,
84    /// (store granularity) generation roots the peer holds for the store, newest-first.
85    #[serde(default, skip_serializing_if = "Option::is_none")]
86    pub roots: Option<Vec<String>>,
87    /// (root/resource granularity) the ciphertext length — lets the caller plan its ranges.
88    #[serde(default, skip_serializing_if = "Option::is_none")]
89    pub total_length: Option<u64>,
90    /// (root/resource granularity) the chunk count.
91    #[serde(default, skip_serializing_if = "Option::is_none")]
92    pub chunk_count: Option<u64>,
93    /// Whether the peer holds the FULL resource/capsule (`true`) or only part (`false`).
94    #[serde(default, skip_serializing_if = "Option::is_none")]
95    pub complete: Option<bool>,
96}
97
98/// The peer's answer to an [`AvailabilityRequest`]: one [`AvailabilityAnswer`] per queried item,
99/// positionally aligned with the request's `items`.
100#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
101#[non_exhaustive]
102pub struct AvailabilityResponse {
103    /// One answer per queried item, in request order.
104    pub items: Vec<AvailabilityAnswer>,
105}
106
107/// A byte-range request (`dig.fetchRange`, L7 spec §9) written at the start of a range-scoped stream.
108/// Identifies a resource (`store_id` + `retrieval_key` [+ `root`]) or a whole capsule
109/// (`capsule: true`, identified by `store_id` [+ `root`]) and the `[offset, offset+length)` range.
110#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
111#[non_exhaustive]
112pub struct RangeRequest {
113    /// The store id (64-hex).
114    pub store_id: String,
115    /// The resource retrieval key (64-hex). Omitted when `capsule` is true.
116    #[serde(default, skip_serializing_if = "Option::is_none")]
117    pub retrieval_key: Option<String>,
118    /// The generation root (64-hex). Optional — defaults to the chain-anchored tip.
119    #[serde(default, skip_serializing_if = "Option::is_none")]
120    pub root: Option<String>,
121    /// Fetch a whole capsule / `.dig` (identified by `store_id` [+ `root`]) rather than one resource.
122    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
123    pub capsule: bool,
124    /// Start offset (bytes) into the resource ciphertext. Default `0`.
125    #[serde(default)]
126    pub offset: u64,
127    /// Length (bytes) to return (widened to whole-chunk boundaries; clamped to the node window).
128    pub length: u64,
129    /// Suppress the resource-scaling layout metadata (`chunk_lens` + `inclusion_proof`) on this
130    /// stream's frames, because the client already holds the commitment for this `root`.
131    ///
132    /// A client that has already read the layout once — a resumed download, a second range of the same
133    /// resource, a parallel fetch from another holder — does not need it again, and for a large
134    /// resource re-sending it costs a paged prologue per stream. Absent or `false` preserves the
135    /// pre-0.13.0 behaviour, so an older holder that ignores this field is never broken by it: it
136    /// simply sends metadata the client discards. The fixed-size identity fields (`root`,
137    /// `total_length`, `chunk_count`, `chunk_index`) are NOT suppressed — they are what detects a
138    /// wrong-generation holder on arrival.
139    #[serde(default, skip_serializing_if = "Option::is_none")]
140    pub skip_layout: Option<bool>,
141}
142
143/// One streamed `dig.fetchRange` frame (L7 spec §8 framing). Frames arrive in ascending `offset`
144/// order and tile the requested range exactly; the caller reassembles by `offset` and stops on
145/// `complete`.
146///
147/// ## The metadata splits into two sets, and which frames carry them differs
148///
149/// Read this before implementing either side — the two halves have different rules, and conforming to
150/// the wrong half produces a verification miss on a multi-frame read rather than a clean failure.
151///
152/// - **Identity — fixed-size, on EVERY frame:** `root`, `total_length`, `chunk_count`, plus
153///   `chunk_index` where the window is chunk-aligned. These are what let a reader reject a
154///   wrong-generation or wrong-layout holder the moment a frame arrives, which is a property the
155///   once-per-stream set can never have. They cost a bounded number of bytes, so carrying them
156///   everywhere is cheap. Set them with [`with_identity`](Self::with_identity) +
157///   [`with_chunk_index`](Self::with_chunk_index).
158/// - **Prologue — resource-scaling, ONCE per range stream:** `chunk_lens` (paged, located by
159///   `chunk_lens_offset`) and `inclusion_proof`. Their size is a function of the RESOURCE rather than of
160///   the frame, so they ride the first frame or a paged prologue and **MUST NOT** be repeated on later
161///   frames. Set them with [`with_chunk_lens_page`](Self::with_chunk_lens_page) +
162///   [`with_inclusion_proof`](Self::with_inclusion_proof), or omit them entirely when the request set
163///   [`RangeRequest::skip_layout`].
164///
165/// Before 0.13.0 every one of these was "first frame only", because the whole layout had to fit one
166/// frame or the range was unservable (#1640). `SPEC.md` §5.1.1 is normative.
167#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
168#[non_exhaustive]
169pub struct RangeFrame {
170    /// This frame's start offset within the requested range.
171    pub offset: u64,
172    /// This frame's byte length.
173    pub length: u64,
174    /// The raw ciphertext bytes. On the wire they are **base64** (`base64_bytes`) — the canonical
175    /// `dig.fetchRange` frame encoding every producer emits.
176    #[serde(with = "base64_bytes")]
177    pub bytes: Vec<u8>,
178    /// Whether this is the final frame of the range.
179    pub complete: bool,
180    /// **(identity — every frame)** the full resource ciphertext length.
181    #[serde(default, skip_serializing_if = "Option::is_none")]
182    pub total_length: Option<u64>,
183    /// **(prologue — once per stream)** per-chunk ciphertext lengths of the whole resource, in order,
184    /// or ONE PAGE of them beginning at `chunk_lens_offset`.
185    ///
186    /// Resource-scaling, so it MUST NOT be repeated on later frames: on a later frame it was only a
187    /// redundant equality check on an array the client already held. Omitted entirely when the request
188    /// set [`RangeRequest::skip_layout`].
189    #[serde(default, skip_serializing_if = "Option::is_none")]
190    pub chunk_lens: Option<Vec<u64>>,
191    /// **(identity — every frame, where the window is chunk-aligned)** index into the resource's
192    /// `chunk_lens` array of the first chunk in THIS frame.
193    ///
194    /// Fixed-size, and about this frame rather than about the resource, so a chunk-aligned continuation
195    /// frame states it too — see [`with_chunk_index`](Self::with_chunk_index), which is deliberately
196    /// separate from the proof setter so stating alignment never costs a repeated proof.
197    #[serde(default, skip_serializing_if = "Option::is_none")]
198    pub chunk_index: Option<u64>,
199    /// **(prologue — once per stream)** merkle inclusion proof of the whole resource vs the generation
200    /// `root` (base64, relayed verbatim); `null`/absent for `capsule: true` (self-verifying on install).
201    ///
202    /// Resource-scaling and capped at [`MAX_INCLUSION_PROOF_B64`]. At 4,096 B it would consume the
203    /// entire remaining frame budget if repeated, which is the concrete reason "once per stream" is a
204    /// MUST NOT rather than a preference.
205    #[serde(default, skip_serializing_if = "Option::is_none")]
206    pub inclusion_proof: Option<String>,
207    /// **(identity — every frame)** the generation root (64-hex) this range is served from, and which
208    /// the inclusion proof is against.
209    #[serde(default, skip_serializing_if = "Option::is_none")]
210    pub root: Option<String>,
211    /// **(identity — every frame)** the resource's TOTAL chunk count — how many entries the
212    /// reassembled `chunk_lens` array has.
213    ///
214    /// Together with `root` and `total_length` it is what lets a reader detect a wrong-generation or
215    /// wrong-layout holder the moment a frame arrives. It is also how a reader sizes the array it is
216    /// paging in, and how it knows the prologue is complete.
217    #[serde(default, skip_serializing_if = "Option::is_none")]
218    pub chunk_count: Option<u64>,
219    /// **(prologue — once per page)** the index into the resource's `chunk_lens` array at which THIS
220    /// frame's `chunk_lens` page begins — how a **paged prologue** is located and reassembled.
221    ///
222    /// A resource whose layout exceeds [`MAX_CHUNK_LENS_PER_FRAME`] cannot state it on one frame, so
223    /// the sender pages it: successive frames each carry up to that many entries, stamped with the
224    /// offset they start at. A reader places each page at its offset and has the whole array once it
225    /// holds `chunk_count` entries. Absent means "this frame's `chunk_lens`, if any, begins at 0" — the
226    /// single-frame layout, which is the pre-0.13.0 shape.
227    #[serde(default, skip_serializing_if = "Option::is_none")]
228    pub chunk_lens_offset: Option<u64>,
229}
230
231/// Serde for [`RangeFrame::bytes`]: **base64** on the wire, raw `Vec<u8>` in Rust.
232///
233/// The `dig.fetchRange` frame is JSON, and the canonical wire type
234/// (`dig_rpc_protocol::types::RangeFrame`, "this window's ciphertext, base64") — and every real
235/// producer, including the dig-node peer serve path — encodes the window as a base64 STRING. Reading
236/// it with `serde_bytes` instead yielded the string's literal characters, so a served window arrived
237/// as its own base64 text and the reassembler rejected the frame (#1586, the read-leg blocker).
238///
239/// Deserialization is tolerant: a base64 string (canonical) OR a byte array (what an older dig-nat
240/// emitted) both decode, so a mixed-version peer is never dropped.
241mod base64_bytes {
242    use base64::Engine as _;
243    use serde::de::{SeqAccess, Visitor};
244    use serde::{Deserializer, Serializer};
245    use std::fmt;
246
247    pub fn serialize<S: Serializer>(bytes: &[u8], s: S) -> Result<S::Ok, S::Error> {
248        s.serialize_str(&base64::engine::general_purpose::STANDARD.encode(bytes))
249    }
250
251    pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<u8>, D::Error> {
252        d.deserialize_any(Base64OrArray)
253    }
254
255    struct Base64OrArray;
256
257    impl<'de> Visitor<'de> for Base64OrArray {
258        type Value = Vec<u8>;
259
260        fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
261            f.write_str("base64-encoded ciphertext (or a legacy byte array)")
262        }
263
264        fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<Self::Value, E> {
265            base64::engine::general_purpose::STANDARD
266                .decode(v)
267                .map_err(|e| E::custom(format!("range frame bytes are not valid base64: {e}")))
268        }
269
270        fn visit_bytes<E: serde::de::Error>(self, v: &[u8]) -> Result<Self::Value, E> {
271            Ok(v.to_vec())
272        }
273
274        fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value, A::Error> {
275            let mut out = Vec::with_capacity(seq.size_hint().unwrap_or_default());
276            while let Some(b) = seq.next_element::<u8>()? {
277                out.push(b);
278            }
279            Ok(out)
280        }
281    }
282}
283
284impl AvailabilityItem {
285    /// A *has_store* query — does the peer hold anything for this store?
286    pub fn store(store_id: impl Into<String>) -> Self {
287        AvailabilityItem {
288            store_id: store_id.into(),
289            root: None,
290            retrieval_key: None,
291        }
292    }
293
294    /// Narrow the query to one generation `root` — a *has_root* query for the capsule `store_id:root`.
295    pub fn with_root(mut self, root: impl Into<String>) -> Self {
296        self.root = Some(root.into());
297        self
298    }
299
300    /// Narrow the query to one resource — a *has_resource* query.
301    pub fn with_retrieval_key(mut self, retrieval_key: impl Into<String>) -> Self {
302        self.retrieval_key = Some(retrieval_key.into());
303        self
304    }
305}
306
307impl AvailabilityRequest {
308    /// A batched availability pre-check over `items`.
309    pub fn new(items: Vec<AvailabilityItem>) -> Self {
310        AvailabilityRequest { items }
311    }
312
313    /// Serialize as a `u32` big-endian length prefix + JSON body (the uniform control framing).
314    pub fn encode(&self) -> io::Result<Vec<u8>> {
315        encode_framed(self)
316    }
317    /// Read + decode an [`AvailabilityRequest`] from `r` (the peer/serving side).
318    pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
319        decode_framed(r).await
320    }
321}
322
323impl AvailabilityAnswer {
324    /// The peer holds the queried item. Describe WHAT it holds with the `with_*` setters.
325    pub fn available() -> Self {
326        AvailabilityAnswer::with_availability(true)
327    }
328
329    /// The peer does not hold the queried item — the whole answer, since no other field is meaningful.
330    pub fn unavailable() -> Self {
331        AvailabilityAnswer::with_availability(false)
332    }
333
334    fn with_availability(available: bool) -> Self {
335        AvailabilityAnswer {
336            available,
337            roots: None,
338            total_length: None,
339            chunk_count: None,
340            complete: None,
341        }
342    }
343
344    /// (store granularity) the generation roots held for the store, newest-first.
345    pub fn with_roots(mut self, roots: Vec<String>) -> Self {
346        self.roots = Some(roots);
347        self
348    }
349
350    /// (root/resource granularity) the resource's ciphertext length, so the caller can plan its ranges.
351    pub fn with_total_length(mut self, total_length: u64) -> Self {
352        self.total_length = Some(total_length);
353        self
354    }
355
356    /// (root/resource granularity) the resource's chunk count.
357    pub fn with_chunk_count(mut self, chunk_count: u64) -> Self {
358        self.chunk_count = Some(chunk_count);
359        self
360    }
361
362    /// Whether the peer holds the FULL resource/capsule rather than only part of it.
363    pub fn with_complete(mut self, complete: bool) -> Self {
364        self.complete = Some(complete);
365        self
366    }
367}
368
369impl AvailabilityResponse {
370    /// The answers to an [`AvailabilityRequest`], positionally aligned with its `items`.
371    pub fn new(items: Vec<AvailabilityAnswer>) -> Self {
372        AvailabilityResponse { items }
373    }
374
375    /// Serialize as a `u32` big-endian length prefix + JSON body.
376    pub fn encode(&self) -> io::Result<Vec<u8>> {
377        encode_framed(self)
378    }
379    /// Read + decode an [`AvailabilityResponse`] from `r` (the requesting side).
380    pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
381        decode_framed(r).await
382    }
383}
384
385impl RangeRequest {
386    /// A range request for a content resource (`store_id` + `retrieval_key`).
387    pub fn resource(
388        store_id: impl Into<String>,
389        retrieval_key: impl Into<String>,
390        offset: u64,
391        length: u64,
392    ) -> Self {
393        RangeRequest::resource_or_capsule(
394            store_id,
395            Some(retrieval_key.into()),
396            false,
397            offset,
398            length,
399        )
400    }
401
402    /// A range request for a whole capsule / `.dig`, identified by `store_id` (plus a `root` if the
403    /// caller pins a generation).
404    pub fn capsule(store_id: impl Into<String>, offset: u64, length: u64) -> Self {
405        RangeRequest::resource_or_capsule(store_id, None, true, offset, length)
406    }
407
408    fn resource_or_capsule(
409        store_id: impl Into<String>,
410        retrieval_key: Option<String>,
411        capsule: bool,
412        offset: u64,
413        length: u64,
414    ) -> Self {
415        RangeRequest {
416            store_id: store_id.into(),
417            retrieval_key,
418            root: None,
419            capsule,
420            offset,
421            length,
422            skip_layout: None,
423        }
424    }
425
426    /// Pin the request to one generation `root` (64-hex) rather than the chain-anchored tip.
427    pub fn with_root(mut self, root: impl Into<String>) -> Self {
428        self.root = Some(root.into());
429        self
430    }
431
432    /// Ask the holder to omit the resource-scaling layout metadata, because this client already holds
433    /// the commitment for this `root`. See [`RangeRequest::skip_layout`].
434    pub fn with_skip_layout(mut self, skip_layout: bool) -> Self {
435        self.skip_layout = Some(skip_layout);
436        self
437    }
438
439    /// Serialize as a `u32` big-endian length prefix + JSON body — the preamble a peer reads to learn
440    /// the resource + range before streaming the frames.
441    pub fn encode(&self) -> io::Result<Vec<u8>> {
442        encode_framed(self)
443    }
444    /// Read + decode a [`RangeRequest`] preamble from `r` (the serving side of a range stream).
445    pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
446        decode_framed(r).await
447    }
448}
449
450impl RangeFrame {
451    /// A **data frame**: `bytes` at `offset`, with no metadata — the shape of every continuation frame.
452    ///
453    /// `length` is derived from `bytes`, since a frame that misdeclares its own payload length is never
454    /// what a serve path wants. The metadata a FIRST frame or a prologue page carries is layered on with
455    /// the `with_*` setters, so a continuation frame cannot accidentally claim a layout it is not
456    /// stating — the two shapes are different call chains, not one call with a pile of `None`s.
457    pub fn data(offset: u64, bytes: Vec<u8>) -> Self {
458        RangeFrame {
459            offset,
460            length: bytes.len() as u64,
461            bytes,
462            complete: false,
463            total_length: None,
464            chunk_lens: None,
465            chunk_index: None,
466            inclusion_proof: None,
467            root: None,
468            chunk_count: None,
469            chunk_lens_offset: None,
470        }
471    }
472
473    /// Mark this as the final frame of the range.
474    pub fn with_complete(mut self, complete: bool) -> Self {
475        self.complete = complete;
476        self
477    }
478
479    /// The fixed-size identity metadata EVERY frame of a range should carry: the generation `root`
480    /// (64-hex) the range is served from, the resource's ciphertext `total_length`, and its
481    /// `chunk_count`.
482    ///
483    /// These three are what let a reader reject a wrong-generation or wrong-layout holder the moment a
484    /// frame arrives — which the resource-scaling metadata never could, since it arrives once. They are
485    /// fixed-size, so carrying them everywhere costs a bounded number of bytes.
486    pub fn with_identity(
487        mut self,
488        root: impl Into<String>,
489        total_length: u64,
490        chunk_count: u64,
491    ) -> Self {
492        self.root = Some(root.into());
493        self.total_length = Some(total_length);
494        self.chunk_count = Some(chunk_count);
495        self
496    }
497
498    /// A page of the resource's `chunk_lens` array, beginning at entry `chunk_lens_offset`.
499    ///
500    /// Call it once with offset `0` for a layout that fits one frame (at most
501    /// [`MAX_CHUNK_LENS_PER_FRAME`] entries), or once per page of a **paged prologue**. `chunk_lens` is
502    /// a DECRYPT input — per-chunk AES-GCM-SIV needs the WHOLE array, and a reader rejects an array
503    /// whose sum differs from `total_length` — so a page is only ever useful as part of a complete set.
504    pub fn with_chunk_lens_page(mut self, chunk_lens_offset: u64, chunk_lens: Vec<u64>) -> Self {
505        self.chunk_lens_offset = Some(chunk_lens_offset);
506        self.chunk_lens = Some(chunk_lens);
507        self
508    }
509
510    /// Where this frame starts in the resource's chunk sequence — the index into `chunk_lens` of its
511    /// first chunk, for a chunk-aligned window.
512    ///
513    /// Fixed-size **identity**, so a chunk-aligned CONTINUATION frame states it too, not only the first
514    /// frame. It has its own setter for exactly that reason: it used to be a parameter of
515    /// [`with_inclusion_proof`](Self::with_inclusion_proof), which meant the one field every reader wants
516    /// on every frame could not be stated without repeating a once-per-stream proof — a `SPEC.md` §5.1.1
517    /// MUST NOT, and 4,096 B per frame against a budget with zero slack.
518    pub fn with_chunk_index(mut self, chunk_index: u64) -> Self {
519        self.chunk_index = Some(chunk_index);
520        self
521    }
522
523    /// The merkle inclusion proof of the whole resource against the generation `root` (base64, relayed
524    /// verbatim).
525    ///
526    /// Resource-scaling, so it rides the first frame or the prologue — once per range stream, never
527    /// repeated. Base64 longer than [`MAX_INCLUSION_PROOF_B64`] is refused by
528    /// [`encode`](Self::encode): such a resource has no conforming range stream at all, and the holder
529    /// must say so rather than stream frames a reader cannot verify.
530    pub fn with_inclusion_proof(mut self, inclusion_proof: impl Into<String>) -> Self {
531        self.inclusion_proof = Some(inclusion_proof.into());
532        self
533    }
534
535    /// Override the declared `length`, which [`data`](Self::data) otherwise derives from `bytes`.
536    ///
537    /// For budget fixtures that hold every scalar at its widest legal value. A serve path has no reason
538    /// to call this: a frame whose `length` disagrees with its payload is a frame the reader distrusts.
539    pub fn with_declared_length(mut self, length: u64) -> Self {
540        self.length = length;
541        self
542    }
543
544    /// Serialize as a `u32` big-endian length prefix + JSON body (one framed frame on the stream).
545    ///
546    /// Fails with [`io::ErrorKind::InvalidData`] if [`bytes`](Self::bytes) exceeds
547    /// [`MAX_RANGE_FRAME_PAYLOAD`], if [`inclusion_proof`](Self::inclusion_proof) exceeds
548    /// [`MAX_INCLUSION_PROOF_B64`], or if the serialized body exceeds [`MAX_FRAMED_BODY`] for any
549    /// other reason (an unusually large `chunk_lens` page). A serving peer therefore cannot
550    /// emit a frame [`decode`](Self::decode) is required to reject: it splits its resource on
551    /// [`MAX_RANGE_FRAME_PAYLOAD`] or it learns about it here, at the send site, with the ceiling
552    /// named in the error.
553    ///
554    /// The payload is checked separately from the body because the body check alone is too weak: a
555    /// payload well over the ceiling still fits in [`MAX_FRAMED_BODY`] once base64'd when the frame
556    /// carries no metadata, so it would encode here and then overflow the moment the same span rode a
557    /// FIRST frame with a chunk table attached — a size-dependent, intermittent failure. One explicit
558    /// ceiling on `bytes` makes the limit the same for every frame. The proof is checked for the same
559    /// reason: it is **premise 2** of [`MAX_FIRST_FRAME_CHUNK_LENS`], and a premise only a test believes
560    /// is not a bound (#1655).
561    pub fn encode(&self) -> io::Result<Vec<u8>> {
562        if self.bytes.len() > MAX_RANGE_FRAME_PAYLOAD {
563            return Err(io::Error::new(
564                io::ErrorKind::InvalidData,
565                format!(
566                    "RangeFrame payload {} exceeds MAX_RANGE_FRAME_PAYLOAD {MAX_RANGE_FRAME_PAYLOAD}; \
567                     split the range into ceiling-sized frames",
568                    self.bytes.len()
569                ),
570            ));
571        }
572        if let Some(proof) = &self.inclusion_proof {
573            if proof.len() > MAX_INCLUSION_PROOF_B64 {
574                return Err(io::Error::new(
575                    io::ErrorKind::InvalidData,
576                    format!(
577                        "RangeFrame inclusion_proof {} exceeds MAX_INCLUSION_PROOF_B64 \
578                         {MAX_INCLUSION_PROOF_B64}; the resource has no conforming range stream, so \
579                         answer RANGE_METADATA_UNREPRESENTABLE rather than streaming frames",
580                        proof.len()
581                    ),
582                ));
583            }
584        }
585        encode_framed(self)
586    }
587    /// Read + decode one [`RangeFrame`] from `r`. Returns `Ok(None)` at clean end-of-stream (the
588    /// reader hit EOF on a frame boundary), so a consumer loops until `None` or `complete`.
589    pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Option<Self>> {
590        decode_framed_opt(r).await
591    }
592}
593
594/// Maximum length-prefixed frame BODY, in bytes — the one number both sides of the framing contract
595/// obey. A decoder rejects a longer body (guarding against a malicious length prefix forcing a huge
596/// allocation) and, since #1640, an encoder refuses to produce one, so a sender can never emit a
597/// frame a conforming receiver is required to reject.
598///
599/// This is a **shared byte-identical wire constant**: any second implementation of DIG peer framing
600/// — including `dig-node`'s own `write_framed`/`read_framed` — MUST use this exact value.
601pub const MAX_FRAMED_BODY: usize = 64 * 1024;
602
603/// Maximum raw [`RangeFrame::bytes`] length, in bytes, that a single frame may carry — the number a
604/// serving peer splits a resource on. **32 KiB.**
605///
606/// ## Why it is not ~48 KiB
607///
608/// `bytes` travels base64 (4 output bytes per 3 input bytes), so [`MAX_FRAMED_BODY`] of body would
609/// hold at most ~48 KiB of raw payload *if the payload were the only thing in the frame*. It is not.
610/// The frame is JSON and, per #1577, the FIRST frame of every range additionally carries `root`,
611/// `total_length`, `chunk_index`, a base64 `inclusion_proof`, and **the entire `chunk_lens` array of
612/// the whole resource** — whose size is driven by the RESOURCE, not by the frame.
613///
614/// So the ceiling is deliberately CONSERVATIVE rather than exact-fit:
615///
616/// | component | budget |
617/// |---|---|
618/// | base64 of 32 KiB of `bytes` | 43,692 B |
619/// | remaining allowance for JSON keys + `chunk_lens` + proof + root | 21,844 B |
620/// | **total** | **65,536 B** = [`MAX_FRAMED_BODY`] |
621///
622/// **Resist tightening it.** An exact-fit constant satisfies a naive round-trip test and then
623/// overflows in production on the first resource with a large chunk table — which is precisely the
624/// class of defect #1640 was.
625///
626/// ## The allowance does NOT cover every permitted resource
627///
628/// It is bounded, and the bound is [`MAX_FIRST_FRAME_CHUNK_LENS`] — read that before assuming this
629/// ceiling makes any legal range answerable. It does not.
630pub const MAX_RANGE_FRAME_PAYLOAD: usize = 32 * 1024;
631
632/// Maximum length of a base64 [`RangeFrame::inclusion_proof`], in bytes: **4,096** — enough for a
633/// 96-level merkle path, and **premise 2 of [`MAX_FIRST_FRAME_CHUNK_LENS`]**.
634///
635/// Published and enforced from 0.13.0. Before that it lived only in a test-local `const` and in prose,
636/// so the word GUARANTEED on the entry bound rested on a cap nothing checked: an 8 KiB proof made a
637/// *sub*-bound resource unsendable, which the published bound says cannot happen (#1655).
638/// [`RangeFrame::encode`] now refuses an over-cap proof at the SENDER, naming this constant.
639///
640/// Widening it is not a local decision: every byte added here comes out of
641/// [`MAX_FIRST_FRAME_CHUNK_LENS`], which must then be re-derived and re-published. A proof that does
642/// not fit is a resource with NO conforming range stream — the holder answers a structured
643/// `RANGE_METADATA_UNREPRESENTABLE` error rather than streaming an unverifiable frame.
644///
645/// This is a **shared byte-identical wire constant**: every implementation of DIG peer framing MUST
646/// use this exact value.
647pub const MAX_INCLUSION_PROOF_B64: usize = 4096;
648
649/// The `chunk_lens` entries a serving peer puts on ONE frame before it starts **paging the prologue**:
650/// **2,048**.
651///
652/// This is the SENDER's threshold, deliberately below the hard arithmetic ceiling
653/// [`MAX_FIRST_FRAME_CHUNK_LENS`]; the gap between the two is margin, and collapsing one into the other
654/// removes the only slack the budget has. A resource whose layout exceeds this is described by a
655/// **paged prologue**: successive frames each carry up to this many entries, and every page states the
656/// [`RangeFrame::chunk_lens_offset`] it begins at, so the reader reassembles one array of
657/// [`RangeFrame::chunk_count`] entries before it decrypts.
658///
659/// This is a **shared byte-identical wire constant**: every implementation of DIG peer framing MUST
660/// use this exact value.
661pub const MAX_CHUNK_LENS_PER_FRAME: usize = 2048;
662
663/// The largest `chunk_lens` array, in entries, that is GUARANTEED to fit on a first frame: **2,486**.
664///
665/// ## The four maxima this is derived against — all of them, simultaneously
666///
667/// Read these before using or changing the number. The same budget has been derived wrong three times,
668/// each time because ONE of these stayed implicit, and each time in the UNSAFE direction — too
669/// generous, which lets a conforming sender emit a frame the receiver must reject, i.e. exactly the
670/// defect #1640 exists to close.
671///
672/// 1. **Payload at its cap** — [`MAX_RANGE_FRAME_PAYLOAD`] (32,768 B raw → 43,692 B base64).
673/// 2. **Inclusion proof at its cap** — `MAX_INCLUSION_PROOF_B64` = 4,096 B. *(The 2,891 figure this
674///    constant replaces held the proof at ~1,400 B, so it only fit a small-proof frame: at 2,891
675///    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
676///    real cap.)*
677/// 3. **Entries at their maximum decimal WIDTH** — a chunk may legally be 256 KiB, and `chunk_lens` is
678///    JSON decimal, so a worst-case entry is SIX digits. *(The 3,373 figure assumed five, from a wrong
679///    premise that the chunker targets 256 KiB; it is FastCDC target **64 KiB**, min 16 KiB, max
680///    256 KiB — `digs crates/digstore-chunker/src/config.rs`.)*
681/// 4. **Every `u64` scalar at max width, including the 0.13.0 ones** — `offset`, `length`,
682///    `total_length`, `chunk_index`,
683///    `chunk_count`, `chunk_lens_offset`, each at 20 digits. Deliberately stricter than the
684///    protocol-tight widths (`length` ≤ 32,768 is 5 digits, `chunk_index` and `chunk_count`
685///    ≤ 1,048,576 are 7). That slack is given up on purpose: a bound that depends on a scalar
686///    happening to be small is a bound with a fifth implicit premise. Holding `chunk_count` at 7
687///    digits rather than 20 is worth 13 B — very nearly two entries — which is exactly the size of
688///    mistake this rule exists to prevent.
689///
690/// Measured on the **real 0.13.0 struct**, `chunk_count` and `chunk_lens_offset` included. It was
691/// pre-derived against those fields before they existed — via a mirror struct, which reserved a measured
692/// 76 B — and the reservation proved exact: the number did NOT move when the fields landed. That is the
693/// technique worth reusing, because this bound has **zero** slack (65,536 B, the cap exactly), so any
694/// further field WILL move it, and only the both-sides pin in `tests/framing_ceiling.rs` will say so.
695///
696/// ## Measured, never argued
697///
698/// Every figure here comes from serializing the real struct. Bodies at the four maxima, by entry
699/// width, on the 0.13.0 field set:
700///
701/// | `chunk_lens` entries | 5-digit entries | 6-digit entries |
702/// |---|---|---|
703/// | 2,048 | 60,422 B | 62,470 B |
704/// | **2,486** | 63,050 B | **65,536 B — the bound, exactly at the cap** |
705/// | 2,487 | 63,056 B | 65,543 B (over by 7) |
706/// | 2,495 | 63,104 B | 65,599 B (over by 63) |
707/// | 2,891 | 65,480 B | 68,371 B (over by 2,835) |
708/// | 4,096 | 72,710 B | 76,806 B |
709///
710/// At five-digit entries 2,900 would fit, but that is a property of the DATA, not of the protocol —
711/// never rely on it.
712///
713/// In resource terms 2,486 entries is about **155 MiB** at the 64 KiB target chunk size and about
714/// **38 MiB** at the 16 KiB minimum.
715///
716/// ## Not the same number as the sender's paging threshold
717///
718/// `MAX_CHUNK_LENS_PER_FRAME` = 2,048 is the 0.13.0 SENDER threshold — where a serve path starts
719/// paging the prologue (62,470 B at these same maxima). This constant is the hard ARITHMETIC ceiling.
720/// Two numbers doing two different jobs, and **the gap between them is deliberate margin**; do not
721/// collapse one into the other.
722///
723/// ## This is a CEILING, not the largest servable resource (#1640, resolved in 0.13.0)
724///
725/// The ecosystem permits resources far past this bound — `digstore-host` sets
726/// `MAX_MODULE_BYTES = 256 MiB` (~4,096 chunks at the 64 KiB target, the last row above) and
727/// dig-download accepts `MAX_MODULE_CHUNK_COUNT = 1,048,576` — and they are all servable, because the
728/// resource-scaling metadata no longer has to fit one frame: it rides a **paged prologue**
729/// ([`MAX_CHUNK_LENS_PER_FRAME`] entries per page, each located by
730/// [`RangeFrame::chunk_lens_offset`]). A 1,048,576-chunk layout costs roughly 7.3 MB across about 512
731/// pages, once per range stream — and nothing at all when the client sets
732/// [`RangeRequest::skip_layout`] because it already holds the layout.
733///
734/// So this constant governs how much layout ONE frame may state, which is why it is still the number a
735/// sender must respect and never the number a resource must fit under. Before 0.13.0 the two were the
736/// same thing, and a resource past the bound simply had no conforming first frame: [`RangeFrame::encode`]
737/// refused it, hard-failing LOUDLY at the sender rather than corrupting a read at the receiver — the
738/// correct half of the trade, but only half. Surrendering the payload entirely bought only so much room:
739/// past **8,727 entries** the metadata ALONE fills the body, which is why no value of
740/// [`MAX_RANGE_FRAME_PAYLOAD`] was ever the fix and the wire shape had to change.
741///
742/// Two workarounds remain forbidden, for reasons that did not go away: raising [`MAX_FRAMED_BODY`] is a
743/// RECEIVER bound (no sender may exceed it until every receiver is deployed, and no finite cap holds a
744/// million entries without abandoning the bounded-allocation property the cap exists for), and
745/// truncating `chunk_lens` is not an option either — it is a DECRYPT input, since per-chunk
746/// AES-256-GCM-SIV needs the whole array and a reader rejects any array whose sum differs from
747/// `total_length`. The ONE remaining unrepresentable input is an `inclusion_proof` over
748/// [`MAX_INCLUSION_PROOF_B64`]; such a resource has no conforming range stream at all, and the holder
749/// says so with a structured `RANGE_METADATA_UNREPRESENTABLE` error instead of streaming frames.
750pub const MAX_FIRST_FRAME_CHUNK_LENS: usize = 2_486;
751
752/// Serialize `value` as a `u32` big-endian length prefix + JSON body — the uniform framing for every
753/// control message on a stream (availability + range preambles, and the range frames themselves).
754///
755/// Fails with [`io::ErrorKind::InvalidData`] if the body would exceed [`MAX_FRAMED_BODY`]. That check
756/// is the whole point: the decoders below MUST reject such a body, so producing one is a bug at the
757/// SENDER and belongs at the sender's call site — not as an opaque `InvalidData` surfacing on some
758/// remote peer's read.
759fn encode_framed<T: Serialize>(value: &T) -> io::Result<Vec<u8>> {
760    let body =
761        serde_json::to_vec(value).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
762    if body.len() > MAX_FRAMED_BODY {
763        return Err(io::Error::new(
764            io::ErrorKind::InvalidData,
765            format!(
766                "framed body {} exceeds MAX_FRAMED_BODY {MAX_FRAMED_BODY}; split the payload on \
767                 MAX_RANGE_FRAME_PAYLOAD ({MAX_RANGE_FRAME_PAYLOAD})",
768                body.len()
769            ),
770        ));
771    }
772    let mut out = Vec::with_capacity(4 + body.len());
773    out.extend_from_slice(&(body.len() as u32).to_be_bytes());
774    out.extend_from_slice(&body);
775    Ok(out)
776}
777
778/// Read + decode a length-prefixed JSON control message from `r`, bounded by [`MAX_FRAMED_BODY`].
779async fn decode_framed<T: for<'de> Deserialize<'de>, R: AsyncRead + Unpin>(
780    r: &mut R,
781) -> io::Result<T> {
782    let mut len_buf = [0u8; 4];
783    r.read_exact(&mut len_buf).await?;
784    let len = u32::from_be_bytes(len_buf) as usize;
785    if len > MAX_FRAMED_BODY {
786        return Err(io::Error::new(
787            io::ErrorKind::InvalidData,
788            "control message too large",
789        ));
790    }
791    let mut body = vec![0u8; len];
792    r.read_exact(&mut body).await?;
793    serde_json::from_slice(&body).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
794}
795
796/// Like [`decode_framed`] but returns `Ok(None)` on a CLEAN end-of-stream at a frame boundary (the
797/// length prefix read hits immediate EOF), so a streaming consumer can loop until the stream ends.
798async fn decode_framed_opt<T: for<'de> Deserialize<'de>, R: AsyncRead + Unpin>(
799    r: &mut R,
800) -> io::Result<Option<T>> {
801    let mut len_buf = [0u8; 4];
802    match r.read_exact(&mut len_buf).await {
803        Ok(_) => {}
804        Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
805        Err(e) => return Err(e),
806    }
807    let len = u32::from_be_bytes(len_buf) as usize;
808    if len > MAX_FRAMED_BODY {
809        return Err(io::Error::new(
810            io::ErrorKind::InvalidData,
811            "control message too large",
812        ));
813    }
814    let mut body = vec![0u8; len];
815    r.read_exact(&mut body).await?;
816    serde_json::from_slice(&body)
817        .map(Some)
818        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
819}
820
821/// One command to the yamux driver task. yamux 0.13 has no `Control` handle, so we drive the
822/// [`yamux::Connection`] in a task and talk to it over this channel.
823enum MuxCommand {
824    /// Open a new outbound stream; the resulting [`yamux::Stream`] (or error) comes back on the sender.
825    OpenOutbound(tokio::sync::oneshot::Sender<Result<yamux::Stream, String>>),
826}
827
828/// A multiplexed session over one peer connection: open many concurrent logical [`PeerStream`]s.
829///
830/// yamux 0.13 exposes a poll-based [`yamux::Connection`] (no `Control` handle), so a background
831/// driver task owns the connection and serves open-stream requests over a command channel; inbound
832/// streams are surfaced on [`Self::inbound_rx`] for a serving node. Dropping the session closes the
833/// command channel, which ends the driver and tears down the underlying byte stream.
834pub struct PeerSession {
835    cmd_tx: tokio::sync::mpsc::Sender<MuxCommand>,
836    /// Inbound streams opened BY the peer (server role / bidirectional use). A pure client can
837    /// ignore this; a serving node reads accepted range-request streams from here.
838    inbound_rx: tokio::sync::mpsc::Receiver<PeerStream>,
839    /// Set + notified when the driver task ends (the underlying byte stream closed — a clean close or
840    /// a transport error). Observed via [`Self::closed_handle`] so fast-connect can detect a transport
841    /// dying and fall back without holding the session lock.
842    closed_flag: Arc<AtomicBool>,
843    closed_notify: Arc<Notify>,
844}
845
846/// A cheap, cloneable observer of a [`PeerSession`]'s underlying byte stream closing (the mux driver
847/// task ending — a clean close OR a transport error). Fast-connect's promotion guard holds one to
848/// detect the active transport dying and fall back to another tier, without locking the session.
849#[derive(Clone)]
850pub struct ClosedHandle {
851    flag: Arc<AtomicBool>,
852    notify: Arc<Notify>,
853}
854
855impl ClosedHandle {
856    /// Whether the session's transport has already closed.
857    pub fn is_closed(&self) -> bool {
858        self.flag.load(Ordering::Acquire)
859    }
860
861    /// Resolve once the session's transport has closed (returns immediately if already closed).
862    pub async fn closed(&self) {
863        loop {
864            if self.flag.load(Ordering::Acquire) {
865                return;
866            }
867            // Arm the wait, then re-check the flag: the driver sets the flag BEFORE
868            // `notify_waiters`, so a close that races this arming is caught by the recheck (no lost
869            // wakeup).
870            let notified = self.notify.notified();
871            if self.flag.load(Ordering::Acquire) {
872                return;
873            }
874            notified.await;
875        }
876    }
877}
878
879impl std::fmt::Debug for PeerSession {
880    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
881        f.debug_struct("PeerSession").finish_non_exhaustive()
882    }
883}
884
885impl PeerSession {
886    /// Wrap an established mTLS byte stream in yamux as the **client** (outbound-stream opener) and
887    /// spawn the driver. `io` is any tokio duplex stream (the mTLS [`tokio_rustls::client::TlsStream`]
888    /// or, in tests, a loopback stream). Returns the session; open streams with
889    /// [`Self::open_stream`] / [`Self::open_range_stream`].
890    pub fn client<S>(io: S) -> Self
891    where
892        S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
893    {
894        Self::new(io, yamux::Mode::Client)
895    }
896
897    /// Wrap an established byte stream in yamux as the **server** (accepts inbound streams). Inbound
898    /// streams the peer opens are delivered via [`Self::accept_stream`]. Provided for symmetry + the
899    /// serving side of tests.
900    pub fn server<S>(io: S) -> Self
901    where
902        S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
903    {
904        Self::new(io, yamux::Mode::Server)
905    }
906
907    fn new<S>(io: S, mode: yamux::Mode) -> Self
908    where
909        S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
910    {
911        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::channel::<MuxCommand>(64);
912        let (inbound_tx, inbound_rx) = tokio::sync::mpsc::channel::<PeerStream>(64);
913        let conn = yamux::Connection::new(io.compat(), yamux::Config::default(), mode);
914        let closed_flag = Arc::new(AtomicBool::new(false));
915        let closed_notify = Arc::new(Notify::new());
916        tokio::spawn(drive_connection(
917            conn,
918            cmd_rx,
919            inbound_tx,
920            Arc::clone(&closed_flag),
921            Arc::clone(&closed_notify),
922        ));
923        PeerSession {
924            cmd_tx,
925            inbound_rx,
926            closed_flag,
927            closed_notify,
928        }
929    }
930
931    /// A cloneable observer of this session's transport closing — see [`ClosedHandle`].
932    pub fn closed_handle(&self) -> ClosedHandle {
933        ClosedHandle {
934            flag: Arc::clone(&self.closed_flag),
935            notify: Arc::clone(&self.closed_notify),
936        }
937    }
938
939    /// Open a new outbound logical stream to the peer. Cheap — open as many as you need to run
940    /// concurrent transfers without head-of-line blocking.
941    pub async fn open_stream(&mut self) -> io::Result<PeerStream> {
942        let (tx, rx) = tokio::sync::oneshot::channel();
943        self.cmd_tx
944            .send(MuxCommand::OpenOutbound(tx))
945            .await
946            .map_err(|_| io::Error::other("mux driver closed"))?;
947        let stream = rx
948            .await
949            .map_err(|_| io::Error::other("mux driver dropped request"))?
950            .map_err(io::Error::other)?;
951        Ok(stream.compat())
952    }
953
954    /// Accept the next inbound logical stream the peer opened (server side). Returns `None` when the
955    /// connection has closed. A pure client never calls this.
956    pub async fn accept_stream(&mut self) -> Option<PeerStream> {
957        self.inbound_rx.recv().await
958    }
959
960    /// Open a `dig.fetchRange` stream for `req`: opens a fresh logical stream, writes the
961    /// [`RangeRequest`] preamble, and returns the stream for the caller to read [`RangeFrame`]s from
962    /// (via [`RangeFrame::decode`]) as they arrive. The building block for multi-source parallel
963    /// range downloads — open one of these per (peer, range) and read them concurrently.
964    pub async fn open_range_stream(&mut self, req: &RangeRequest) -> io::Result<PeerStream> {
965        let mut stream = self.open_stream().await?;
966        stream.write_all(&req.encode()?).await?;
967        stream.flush().await?;
968        Ok(stream)
969    }
970
971    /// **Availability pre-check** (`dig.getAvailability`) — ask the peer which of `items` it holds,
972    /// BEFORE opening any range streams. Opens a short-lived control stream, writes the batched
973    /// [`AvailabilityRequest`], reads the [`AvailabilityResponse`]. A multi-source downloader runs
974    /// this against candidate peers and only range-fetches from holders — the normative flow is:
975    /// discover peers → `query_availability` (batch) → fan byte-ranges across holders → verify each
976    /// vs the chain-anchored root → retry a bad range from another holder → reassemble.
977    pub async fn query_availability(
978        &mut self,
979        items: Vec<AvailabilityItem>,
980    ) -> io::Result<AvailabilityResponse> {
981        let req = AvailabilityRequest { items };
982        let mut stream = self.open_stream().await?;
983        stream.write_all(&req.encode()?).await?;
984        stream.flush().await?;
985        AvailabilityResponse::decode(&mut stream).await
986    }
987}
988
989/// Drive one yamux [`Connection`](yamux::Connection): concurrently service open-outbound commands
990/// and surface inbound streams, until the command channel closes (session dropped) or the connection
991/// errors. This is the task that replaces yamux 0.12's `Control`.
992///
993/// `T` is the futures-io view of the byte stream (a `tokio-util` [`Compat`] of the tokio mTLS
994/// stream), since yamux operates on `futures::AsyncRead + AsyncWrite`.
995async fn drive_connection<T>(
996    mut conn: yamux::Connection<T>,
997    mut cmd_rx: tokio::sync::mpsc::Receiver<MuxCommand>,
998    inbound_tx: tokio::sync::mpsc::Sender<PeerStream>,
999    closed_flag: Arc<AtomicBool>,
1000    closed_notify: Arc<Notify>,
1001) where
1002    T: futures::AsyncRead + futures::AsyncWrite + Send + Unpin + 'static,
1003{
1004    use std::future::poll_fn;
1005
1006    loop {
1007        tokio::select! {
1008            // An open-outbound request from the session.
1009            cmd = cmd_rx.recv() => {
1010                match cmd {
1011                    Some(MuxCommand::OpenOutbound(reply)) => {
1012                        let res = poll_fn(|cx| conn.poll_new_outbound(cx)).await;
1013                        let _ = reply.send(res.map_err(|e| e.to_string()));
1014                    }
1015                    None => {
1016                        // Session dropped — close the connection and end the driver.
1017                        let _ = poll_fn(|cx| conn.poll_close(cx)).await;
1018                        break;
1019                    }
1020                }
1021            }
1022            // An inbound stream opened by the peer.
1023            inbound = poll_fn(|cx| conn.poll_next_inbound(cx)) => {
1024                match inbound {
1025                    Some(Ok(stream)) => {
1026                        // Deliver to a serving node; if no one is accepting, the stream is dropped.
1027                        let _ = inbound_tx.try_send(stream.compat());
1028                    }
1029                    Some(Err(_)) | None => {
1030                        // Connection closed / errored — end the driver.
1031                        break;
1032                    }
1033                }
1034            }
1035        }
1036    }
1037
1038    // The transport is gone: publish closure (flag BEFORE notify, so `ClosedHandle::closed`'s
1039    // arm-then-recheck can never miss the wakeup).
1040    closed_flag.store(true, Ordering::Release);
1041    closed_notify.notify_waiters();
1042}