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}`, first frame adding `total_length` + `chunk_lens` +
29//! `chunk_index` + `inclusion_proof` + `root`). Per-chunk integrity (split by `chunk_lens`, verify
30//! the whole-resource inclusion proof vs the chain-anchored `root`, AES-256-GCM-SIV-open) is done by
31//! the CONTENT layer above dig-nat; dig-nat carries these frames faithfully over the mux transport.
32
33use std::io;
34use std::sync::atomic::{AtomicBool, Ordering};
35use std::sync::Arc;
36
37use serde::{Deserialize, Serialize};
38use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
39use tokio::sync::Notify;
40use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
41
42/// One logical, bidirectional stream to the peer — a tokio [`AsyncRead`] + [`AsyncWrite`]. Reads
43/// deliver bytes incrementally as they arrive (streaming, with yamux-window backpressure); many
44/// [`PeerStream`]s coexist on one [`PeerSession`] without head-of-line blocking.
45///
46/// yamux streams are `futures` streams; this is the tokio-trait view via `tokio-util` compat.
47pub type PeerStream = Compat<yamux::Stream>;
48
49/// One item in a [`dig.getAvailability`](AvailabilityRequest) batch — a resource key at store, root,
50/// or capsule/resource granularity (inferred from which fields are present, per the L7 spec §9):
51/// `store_id` only → *has_store*; `+ root` → *has_root* (the capsule `store_id:root`); `+
52/// retrieval_key` → *has_resource*. Hashes are 64-hex.
53#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
54pub struct AvailabilityItem {
55    /// The store id (64-hex). Always present.
56    pub store_id: String,
57    /// The generation root (64-hex). Present for root/resource granularity.
58    #[serde(default, skip_serializing_if = "Option::is_none")]
59    pub root: Option<String>,
60    /// The resource retrieval key (64-hex). Present for resource granularity.
61    #[serde(default, skip_serializing_if = "Option::is_none")]
62    pub retrieval_key: Option<String>,
63}
64
65/// The **availability pre-check** (`dig.getAvailability`, L7 spec §9) — asked BEFORE any range fetch.
66/// A multi-source download batches candidate peers × items in one call each and only fans byte-range
67/// requests at peers that answer *available* — never opening range streams to peers that may not hold
68/// the content. A message-style control call over the mux'd mTLS connection.
69#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
70pub struct AvailabilityRequest {
71    /// The items to check, batched.
72    pub items: Vec<AvailabilityItem>,
73}
74
75/// One answer in an [`AvailabilityResponse`], positionally aligned with the request `items`.
76#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
77pub struct AvailabilityAnswer {
78    /// Whether the peer holds the queried item. Always present.
79    pub available: bool,
80    /// (store granularity) generation roots the peer holds for the store, newest-first.
81    #[serde(default, skip_serializing_if = "Option::is_none")]
82    pub roots: Option<Vec<String>>,
83    /// (root/resource granularity) the ciphertext length — lets the caller plan its ranges.
84    #[serde(default, skip_serializing_if = "Option::is_none")]
85    pub total_length: Option<u64>,
86    /// (root/resource granularity) the chunk count.
87    #[serde(default, skip_serializing_if = "Option::is_none")]
88    pub chunk_count: Option<u64>,
89    /// Whether the peer holds the FULL resource/capsule (`true`) or only part (`false`).
90    #[serde(default, skip_serializing_if = "Option::is_none")]
91    pub complete: Option<bool>,
92}
93
94/// The peer's answer to an [`AvailabilityRequest`]: one [`AvailabilityAnswer`] per queried item,
95/// positionally aligned with the request's `items`.
96#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
97pub struct AvailabilityResponse {
98    /// One answer per queried item, in request order.
99    pub items: Vec<AvailabilityAnswer>,
100}
101
102/// A byte-range request (`dig.fetchRange`, L7 spec §9) written at the start of a range-scoped stream.
103/// Identifies a resource (`store_id` + `retrieval_key` [+ `root`]) or a whole capsule
104/// (`capsule: true`, identified by `store_id` [+ `root`]) and the `[offset, offset+length)` range.
105#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
106pub struct RangeRequest {
107    /// The store id (64-hex).
108    pub store_id: String,
109    /// The resource retrieval key (64-hex). Omitted when `capsule` is true.
110    #[serde(default, skip_serializing_if = "Option::is_none")]
111    pub retrieval_key: Option<String>,
112    /// The generation root (64-hex). Optional — defaults to the chain-anchored tip.
113    #[serde(default, skip_serializing_if = "Option::is_none")]
114    pub root: Option<String>,
115    /// Fetch a whole capsule / `.dig` (identified by `store_id` [+ `root`]) rather than one resource.
116    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
117    pub capsule: bool,
118    /// Start offset (bytes) into the resource ciphertext. Default `0`.
119    #[serde(default)]
120    pub offset: u64,
121    /// Length (bytes) to return (widened to whole-chunk boundaries; clamped to the node window).
122    pub length: u64,
123}
124
125/// One streamed `dig.fetchRange` frame (L7 spec §8 framing). Frames arrive in ascending `offset`
126/// order and tile the requested range exactly; the caller reassembles by `offset` and stops on
127/// `complete`. The **first frame** additionally carries the per-range verification metadata so a
128/// single-peer range is independently verifiable against the chain-anchored `root`.
129#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
130pub struct RangeFrame {
131    /// This frame's start offset within the requested range.
132    pub offset: u64,
133    /// This frame's byte length.
134    pub length: u64,
135    /// The raw ciphertext bytes. On the wire they are **base64** (`base64_bytes`) — the canonical
136    /// `dig.fetchRange` frame encoding every producer emits.
137    #[serde(with = "base64_bytes")]
138    pub bytes: Vec<u8>,
139    /// Whether this is the final frame of the range.
140    pub complete: bool,
141    /// (first frame only) the full resource ciphertext length.
142    #[serde(default, skip_serializing_if = "Option::is_none")]
143    pub total_length: Option<u64>,
144    /// (first frame only) per-chunk ciphertext lengths of the whole resource, in order.
145    #[serde(default, skip_serializing_if = "Option::is_none")]
146    pub chunk_lens: Option<Vec<u64>>,
147    /// (first frame only) index into `chunk_lens` of the first chunk in this frame.
148    #[serde(default, skip_serializing_if = "Option::is_none")]
149    pub chunk_index: Option<u64>,
150    /// (first frame only) merkle inclusion proof of the whole resource vs the generation `root`
151    /// (base64, relayed verbatim); `null`/absent for `capsule: true` (self-verifying on install).
152    #[serde(default, skip_serializing_if = "Option::is_none")]
153    pub inclusion_proof: Option<String>,
154    /// (first frame only) the generation root (64-hex) the inclusion proof is against.
155    #[serde(default, skip_serializing_if = "Option::is_none")]
156    pub root: Option<String>,
157}
158
159/// Serde for [`RangeFrame::bytes`]: **base64** on the wire, raw `Vec<u8>` in Rust.
160///
161/// The `dig.fetchRange` frame is JSON, and the canonical wire type
162/// (`dig_rpc_protocol::types::RangeFrame`, "this window's ciphertext, base64") — and every real
163/// producer, including the dig-node peer serve path — encodes the window as a base64 STRING. Reading
164/// it with `serde_bytes` instead yielded the string's literal characters, so a served window arrived
165/// as its own base64 text and the reassembler rejected the frame (#1586, the read-leg blocker).
166///
167/// Deserialization is tolerant: a base64 string (canonical) OR a byte array (what an older dig-nat
168/// emitted) both decode, so a mixed-version peer is never dropped.
169mod base64_bytes {
170    use base64::Engine as _;
171    use serde::de::{SeqAccess, Visitor};
172    use serde::{Deserializer, Serializer};
173    use std::fmt;
174
175    pub fn serialize<S: Serializer>(bytes: &[u8], s: S) -> Result<S::Ok, S::Error> {
176        s.serialize_str(&base64::engine::general_purpose::STANDARD.encode(bytes))
177    }
178
179    pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<u8>, D::Error> {
180        d.deserialize_any(Base64OrArray)
181    }
182
183    struct Base64OrArray;
184
185    impl<'de> Visitor<'de> for Base64OrArray {
186        type Value = Vec<u8>;
187
188        fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
189            f.write_str("base64-encoded ciphertext (or a legacy byte array)")
190        }
191
192        fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<Self::Value, E> {
193            base64::engine::general_purpose::STANDARD
194                .decode(v)
195                .map_err(|e| E::custom(format!("range frame bytes are not valid base64: {e}")))
196        }
197
198        fn visit_bytes<E: serde::de::Error>(self, v: &[u8]) -> Result<Self::Value, E> {
199            Ok(v.to_vec())
200        }
201
202        fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value, A::Error> {
203            let mut out = Vec::with_capacity(seq.size_hint().unwrap_or_default());
204            while let Some(b) = seq.next_element::<u8>()? {
205                out.push(b);
206            }
207            Ok(out)
208        }
209    }
210}
211
212impl AvailabilityRequest {
213    /// Serialize as a `u32` big-endian length prefix + JSON body (the uniform control framing).
214    pub fn encode(&self) -> Vec<u8> {
215        encode_framed(self)
216    }
217    /// Read + decode an [`AvailabilityRequest`] from `r` (the peer/serving side).
218    pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
219        decode_framed(r).await
220    }
221}
222
223impl AvailabilityResponse {
224    /// Serialize as a `u32` big-endian length prefix + JSON body.
225    pub fn encode(&self) -> Vec<u8> {
226        encode_framed(self)
227    }
228    /// Read + decode an [`AvailabilityResponse`] from `r` (the requesting side).
229    pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
230        decode_framed(r).await
231    }
232}
233
234impl RangeRequest {
235    /// A range request for a content resource (`store_id` + `retrieval_key`).
236    pub fn resource(
237        store_id: impl Into<String>,
238        retrieval_key: impl Into<String>,
239        offset: u64,
240        length: u64,
241    ) -> Self {
242        RangeRequest {
243            store_id: store_id.into(),
244            retrieval_key: Some(retrieval_key.into()),
245            root: None,
246            capsule: false,
247            offset,
248            length,
249        }
250    }
251
252    /// Serialize as a `u32` big-endian length prefix + JSON body — the preamble a peer reads to learn
253    /// the resource + range before streaming the frames.
254    pub fn encode(&self) -> Vec<u8> {
255        encode_framed(self)
256    }
257    /// Read + decode a [`RangeRequest`] preamble from `r` (the serving side of a range stream).
258    pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
259        decode_framed(r).await
260    }
261}
262
263impl RangeFrame {
264    /// Serialize as a `u32` big-endian length prefix + JSON body (one framed frame on the stream).
265    pub fn encode(&self) -> Vec<u8> {
266        encode_framed(self)
267    }
268    /// Read + decode one [`RangeFrame`] from `r`. Returns `Ok(None)` at clean end-of-stream (the
269    /// reader hit EOF on a frame boundary), so a consumer loops until `None` or `complete`.
270    pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Option<Self>> {
271        decode_framed_opt(r).await
272    }
273}
274
275/// Maximum length-prefixed control/preamble body — guards against a malicious length prefix forcing
276/// a huge allocation.
277const MAX_FRAMED_BODY: usize = 64 * 1024;
278
279/// Serialize `value` as a `u32` big-endian length prefix + JSON body — the uniform framing for every
280/// small control message on a stream (availability + range preambles).
281fn encode_framed<T: Serialize>(value: &T) -> Vec<u8> {
282    let body = serde_json::to_vec(value).expect("control message serializes");
283    let mut out = Vec::with_capacity(4 + body.len());
284    out.extend_from_slice(&(body.len() as u32).to_be_bytes());
285    out.extend_from_slice(&body);
286    out
287}
288
289/// Read + decode a length-prefixed JSON control message from `r`, bounded by [`MAX_FRAMED_BODY`].
290async fn decode_framed<T: for<'de> Deserialize<'de>, R: AsyncRead + Unpin>(
291    r: &mut R,
292) -> io::Result<T> {
293    let mut len_buf = [0u8; 4];
294    r.read_exact(&mut len_buf).await?;
295    let len = u32::from_be_bytes(len_buf) as usize;
296    if len > MAX_FRAMED_BODY {
297        return Err(io::Error::new(
298            io::ErrorKind::InvalidData,
299            "control message too large",
300        ));
301    }
302    let mut body = vec![0u8; len];
303    r.read_exact(&mut body).await?;
304    serde_json::from_slice(&body).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
305}
306
307/// Like [`decode_framed`] but returns `Ok(None)` on a CLEAN end-of-stream at a frame boundary (the
308/// length prefix read hits immediate EOF), so a streaming consumer can loop until the stream ends.
309async fn decode_framed_opt<T: for<'de> Deserialize<'de>, R: AsyncRead + Unpin>(
310    r: &mut R,
311) -> io::Result<Option<T>> {
312    let mut len_buf = [0u8; 4];
313    match r.read_exact(&mut len_buf).await {
314        Ok(_) => {}
315        Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
316        Err(e) => return Err(e),
317    }
318    let len = u32::from_be_bytes(len_buf) as usize;
319    if len > MAX_FRAMED_BODY {
320        return Err(io::Error::new(
321            io::ErrorKind::InvalidData,
322            "control message too large",
323        ));
324    }
325    let mut body = vec![0u8; len];
326    r.read_exact(&mut body).await?;
327    serde_json::from_slice(&body)
328        .map(Some)
329        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
330}
331
332/// One command to the yamux driver task. yamux 0.13 has no `Control` handle, so we drive the
333/// [`yamux::Connection`] in a task and talk to it over this channel.
334enum MuxCommand {
335    /// Open a new outbound stream; the resulting [`yamux::Stream`] (or error) comes back on the sender.
336    OpenOutbound(tokio::sync::oneshot::Sender<Result<yamux::Stream, String>>),
337}
338
339/// A multiplexed session over one peer connection: open many concurrent logical [`PeerStream`]s.
340///
341/// yamux 0.13 exposes a poll-based [`yamux::Connection`] (no `Control` handle), so a background
342/// driver task owns the connection and serves open-stream requests over a command channel; inbound
343/// streams are surfaced on [`Self::inbound_rx`] for a serving node. Dropping the session closes the
344/// command channel, which ends the driver and tears down the underlying byte stream.
345pub struct PeerSession {
346    cmd_tx: tokio::sync::mpsc::Sender<MuxCommand>,
347    /// Inbound streams opened BY the peer (server role / bidirectional use). A pure client can
348    /// ignore this; a serving node reads accepted range-request streams from here.
349    inbound_rx: tokio::sync::mpsc::Receiver<PeerStream>,
350    /// Set + notified when the driver task ends (the underlying byte stream closed — a clean close or
351    /// a transport error). Observed via [`Self::closed_handle`] so fast-connect can detect a transport
352    /// dying and fall back without holding the session lock.
353    closed_flag: Arc<AtomicBool>,
354    closed_notify: Arc<Notify>,
355}
356
357/// A cheap, cloneable observer of a [`PeerSession`]'s underlying byte stream closing (the mux driver
358/// task ending — a clean close OR a transport error). Fast-connect's promotion guard holds one to
359/// detect the active transport dying and fall back to another tier, without locking the session.
360#[derive(Clone)]
361pub struct ClosedHandle {
362    flag: Arc<AtomicBool>,
363    notify: Arc<Notify>,
364}
365
366impl ClosedHandle {
367    /// Whether the session's transport has already closed.
368    pub fn is_closed(&self) -> bool {
369        self.flag.load(Ordering::Acquire)
370    }
371
372    /// Resolve once the session's transport has closed (returns immediately if already closed).
373    pub async fn closed(&self) {
374        loop {
375            if self.flag.load(Ordering::Acquire) {
376                return;
377            }
378            // Arm the wait, then re-check the flag: the driver sets the flag BEFORE
379            // `notify_waiters`, so a close that races this arming is caught by the recheck (no lost
380            // wakeup).
381            let notified = self.notify.notified();
382            if self.flag.load(Ordering::Acquire) {
383                return;
384            }
385            notified.await;
386        }
387    }
388}
389
390impl std::fmt::Debug for PeerSession {
391    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
392        f.debug_struct("PeerSession").finish_non_exhaustive()
393    }
394}
395
396impl PeerSession {
397    /// Wrap an established mTLS byte stream in yamux as the **client** (outbound-stream opener) and
398    /// spawn the driver. `io` is any tokio duplex stream (the mTLS [`tokio_rustls::client::TlsStream`]
399    /// or, in tests, a loopback stream). Returns the session; open streams with
400    /// [`Self::open_stream`] / [`Self::open_range_stream`].
401    pub fn client<S>(io: S) -> Self
402    where
403        S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
404    {
405        Self::new(io, yamux::Mode::Client)
406    }
407
408    /// Wrap an established byte stream in yamux as the **server** (accepts inbound streams). Inbound
409    /// streams the peer opens are delivered via [`Self::accept_stream`]. Provided for symmetry + the
410    /// serving side of tests.
411    pub fn server<S>(io: S) -> Self
412    where
413        S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
414    {
415        Self::new(io, yamux::Mode::Server)
416    }
417
418    fn new<S>(io: S, mode: yamux::Mode) -> Self
419    where
420        S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
421    {
422        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::channel::<MuxCommand>(64);
423        let (inbound_tx, inbound_rx) = tokio::sync::mpsc::channel::<PeerStream>(64);
424        let conn = yamux::Connection::new(io.compat(), yamux::Config::default(), mode);
425        let closed_flag = Arc::new(AtomicBool::new(false));
426        let closed_notify = Arc::new(Notify::new());
427        tokio::spawn(drive_connection(
428            conn,
429            cmd_rx,
430            inbound_tx,
431            Arc::clone(&closed_flag),
432            Arc::clone(&closed_notify),
433        ));
434        PeerSession {
435            cmd_tx,
436            inbound_rx,
437            closed_flag,
438            closed_notify,
439        }
440    }
441
442    /// A cloneable observer of this session's transport closing — see [`ClosedHandle`].
443    pub fn closed_handle(&self) -> ClosedHandle {
444        ClosedHandle {
445            flag: Arc::clone(&self.closed_flag),
446            notify: Arc::clone(&self.closed_notify),
447        }
448    }
449
450    /// Open a new outbound logical stream to the peer. Cheap — open as many as you need to run
451    /// concurrent transfers without head-of-line blocking.
452    pub async fn open_stream(&mut self) -> io::Result<PeerStream> {
453        let (tx, rx) = tokio::sync::oneshot::channel();
454        self.cmd_tx
455            .send(MuxCommand::OpenOutbound(tx))
456            .await
457            .map_err(|_| io::Error::other("mux driver closed"))?;
458        let stream = rx
459            .await
460            .map_err(|_| io::Error::other("mux driver dropped request"))?
461            .map_err(io::Error::other)?;
462        Ok(stream.compat())
463    }
464
465    /// Accept the next inbound logical stream the peer opened (server side). Returns `None` when the
466    /// connection has closed. A pure client never calls this.
467    pub async fn accept_stream(&mut self) -> Option<PeerStream> {
468        self.inbound_rx.recv().await
469    }
470
471    /// Open a `dig.fetchRange` stream for `req`: opens a fresh logical stream, writes the
472    /// [`RangeRequest`] preamble, and returns the stream for the caller to read [`RangeFrame`]s from
473    /// (via [`RangeFrame::decode`]) as they arrive. The building block for multi-source parallel
474    /// range downloads — open one of these per (peer, range) and read them concurrently.
475    pub async fn open_range_stream(&mut self, req: &RangeRequest) -> io::Result<PeerStream> {
476        let mut stream = self.open_stream().await?;
477        stream.write_all(&req.encode()).await?;
478        stream.flush().await?;
479        Ok(stream)
480    }
481
482    /// **Availability pre-check** (`dig.getAvailability`) — ask the peer which of `items` it holds,
483    /// BEFORE opening any range streams. Opens a short-lived control stream, writes the batched
484    /// [`AvailabilityRequest`], reads the [`AvailabilityResponse`]. A multi-source downloader runs
485    /// this against candidate peers and only range-fetches from holders — the normative flow is:
486    /// discover peers → `query_availability` (batch) → fan byte-ranges across holders → verify each
487    /// vs the chain-anchored root → retry a bad range from another holder → reassemble.
488    pub async fn query_availability(
489        &mut self,
490        items: Vec<AvailabilityItem>,
491    ) -> io::Result<AvailabilityResponse> {
492        let req = AvailabilityRequest { items };
493        let mut stream = self.open_stream().await?;
494        stream.write_all(&req.encode()).await?;
495        stream.flush().await?;
496        AvailabilityResponse::decode(&mut stream).await
497    }
498}
499
500/// Drive one yamux [`Connection`](yamux::Connection): concurrently service open-outbound commands
501/// and surface inbound streams, until the command channel closes (session dropped) or the connection
502/// errors. This is the task that replaces yamux 0.12's `Control`.
503///
504/// `T` is the futures-io view of the byte stream (a `tokio-util` [`Compat`] of the tokio mTLS
505/// stream), since yamux operates on `futures::AsyncRead + AsyncWrite`.
506async fn drive_connection<T>(
507    mut conn: yamux::Connection<T>,
508    mut cmd_rx: tokio::sync::mpsc::Receiver<MuxCommand>,
509    inbound_tx: tokio::sync::mpsc::Sender<PeerStream>,
510    closed_flag: Arc<AtomicBool>,
511    closed_notify: Arc<Notify>,
512) where
513    T: futures::AsyncRead + futures::AsyncWrite + Send + Unpin + 'static,
514{
515    use std::future::poll_fn;
516
517    loop {
518        tokio::select! {
519            // An open-outbound request from the session.
520            cmd = cmd_rx.recv() => {
521                match cmd {
522                    Some(MuxCommand::OpenOutbound(reply)) => {
523                        let res = poll_fn(|cx| conn.poll_new_outbound(cx)).await;
524                        let _ = reply.send(res.map_err(|e| e.to_string()));
525                    }
526                    None => {
527                        // Session dropped — close the connection and end the driver.
528                        let _ = poll_fn(|cx| conn.poll_close(cx)).await;
529                        break;
530                    }
531                }
532            }
533            // An inbound stream opened by the peer.
534            inbound = poll_fn(|cx| conn.poll_next_inbound(cx)) => {
535                match inbound {
536                    Some(Ok(stream)) => {
537                        // Deliver to a serving node; if no one is accepting, the stream is dropped.
538                        let _ = inbound_tx.try_send(stream.compat());
539                    }
540                    Some(Err(_)) | None => {
541                        // Connection closed / errored — end the driver.
542                        break;
543                    }
544                }
545            }
546        }
547    }
548
549    // The transport is gone: publish closure (flag BEFORE notify, so `ClosedHandle::closed`'s
550    // arm-then-recheck can never miss the wakeup).
551    closed_flag.store(true, Ordering::Release);
552    closed_notify.notify_waiters();
553}