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}