hatchery_node/

restream_session.rs

//! Node end of a live restreamer session (A2 / A3).
//!
//! The node dials C2's uplink, OPENs a channel, then runs the inner Noise_NK
//! handshake toward the relayed recipient key as ordinary DATA frames:
//! msg1 goes up through C2 into the core drain, msg2 comes back down the same
//! stream. After that every request and reply is Noise transport ciphertext,
//! wrapped in C2's outer layer on the way up. C2 only ever sees opaque bytes,
//! and the node never learns what sits behind C2.
//!
//! A node may open one TCP per role (g4a/PTY bridge, dedicated m4a). C2 routes
//! downlink by `(node_id, channel_id)` to the TCP that OPENed that channel.
//! On a shared TCP, a single FIFO send queue would let browser frames or a stuck
//! mail channel wedge PTY control (head-of-line blocking). Outbound traffic is
//! therefore queued per-channel with priority scheduling, browser frames are
//! dropped *before* Noise encrypt under pressure (never dequeue already-encrypted
//! DATA — that desyncs nonces and whitescreens HQ), and a write timeout tears
//! the socket down so the bridge reconnects.

use crate::uplink_journal::UplinkJournal;
use hatchery_core::inner::{decode_inner, encode_inner, InnerKind};
use hatchery_core::nested::{node_seal_outer, NoiseInitiator, NoiseTransport, PublicKeyBytes};
use hatchery_core::{encode_udp_data, grow_read_buf, ChannelKind, Frame, FrameError, FrameType, OpenPayload};
use serde_json::{json, Value};
use std::collections::{BTreeMap, VecDeque};
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::sync::{mpsc, Notify};
use tokio::task::JoinHandle;

/// App-level ping on an idle node→C2 uplink (C2 answers Pong).
const UPLINK_PING_EVERY: Duration = Duration::from_secs(10);
/// No inbound bytes (pong or data) for this long => path is dead; reconnect.
const UPLINK_DEAD_AFTER: Duration = Duration::from_secs(35);

/// Keys C2 hands a node at enrollment: its own outer key and the upstream recipient key.
#[derive(Clone, Copy, Debug)]
pub struct UpstreamKeys {
    pub c2_public: PublicKeyBytes,
    pub recipient_public: PublicKeyBytes,
}

/// Outbound priority on the shared uplink. Lower ordinal = sent first.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum SendPrio {
    /// OPEN, Noise handshake, journal replies, g4a request/response (incl. PTY input path).
    Control = 0,
    /// g4a terminal / event frames.
    Pty = 1,
    /// m4a mail/task calls.
    Mail = 2,
    /// dig2browser frames — droppable / coalesced under pressure.
    Browser = 3,
}

const WRITE_TIMEOUT: Duration = Duration::from_secs(5);

fn udp_pty_enabled() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| {
        // Default on (0.2.53): only sessions that opted in via `set_udp` use it.
        !matches!(
            std::env::var("HATCHERY_PTY_UDP").ok().as_deref(),
            Some("0") | Some("false") | Some("FALSE") | Some("no")
        )
    })
}

const CHAN_CAP_CONTROL: usize = 64;
const CHAN_CAP_PTY: usize = 128;
const CHAN_CAP_MAIL: usize = 16;
const CHAN_CAP_BROWSER: usize = 8;

struct Queued {
    prio: SendPrio,
    frame: Frame,
}

struct ChanQ {
    q: VecDeque<Queued>,
    /// Remaining WINDOW credit for DATA frames; None until the first WINDOW.
    credit: Option<u32>,
}

impl ChanQ {
    fn new() -> Self {
        Self { q: VecDeque::new(), credit: None }
    }

    fn cap_for(prio: SendPrio) -> usize {
        match prio {
            SendPrio::Control => CHAN_CAP_CONTROL,
            SendPrio::Pty => CHAN_CAP_PTY,
            SendPrio::Mail => CHAN_CAP_MAIL,
            SendPrio::Browser => CHAN_CAP_BROWSER,
        }
    }
}

struct OutState {
    chans: BTreeMap<u32, ChanQ>,
}

/// One TCP uplink to C2 carrying any number of channels.
/// Optional UDP socket on the same peer for PTY bulk (SendPrio::Pty DATA).
pub struct NodeUplink {
    enqueue: mpsc::UnboundedSender<(u32, SendPrio, Frame)>,
    /// Shared with the ingress task so we can refuse Browser encrypt when congested.
    out_state: Arc<Mutex<OutState>>,
    routes: Arc<Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>>,
    /// Last WINDOW credit per channel (updated by the reader).
    credits: Arc<Mutex<BTreeMap<u32, u32>>>,
    credit_pulse: Arc<Notify>,
    journal: Option<Arc<UplinkJournal>>,
    alive: Arc<AtomicBool>,
    dead: Arc<Notify>,
    tasks: Vec<JoinHandle<()>>,
    /// Peer address for the UDP data plane (same as TCP uplink).
    udp_peer: Option<std::net::SocketAddr>,
    /// Fallback OS UDP socket when userspace WG is not carrying the peer
    /// (kernel WireGuard / no-wg builds). Prefer kit_wireguard::udp_send.
    udp: Option<Arc<std::net::UdpSocket>>,
    udp_node_id: Arc<Mutex<Option<String>>>,
    udp_sent: AtomicU64,
}

impl Drop for NodeUplink {
    fn drop(&mut self) {
        self.alive.store(false, Ordering::SeqCst);
        for t in &self.tasks {
            t.abort();
        }
    }
}

fn send_udp_datagram(
    peer: std::net::SocketAddr,
    dg: &[u8],
    os: Option<&std::net::UdpSocket>,
) -> Result<(), String> {
    #[cfg(feature = "wireguard")]
    {
        match gate4agent_node_core::kit_wireguard::udp_send(peer, dg) {
            Ok(()) => return Ok(()),
            Err(e) if e.kind() == std::io::ErrorKind::InvalidInput => {
                return Err(format!("udp pty too large for WG MTU: {e}"));
            }
            Err(e) => {
                // Fall through to OS socket (kernel WG).
                tracing::debug!(error = %e, "userspace udp_send failed; trying OS");
            }
        }
    }
    let Some(udp) = os else {
        return Err("udp data plane unavailable".into());
    };
    udp.send(dg).map_err(|e| format!("udp pty send: {e}"))?;
    Ok(())
}

/// Dial the C2 uplink. With the `wireguard` feature, an address served by an
/// in-process WireGuard session (Windows) is carried inside it; otherwise this
/// is a plain TCP connect.
fn dial_uplink(addr: SocketAddr) -> std::io::Result<std::net::TcpStream> {
    #[cfg(feature = "wireguard")]
    {
        gate4agent_node_core::kit_wireguard::tcp_connect(addr)
    }
    #[cfg(not(feature = "wireguard"))]
    {
        std::net::TcpStream::connect(addr)
    }
}

impl NodeUplink {
    pub async fn connect(uplink: SocketAddr, journal: Option<Arc<UplinkJournal>>) -> Result<Arc<Self>, String> {
        let dial = uplink;
        let started = std::time::Instant::now();
        crate::node_log::info(
            "uplink_connect_attempt",
            &[("addr_port", dial.port().to_string())],
        );
        let std_stream = match tokio::task::spawn_blocking(move || dial_uplink(dial)).await {
            Ok(Ok(s)) => {
                crate::node_log::info(
                    "uplink_connect_ok",
                    &[
                        ("addr_port", dial.port().to_string()),
                        ("ms", started.elapsed().as_millis().to_string()),
                    ],
                );
                s
            }
            Ok(Err(e)) => {
                crate::node_log::warn(
                    "uplink_connect_fail",
                    &[
                        ("addr_port", dial.port().to_string()),
                        ("error", e.to_string()),
                        ("ms", started.elapsed().as_millis().to_string()),
                    ],
                );
                return Err(e.to_string());
            }
            Err(e) => {
                crate::node_log::warn(
                    "uplink_connect_fail",
                    &[
                        ("addr_port", dial.port().to_string()),
                        ("error", e.to_string()),
                        ("ms", started.elapsed().as_millis().to_string()),
                    ],
                );
                return Err(e.to_string());
            }
        };
        let _ = std_stream.set_nodelay(true);
        // OS keepalive backs up the app-level ping. On the in-process
        // WireGuard path this socket is loopback and the inner tunnel socket
        // carries its own keepalive (gate4agent-node-core).
        if let Err(error) = crate::server::enable_tcp_keepalive(&std_stream) {
            tracing::warn!(cause = %error, "could not enable TCP keepalive on the uplink");
        }
        // current_thread runtime. A blocking socket here stalls the only
        // worker inside the first read, so the OPEN writer and the handshake
        // timeout never run. Windows does not flip this in `from_std`.
        std_stream.set_nonblocking(true).map_err(|e| e.to_string())?;
        let stream = TcpStream::from_std(std_stream).map_err(|e| e.to_string())?;
        // UDP data plane (same peer): PTY bulk. Prefer userspace-WG udp_send
        // (plain OS UDP never enters the in-process tunnel). OS socket is a
        // fallback for kernel WireGuard / non-wg builds.
        let udp_peer = Some(dial);
        let udp = match std::net::UdpSocket::bind("0.0.0.0:0") {
            Ok(sock) => match sock.connect(dial) {
                Ok(()) => {
                    let _ = sock.set_nonblocking(true);
                    Some(Arc::new(sock))
                }
                Err(_) => None,
            },
            Err(_) => None,
        };
        tracing::info!(%dial, via = "wireguard-or-os", "uplink UDP data plane ready");
        let udp_node_id = Arc::new(Mutex::new(None));
        let (mut rd, mut wr) = stream.into_split();
        let (enq_tx, mut enq_rx) = mpsc::unbounded_channel::<(u32, SendPrio, Frame)>();
        let routes: Arc<Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>> = Arc::default();
        let credits: Arc<Mutex<BTreeMap<u32, u32>>> = Arc::default();
        let credit_pulse = Arc::new(Notify::new());
        let alive = Arc::new(AtomicBool::new(true));
        let dead = Arc::new(Notify::new());
        let wake = Arc::new(Notify::new());

        let out_state = Arc::new(Mutex::new(OutState { chans: BTreeMap::new() }));
        // Last inbound byte (ms since `born`). The keepalive below pings C2 on
        // an idle uplink so NAT/WG/conntrack state stays warm and a half-open
        // path is detected in seconds, not discovered by the next HQ spawn
        // (owner spawn after ~28 min idle sat ~23 s before reaching the node).
        let born = std::time::Instant::now();
        let last_rx = Arc::new(AtomicU64::new(0));

        // Ingress into per-channel queues (never blocks the producer on TCP).
        let state_in = Arc::clone(&out_state);
        let wake_in = Arc::clone(&wake);
        let alive_in = Arc::clone(&alive);
        let ingress = tokio::spawn(async move {
            while let Some((ch, prio, frame)) = enq_rx.recv().await {
                if !alive_in.load(Ordering::SeqCst) {
                    break;
                }
                let mut st = state_in.lock().expect("out");
                let q = st.chans.entry(ch).or_insert_with(ChanQ::new);
                let cap = ChanQ::cap_for(prio);
                // Noise transport nonces are assigned at encrypt time (before enqueue).
                // Never dequeue ciphertexts. Pre-encrypt enqueue_ok should prevent this;
                // if a race still delivers an over-cap frame, tear the uplink down.
                if q.q.len() >= cap {
                    tracing::error!(?prio, cap, qlen = q.q.len(), "ciphertext dropped after encrypt (nonce gap); reconnecting uplink");
                    drop(st);
                    alive_in.store(false, Ordering::SeqCst);
                    break;
                }
                if prio == SendPrio::Browser
                    && q.q.iter().any(|x| x.prio == SendPrio::Browser)
                {
                    // Pre-encrypt gate should prevent this; a race burns a nonce → desync.
                    tracing::error!("browser ciphertext dropped after encrypt (nonce gap); reconnecting uplink");
                    drop(st);
                    alive_in.store(false, Ordering::SeqCst);
                    break;
                }
                q.q.push_back(Queued { prio, frame });
                drop(st);
                wake_in.notify_one();
            }
        });

        // Fair/priority writer: always pick the globally highest-priority waiting frame
        // whose channel has credit (DATA) or needs no credit (OPEN/control non-DATA).
        let state_wr = Arc::clone(&out_state);
        let wake_wr = Arc::clone(&wake);
        let alive_wr = Arc::clone(&alive);
        let dead_wr = Arc::clone(&dead);
        let routes_wr = Arc::clone(&routes);
        let writer = tokio::spawn(async move {
            loop {
                if !alive_wr.load(Ordering::SeqCst) {
                    break;
                }
                let next = {
                    let mut st = state_wr.lock().expect("out");
                    pick_next(&mut st)
                };
                let Some((ch, is_data, frame)) = next else {
                    wake_wr.notified().await;
                    continue;
                };
                let Ok(bytes) = frame.encode() else {
                    kill(&alive_wr, &dead_wr, &routes_wr);
                    break;
                };
                match tokio::time::timeout(WRITE_TIMEOUT, wr.write_all(&bytes)).await {
                    Ok(Ok(())) => {
                        if is_data {
                            let mut st = state_wr.lock().expect("out");
                            if let Some(q) = st.chans.get_mut(&ch) {
                                if let Some(c) = q.credit.as_mut() {
                                    *c = c.saturating_sub(1);
                                }
                            }
                        }
                    }
                    Ok(Err(_)) | Err(_) => {
                        tracing::warn!(%ch, "uplink write timed out or failed; reconnecting");
                        kill(&alive_wr, &dead_wr, &routes_wr);
                        break;
                    }
                }
            }
        });

        let enq_reader = enq_tx.clone();
        let reader_journal = journal.clone();
        let reader_routes = Arc::clone(&routes);
        let alive_rd = Arc::clone(&alive);
        let dead_rd = Arc::clone(&dead);
        let state_rd = Arc::clone(&out_state);
        let wake_rd = Arc::clone(&wake);
        let credits_rd = Arc::clone(&credits);
        let credit_pulse_rd = Arc::clone(&credit_pulse);
        let last_rx_rd = Arc::clone(&last_rx);
        let reader = tokio::spawn(async move {
            let mut buf = vec![0u8; 128 * 1024];
            let mut filled = 0usize;
            loop {
                let Ok(n) = rd.read(&mut buf[filled..]).await else { break };
                if n == 0 {
                    break;
                }
                filled += n;
                last_rx_rd.store(born.elapsed().as_millis() as u64, Ordering::SeqCst);
                loop {
                    match Frame::decode(&buf[..filled]) {
                        Ok((frame, used)) => {
                            buf.copy_within(used..filled, 0);
                            filled -= used;
                            match frame.frame_type {
                                FrameType::Data => {
                                    let tx = reader_routes.lock().expect("routes").get(&frame.channel_id).cloned();
                                    if let Some(tx) = tx {
                                        let _ = tx.send(frame.payload);
                                    }
                                }
                                FrameType::Window => {
                                    let credit = frame.seq.min(u64::from(u32::MAX)) as u32;
                                    {
                                        let mut st = state_rd.lock().expect("out");
                                        st.chans.entry(frame.channel_id).or_insert_with(ChanQ::new).credit = Some(credit);
                                    }
                                    credits_rd.lock().expect("credits").insert(frame.channel_id, credit);
                                    credit_pulse_rd.notify_waiters();
                                    wake_rd.notify_one();
                                }
                                FrameType::JournalReq => {
                                    let Some(j) = reader_journal.as_ref() else { continue };
                                    let limit = frame.journal_req_limit().unwrap_or(64) as usize;
                                    if let Ok(body) = j.journal().reply_body(frame.channel_id, frame.seq, limit) {
                                        if let Ok(reply) = Frame::journal_reply(frame.channel_id, &body) {
                                            let _ = enq_reader.send((frame.channel_id, SendPrio::Control, reply));
                                        }
                                    }
                                }
                                FrameType::GoAway | FrameType::Close => {
                                    kill(&alive_rd, &dead_rd, &reader_routes);
                                    return;
                                }
                                _ => {}
                            }
                        }
                        Err(FrameError::Truncated) => break,
                        Err(_) => {
                            kill(&alive_rd, &dead_rd, &reader_routes);
                            return;
                        }
                    }
                }
                if filled == buf.len() {
                    // Same grow path as C2: do not tear the bridge when a
                    // single frame exceeds the initial 128 KiB read buf.
                    if grow_read_buf(&mut buf, filled).is_err() {
                        break;
                    }
                }
            }
            kill(&alive_rd, &dead_rd, &reader_routes);
        });

        let ka_tx = enq_tx.clone();
        let alive_ka = Arc::clone(&alive);
        let dead_ka = Arc::clone(&dead);
        let routes_ka = Arc::clone(&routes);
        let keepalive = tokio::spawn(async move {
            let mut nonce = 0u64;
            let mut tick = tokio::time::interval(UPLINK_PING_EVERY);
            tick.tick().await;
            loop {
                tick.tick().await;
                if !alive_ka.load(Ordering::SeqCst) {
                    break;
                }
                let now = born.elapsed().as_millis() as u64;
                let idle = now.saturating_sub(last_rx.load(Ordering::SeqCst));
                if idle > UPLINK_DEAD_AFTER.as_millis() as u64 {
                    tracing::warn!(idle_ms = idle, "uplink silent past keepalive deadline; reconnecting");
                    kill(&alive_ka, &dead_ka, &routes_ka);
                    break;
                }
                nonce = nonce.wrapping_add(1);
                if ka_tx.send((0, SendPrio::Control, Frame::ping(nonce))).is_err() {
                    break;
                }
            }
        });

        Ok(Arc::new(Self {
            enqueue: enq_tx,
            out_state: Arc::clone(&out_state),
            routes,
            credits,
            credit_pulse,
            journal,
            alive,
            dead,
            tasks: vec![ingress, writer, reader, keepalive],
            udp_peer,
            udp,
            udp_node_id,
            udp_sent: AtomicU64::new(0),
        }))
    }

    /// Datagrams sent on the UDP data plane by this uplink.
    pub fn udp_sent(&self) -> u64 {
        self.udp_sent.load(Ordering::Relaxed)
    }

    pub fn is_alive(&self) -> bool {
        self.alive.load(Ordering::SeqCst)
    }

    /// Wait until the uplink is marked dead (write timeout / peer close).
    pub async fn closed(&self) {
        if !self.is_alive() {
            return;
        }
        self.dead.notified().await;
    }

    /// True if a new ciphertext can be enqueued without dropping anything.
    /// Call *before* Noise encrypt so a refused frame does not burn a nonce.
    fn enqueue_ok(&self, channel_id: u32, prio: SendPrio) -> bool {
        let st = self.out_state.lock().expect("out");
        let Some(q) = st.chans.get(&channel_id) else {
            return true;
        };
        if q.q.len() >= ChanQ::cap_for(prio) {
            return false;
        }
        if prio == SendPrio::Browser && q.q.iter().any(|x| x.prio == SendPrio::Browser) {
            return false;
        }
        true
    }

    fn push(&self, channel_id: u32, prio: SendPrio, frame: Frame) -> Result<(), String> {
        if !self.is_alive() {
            return Err("uplink closed".into());
        }
        self.enqueue
            .send((channel_id, prio, frame))
            .map_err(|_| "uplink closed".to_owned())
    }

    /// OPEN `channel_id`, wait for WINDOW credit, then complete Noise_NK.
    pub async fn open(
        self: &Arc<Self>,
        node_id: &str,
        channel_id: u32,
        kind: ChannelKind,
        keys: UpstreamKeys,
        handshake_timeout: Duration,
    ) -> Result<NodeSession, String> {
        let (down_tx, mut down_rx) = mpsc::unbounded_channel::<Vec<u8>>();
        self.routes.lock().expect("routes").insert(channel_id, down_tx);
        self.credits.lock().expect("credits").remove(&channel_id);
        {
            let mut g = self.udp_node_id.lock().expect("udp_node_id");
            if g.is_none() {
                *g = Some(node_id.to_owned());
            }
        }
        let open = Frame::open(channel_id, &OpenPayload { channel_kind: kind, node_id: node_id.to_owned() })
            .map_err(|e| e.to_string())?;
        self.push(channel_id, SendPrio::Control, open)?;

        // Wait until C2 grants WINDOW credit after OPEN (avoids sending Handshake1
        // into a stream that still has credit 0 from a previous session).
        tokio::time::timeout(handshake_timeout, async {
            loop {
                if !self.is_alive() {
                    return Err("uplink closed".to_owned());
                }
                if self.credits.lock().expect("credits").get(&channel_id).copied().unwrap_or(0) > 0 {
                    return Ok(());
                }
                tokio::select! {
                    _ = self.credit_pulse.notified() => {}
                    _ = self.closed() => return Err("uplink closed".to_owned()),
                }
            }
        })
        .await
        .map_err(|_| "window credit timeout after OPEN".to_owned())??;

        let mut initiator = NoiseInitiator::new(keys.recipient_public).map_err(|e| e.to_string())?;
        let msg1 = initiator.write_message1().map_err(|e| e.to_string())?;
        self.send_inner(channel_id, keys, SendPrio::Control, &encode_inner(InnerKind::Handshake1, &msg1), false)?;

        let msg2 = tokio::time::timeout(handshake_timeout, async {
            loop {
                let Some(bytes) = down_rx.recv().await else { return Err("uplink closed".to_owned()) };
                if let Some((InnerKind::Handshake2, body)) = decode_inner(&bytes) {
                    return Ok(body.to_vec());
                }
            }
        })
        .await
        .map_err(|_| "handshake timeout".to_owned())??;
        let transport = initiator.read_message2(&msg2).map_err(|e| e.to_string())?;
        Ok(NodeSession {
            channel_id,
            keys,
            default_prio: match kind {
                ChannelKind::M4a => SendPrio::Mail,
                ChannelKind::G4aControl | ChannelKind::G4aPty | ChannelKind::G4aEvents => SendPrio::Pty,
                ChannelKind::Vault | ChannelKind::Observation => SendPrio::Mail,
            },
            down: down_rx,
            transport,
            journal: self.journal.clone(),
            uplink: Arc::clone(self),
            udp: false,
        })
    }

    fn send_inner(&self, channel_id: u32, keys: UpstreamKeys, prio: SendPrio, inner: &[u8], udp_on: bool) -> Result<(), String> {
        if let Some(j) = &self.journal {
            j.append(channel_id, inner)?;
        }
        let sealed = node_seal_outer(keys.c2_public, inner).map_err(|e| e.to_string())?;
        let frame = Frame::data(channel_id, 0, sealed).map_err(|e| e.to_string())?;
        // PTY bulk over UDP (opt-in): same Noise ciphertext as TCP would carry, so
        // HATCHERY_PTY_UDP=1 must send *all* ch3 DATA via UDP (no mix) to keep nonces
        // ordered. Default stays TCP on ch3 — still isolates WINDOW credit from control.
        // Only a session that opted in (dedicated ch3) uses UDP: one channel never mixes paths.
        if udp_on && prio == SendPrio::Pty && udp_pty_enabled() {
            if let (Some(peer), Some(nid)) = (
                self.udp_peer,
                self.udp_node_id.lock().expect("udp_node_id").clone(),
            ) {
                match encode_udp_data(&nid, &frame) {
                    Ok(dg) => {
                        // No TCP mix: would break Noise nonces on ch3.
                        send_udp_datagram(peer, &dg, self.udp.as_deref())?;
                        self.udp_sent.fetch_add(1, Ordering::Relaxed);
                        return Ok(());
                    }
                    Err(e) => return Err(format!("udp encode: {e}")),
                }
            }
        }
        self.push(channel_id, prio, frame)
    }
}

fn kill(alive: &AtomicBool, dead: &Notify, routes: &Mutex<BTreeMap<u32, mpsc::UnboundedSender<Vec<u8>>>>) {
    if alive.swap(false, Ordering::SeqCst) {
        routes.lock().expect("routes").clear();
        dead.notify_waiters();
    }
}

/// Pick the highest-priority queued frame that may be sent now.
fn pick_next(st: &mut OutState) -> Option<(u32, bool, Frame)> {
    // Candidates: (prio, channel, index). DATA frames of one channel are Noise ciphertexts
    // encrypted in queue order, so only the channel's first DATA frame may go next
    // (overtaking breaks the receiver's nonce sequence); non-DATA frames may go any time.
    let mut best: Option<(SendPrio, u32, usize)> = None;
    for (&ch, cq) in &st.chans {
        let first_data = cq.q.iter().position(|it| matches!(it.frame.frame_type, FrameType::Data));
        for (i, item) in cq.q.iter().enumerate() {
            let needs_credit = matches!(item.frame.frame_type, FrameType::Data);
            if needs_credit && Some(i) != first_data {
                continue;
            }
            if needs_credit {
                match cq.credit {
                    Some(0) => continue,
                    None => {
                        // No WINDOW yet: allow Control DATA (handshake) so OPEN→Handshake1 works;
                        // hold back bulk traffic.
                        if item.prio > SendPrio::Control {
                            continue;
                        }
                    }
                    Some(_) => {}
                }
            }
            let cand = (item.prio, ch, i);
            if best.is_none_or(|b| cand.0 < b.0 || (cand.0 == b.0 && cand.1 < b.1)) {
                best = Some(cand);
            }
        }
    }
    let (prio, ch, idx) = best?;
    let cq = st.chans.get_mut(&ch)?;
    let item = cq.q.remove(idx)?;
    debug_assert_eq!(item.prio, prio);
    let is_data = matches!(item.frame.frame_type, FrameType::Data);
    Some((ch, is_data, item.frame))
}

pub struct NodeSession {
    channel_id: u32,
    keys: UpstreamKeys,
    default_prio: SendPrio,
    down: mpsc::UnboundedReceiver<Vec<u8>>,
    transport: NoiseTransport,
    journal: Option<Arc<UplinkJournal>>,
    uplink: Arc<NodeUplink>,
    udp: bool,
}

impl NodeSession {
    /// Whether this uplink has a connected UDP socket and the data plane is enabled.
    pub fn udp_available(&self) -> bool {
        self.uplink.udp_peer.is_some() && udp_pty_enabled()
    }

    /// Route this session's `SendPrio::Pty` DATA over UDP. Set once, right after open,
    /// before any transport frame: switching later would reorder Noise nonces.
    pub fn set_udp(&mut self, on: bool) {
        self.udp = on && self.udp_available();
    }

    pub fn udp(&self) -> bool {
        self.udp
    }

    pub async fn connect(
        uplink: SocketAddr,
        node_id: &str,
        channel_id: u32,
        kind: ChannelKind,
        keys: UpstreamKeys,
        journal: Option<Arc<UplinkJournal>>,
        handshake_timeout: Duration,
    ) -> Result<Self, String> {
        let link = NodeUplink::connect(uplink, journal).await?;
        link.open(node_id, channel_id, kind, keys, handshake_timeout).await
    }

    pub fn uplink(&self) -> Arc<NodeUplink> {
        Arc::clone(&self.uplink)
    }

    pub async fn send(&mut self, plaintext: &[u8]) -> Result<(), String> {
        self.send_prio(plaintext, self.default_prio).await
    }

    pub async fn send_prio(&mut self, plaintext: &[u8], prio: SendPrio) -> Result<(), String> {
        if !self.uplink.is_alive() {
            return Err("uplink closed".into());
        }
        // Refuse *before* encrypt when the channel queue cannot take this frame.
        // Encrypt-then-drop desyncs Noise nonces (HQ white canvas / bridge decrypt errors).
        if !self.uplink.enqueue_ok(self.channel_id, prio) {
            if prio == SendPrio::Browser {
                return Ok(()); // droppable
            }
            return Err("uplink send queue full".into());
        }
        let ct = self.transport.encrypt(plaintext).map_err(|e| e.to_string())?;
        let inner = encode_inner(InnerKind::Transport, &ct);
        self.uplink.send_inner(self.channel_id, self.keys, prio, &inner, self.udp)
    }

    pub async fn recv(&mut self, timeout: Duration) -> Result<Vec<u8>, String> {
        tokio::time::timeout(timeout, async {
            loop {
                tokio::select! {
                    bytes = self.down.recv() => {
                        let Some(bytes) = bytes else { return Err("uplink closed".to_owned()) };
                        if let Some((InnerKind::Transport, ct)) = decode_inner(&bytes) {
                            return self.transport.decrypt(ct).map_err(|e| e.to_string());
                        }
                    }
                    _ = self.uplink.closed() => return Err("uplink closed".to_owned()),
                }
            }
        })
        .await
        .map_err(|_| "reply timeout".to_owned())?
    }

    pub async fn call(&mut self, op: &str, payload: &Value, timeout: Duration) -> Result<Value, String> {
        let req = serde_json::to_vec(&json!({ "op": op, "payload": payload })).map_err(|e| e.to_string())?;
        self.send_prio(&req, SendPrio::Mail).await?;
        let reply = self.recv(timeout).await?;
        let reply: Value = serde_json::from_slice(&reply).map_err(|e| e.to_string())?;
        let status = reply.get("status").and_then(Value::as_u64).unwrap_or(500);
        let body = reply.get("body").cloned().unwrap_or(Value::Null);
        if status == 200 {
            Ok(body)
        } else {
            let msg = body.get("error").and_then(Value::as_str).unwrap_or("error");
            Err(format!("{status}: {msg}"))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pick_prefers_control_across_channels_but_keeps_data_order_within_one() {
        let mut st = OutState { chans: BTreeMap::new() };
        let a = st.chans.entry(1).or_insert_with(ChanQ::new);
        a.credit = Some(10);
        a.q.push_back(Queued { prio: SendPrio::Browser, frame: Frame::data(1, 0, vec![1]).unwrap() });
        a.q.push_back(Queued { prio: SendPrio::Control, frame: Frame::data(1, 0, vec![2]).unwrap() });
        let b = st.chans.entry(2).or_insert_with(ChanQ::new);
        b.credit = Some(10);
        b.q.push_back(Queued { prio: SendPrio::Pty, frame: Frame::data(2, 0, vec![3]).unwrap() });
        // Channel 2's Pty beats channel 1's head (Browser); channel 1 stays in order.
        let order: Vec<u8> = std::iter::from_fn(|| pick_next(&mut st).map(|(_, _, f)| f.payload[0])).collect();
        assert_eq!(order, vec![3, 1, 2]);
    }

    #[test]
    fn pick_holds_bulk_until_window() {
        let mut st = OutState { chans: BTreeMap::new() };
        let q = st.chans.entry(2).or_insert_with(ChanQ::new);
        q.credit = None;
        q.q.push_back(Queued {
            prio: SendPrio::Control,
            frame: Frame::data(2, 0, vec![1]).unwrap(),
        });
        q.q.push_back(Queued {
            prio: SendPrio::Mail,
            frame: Frame::data(2, 0, vec![9]).unwrap(),
        });
        let (_, _, f) = pick_next(&mut st).unwrap();
        assert_eq!(f.payload, vec![1]);
        assert!(pick_next(&mut st).is_none());
        st.chans.get_mut(&2).unwrap().credit = Some(4);
        let (_, _, f) = pick_next(&mut st).unwrap();
        assert_eq!(f.payload, vec![9]);
    }
}