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 started = std::time::Instant::now();
132        crate::node_log::info(
133            "uplink_connect_attempt",
134            &[("addr_port", dial.port().to_string())],
135        );
136        let std_stream = match tokio::task::spawn_blocking(move || dial_uplink(dial)).await {
137            Ok(Ok(s)) => {
138                crate::node_log::info(
139                    "uplink_connect_ok",
140                    &[
141                        ("addr_port", dial.port().to_string()),
142                        ("ms", started.elapsed().as_millis().to_string()),
143                    ],
144                );
145                s
146            }
147            Ok(Err(e)) => {
148                crate::node_log::warn(
149                    "uplink_connect_fail",
150                    &[
151                        ("addr_port", dial.port().to_string()),
152                        ("error", e.to_string()),
153                        ("ms", started.elapsed().as_millis().to_string()),
154                    ],
155                );
156                return Err(e.to_string());
157            }
158            Err(e) => {
159                crate::node_log::warn(
160                    "uplink_connect_fail",
161                    &[
162                        ("addr_port", dial.port().to_string()),
163                        ("error", e.to_string()),
164                        ("ms", started.elapsed().as_millis().to_string()),
165                    ],
166                );
167                return Err(e.to_string());
168            }
169        };
170        let _ = std_stream.set_nodelay(true);
171        // OS keepalive backs up the app-level ping. On the in-process
172        // WireGuard path this socket is loopback and the inner tunnel socket
173        // carries its own keepalive (gate4agent-node-core).
174        if let Err(error) = crate::server::enable_tcp_keepalive(&std_stream) {
175            tracing::warn!(cause = %error, "could not enable TCP keepalive on the uplink");
176        }
177        // current_thread runtime. A blocking socket here stalls the only
178        // worker inside the first read, so the OPEN writer and the handshake
179        // timeout never run. Windows does not flip this in `from_std`.
180        std_stream.set_nonblocking(true).map_err(|e| e.to_string())?;
181        let stream = TcpStream::from_std(std_stream).map_err(|e| e.to_string())?;
182        let (mut rd, mut wr) = stream.into_split();
183        let (enq_tx, mut enq_rx) = mpsc::unbounded_channel::<(u32, SendPrio, Frame)>();
184        let routes: Arc<Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>> = Arc::default();
185        let credits: Arc<Mutex<BTreeMap<u32, u32>>> = Arc::default();
186        let credit_pulse = Arc::new(Notify::new());
187        let alive = Arc::new(AtomicBool::new(true));
188        let dead = Arc::new(Notify::new());
189        let wake = Arc::new(Notify::new());
190
191        let out_state = Arc::new(Mutex::new(OutState { chans: BTreeMap::new() }));
192        // Last inbound byte (ms since `born`). The keepalive below pings C2 on
193        // an idle uplink so NAT/WG/conntrack state stays warm and a half-open
194        // path is detected in seconds, not discovered by the next HQ spawn
195        // (owner spawn after ~28 min idle sat ~23 s before reaching the node).
196        let born = std::time::Instant::now();
197        let last_rx = Arc::new(AtomicU64::new(0));
198
199        // Ingress into per-channel queues (never blocks the producer on TCP).
200        let state_in = Arc::clone(&out_state);
201        let wake_in = Arc::clone(&wake);
202        let alive_in = Arc::clone(&alive);
203        let ingress = tokio::spawn(async move {
204            while let Some((ch, prio, frame)) = enq_rx.recv().await {
205                if !alive_in.load(Ordering::SeqCst) {
206                    break;
207                }
208                let mut st = state_in.lock().expect("out");
209                let q = st.chans.entry(ch).or_insert_with(ChanQ::new);
210                let cap = ChanQ::cap_for(prio);
211                if prio == SendPrio::Browser {
212                    // Coalesce: keep only the newest browser frame under pressure.
213                    q.q.retain(|x| x.prio != SendPrio::Browser);
214                    if q.q.len() >= cap {
215                        // Still full of higher-prio work: drop this browser frame.
216                        continue;
217                    }
218                } else {
219                    while q.q.len() >= cap {
220                        // Drop lowest-priority (browser first, then mail) before refusing control/PTY.
221                        if let Some(i) = q
222                            .q
223                            .iter()
224                            .enumerate()
225                            .filter(|(_, x)| x.prio >= SendPrio::Mail)
226                            .max_by_key(|(_, x)| x.prio)
227                            .map(|(i, _)| i)
228                        {
229                            q.q.remove(i);
230                        } else if prio <= SendPrio::Pty {
231                            // Control/PTY: drop oldest same-or-lower in this channel.
232                            q.q.pop_front();
233                        } else {
234                            break;
235                        }
236                    }
237                    if q.q.len() >= cap {
238                        continue;
239                    }
240                }
241                q.q.push_back(Queued { prio, frame });
242                drop(st);
243                wake_in.notify_one();
244            }
245        });
246
247        // Fair/priority writer: always pick the globally highest-priority waiting frame
248        // whose channel has credit (DATA) or needs no credit (OPEN/control non-DATA).
249        let state_wr = Arc::clone(&out_state);
250        let wake_wr = Arc::clone(&wake);
251        let alive_wr = Arc::clone(&alive);
252        let dead_wr = Arc::clone(&dead);
253        let routes_wr = Arc::clone(&routes);
254        let writer = tokio::spawn(async move {
255            loop {
256                if !alive_wr.load(Ordering::SeqCst) {
257                    break;
258                }
259                let next = {
260                    let mut st = state_wr.lock().expect("out");
261                    pick_next(&mut st)
262                };
263                let Some((ch, is_data, frame)) = next else {
264                    wake_wr.notified().await;
265                    continue;
266                };
267                let Ok(bytes) = frame.encode() else {
268                    kill(&alive_wr, &dead_wr, &routes_wr);
269                    break;
270                };
271                match tokio::time::timeout(WRITE_TIMEOUT, wr.write_all(&bytes)).await {
272                    Ok(Ok(())) => {
273                        if is_data {
274                            let mut st = state_wr.lock().expect("out");
275                            if let Some(q) = st.chans.get_mut(&ch) {
276                                if let Some(c) = q.credit.as_mut() {
277                                    *c = c.saturating_sub(1);
278                                }
279                            }
280                        }
281                    }
282                    Ok(Err(_)) | Err(_) => {
283                        tracing::warn!(%ch, "uplink write timed out or failed; reconnecting");
284                        kill(&alive_wr, &dead_wr, &routes_wr);
285                        break;
286                    }
287                }
288            }
289        });
290
291        let enq_reader = enq_tx.clone();
292        let reader_journal = journal.clone();
293        let reader_routes = Arc::clone(&routes);
294        let alive_rd = Arc::clone(&alive);
295        let dead_rd = Arc::clone(&dead);
296        let state_rd = Arc::clone(&out_state);
297        let wake_rd = Arc::clone(&wake);
298        let credits_rd = Arc::clone(&credits);
299        let credit_pulse_rd = Arc::clone(&credit_pulse);
300        let last_rx_rd = Arc::clone(&last_rx);
301        let reader = tokio::spawn(async move {
302            let mut buf = vec![0u8; 128 * 1024];
303            let mut filled = 0usize;
304            loop {
305                let Ok(n) = rd.read(&mut buf[filled..]).await else { break };
306                if n == 0 {
307                    break;
308                }
309                filled += n;
310                last_rx_rd.store(born.elapsed().as_millis() as u64, Ordering::SeqCst);
311                loop {
312                    match Frame::decode(&buf[..filled]) {
313                        Ok((frame, used)) => {
314                            buf.copy_within(used..filled, 0);
315                            filled -= used;
316                            match frame.frame_type {
317                                FrameType::Data => {
318                                    let tx = reader_routes.lock().expect("routes").get(&frame.channel_id).cloned();
319                                    if let Some(tx) = tx {
320                                        let _ = tx.send(frame.payload);
321                                    }
322                                }
323                                FrameType::Window => {
324                                    let credit = frame.seq.min(u64::from(u32::MAX)) as u32;
325                                    {
326                                        let mut st = state_rd.lock().expect("out");
327                                        st.chans.entry(frame.channel_id).or_insert_with(ChanQ::new).credit = Some(credit);
328                                    }
329                                    credits_rd.lock().expect("credits").insert(frame.channel_id, credit);
330                                    credit_pulse_rd.notify_waiters();
331                                    wake_rd.notify_one();
332                                }
333                                FrameType::JournalReq => {
334                                    let Some(j) = reader_journal.as_ref() else { continue };
335                                    let limit = frame.journal_req_limit().unwrap_or(64) as usize;
336                                    if let Ok(body) = j.journal().reply_body(frame.channel_id, frame.seq, limit) {
337                                        if let Ok(reply) = Frame::journal_reply(frame.channel_id, &body) {
338                                            let _ = enq_reader.send((frame.channel_id, SendPrio::Control, reply));
339                                        }
340                                    }
341                                }
342                                FrameType::GoAway | FrameType::Close => {
343                                    kill(&alive_rd, &dead_rd, &reader_routes);
344                                    return;
345                                }
346                                _ => {}
347                            }
348                        }
349                        Err(FrameError::Truncated) => break,
350                        Err(_) => {
351                            kill(&alive_rd, &dead_rd, &reader_routes);
352                            return;
353                        }
354                    }
355                }
356                if filled == buf.len() {
357                    // Same grow path as C2: do not tear the bridge when a
358                    // single frame exceeds the initial 128 KiB read buf.
359                    if grow_read_buf(&mut buf, filled).is_err() {
360                        break;
361                    }
362                }
363            }
364            kill(&alive_rd, &dead_rd, &reader_routes);
365        });
366
367        let ka_tx = enq_tx.clone();
368        let alive_ka = Arc::clone(&alive);
369        let dead_ka = Arc::clone(&dead);
370        let routes_ka = Arc::clone(&routes);
371        let keepalive = tokio::spawn(async move {
372            let mut nonce = 0u64;
373            let mut tick = tokio::time::interval(UPLINK_PING_EVERY);
374            tick.tick().await;
375            loop {
376                tick.tick().await;
377                if !alive_ka.load(Ordering::SeqCst) {
378                    break;
379                }
380                let now = born.elapsed().as_millis() as u64;
381                let idle = now.saturating_sub(last_rx.load(Ordering::SeqCst));
382                if idle > UPLINK_DEAD_AFTER.as_millis() as u64 {
383                    tracing::warn!(idle_ms = idle, "uplink silent past keepalive deadline; reconnecting");
384                    kill(&alive_ka, &dead_ka, &routes_ka);
385                    break;
386                }
387                nonce = nonce.wrapping_add(1);
388                if ka_tx.send((0, SendPrio::Control, Frame::ping(nonce))).is_err() {
389                    break;
390                }
391            }
392        });
393
394        Ok(Arc::new(Self {
395            enqueue: enq_tx,
396            routes,
397            credits,
398            credit_pulse,
399            journal,
400            alive,
401            dead,
402            tasks: vec![ingress, writer, reader, keepalive],
403        }))
404    }
405
406    pub fn is_alive(&self) -> bool {
407        self.alive.load(Ordering::SeqCst)
408    }
409
410    /// Wait until the uplink is marked dead (write timeout / peer close).
411    pub async fn closed(&self) {
412        if !self.is_alive() {
413            return;
414        }
415        self.dead.notified().await;
416    }
417
418    fn push(&self, channel_id: u32, prio: SendPrio, frame: Frame) -> Result<(), String> {
419        if !self.is_alive() {
420            return Err("uplink closed".into());
421        }
422        self.enqueue
423            .send((channel_id, prio, frame))
424            .map_err(|_| "uplink closed".to_owned())
425    }
426
427    /// OPEN `channel_id`, wait for WINDOW credit, then complete Noise_NK.
428    pub async fn open(
429        self: &Arc<Self>,
430        node_id: &str,
431        channel_id: u32,
432        kind: ChannelKind,
433        keys: UpstreamKeys,
434        handshake_timeout: Duration,
435    ) -> Result<NodeSession, String> {
436        let (down_tx, mut down_rx) = mpsc::unbounded_channel::<Vec<u8>>();
437        self.routes.lock().expect("routes").insert(channel_id, down_tx);
438        self.credits.lock().expect("credits").remove(&channel_id);
439        let open = Frame::open(channel_id, &OpenPayload { channel_kind: kind, node_id: node_id.to_owned() })
440            .map_err(|e| e.to_string())?;
441        self.push(channel_id, SendPrio::Control, open)?;
442
443        // Wait until C2 grants WINDOW credit after OPEN (avoids sending Handshake1
444        // into a stream that still has credit 0 from a previous session).
445        tokio::time::timeout(handshake_timeout, async {
446            loop {
447                if !self.is_alive() {
448                    return Err("uplink closed".to_owned());
449                }
450                if self.credits.lock().expect("credits").get(&channel_id).copied().unwrap_or(0) > 0 {
451                    return Ok(());
452                }
453                tokio::select! {
454                    _ = self.credit_pulse.notified() => {}
455                    _ = self.closed() => return Err("uplink closed".to_owned()),
456                }
457            }
458        })
459        .await
460        .map_err(|_| "window credit timeout after OPEN".to_owned())??;
461
462        let mut initiator = NoiseInitiator::new(keys.recipient_public).map_err(|e| e.to_string())?;
463        let msg1 = initiator.write_message1().map_err(|e| e.to_string())?;
464        self.send_inner(channel_id, keys, SendPrio::Control, &encode_inner(InnerKind::Handshake1, &msg1))?;
465
466        let msg2 = tokio::time::timeout(handshake_timeout, async {
467            loop {
468                let Some(bytes) = down_rx.recv().await else { return Err("uplink closed".to_owned()) };
469                if let Some((InnerKind::Handshake2, body)) = decode_inner(&bytes) {
470                    return Ok(body.to_vec());
471                }
472            }
473        })
474        .await
475        .map_err(|_| "handshake timeout".to_owned())??;
476        let transport = initiator.read_message2(&msg2).map_err(|e| e.to_string())?;
477        Ok(NodeSession {
478            channel_id,
479            keys,
480            default_prio: match kind {
481                ChannelKind::M4a => SendPrio::Mail,
482                ChannelKind::G4aControl | ChannelKind::G4aPty | ChannelKind::G4aEvents => SendPrio::Pty,
483                ChannelKind::Vault | ChannelKind::Observation => SendPrio::Mail,
484            },
485            down: down_rx,
486            transport,
487            journal: self.journal.clone(),
488            uplink: Arc::clone(self),
489        })
490    }
491
492    fn send_inner(&self, channel_id: u32, keys: UpstreamKeys, prio: SendPrio, inner: &[u8]) -> Result<(), String> {
493        if let Some(j) = &self.journal {
494            j.append(channel_id, inner)?;
495        }
496        let sealed = node_seal_outer(keys.c2_public, inner).map_err(|e| e.to_string())?;
497        let frame = Frame::data(channel_id, 0, sealed).map_err(|e| e.to_string())?;
498        self.push(channel_id, prio, frame)
499    }
500}
501
502fn kill(alive: &AtomicBool, dead: &Notify, routes: &Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>) {
503    if alive.swap(false, Ordering::SeqCst) {
504        routes.lock().expect("routes").clear();
505        dead.notify_waiters();
506    }
507}
508
509/// Pick the highest-priority queued frame that may be sent now.
510fn pick_next(st: &mut OutState) -> Option<(u32, bool, Frame)> {
511    // Candidates: (prio, channel, index). DATA frames of one channel are Noise ciphertexts
512    // encrypted in queue order, so only the channel's first DATA frame may go next
513    // (overtaking breaks the receiver's nonce sequence); non-DATA frames may go any time.
514    let mut best: Option<(SendPrio, u32, usize)> = None;
515    for (&ch, cq) in &st.chans {
516        let first_data = cq.q.iter().position(|it| matches!(it.frame.frame_type, FrameType::Data));
517        for (i, item) in cq.q.iter().enumerate() {
518            let needs_credit = matches!(item.frame.frame_type, FrameType::Data);
519            if needs_credit && Some(i) != first_data {
520                continue;
521            }
522            if needs_credit {
523                match cq.credit {
524                    Some(0) => continue,
525                    None => {
526                        // No WINDOW yet: allow Control DATA (handshake) so OPEN→Handshake1 works;
527                        // hold back bulk traffic.
528                        if item.prio > SendPrio::Control {
529                            continue;
530                        }
531                    }
532                    Some(_) => {}
533                }
534            }
535            let cand = (item.prio, ch, i);
536            if best.is_none_or(|b| cand.0 < b.0 || (cand.0 == b.0 && cand.1 < b.1)) {
537                best = Some(cand);
538            }
539        }
540    }
541    let (prio, ch, idx) = best?;
542    let cq = st.chans.get_mut(&ch)?;
543    let item = cq.q.remove(idx)?;
544    debug_assert_eq!(item.prio, prio);
545    let is_data = matches!(item.frame.frame_type, FrameType::Data);
546    Some((ch, is_data, item.frame))
547}
548
549pub struct NodeSession {
550    channel_id: u32,
551    keys: UpstreamKeys,
552    default_prio: SendPrio,
553    down: mpsc::UnboundedReceiver<Vec<u8>>,
554    transport: NoiseTransport,
555    journal: Option<Arc<UplinkJournal>>,
556    uplink: Arc<NodeUplink>,
557}
558
559impl NodeSession {
560    pub async fn connect(
561        uplink: SocketAddr,
562        node_id: &str,
563        channel_id: u32,
564        kind: ChannelKind,
565        keys: UpstreamKeys,
566        journal: Option<Arc<UplinkJournal>>,
567        handshake_timeout: Duration,
568    ) -> Result<Self, String> {
569        let link = NodeUplink::connect(uplink, journal).await?;
570        link.open(node_id, channel_id, kind, keys, handshake_timeout).await
571    }
572
573    pub fn uplink(&self) -> Arc<NodeUplink> {
574        Arc::clone(&self.uplink)
575    }
576
577    pub async fn send(&mut self, plaintext: &[u8]) -> Result<(), String> {
578        self.send_prio(plaintext, self.default_prio).await
579    }
580
581    pub async fn send_prio(&mut self, plaintext: &[u8], prio: SendPrio) -> Result<(), String> {
582        if !self.uplink.is_alive() {
583            return Err("uplink closed".into());
584        }
585        let ct = self.transport.encrypt(plaintext).map_err(|e| e.to_string())?;
586        let inner = encode_inner(InnerKind::Transport, &ct);
587        self.uplink.send_inner(self.channel_id, self.keys, prio, &inner)
588    }
589
590    pub async fn recv(&mut self, timeout: Duration) -> Result<Vec<u8>, String> {
591        tokio::time::timeout(timeout, async {
592            loop {
593                tokio::select! {
594                    bytes = self.down.recv() => {
595                        let Some(bytes) = bytes else { return Err("uplink closed".to_owned()) };
596                        if let Some((InnerKind::Transport, ct)) = decode_inner(&bytes) {
597                            return self.transport.decrypt(ct).map_err(|e| e.to_string());
598                        }
599                    }
600                    _ = self.uplink.closed() => return Err("uplink closed".to_owned()),
601                }
602            }
603        })
604        .await
605        .map_err(|_| "reply timeout".to_owned())?
606    }
607
608    pub async fn call(&mut self, op: &str, payload: &Value, timeout: Duration) -> Result<Value, String> {
609        let req = serde_json::to_vec(&json!({ "op": op, "payload": payload })).map_err(|e| e.to_string())?;
610        self.send_prio(&req, SendPrio::Mail).await?;
611        let reply = self.recv(timeout).await?;
612        let reply: Value = serde_json::from_slice(&reply).map_err(|e| e.to_string())?;
613        let status = reply.get("status").and_then(Value::as_u64).unwrap_or(500);
614        let body = reply.get("body").cloned().unwrap_or(Value::Null);
615        if status == 200 {
616            Ok(body)
617        } else {
618            let msg = body.get("error").and_then(Value::as_str).unwrap_or("error");
619            Err(format!("{status}: {msg}"))
620        }
621    }
622}
623
624#[cfg(test)]
625mod tests {
626    use super::*;
627
628    #[test]
629    fn pick_prefers_control_across_channels_but_keeps_data_order_within_one() {
630        let mut st = OutState { chans: BTreeMap::new() };
631        let a = st.chans.entry(1).or_insert_with(ChanQ::new);
632        a.credit = Some(10);
633        a.q.push_back(Queued { prio: SendPrio::Browser, frame: Frame::data(1, 0, vec![1]).unwrap() });
634        a.q.push_back(Queued { prio: SendPrio::Control, frame: Frame::data(1, 0, vec![2]).unwrap() });
635        let b = st.chans.entry(2).or_insert_with(ChanQ::new);
636        b.credit = Some(10);
637        b.q.push_back(Queued { prio: SendPrio::Pty, frame: Frame::data(2, 0, vec![3]).unwrap() });
638        // Channel 2's Pty beats channel 1's head (Browser); channel 1 stays in order.
639        let order: Vec<u8> = std::iter::from_fn(|| pick_next(&mut st).map(|(_, _, f)| f.payload[0])).collect();
640        assert_eq!(order, vec![3, 1, 2]);
641    }
642
643    #[test]
644    fn pick_holds_bulk_until_window() {
645        let mut st = OutState { chans: BTreeMap::new() };
646        let q = st.chans.entry(2).or_insert_with(ChanQ::new);
647        q.credit = None;
648        q.q.push_back(Queued {
649            prio: SendPrio::Control,
650            frame: Frame::data(2, 0, vec![1]).unwrap(),
651        });
652        q.q.push_back(Queued {
653            prio: SendPrio::Mail,
654            frame: Frame::data(2, 0, vec![9]).unwrap(),
655        });
656        let (_, _, f) = pick_next(&mut st).unwrap();
657        assert_eq!(f.payload, vec![1]);
658        assert!(pick_next(&mut st).is_none());
659        st.chans.get_mut(&2).unwrap().credit = Some(4);
660        let (_, _, f) = pick_next(&mut st).unwrap();
661        assert_eq!(f.payload, vec![9]);
662    }
663}
hatchery_node/restream_session.rs

hatchery_node/
restream_session.rs
