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, 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/// Keys C2 hands a node at enrollment: its own outer key and the upstream recipient key.
32#[derive(Clone, Copy, Debug)]
33pub struct UpstreamKeys {
34    pub c2_public: PublicKeyBytes,
35    pub recipient_public: PublicKeyBytes,
36}
37
38/// Outbound priority on the shared uplink. Lower ordinal = sent first.
39#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
40pub enum SendPrio {
41    /// OPEN, Noise handshake, journal replies, g4a request/response (incl. PTY input path).
42    Control = 0,
43    /// g4a terminal / event frames.
44    Pty = 1,
45    /// m4a mail/task calls.
46    Mail = 2,
47    /// dig2browser frames — droppable / coalesced under pressure.
48    Browser = 3,
49}
50
51const WRITE_TIMEOUT: Duration = Duration::from_secs(5);
52const CHAN_CAP_CONTROL: usize = 64;
53const CHAN_CAP_PTY: usize = 128;
54const CHAN_CAP_MAIL: usize = 16;
55const CHAN_CAP_BROWSER: usize = 8;
56
57struct Queued {
58    prio: SendPrio,
59    frame: Frame,
60}
61
62struct ChanQ {
63    q: VecDeque<Queued>,
64    /// Remaining WINDOW credit for DATA frames; None until the first WINDOW.
65    credit: Option<u32>,
66}
67
68impl ChanQ {
69    fn new() -> Self {
70        Self { q: VecDeque::new(), credit: None }
71    }
72
73    fn cap_for(prio: SendPrio) -> usize {
74        match prio {
75            SendPrio::Control => CHAN_CAP_CONTROL,
76            SendPrio::Pty => CHAN_CAP_PTY,
77            SendPrio::Mail => CHAN_CAP_MAIL,
78            SendPrio::Browser => CHAN_CAP_BROWSER,
79        }
80    }
81}
82
83struct OutState {
84    chans: BTreeMap<u32, ChanQ>,
85}
86
87/// One TCP uplink to C2 carrying any number of channels.
88pub struct NodeUplink {
89    enqueue: mpsc::UnboundedSender<(u32, SendPrio, Frame)>,
90    routes: Arc<Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>>,
91    /// Last WINDOW credit per channel (updated by the reader).
92    credits: Arc<Mutex<BTreeMap<u32, u32>>>,
93    credit_pulse: Arc<Notify>,
94    journal: Option<Arc<UplinkJournal>>,
95    alive: Arc<AtomicBool>,
96    dead: Arc<Notify>,
97    tasks: Vec<JoinHandle<()>>,
98}
99
100impl Drop for NodeUplink {
101    fn drop(&mut self) {
102        self.alive.store(false, Ordering::SeqCst);
103        for t in &self.tasks {
104            t.abort();
105        }
106    }
107}
108
109impl NodeUplink {
110    pub async fn connect(uplink: SocketAddr, journal: Option<Arc<UplinkJournal>>) -> Result<Arc<Self>, String> {
111        let stream = TcpStream::connect(uplink).await.map_err(|e| e.to_string())?;
112        let _ = stream.set_nodelay(true);
113        let (mut rd, mut wr) = stream.into_split();
114        let (enq_tx, mut enq_rx) = mpsc::unbounded_channel::<(u32, SendPrio, Frame)>();
115        let routes: Arc<Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>> = Arc::default();
116        let credits: Arc<Mutex<BTreeMap<u32, u32>>> = Arc::default();
117        let credit_pulse = Arc::new(Notify::new());
118        let alive = Arc::new(AtomicBool::new(true));
119        let dead = Arc::new(Notify::new());
120        let wake = Arc::new(Notify::new());
121
122        let out_state = Arc::new(Mutex::new(OutState { chans: BTreeMap::new() }));
123
124        // Ingress into per-channel queues (never blocks the producer on TCP).
125        let state_in = Arc::clone(&out_state);
126        let wake_in = Arc::clone(&wake);
127        let alive_in = Arc::clone(&alive);
128        let ingress = tokio::spawn(async move {
129            while let Some((ch, prio, frame)) = enq_rx.recv().await {
130                if !alive_in.load(Ordering::SeqCst) {
131                    break;
132                }
133                let mut st = state_in.lock().expect("out");
134                let q = st.chans.entry(ch).or_insert_with(ChanQ::new);
135                let cap = ChanQ::cap_for(prio);
136                if prio == SendPrio::Browser {
137                    // Coalesce: keep only the newest browser frame under pressure.
138                    q.q.retain(|x| x.prio != SendPrio::Browser);
139                    if q.q.len() >= cap {
140                        // Still full of higher-prio work: drop this browser frame.
141                        continue;
142                    }
143                } else {
144                    while q.q.len() >= cap {
145                        // Drop lowest-priority (browser first, then mail) before refusing control/PTY.
146                        if let Some(i) = q
147                            .q
148                            .iter()
149                            .enumerate()
150                            .filter(|(_, x)| x.prio >= SendPrio::Mail)
151                            .max_by_key(|(_, x)| x.prio)
152                            .map(|(i, _)| i)
153                        {
154                            q.q.remove(i);
155                        } else if prio <= SendPrio::Pty {
156                            // Control/PTY: drop oldest same-or-lower in this channel.
157                            q.q.pop_front();
158                        } else {
159                            break;
160                        }
161                    }
162                    if q.q.len() >= cap {
163                        continue;
164                    }
165                }
166                q.q.push_back(Queued { prio, frame });
167                drop(st);
168                wake_in.notify_one();
169            }
170        });
171
172        // Fair/priority writer: always pick the globally highest-priority waiting frame
173        // whose channel has credit (DATA) or needs no credit (OPEN/control non-DATA).
174        let state_wr = Arc::clone(&out_state);
175        let wake_wr = Arc::clone(&wake);
176        let alive_wr = Arc::clone(&alive);
177        let dead_wr = Arc::clone(&dead);
178        let routes_wr = Arc::clone(&routes);
179        let writer = tokio::spawn(async move {
180            loop {
181                if !alive_wr.load(Ordering::SeqCst) {
182                    break;
183                }
184                let next = {
185                    let mut st = state_wr.lock().expect("out");
186                    pick_next(&mut st)
187                };
188                let Some((ch, is_data, frame)) = next else {
189                    wake_wr.notified().await;
190                    continue;
191                };
192                let Ok(bytes) = frame.encode() else {
193                    kill(&alive_wr, &dead_wr, &routes_wr);
194                    break;
195                };
196                match tokio::time::timeout(WRITE_TIMEOUT, wr.write_all(&bytes)).await {
197                    Ok(Ok(())) => {
198                        if is_data {
199                            let mut st = state_wr.lock().expect("out");
200                            if let Some(q) = st.chans.get_mut(&ch) {
201                                if let Some(c) = q.credit.as_mut() {
202                                    *c = c.saturating_sub(1);
203                                }
204                            }
205                        }
206                    }
207                    Ok(Err(_)) | Err(_) => {
208                        tracing::warn!(%ch, "uplink write timed out or failed; reconnecting");
209                        kill(&alive_wr, &dead_wr, &routes_wr);
210                        break;
211                    }
212                }
213            }
214        });
215
216        let enq_reader = enq_tx.clone();
217        let reader_journal = journal.clone();
218        let reader_routes = Arc::clone(&routes);
219        let alive_rd = Arc::clone(&alive);
220        let dead_rd = Arc::clone(&dead);
221        let state_rd = Arc::clone(&out_state);
222        let wake_rd = Arc::clone(&wake);
223        let credits_rd = Arc::clone(&credits);
224        let credit_pulse_rd = Arc::clone(&credit_pulse);
225        let reader = tokio::spawn(async move {
226            let mut buf = vec![0u8; 128 * 1024];
227            let mut filled = 0usize;
228            loop {
229                let Ok(n) = rd.read(&mut buf[filled..]).await else { break };
230                if n == 0 {
231                    break;
232                }
233                filled += n;
234                loop {
235                    match Frame::decode(&buf[..filled]) {
236                        Ok((frame, used)) => {
237                            buf.copy_within(used..filled, 0);
238                            filled -= used;
239                            match frame.frame_type {
240                                FrameType::Data => {
241                                    let tx = reader_routes.lock().expect("routes").get(&frame.channel_id).cloned();
242                                    if let Some(tx) = tx {
243                                        let _ = tx.send(frame.payload);
244                                    }
245                                }
246                                FrameType::Window => {
247                                    let credit = frame.seq.min(u64::from(u32::MAX)) as u32;
248                                    {
249                                        let mut st = state_rd.lock().expect("out");
250                                        st.chans.entry(frame.channel_id).or_insert_with(ChanQ::new).credit = Some(credit);
251                                    }
252                                    credits_rd.lock().expect("credits").insert(frame.channel_id, credit);
253                                    credit_pulse_rd.notify_waiters();
254                                    wake_rd.notify_one();
255                                }
256                                FrameType::JournalReq => {
257                                    let Some(j) = reader_journal.as_ref() else { continue };
258                                    let limit = frame.journal_req_limit().unwrap_or(64) as usize;
259                                    if let Ok(body) = j.journal().reply_body(frame.channel_id, frame.seq, limit) {
260                                        if let Ok(reply) = Frame::journal_reply(frame.channel_id, &body) {
261                                            let _ = enq_reader.send((frame.channel_id, SendPrio::Control, reply));
262                                        }
263                                    }
264                                }
265                                FrameType::GoAway | FrameType::Close => {
266                                    kill(&alive_rd, &dead_rd, &reader_routes);
267                                    return;
268                                }
269                                _ => {}
270                            }
271                        }
272                        Err(FrameError::Truncated) => break,
273                        Err(_) => {
274                            kill(&alive_rd, &dead_rd, &reader_routes);
275                            return;
276                        }
277                    }
278                }
279                if filled == buf.len() {
280                    // Same grow path as C2: do not tear the bridge when a
281                    // single frame exceeds the initial 128 KiB read buf.
282                    if grow_read_buf(&mut buf, filled).is_err() {
283                        break;
284                    }
285                }
286            }
287            kill(&alive_rd, &dead_rd, &reader_routes);
288        });
289
290        Ok(Arc::new(Self {
291            enqueue: enq_tx,
292            routes,
293            credits,
294            credit_pulse,
295            journal,
296            alive,
297            dead,
298            tasks: vec![ingress, writer, reader],
299        }))
300    }
301
302    pub fn is_alive(&self) -> bool {
303        self.alive.load(Ordering::SeqCst)
304    }
305
306    /// Wait until the uplink is marked dead (write timeout / peer close).
307    pub async fn closed(&self) {
308        if !self.is_alive() {
309            return;
310        }
311        self.dead.notified().await;
312    }
313
314    fn push(&self, channel_id: u32, prio: SendPrio, frame: Frame) -> Result<(), String> {
315        if !self.is_alive() {
316            return Err("uplink closed".into());
317        }
318        self.enqueue
319            .send((channel_id, prio, frame))
320            .map_err(|_| "uplink closed".to_owned())
321    }
322
323    /// OPEN `channel_id`, wait for WINDOW credit, then complete Noise_NK.
324    pub async fn open(
325        self: &Arc<Self>,
326        node_id: &str,
327        channel_id: u32,
328        kind: ChannelKind,
329        keys: UpstreamKeys,
330        handshake_timeout: Duration,
331    ) -> Result<NodeSession, String> {
332        let (down_tx, mut down_rx) = mpsc::unbounded_channel::<Vec<u8>>();
333        self.routes.lock().expect("routes").insert(channel_id, down_tx);
334        self.credits.lock().expect("credits").remove(&channel_id);
335        let open = Frame::open(channel_id, &OpenPayload { channel_kind: kind, node_id: node_id.to_owned() })
336            .map_err(|e| e.to_string())?;
337        self.push(channel_id, SendPrio::Control, open)?;
338
339        // Wait until C2 grants WINDOW credit after OPEN (avoids sending Handshake1
340        // into a stream that still has credit 0 from a previous session).
341        tokio::time::timeout(handshake_timeout, async {
342            loop {
343                if !self.is_alive() {
344                    return Err("uplink closed".to_owned());
345                }
346                if self.credits.lock().expect("credits").get(&channel_id).copied().unwrap_or(0) > 0 {
347                    return Ok(());
348                }
349                tokio::select! {
350                    _ = self.credit_pulse.notified() => {}
351                    _ = self.closed() => return Err("uplink closed".to_owned()),
352                }
353            }
354        })
355        .await
356        .map_err(|_| "window credit timeout after OPEN".to_owned())??;
357
358        let mut initiator = NoiseInitiator::new(keys.recipient_public).map_err(|e| e.to_string())?;
359        let msg1 = initiator.write_message1().map_err(|e| e.to_string())?;
360        self.send_inner(channel_id, keys, SendPrio::Control, &encode_inner(InnerKind::Handshake1, &msg1))?;
361
362        let msg2 = tokio::time::timeout(handshake_timeout, async {
363            loop {
364                let Some(bytes) = down_rx.recv().await else { return Err("uplink closed".to_owned()) };
365                if let Some((InnerKind::Handshake2, body)) = decode_inner(&bytes) {
366                    return Ok(body.to_vec());
367                }
368            }
369        })
370        .await
371        .map_err(|_| "handshake timeout".to_owned())??;
372        let transport = initiator.read_message2(&msg2).map_err(|e| e.to_string())?;
373        Ok(NodeSession {
374            channel_id,
375            keys,
376            default_prio: match kind {
377                ChannelKind::M4a => SendPrio::Mail,
378                ChannelKind::G4aControl | ChannelKind::G4aPty | ChannelKind::G4aEvents => SendPrio::Pty,
379                ChannelKind::Vault | ChannelKind::Observation => SendPrio::Mail,
380            },
381            down: down_rx,
382            transport,
383            journal: self.journal.clone(),
384            uplink: Arc::clone(self),
385        })
386    }
387
388    fn send_inner(&self, channel_id: u32, keys: UpstreamKeys, prio: SendPrio, inner: &[u8]) -> Result<(), String> {
389        if let Some(j) = &self.journal {
390            j.append(channel_id, inner)?;
391        }
392        let sealed = node_seal_outer(keys.c2_public, inner).map_err(|e| e.to_string())?;
393        let frame = Frame::data(channel_id, 0, sealed).map_err(|e| e.to_string())?;
394        self.push(channel_id, prio, frame)
395    }
396}
397
398fn kill(alive: &AtomicBool, dead: &Notify, routes: &Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>) {
399    if alive.swap(false, Ordering::SeqCst) {
400        routes.lock().expect("routes").clear();
401        dead.notify_waiters();
402    }
403}
404
405/// Pick the highest-priority queued frame that may be sent now.
406fn pick_next(st: &mut OutState) -> Option<(u32, bool, Frame)> {
407    // Flatten candidates: (prio, channel, index)
408    let mut best: Option<(SendPrio, u32, usize)> = None;
409    for (&ch, cq) in &st.chans {
410        for (i, item) in cq.q.iter().enumerate() {
411            let needs_credit = matches!(item.frame.frame_type, FrameType::Data);
412            if needs_credit {
413                match cq.credit {
414                    Some(0) => continue,
415                    None => {
416                        // No WINDOW yet: allow Control DATA (handshake) so OPEN→Handshake1 works;
417                        // hold back bulk traffic.
418                        if item.prio > SendPrio::Control {
419                            continue;
420                        }
421                    }
422                    Some(_) => {}
423                }
424            }
425            let cand = (item.prio, ch, i);
426            if best.is_none_or(|b| cand.0 < b.0 || (cand.0 == b.0 && cand.1 < b.1)) {
427                best = Some(cand);
428            }
429        }
430    }
431    let (prio, ch, idx) = best?;
432    let cq = st.chans.get_mut(&ch)?;
433    let item = cq.q.remove(idx)?;
434    debug_assert_eq!(item.prio, prio);
435    let is_data = matches!(item.frame.frame_type, FrameType::Data);
436    Some((ch, is_data, item.frame))
437}
438
439pub struct NodeSession {
440    channel_id: u32,
441    keys: UpstreamKeys,
442    default_prio: SendPrio,
443    down: mpsc::UnboundedReceiver<Vec<u8>>,
444    transport: NoiseTransport,
445    journal: Option<Arc<UplinkJournal>>,
446    uplink: Arc<NodeUplink>,
447}
448
449impl NodeSession {
450    pub async fn connect(
451        uplink: SocketAddr,
452        node_id: &str,
453        channel_id: u32,
454        kind: ChannelKind,
455        keys: UpstreamKeys,
456        journal: Option<Arc<UplinkJournal>>,
457        handshake_timeout: Duration,
458    ) -> Result<Self, String> {
459        let link = NodeUplink::connect(uplink, journal).await?;
460        link.open(node_id, channel_id, kind, keys, handshake_timeout).await
461    }
462
463    pub fn uplink(&self) -> Arc<NodeUplink> {
464        Arc::clone(&self.uplink)
465    }
466
467    pub async fn send(&mut self, plaintext: &[u8]) -> Result<(), String> {
468        self.send_prio(plaintext, self.default_prio).await
469    }
470
471    pub async fn send_prio(&mut self, plaintext: &[u8], prio: SendPrio) -> Result<(), String> {
472        if !self.uplink.is_alive() {
473            return Err("uplink closed".into());
474        }
475        let ct = self.transport.encrypt(plaintext).map_err(|e| e.to_string())?;
476        let inner = encode_inner(InnerKind::Transport, &ct);
477        self.uplink.send_inner(self.channel_id, self.keys, prio, &inner)
478    }
479
480    pub async fn recv(&mut self, timeout: Duration) -> Result<Vec<u8>, String> {
481        tokio::time::timeout(timeout, async {
482            loop {
483                tokio::select! {
484                    bytes = self.down.recv() => {
485                        let Some(bytes) = bytes else { return Err("uplink closed".to_owned()) };
486                        if let Some((InnerKind::Transport, ct)) = decode_inner(&bytes) {
487                            return self.transport.decrypt(ct).map_err(|e| e.to_string());
488                        }
489                    }
490                    _ = self.uplink.closed() => return Err("uplink closed".to_owned()),
491                }
492            }
493        })
494        .await
495        .map_err(|_| "reply timeout".to_owned())?
496    }
497
498    pub async fn call(&mut self, op: &str, payload: &Value, timeout: Duration) -> Result<Value, String> {
499        let req = serde_json::to_vec(&json!({ "op": op, "payload": payload })).map_err(|e| e.to_string())?;
500        self.send_prio(&req, SendPrio::Mail).await?;
501        let reply = self.recv(timeout).await?;
502        let reply: Value = serde_json::from_slice(&reply).map_err(|e| e.to_string())?;
503        let status = reply.get("status").and_then(Value::as_u64).unwrap_or(500);
504        let body = reply.get("body").cloned().unwrap_or(Value::Null);
505        if status == 200 {
506            Ok(body)
507        } else {
508            let msg = body.get("error").and_then(Value::as_str).unwrap_or("error");
509            Err(format!("{status}: {msg}"))
510        }
511    }
512}
513
514#[cfg(test)]
515mod tests {
516    use super::*;
517
518    #[test]
519    fn pick_prefers_control_over_browser() {
520        let mut st = OutState { chans: BTreeMap::new() };
521        let q = st.chans.entry(1).or_insert_with(ChanQ::new);
522        q.credit = Some(10);
523        q.q.push_back(Queued {
524            prio: SendPrio::Browser,
525            frame: Frame::data(1, 0, vec![1]).unwrap(),
526        });
527        q.q.push_back(Queued {
528            prio: SendPrio::Control,
529            frame: Frame::data(1, 0, vec![2]).unwrap(),
530        });
531        q.q.push_back(Queued {
532            prio: SendPrio::Pty,
533            frame: Frame::data(1, 0, vec![3]).unwrap(),
534        });
535        let (_, _, f) = pick_next(&mut st).unwrap();
536        assert_eq!(f.payload, vec![2]);
537        let (_, _, f) = pick_next(&mut st).unwrap();
538        assert_eq!(f.payload, vec![3]);
539        let (_, _, f) = pick_next(&mut st).unwrap();
540        assert_eq!(f.payload, vec![1]);
541    }
542
543    #[test]
544    fn pick_holds_bulk_until_window() {
545        let mut st = OutState { chans: BTreeMap::new() };
546        let q = st.chans.entry(2).or_insert_with(ChanQ::new);
547        q.credit = None;
548        q.q.push_back(Queued {
549            prio: SendPrio::Mail,
550            frame: Frame::data(2, 0, vec![9]).unwrap(),
551        });
552        q.q.push_back(Queued {
553            prio: SendPrio::Control,
554            frame: Frame::data(2, 0, vec![1]).unwrap(),
555        });
556        let (_, _, f) = pick_next(&mut st).unwrap();
557        assert_eq!(f.payload, vec![1]);
558        assert!(pick_next(&mut st).is_none());
559        st.chans.get_mut(&2).unwrap().credit = Some(4);
560        let (_, _, f) = pick_next(&mut st).unwrap();
561        assert_eq!(f.payload, vec![9]);
562    }
563}
hatchery_node/restream_session.rs

hatchery_node/
restream_session.rs
