hatchery_node/
restream_session.rs

1//! Node end of a live restreamer session (A2 / A3).
2//!
3//! The node dials C2's uplink, OPENs a channel, then runs the inner Noise_NK
4//! handshake toward the relayed recipient key as ordinary DATA frames:
5//! msg1 goes up through C2 into the core drain, msg2 comes back down the same
6//! stream. After that every request and reply is Noise transport ciphertext,
7//! wrapped in C2's outer layer on the way up. C2 only ever sees opaque bytes,
8//! and the node never learns what sits behind C2.
9//!
10//! One TCP carries every channel (C2 routes downlink by node id). A single
11//! FIFO send queue would let browser frames or a stuck mail channel wedge PTY
12//! control (head-of-line blocking). Outbound traffic is therefore queued
13//! per-channel with priority scheduling, browser frames are coalesced under
14//! pressure, and a write timeout tears the socket down so the bridge reconnects.
15
16use crate::uplink_journal::UplinkJournal;
17use hatchery_core::inner::{decode_inner, encode_inner, InnerKind};
18use hatchery_core::nested::{node_seal_outer, NoiseInitiator, NoiseTransport, PublicKeyBytes};
19use hatchery_core::{grow_read_buf, ChannelKind, Frame, FrameError, FrameType, OpenPayload};
20use serde_json::{json, Value};
21use std::collections::{BTreeMap, VecDeque};
22use std::net::SocketAddr;
23use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
24use std::sync::{Arc, Mutex};
25use std::time::Duration;
26use tokio::io::{AsyncReadExt, AsyncWriteExt};
27use tokio::net::TcpStream;
28use tokio::sync::{mpsc, Notify};
29use tokio::task::JoinHandle;
30
31/// App-level ping on an idle node→C2 uplink (C2 answers Pong).
32const UPLINK_PING_EVERY: Duration = Duration::from_secs(10);
33/// No inbound bytes (pong or data) for this long => path is dead; reconnect.
34const UPLINK_DEAD_AFTER: Duration = Duration::from_secs(35);
35
36/// Keys C2 hands a node at enrollment: its own outer key and the upstream recipient key.
37#[derive(Clone, Copy, Debug)]
38pub struct UpstreamKeys {
39    pub c2_public: PublicKeyBytes,
40    pub recipient_public: PublicKeyBytes,
41}
42
43/// Outbound priority on the shared uplink. Lower ordinal = sent first.
44#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
45pub enum SendPrio {
46    /// OPEN, Noise handshake, journal replies, g4a request/response (incl. PTY input path).
47    Control = 0,
48    /// g4a terminal / event frames.
49    Pty = 1,
50    /// m4a mail/task calls.
51    Mail = 2,
52    /// dig2browser frames — droppable / coalesced under pressure.
53    Browser = 3,
54}
55
56const WRITE_TIMEOUT: Duration = Duration::from_secs(5);
57const CHAN_CAP_CONTROL: usize = 64;
58const CHAN_CAP_PTY: usize = 128;
59const CHAN_CAP_MAIL: usize = 16;
60const CHAN_CAP_BROWSER: usize = 8;
61
62struct Queued {
63    prio: SendPrio,
64    frame: Frame,
65}
66
67struct ChanQ {
68    q: VecDeque<Queued>,
69    /// Remaining WINDOW credit for DATA frames; None until the first WINDOW.
70    credit: Option<u32>,
71}
72
73impl ChanQ {
74    fn new() -> Self {
75        Self { q: VecDeque::new(), credit: None }
76    }
77
78    fn cap_for(prio: SendPrio) -> usize {
79        match prio {
80            SendPrio::Control => CHAN_CAP_CONTROL,
81            SendPrio::Pty => CHAN_CAP_PTY,
82            SendPrio::Mail => CHAN_CAP_MAIL,
83            SendPrio::Browser => CHAN_CAP_BROWSER,
84        }
85    }
86}
87
88struct OutState {
89    chans: BTreeMap<u32, ChanQ>,
90}
91
92/// One TCP uplink to C2 carrying any number of channels.
93pub struct NodeUplink {
94    enqueue: mpsc::UnboundedSender<(u32, SendPrio, Frame)>,
95    routes: Arc<Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>>,
96    /// Last WINDOW credit per channel (updated by the reader).
97    credits: Arc<Mutex<BTreeMap<u32, u32>>>,
98    credit_pulse: Arc<Notify>,
99    journal: Option<Arc<UplinkJournal>>,
100    alive: Arc<AtomicBool>,
101    dead: Arc<Notify>,
102    tasks: Vec<JoinHandle<()>>,
103}
104
105impl Drop for NodeUplink {
106    fn drop(&mut self) {
107        self.alive.store(false, Ordering::SeqCst);
108        for t in &self.tasks {
109            t.abort();
110        }
111    }
112}
113
114/// Dial the C2 uplink. With the `wireguard` feature, an address served by an
115/// in-process WireGuard session (Windows) is carried inside it; otherwise this
116/// is a plain TCP connect.
117fn dial_uplink(addr: SocketAddr) -> std::io::Result<std::net::TcpStream> {
118    #[cfg(feature = "wireguard")]
119    {
120        gate4agent_node_core::kit_wireguard::tcp_connect(addr)
121    }
122    #[cfg(not(feature = "wireguard"))]
123    {
124        std::net::TcpStream::connect(addr)
125    }
126}
127
128impl NodeUplink {
129    pub async fn connect(uplink: SocketAddr, journal: Option<Arc<UplinkJournal>>) -> Result<Arc<Self>, String> {
130        let dial = uplink;
131        let std_stream = tokio::task::spawn_blocking(move || dial_uplink(dial))
132        .await
133        .map_err(|e| e.to_string())?
134        .map_err(|e| e.to_string())?;
135        let _ = std_stream.set_nodelay(true);
136        // OS keepalive backs up the app-level ping. On the in-process
137        // WireGuard path this socket is loopback and the inner tunnel socket
138        // carries its own keepalive (gate4agent-node-core).
139        if let Err(error) = crate::server::enable_tcp_keepalive(&std_stream) {
140            tracing::warn!(cause = %error, "could not enable TCP keepalive on the uplink");
141        }
142        // current_thread runtime. A blocking socket here stalls the only
143        // worker inside the first read, so the OPEN writer and the handshake
144        // timeout never run. Windows does not flip this in `from_std`.
145        std_stream.set_nonblocking(true).map_err(|e| e.to_string())?;
146        let stream = TcpStream::from_std(std_stream).map_err(|e| e.to_string())?;
147        let (mut rd, mut wr) = stream.into_split();
148        let (enq_tx, mut enq_rx) = mpsc::unbounded_channel::<(u32, SendPrio, Frame)>();
149        let routes: Arc<Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>> = Arc::default();
150        let credits: Arc<Mutex<BTreeMap<u32, u32>>> = Arc::default();
151        let credit_pulse = Arc::new(Notify::new());
152        let alive = Arc::new(AtomicBool::new(true));
153        let dead = Arc::new(Notify::new());
154        let wake = Arc::new(Notify::new());
155
156        let out_state = Arc::new(Mutex::new(OutState { chans: BTreeMap::new() }));
157        // Last inbound byte (ms since `born`). The keepalive below pings C2 on
158        // an idle uplink so NAT/WG/conntrack state stays warm and a half-open
159        // path is detected in seconds, not discovered by the next HQ spawn
160        // (owner spawn after ~28 min idle sat ~23 s before reaching the node).
161        let born = std::time::Instant::now();
162        let last_rx = Arc::new(AtomicU64::new(0));
163
164        // Ingress into per-channel queues (never blocks the producer on TCP).
165        let state_in = Arc::clone(&out_state);
166        let wake_in = Arc::clone(&wake);
167        let alive_in = Arc::clone(&alive);
168        let ingress = tokio::spawn(async move {
169            while let Some((ch, prio, frame)) = enq_rx.recv().await {
170                if !alive_in.load(Ordering::SeqCst) {
171                    break;
172                }
173                let mut st = state_in.lock().expect("out");
174                let q = st.chans.entry(ch).or_insert_with(ChanQ::new);
175                let cap = ChanQ::cap_for(prio);
176                if prio == SendPrio::Browser {
177                    // Coalesce: keep only the newest browser frame under pressure.
178                    q.q.retain(|x| x.prio != SendPrio::Browser);
179                    if q.q.len() >= cap {
180                        // Still full of higher-prio work: drop this browser frame.
181                        continue;
182                    }
183                } else {
184                    while q.q.len() >= cap {
185                        // Drop lowest-priority (browser first, then mail) before refusing control/PTY.
186                        if let Some(i) = q
187                            .q
188                            .iter()
189                            .enumerate()
190                            .filter(|(_, x)| x.prio >= SendPrio::Mail)
191                            .max_by_key(|(_, x)| x.prio)
192                            .map(|(i, _)| i)
193                        {
194                            q.q.remove(i);
195                        } else if prio <= SendPrio::Pty {
196                            // Control/PTY: drop oldest same-or-lower in this channel.
197                            q.q.pop_front();
198                        } else {
199                            break;
200                        }
201                    }
202                    if q.q.len() >= cap {
203                        continue;
204                    }
205                }
206                q.q.push_back(Queued { prio, frame });
207                drop(st);
208                wake_in.notify_one();
209            }
210        });
211
212        // Fair/priority writer: always pick the globally highest-priority waiting frame
213        // whose channel has credit (DATA) or needs no credit (OPEN/control non-DATA).
214        let state_wr = Arc::clone(&out_state);
215        let wake_wr = Arc::clone(&wake);
216        let alive_wr = Arc::clone(&alive);
217        let dead_wr = Arc::clone(&dead);
218        let routes_wr = Arc::clone(&routes);
219        let writer = tokio::spawn(async move {
220            loop {
221                if !alive_wr.load(Ordering::SeqCst) {
222                    break;
223                }
224                let next = {
225                    let mut st = state_wr.lock().expect("out");
226                    pick_next(&mut st)
227                };
228                let Some((ch, is_data, frame)) = next else {
229                    wake_wr.notified().await;
230                    continue;
231                };
232                let Ok(bytes) = frame.encode() else {
233                    kill(&alive_wr, &dead_wr, &routes_wr);
234                    break;
235                };
236                match tokio::time::timeout(WRITE_TIMEOUT, wr.write_all(&bytes)).await {
237                    Ok(Ok(())) => {
238                        if is_data {
239                            let mut st = state_wr.lock().expect("out");
240                            if let Some(q) = st.chans.get_mut(&ch) {
241                                if let Some(c) = q.credit.as_mut() {
242                                    *c = c.saturating_sub(1);
243                                }
244                            }
245                        }
246                    }
247                    Ok(Err(_)) | Err(_) => {
248                        tracing::warn!(%ch, "uplink write timed out or failed; reconnecting");
249                        kill(&alive_wr, &dead_wr, &routes_wr);
250                        break;
251                    }
252                }
253            }
254        });
255
256        let enq_reader = enq_tx.clone();
257        let reader_journal = journal.clone();
258        let reader_routes = Arc::clone(&routes);
259        let alive_rd = Arc::clone(&alive);
260        let dead_rd = Arc::clone(&dead);
261        let state_rd = Arc::clone(&out_state);
262        let wake_rd = Arc::clone(&wake);
263        let credits_rd = Arc::clone(&credits);
264        let credit_pulse_rd = Arc::clone(&credit_pulse);
265        let last_rx_rd = Arc::clone(&last_rx);
266        let reader = tokio::spawn(async move {
267            let mut buf = vec![0u8; 128 * 1024];
268            let mut filled = 0usize;
269            loop {
270                let Ok(n) = rd.read(&mut buf[filled..]).await else { break };
271                if n == 0 {
272                    break;
273                }
274                filled += n;
275                last_rx_rd.store(born.elapsed().as_millis() as u64, Ordering::SeqCst);
276                loop {
277                    match Frame::decode(&buf[..filled]) {
278                        Ok((frame, used)) => {
279                            buf.copy_within(used..filled, 0);
280                            filled -= used;
281                            match frame.frame_type {
282                                FrameType::Data => {
283                                    let tx = reader_routes.lock().expect("routes").get(&frame.channel_id).cloned();
284                                    if let Some(tx) = tx {
285                                        let _ = tx.send(frame.payload);
286                                    }
287                                }
288                                FrameType::Window => {
289                                    let credit = frame.seq.min(u64::from(u32::MAX)) as u32;
290                                    {
291                                        let mut st = state_rd.lock().expect("out");
292                                        st.chans.entry(frame.channel_id).or_insert_with(ChanQ::new).credit = Some(credit);
293                                    }
294                                    credits_rd.lock().expect("credits").insert(frame.channel_id, credit);
295                                    credit_pulse_rd.notify_waiters();
296                                    wake_rd.notify_one();
297                                }
298                                FrameType::JournalReq => {
299                                    let Some(j) = reader_journal.as_ref() else { continue };
300                                    let limit = frame.journal_req_limit().unwrap_or(64) as usize;
301                                    if let Ok(body) = j.journal().reply_body(frame.channel_id, frame.seq, limit) {
302                                        if let Ok(reply) = Frame::journal_reply(frame.channel_id, &body) {
303                                            let _ = enq_reader.send((frame.channel_id, SendPrio::Control, reply));
304                                        }
305                                    }
306                                }
307                                FrameType::GoAway | FrameType::Close => {
308                                    kill(&alive_rd, &dead_rd, &reader_routes);
309                                    return;
310                                }
311                                _ => {}
312                            }
313                        }
314                        Err(FrameError::Truncated) => break,
315                        Err(_) => {
316                            kill(&alive_rd, &dead_rd, &reader_routes);
317                            return;
318                        }
319                    }
320                }
321                if filled == buf.len() {
322                    // Same grow path as C2: do not tear the bridge when a
323                    // single frame exceeds the initial 128 KiB read buf.
324                    if grow_read_buf(&mut buf, filled).is_err() {
325                        break;
326                    }
327                }
328            }
329            kill(&alive_rd, &dead_rd, &reader_routes);
330        });
331
332        let ka_tx = enq_tx.clone();
333        let alive_ka = Arc::clone(&alive);
334        let dead_ka = Arc::clone(&dead);
335        let routes_ka = Arc::clone(&routes);
336        let keepalive = tokio::spawn(async move {
337            let mut nonce = 0u64;
338            let mut tick = tokio::time::interval(UPLINK_PING_EVERY);
339            tick.tick().await;
340            loop {
341                tick.tick().await;
342                if !alive_ka.load(Ordering::SeqCst) {
343                    break;
344                }
345                let now = born.elapsed().as_millis() as u64;
346                let idle = now.saturating_sub(last_rx.load(Ordering::SeqCst));
347                if idle > UPLINK_DEAD_AFTER.as_millis() as u64 {
348                    tracing::warn!(idle_ms = idle, "uplink silent past keepalive deadline; reconnecting");
349                    kill(&alive_ka, &dead_ka, &routes_ka);
350                    break;
351                }
352                nonce = nonce.wrapping_add(1);
353                if ka_tx.send((0, SendPrio::Control, Frame::ping(nonce))).is_err() {
354                    break;
355                }
356            }
357        });
358
359        Ok(Arc::new(Self {
360            enqueue: enq_tx,
361            routes,
362            credits,
363            credit_pulse,
364            journal,
365            alive,
366            dead,
367            tasks: vec![ingress, writer, reader, keepalive],
368        }))
369    }
370
371    pub fn is_alive(&self) -> bool {
372        self.alive.load(Ordering::SeqCst)
373    }
374
375    /// Wait until the uplink is marked dead (write timeout / peer close).
376    pub async fn closed(&self) {
377        if !self.is_alive() {
378            return;
379        }
380        self.dead.notified().await;
381    }
382
383    fn push(&self, channel_id: u32, prio: SendPrio, frame: Frame) -> Result<(), String> {
384        if !self.is_alive() {
385            return Err("uplink closed".into());
386        }
387        self.enqueue
388            .send((channel_id, prio, frame))
389            .map_err(|_| "uplink closed".to_owned())
390    }
391
392    /// OPEN `channel_id`, wait for WINDOW credit, then complete Noise_NK.
393    pub async fn open(
394        self: &Arc<Self>,
395        node_id: &str,
396        channel_id: u32,
397        kind: ChannelKind,
398        keys: UpstreamKeys,
399        handshake_timeout: Duration,
400    ) -> Result<NodeSession, String> {
401        let (down_tx, mut down_rx) = mpsc::unbounded_channel::<Vec<u8>>();
402        self.routes.lock().expect("routes").insert(channel_id, down_tx);
403        self.credits.lock().expect("credits").remove(&channel_id);
404        let open = Frame::open(channel_id, &OpenPayload { channel_kind: kind, node_id: node_id.to_owned() })
405            .map_err(|e| e.to_string())?;
406        self.push(channel_id, SendPrio::Control, open)?;
407
408        // Wait until C2 grants WINDOW credit after OPEN (avoids sending Handshake1
409        // into a stream that still has credit 0 from a previous session).
410        tokio::time::timeout(handshake_timeout, async {
411            loop {
412                if !self.is_alive() {
413                    return Err("uplink closed".to_owned());
414                }
415                if self.credits.lock().expect("credits").get(&channel_id).copied().unwrap_or(0) > 0 {
416                    return Ok(());
417                }
418                tokio::select! {
419                    _ = self.credit_pulse.notified() => {}
420                    _ = self.closed() => return Err("uplink closed".to_owned()),
421                }
422            }
423        })
424        .await
425        .map_err(|_| "window credit timeout after OPEN".to_owned())??;
426
427        let mut initiator = NoiseInitiator::new(keys.recipient_public).map_err(|e| e.to_string())?;
428        let msg1 = initiator.write_message1().map_err(|e| e.to_string())?;
429        self.send_inner(channel_id, keys, SendPrio::Control, &encode_inner(InnerKind::Handshake1, &msg1))?;
430
431        let msg2 = tokio::time::timeout(handshake_timeout, async {
432            loop {
433                let Some(bytes) = down_rx.recv().await else { return Err("uplink closed".to_owned()) };
434                if let Some((InnerKind::Handshake2, body)) = decode_inner(&bytes) {
435                    return Ok(body.to_vec());
436                }
437            }
438        })
439        .await
440        .map_err(|_| "handshake timeout".to_owned())??;
441        let transport = initiator.read_message2(&msg2).map_err(|e| e.to_string())?;
442        Ok(NodeSession {
443            channel_id,
444            keys,
445            default_prio: match kind {
446                ChannelKind::M4a => SendPrio::Mail,
447                ChannelKind::G4aControl | ChannelKind::G4aPty | ChannelKind::G4aEvents => SendPrio::Pty,
448                ChannelKind::Vault | ChannelKind::Observation => SendPrio::Mail,
449            },
450            down: down_rx,
451            transport,
452            journal: self.journal.clone(),
453            uplink: Arc::clone(self),
454        })
455    }
456
457    fn send_inner(&self, channel_id: u32, keys: UpstreamKeys, prio: SendPrio, inner: &[u8]) -> Result<(), String> {
458        if let Some(j) = &self.journal {
459            j.append(channel_id, inner)?;
460        }
461        let sealed = node_seal_outer(keys.c2_public, inner).map_err(|e| e.to_string())?;
462        let frame = Frame::data(channel_id, 0, sealed).map_err(|e| e.to_string())?;
463        self.push(channel_id, prio, frame)
464    }
465}
466
467fn kill(alive: &AtomicBool, dead: &Notify, routes: &Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>) {
468    if alive.swap(false, Ordering::SeqCst) {
469        routes.lock().expect("routes").clear();
470        dead.notify_waiters();
471    }
472}
473
474/// Pick the highest-priority queued frame that may be sent now.
475fn pick_next(st: &mut OutState) -> Option<(u32, bool, Frame)> {
476    // Candidates: (prio, channel, index). DATA frames of one channel are Noise ciphertexts
477    // encrypted in queue order, so only the channel's first DATA frame may go next
478    // (overtaking breaks the receiver's nonce sequence); non-DATA frames may go any time.
479    let mut best: Option<(SendPrio, u32, usize)> = None;
480    for (&ch, cq) in &st.chans {
481        let first_data = cq.q.iter().position(|it| matches!(it.frame.frame_type, FrameType::Data));
482        for (i, item) in cq.q.iter().enumerate() {
483            let needs_credit = matches!(item.frame.frame_type, FrameType::Data);
484            if needs_credit && Some(i) != first_data {
485                continue;
486            }
487            if needs_credit {
488                match cq.credit {
489                    Some(0) => continue,
490                    None => {
491                        // No WINDOW yet: allow Control DATA (handshake) so OPEN→Handshake1 works;
492                        // hold back bulk traffic.
493                        if item.prio > SendPrio::Control {
494                            continue;
495                        }
496                    }
497                    Some(_) => {}
498                }
499            }
500            let cand = (item.prio, ch, i);
501            if best.is_none_or(|b| cand.0 < b.0 || (cand.0 == b.0 && cand.1 < b.1)) {
502                best = Some(cand);
503            }
504        }
505    }
506    let (prio, ch, idx) = best?;
507    let cq = st.chans.get_mut(&ch)?;
508    let item = cq.q.remove(idx)?;
509    debug_assert_eq!(item.prio, prio);
510    let is_data = matches!(item.frame.frame_type, FrameType::Data);
511    Some((ch, is_data, item.frame))
512}
513
514pub struct NodeSession {
515    channel_id: u32,
516    keys: UpstreamKeys,
517    default_prio: SendPrio,
518    down: mpsc::UnboundedReceiver<Vec<u8>>,
519    transport: NoiseTransport,
520    journal: Option<Arc<UplinkJournal>>,
521    uplink: Arc<NodeUplink>,
522}
523
524impl NodeSession {
525    pub async fn connect(
526        uplink: SocketAddr,
527        node_id: &str,
528        channel_id: u32,
529        kind: ChannelKind,
530        keys: UpstreamKeys,
531        journal: Option<Arc<UplinkJournal>>,
532        handshake_timeout: Duration,
533    ) -> Result<Self, String> {
534        let link = NodeUplink::connect(uplink, journal).await?;
535        link.open(node_id, channel_id, kind, keys, handshake_timeout).await
536    }
537
538    pub fn uplink(&self) -> Arc<NodeUplink> {
539        Arc::clone(&self.uplink)
540    }
541
542    pub async fn send(&mut self, plaintext: &[u8]) -> Result<(), String> {
543        self.send_prio(plaintext, self.default_prio).await
544    }
545
546    pub async fn send_prio(&mut self, plaintext: &[u8], prio: SendPrio) -> Result<(), String> {
547        if !self.uplink.is_alive() {
548            return Err("uplink closed".into());
549        }
550        let ct = self.transport.encrypt(plaintext).map_err(|e| e.to_string())?;
551        let inner = encode_inner(InnerKind::Transport, &ct);
552        self.uplink.send_inner(self.channel_id, self.keys, prio, &inner)
553    }
554
555    pub async fn recv(&mut self, timeout: Duration) -> Result<Vec<u8>, String> {
556        tokio::time::timeout(timeout, async {
557            loop {
558                tokio::select! {
559                    bytes = self.down.recv() => {
560                        let Some(bytes) = bytes else { return Err("uplink closed".to_owned()) };
561                        if let Some((InnerKind::Transport, ct)) = decode_inner(&bytes) {
562                            return self.transport.decrypt(ct).map_err(|e| e.to_string());
563                        }
564                    }
565                    _ = self.uplink.closed() => return Err("uplink closed".to_owned()),
566                }
567            }
568        })
569        .await
570        .map_err(|_| "reply timeout".to_owned())?
571    }
572
573    pub async fn call(&mut self, op: &str, payload: &Value, timeout: Duration) -> Result<Value, String> {
574        let req = serde_json::to_vec(&json!({ "op": op, "payload": payload })).map_err(|e| e.to_string())?;
575        self.send_prio(&req, SendPrio::Mail).await?;
576        let reply = self.recv(timeout).await?;
577        let reply: Value = serde_json::from_slice(&reply).map_err(|e| e.to_string())?;
578        let status = reply.get("status").and_then(Value::as_u64).unwrap_or(500);
579        let body = reply.get("body").cloned().unwrap_or(Value::Null);
580        if status == 200 {
581            Ok(body)
582        } else {
583            let msg = body.get("error").and_then(Value::as_str).unwrap_or("error");
584            Err(format!("{status}: {msg}"))
585        }
586    }
587}
588
589#[cfg(test)]
590mod tests {
591    use super::*;
592
593    #[test]
594    fn pick_prefers_control_across_channels_but_keeps_data_order_within_one() {
595        let mut st = OutState { chans: BTreeMap::new() };
596        let a = st.chans.entry(1).or_insert_with(ChanQ::new);
597        a.credit = Some(10);
598        a.q.push_back(Queued { prio: SendPrio::Browser, frame: Frame::data(1, 0, vec![1]).unwrap() });
599        a.q.push_back(Queued { prio: SendPrio::Control, frame: Frame::data(1, 0, vec![2]).unwrap() });
600        let b = st.chans.entry(2).or_insert_with(ChanQ::new);
601        b.credit = Some(10);
602        b.q.push_back(Queued { prio: SendPrio::Pty, frame: Frame::data(2, 0, vec![3]).unwrap() });
603        // Channel 2's Pty beats channel 1's head (Browser); channel 1 stays in order.
604        let order: Vec<u8> = std::iter::from_fn(|| pick_next(&mut st).map(|(_, _, f)| f.payload[0])).collect();
605        assert_eq!(order, vec![3, 1, 2]);
606    }
607
608    #[test]
609    fn pick_holds_bulk_until_window() {
610        let mut st = OutState { chans: BTreeMap::new() };
611        let q = st.chans.entry(2).or_insert_with(ChanQ::new);
612        q.credit = None;
613        q.q.push_back(Queued {
614            prio: SendPrio::Control,
615            frame: Frame::data(2, 0, vec![1]).unwrap(),
616        });
617        q.q.push_back(Queued {
618            prio: SendPrio::Mail,
619            frame: Frame::data(2, 0, vec![9]).unwrap(),
620        });
621        let (_, _, f) = pick_next(&mut st).unwrap();
622        assert_eq!(f.payload, vec![1]);
623        assert!(pick_next(&mut st).is_none());
624        st.chans.get_mut(&2).unwrap().credit = Some(4);
625        let (_, _, f) = pick_next(&mut st).unwrap();
626        assert_eq!(f.payload, vec![9]);
627    }
628}
hatchery_node/restream_session.rs

hatchery_node/
restream_session.rs
