supermachine 0.7.4

// Status: minimal — connection table + REQUEST→RESPONSE/RST routing
// + TSI control DGRAM (port 1024..1031) handling.
//
// What we DON'T have yet (later commits):
//   - tsi_stream::TsiStreamProxy (TCP backend) — TSI_LISTEN
//     creates a stub Proxy that always replies success but doesn't
//     actually bind a host socket
//   - epoll-driven proxy event handling — without TCP backend, no
//     events to drive
//   - reaper thread — proxies just live until VM shutdown
//   - DGRAM (UDP) proxies via tsi_dgram

#![allow(dead_code)]

use std::collections::HashMap;
use std::net::{SocketAddr, TcpStream, UdpSocket};
use std::os::unix::net::UnixStream;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{mpsc, Arc, Mutex};
use std::time::Instant;

use super::mux_profile::{self, Stage};
use super::muxer_rxq::MuxerRxQ;
use super::muxer_thread::{self, MuxerCmd, MuxerStream};
use super::packet::{
    Header, RxPacket, VSOCK_HOST_CID, VSOCK_OP_CREDIT_REQUEST, VSOCK_OP_CREDIT_UPDATE,
    VSOCK_OP_REQUEST, VSOCK_OP_RESPONSE, VSOCK_OP_RST, VSOCK_OP_RW, VSOCK_OP_SHUTDOWN,
    VSOCK_TYPE_DGRAM, VSOCK_TYPE_STREAM,
};
use super::proxy::{proxy_key, Proxy};
use super::tsi_stream::TsiListener;

/// AF_TSI control op opcodes, encoded as the dst_port of a DGRAM
/// packet sent by the guest kernel's TSI patches. See the patch
/// series under `kernel-build/patches/af-tsi/`.
pub const TSI_PROXY_PORT: u32 = 620;
pub const VSOCK_ENV_PORT: u32 = 1026;
pub const TSI_PROXY_CREATE: u32 = 1024;
pub const TSI_CONNECT: u32 = 1025;
pub const TSI_GETNAME: u32 = 1026;
pub const TSI_SENDTO_ADDR: u32 = 1027;
pub const TSI_SENDTO_DATA: u32 = 1028;
pub const TSI_LISTEN: u32 = 1029;
pub const TSI_ACCEPT: u32 = 1030;
pub const TSI_PROXY_RELEASE: u32 = 1031;

/// Snapshot record for one TSI listener: the (cid, peer_port, vm_port,
/// family, socktype) tuple needed to re-bind a host TcpListener on
/// restore. The host listener itself can't be serialized (OS socket
/// fd) — it gets rebuilt on restore via `restore_tsi_listeners`.
#[derive(Clone, Debug)]
pub struct TsiListenerSnapshot {
    pub cid: u64,
    pub peer_port: u32,
    pub vm_port: u32,
    pub family: u16,
    pub socktype: u16,
}

/// State for one TSI listener proxy. Created on TSI_PROXY_CREATE,
/// activated on TSI_LISTEN. We don't yet bind a host socket (no
/// tsi_stream port); we just remember the (local_port, peer_port,
/// vm_port) so subsequent PROXY_CREATE/LISTEN/ACCEPT calls can be
/// answered.
struct TsiState {
    family: u16,
    socktype: u16,
    /// vm_port from TSI_LISTEN — guest's listening port.
    vm_port: Option<u32>,
    /// Host TCP listener bound on TSI_LISTEN. Drops on
    /// PROXY_RELEASE (or VM shutdown).
    listener: Option<TsiListener>,
}

/// In-flight inbound TCP conn that's been accepted on the host
/// side but the guest hasn't sent RESPONSE yet.
pub struct PendingInbound {
    pub cid: u64,
    pub host_src_port: u32,
    pub vm_port: u32,
    pub stream: Option<MuxerStream>,
}

/// In-flight outbound: guest sent TSI_CONNECT(peer_port, target);
/// we opened TCP to target and now park it until the guest's
/// follow-up VSOCK_OP_REQUEST(src=peer_port, dst=TSI_PROXY_PORT)
/// arrives. Keyed by (cid, peer_port).
pub struct PendingOutbound {
    pub cid: u64,
    pub peer_port: u32,
    pub tcp: Option<TcpStream>,
}

/// Per-conn flow-control accounting. Two directions:
///
/// 1. Guest → host (the bytes the guest writes to us, that we
///    forward to host TCP). The vsock peer sees us as having
///    `OUR_BUF_ALLOC` of receive buffer; every byte they send to
///    us must be ACKed via `our_fwd_cnt` or they stop sending
///    once they've used (peer_buf_alloc - 0) = OUR_BUF_ALLOC
///    bytes. We bump `our_fwd_cnt` as we consume (which is
///    immediate, since we forward straight to TCP), and emit
///    CREDIT_UPDATE when half of the buffer has been consumed
///    since last update.
///
/// 2. Host → guest (the bytes we read from host TCP, that we
///    push as RW packets into the guest's RX queue). The guest
///    advertises its RX buffer via `buf_alloc` on every header
///    and ACKs consumption via `fwd_cnt`. We must respect
///    `available = peer_buf_alloc - (tx_cnt - peer_fwd_cnt)` —
///    when that drops to zero, we MUST pause reading from host
///    TCP, or the guest's vsock RX socket buffer fills and
///    bytes are silently dropped (see issue: 4 MiB GET
///    truncating at ~1.5 MiB). When peer_fwd_cnt advances via
///    a subsequent RW or CREDIT_UPDATE from the guest, we
///    resume.
///
/// Tracked under `Arc<Mutex<...>>` so both the dispatcher path
/// (handle_rw / handle_credit) and the on_data closure invoked
/// from the muxer-io thread share a single view of credit.
pub struct ProxyConnState {
    /// Bytes WE'VE consumed from the guest's TX. Advances on
    /// each handle_rw call as we forward to host TCP.
    pub our_fwd_cnt: u32,
    /// Snapshot of our_fwd_cnt at the last CREDIT_UPDATE we sent
    /// to the guest. Used to decide when to emit the next one.
    pub last_credit_fwd_cnt: u32,
    /// Guest socket's local port — the `dst_port` of every RW
    /// we push toward the guest.
    pub guest_dst_port: u32,
    /// Last buf_alloc the guest advertised. Updated on EVERY
    /// inbound packet from the guest (RW + CREDIT_UPDATE +
    /// CREDIT_REQUEST all carry credit fields in their header).
    pub peer_buf_alloc: u32,
    /// Last fwd_cnt the guest reported — how many bytes of what
    /// we sent it has consumed.
    pub peer_fwd_cnt: u32,
    /// Cumulative bytes we've sent toward the guest on this
    /// conn. Wraps mod 2^32 (vsock uses u32 wrapping arithmetic).
    pub tx_cnt: u32,
    /// host_src_port — the key the muxer-io thread uses to
    /// route Write/Close/PauseReadable/ResumeReadable commands
    /// for this conn.
    pub host_src_port: u32,
    /// True when we asked muxer-io to pause READABLE on the
    /// host TCP because we ran out of credit. Cleared by
    /// handle_rw / handle_credit when credit returns and we
    /// send ResumeReadable.
    pub paused_readable: bool,
    /// Shared flag with the muxer-io thread's `Conn` so the
    /// per-conn read loop can break out of `loop { read }`
    /// inline (without round-tripping a PauseReadable command,
    /// which would be processed only on the next poll wakeup
    /// — by which time many more bytes might have been read).
    /// Set by `try_reserve_credit` when credit hits the low
    /// watermark; cleared when ResumeReadable is processed.
    pub pause_flag: Arc<AtomicBool>,
}

/// Low watermark below which the on_data path stops pushing
/// new bytes and trips the pause flag. Set generously above
/// zero so the in-flight packet that triggers the pause still
/// fits within `peer_buf_alloc`. Linux's vsock_stream tends to
/// advertise 256 KiB; 4 KiB is the smallest individual packet
/// we'd ever push (MAX_PKT_PAYLOAD in muxer_thread), so 4 KiB
/// is the absolute floor — any packet pushed when remaining
/// credit ≥ 4 KiB still fits.
pub const CREDIT_LOW_WATERMARK: u32 = 4096;

impl ProxyConnState {
    pub fn new(guest_dst_port: u32, host_src_port: u32) -> Self {
        Self {
            our_fwd_cnt: 0,
            last_credit_fwd_cnt: 0,
            guest_dst_port,
            peer_buf_alloc: 0,
            peer_fwd_cnt: 0,
            tx_cnt: 0,
            host_src_port,
            paused_readable: false,
            pause_flag: Arc::new(AtomicBool::new(false)),
        }
    }

    /// Available credit toward the guest: how many more bytes
    /// we can push before the guest's RX buffer is conceptually
    /// full. Uses u32 wrapping arithmetic (vsock spec) — if
    /// peer_buf_alloc is zero the result is zero.
    pub fn available_credit(&self) -> u32 {
        let in_flight = self.tx_cnt.wrapping_sub(self.peer_fwd_cnt);
        self.peer_buf_alloc.saturating_sub(in_flight)
    }

    /// Update peer credit fields from an incoming header. Called
    /// on every RW / CREDIT_UPDATE / CREDIT_REQUEST. Returns
    /// true if we were paused AND now have credit above the
    /// watermark (caller should send ResumeReadable).
    pub fn observe_peer(&mut self, peer_buf_alloc: u32, peer_fwd_cnt: u32) -> bool {
        self.peer_buf_alloc = peer_buf_alloc;
        self.peer_fwd_cnt = peer_fwd_cnt;
        if self.paused_readable && self.available_credit() > CREDIT_LOW_WATERMARK {
            self.paused_readable = false;
            self.pause_flag.store(false, Ordering::Release);
            true
        } else {
            false
        }
    }
}

/// Buffer size we advertise to the guest. We forward bytes to
/// host TCP synchronously, so we have effectively-unlimited
/// receive capacity from the guest's perspective. Set to
/// (u32::MAX / 2) so the guest's flow-control math never sees
/// us as "full"; we still emit periodic CREDIT_UPDATE for
/// correctness but guest never blocks waiting for it.
pub const OUR_BUF_ALLOC: u32 = u32::MAX / 2;
const CREDIT_UPDATE_THRESHOLD: u32 = OUR_BUF_ALLOC / 2;

fn body_le_u16(body: &[u8], offset: usize) -> Option<u16> {
    let bytes = body.get(offset..offset + 2)?.try_into().ok()?;
    Some(u16::from_le_bytes(bytes))
}

fn body_be_u16(body: &[u8], offset: usize) -> Option<u16> {
    let bytes = body.get(offset..offset + 2)?.try_into().ok()?;
    Some(u16::from_be_bytes(bytes))
}

fn body_le_u32(body: &[u8], offset: usize) -> Option<u32> {
    let bytes = body.get(offset..offset + 4)?.try_into().ok()?;
    Some(u32::from_le_bytes(bytes))
}

/// Guest connects to (HOST_CID, VSOCK_ENV_PORT) via plain AF_VSOCK to
/// fetch its env JSON. Set at startup by `--env` / `--env-file`. The
/// blob is delivered as a single VSOCK_OP_RW followed by SHUTDOWN.
/// Format: `{"env":{"K":"V",...},"secrets":{"K":"V",...}}`.
pub static VSOCK_ENV_JSON: Mutex<Option<String>> = Mutex::new(None);

pub fn set_env_json(json: String) {
    *VSOCK_ENV_JSON.lock().unwrap() = Some(json);
}

/// Cached `$SUPERMACHINE_VSOCK_TRACE` flag. Calling
/// `std::env::var_os(...)` per packet path is a real perf hit:
/// `getenv` takes a libc global lock + linear scan of `environ`.
/// At nginx c=64 it dominated muxer-io non-syscall CPU. Read
/// once on first use, cache. Toggling at runtime no longer
/// supported — set the env before the worker starts.
#[inline]
pub fn vsock_trace_enabled() -> bool {
    use std::sync::atomic::{AtomicU8, Ordering};
    static CACHED: AtomicU8 = AtomicU8::new(0);
    let v = CACHED.load(Ordering::Relaxed);
    if v != 0 {
        return v == 2;
    }
    // The actual env lookup. Inside the cached helper we call the
    // libc path directly; the sed-rewrite that replaced sites in
    // this file would otherwise recurse here.
    let on = std::env::var_os("SUPERMACHINE_VSOCK_TRACE").is_some();
    CACHED.store(if on { 2 } else { 1 }, Ordering::Relaxed);
    on
}

pub struct VsockMuxer {
    cid: u64,
    /// (local_port, peer_port) → Proxy.
    proxies: Mutex<HashMap<u64, Box<dyn Proxy>>>,
    /// (cid, host_src_port) → PendingInbound waiting for guest's
    /// RESPONSE to a host-initiated REQUEST.
    pending_inbound: Arc<Mutex<HashMap<(u64, u32), PendingInbound>>>,
    /// (cid, peer_port) → PendingOutbound. Set by TSI_CONNECT,
    /// consumed by VSOCK_OP_REQUEST(dst=TSI_PROXY_PORT) from guest.
    pending_outbound: Mutex<HashMap<(u64, u32), PendingOutbound>>,
    /// (cid, peer_port) → staged dst SocketAddr from TSI_SENDTO_ADDR,
    /// consumed by the next TSI_SENDTO_DATA.
    udp_dst: Mutex<HashMap<(u64, u32), SocketAddr>>,
    /// (cid, peer_port) → persistent UdpSocket. Send fd shared with
    /// the muxer-io thread (which holds a try_clone()'d recv fd).
    udp_sockets: Mutex<HashMap<(u64, u32), Arc<UdpSocket>>>,
    /// (cid, host_src_port) → ProxyConnState. The TCP stream
    /// lives in the muxer-io thread; we keep flow-control
    /// accounting (both directions) here so handle_rw can emit
    /// CREDIT_UPDATE packets and the on_data closure (called
    /// from the muxer-io thread) can decrement available credit
    /// and trip the pause flag. The inner Arc<Mutex<...>> is
    /// cloned into the on_data closure on conn establishment.
    inbound_conns: Mutex<HashMap<(u64, u32), Arc<Mutex<ProxyConnState>>>>,
    /// Sender to the single muxer I/O thread. Owns ALL inbound
    /// TcpStreams, replaces per-conn reader+writer threads.
    io_tx: mpsc::Sender<MuxerCmd>,
    /// Waker to nudge the I/O thread after sending a command.
    io_waker: Arc<mio::Waker>,
    /// (peer_cid, peer_port) → TsiState. Indexed by the guest CID
    /// (which is always our `cid` for in-guest TSI) and the
    /// guest's chosen `peer_port` from the TSI control DGRAMs.
    tsi: Mutex<HashMap<(u64, u32), TsiState>>,
    /// Outbound RX queue (intents waiting for guest descriptors).
    /// The vsock device drains this on every notify() and pushes
    /// the packets into RX virtq. Anyone (TsiListener accept
    /// thread, future TsiConn pump threads) can push to it.
    pub rxq: Arc<Mutex<MuxerRxQ>>,
    /// Callback the muxer invokes after pushing to rxq, to wake
    /// the vsock device so it drains. Set by the device after
    /// construction (it has the IRQ-raise closure).
    kick_device: Mutex<Option<Arc<dyn Fn() + Send + Sync>>>,
    /// Expected TSI control-channel auth token (32 bytes). When
    /// `Some`, every TSI control message dispatched in
    /// `handle_tsi_control` must be prefixed with these bytes —
    /// mismatches are silently dropped. The kernel patch 0014
    /// produces matching prefixes; guest userspace can't read the
    /// token (it's static module BSS, scrubbed from /proc/cmdline)
    /// so it can't forge authenticated messages.
    ///
    /// `None` means legacy mode — accept all messages. Used by
    /// pre-0.6.0 snapshots that don't have a baked-in token, and
    /// by unit tests that exercise the dispatcher directly.
    pub(crate) tsi_token: Option<[u8; TSI_TOKEN_LEN]>,
}

/// Length of the TSI control-channel auth token in bytes. Matches
/// the kernel-side `TSI_TOKEN_LEN` in `af_tsi.c`. See
/// `kernel-build/patches/af-tsi/0014-*.patch` for the threat model.
pub const TSI_TOKEN_LEN: usize = 32;

/// Constant-time byte-slice compare. Returns true iff both slices
/// have the same length AND every byte is equal. Visits every byte
/// of `b` so the runtime is independent of where the first mismatch
/// is — closing the timing oracle that a naïve `==` would open. We
/// roll our own instead of pulling in `subtle` because it's three
/// lines and we want zero new deps in the auth path.
#[inline]
fn ct_eq(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    let mut diff: u8 = 0;
    for i in 0..a.len() {
        diff |= a[i] ^ b[i];
    }
    diff == 0
}

impl VsockMuxer {
    pub fn new(cid: u64) -> Result<Self, muxer_thread::StartError> {
        Self::with_tsi_token(cid, None)
    }

    /// Construct a muxer that enforces TSI control-channel auth.
    /// `token` is the 32-byte secret baked into the booting kernel's
    /// `supermachine.tsi_token=<hex>` cmdline arg. Pass `None` to
    /// disable enforcement (legacy snapshots, tests).
    pub fn with_tsi_token(
        cid: u64,
        token: Option<[u8; TSI_TOKEN_LEN]>,
    ) -> Result<Self, muxer_thread::StartError> {
        let (io_tx, io_waker) = muxer_thread::spawn()?;
        Ok(Self {
            cid,
            proxies: Mutex::new(HashMap::new()),
            pending_inbound: Arc::new(Mutex::new(HashMap::new())),
            pending_outbound: Mutex::new(HashMap::new()),
            udp_dst: Mutex::new(HashMap::new()),
            udp_sockets: Mutex::new(HashMap::new()),
            inbound_conns: Mutex::new(HashMap::new()),
            tsi: Mutex::new(HashMap::new()),
            rxq: Arc::new(Mutex::new(MuxerRxQ::new())),
            kick_device: Mutex::new(None),
            io_tx,
            io_waker,
            tsi_token: token,
        })
    }

    pub fn set_kick(&self, kick: Arc<dyn Fn() + Send + Sync>) {
        *self.kick_device.lock().unwrap() = Some(kick);
    }

    /// True when no per-request transport state is still in flight.
    ///
    /// Persistent TSI listeners and cached UDP sockets are intentionally
    /// allowed to remain alive here: the pool can snapshot an idle listening
    /// service, but must not interrupt active TCP streams, queued packets, or
    /// pending outbound connects.
    pub fn is_transport_idle(&self) -> bool {
        self.pending_inbound.lock().unwrap().is_empty()
            && self.pending_outbound.lock().unwrap().is_empty()
            && self.inbound_conns.lock().unwrap().is_empty()
            && self.rxq.lock().unwrap().is_empty()
    }

    /// Drain all per-dispatch state so the muxer is fresh for the
    /// next RESTORE. Pool-worker mode calls this between dispatches.
    /// Concrete steps:
    /// - Drop every TsiState (their TsiListener::Drop closes the
    ///   listener fd + joins the accept thread → no leaked threads).
    /// - Clear all in-flight bookkeeping (proxies, pending_*, udp_*,
    ///   inbound_conns, rxq).
    /// - Send MuxerCmd::Reset to the io thread to drop its tcp/udp
    ///   registrations.
    /// - Keep the io thread + waker + io_tx alive (re-used).
    pub fn reset(&self) {
        // Order matters: drop TSI listeners FIRST so accept threads
        // can't push fresh connections while we're clearing maps.
        self.tsi.lock().unwrap().clear();
        // Reset io thread state and wait for the boundary before
        // draining queues below. Without the ack, a rapid restore can
        // register a fresh stream before the previous async Reset is
        // processed, then have that stale Reset drop the new stream.
        // Also, the io thread can enqueue EOF/SHUTDOWN callbacks while
        // dropping old streams, so rxq must be drained after this ack.
        let (done_tx, done_rx) = mpsc::channel();
        if self.io_tx.send(MuxerCmd::Reset { done: done_tx }).is_ok() {
            let _ = self.io_waker.wake();
            let _ = done_rx.recv_timeout(std::time::Duration::from_millis(50));
        }
        self.proxies.lock().unwrap().clear();
        self.pending_inbound.lock().unwrap().clear();
        self.pending_outbound.lock().unwrap().clear();
        self.udp_dst.lock().unwrap().clear();
        self.udp_sockets.lock().unwrap().clear();
        self.inbound_conns.lock().unwrap().clear();
        // Drain any stale packets, including callbacks the io thread
        // emitted while closing old streams during Reset.
        self.rxq.lock().unwrap().drain();
    }

    /// Look up the host-side TCP port currently bound for the
    /// listener serving guest `vm_port`. Returns None if no listener
    /// is registered yet (e.g. guest hasn't done TSI_LISTEN).
    ///
    /// Caveat: the guest's TSI driver assigns RANDOM vm_ports per
    /// bind — they don't match the AF_INET port the guest userspace
    /// asked for (e.g. http-echo's bind(8000) becomes vm_port=
    /// 0x72b58b71 or similar). So this is rarely what callers want;
    /// `first_host_port` is usually what you reach for.
    pub fn host_port_for_vm_port(&self, vm_port: u32) -> Option<u16> {
        let s = self.tsi.lock().unwrap();
        s.values().find_map(|st| {
            if st.vm_port == Some(vm_port) {
                st.listener.as_ref().map(|l| l.host_addr.port())
            } else {
                None
            }
        })
    }

    /// Return the host TCP port of the first TSI listener registered.
    /// Used by the TLS terminator: typical guest workloads have one
    /// HTTP server, so "the first listener" is unambiguous. Returns
    /// None if no listener has been bound yet.
    pub fn first_host_port(&self) -> Option<u16> {
        let s = self.tsi.lock().unwrap();
        // Prefer the AF_INET TSI listener. Many production images
        // bind both 0.0.0.0 and ::; AF_TSI6 still has a separate,
        // less-tested guest-kernel data path, while AF_TSI is the
        // path we know carries bytes correctly.
        s.values()
            .find_map(|st| {
                if st.family == libc::AF_INET as u16 {
                    st.listener.as_ref().map(|l| l.host_addr.port())
                } else {
                    None
                }
            })
            .or_else(|| {
                s.values()
                    .find_map(|st| st.listener.as_ref().map(|l| l.host_addr.port()))
            })
    }

    fn first_listener_endpoint(&self, vm_port: Option<u32>) -> Option<(u64, u32)> {
        let s = self.tsi.lock().unwrap();
        let endpoint = |cid: u64, st: &TsiState| {
            let listener_vm_port = st.vm_port?;
            if vm_port.is_some_and(|want| want != listener_vm_port) {
                return None;
            }
            Some((cid, listener_vm_port))
        };
        s.iter()
            .find_map(|((cid, _), st)| {
                if st.family == libc::AF_INET as u16 {
                    endpoint(*cid, st)
                } else {
                    None
                }
            })
            .or_else(|| s.iter().find_map(|((cid, _), st)| endpoint(*cid, st)))
    }

    fn open_stream_to_guest(
        &self,
        stream: MuxerStream,
        vm_port: Option<u32>,
    ) -> std::io::Result<()> {
        let Some((cid, vm_port)) = self.first_listener_endpoint(vm_port) else {
            return Err(std::io::Error::new(
                std::io::ErrorKind::NotConnected,
                "no TSI listener",
            ));
        };
        let host_src_port = crate::devices::virtio::vsock::tsi_stream::alloc_host_src_port();
        self.pending_inbound.lock().unwrap().insert(
            (cid, host_src_port),
            PendingInbound {
                cid,
                host_src_port,
                vm_port,
                stream: Some(stream),
            },
        );
        let req = RxPacket {
            hdr: Header {
                src_cid: VSOCK_HOST_CID,
                dst_cid: cid,
                src_port: host_src_port,
                dst_port: vm_port,
                len: 0,
                type_: VSOCK_TYPE_STREAM,
                op: VSOCK_OP_REQUEST,
                flags: 0,
                buf_alloc: OUR_BUF_ALLOC,
                fwd_cnt: 0,
            },
            data: Vec::new(),
        };
        let kick = self.kick_device.lock().unwrap().clone();
        push_rxq_and_kick(&self.rxq, &kick, req);
        Ok(())
    }

    pub fn open_unix_to_guest(
        &self,
        unix: UnixStream,
        vm_port: Option<u32>,
    ) -> std::io::Result<()> {
        self.open_stream_to_guest(MuxerStream::Unix(unix), vm_port)
    }

    /// Host→guest *native* AF_VSOCK connect. Used by the
    /// `<vsock_mux>-exec.sock` frontend to reach the in-guest exec
    /// agent without going through TSI's TCP-emulation path.
    ///
    /// Differences vs [`open_unix_to_guest`] / [`open_tcp_to_guest`]:
    ///
    /// - No TSI listener registry lookup. The destination is just
    ///   `(self.cid, guest_port)` — caller-specified guest port,
    ///   muxer-owned guest CID.
    /// - The guest kernel routes the resulting `OP_REQUEST` to
    ///   whichever vsock socket is bound to `guest_port` (the
    ///   `socket(AF_VSOCK, ...) + bind(...) + listen()` path on the
    ///   guest), independent of TSI.
    /// - `handle_response` matches by `(cid, host_src_port)` from
    ///   `pending_inbound` so the byte bridge is identical to TSI.
    ///
    /// Return value: Ok if the request was queued. The actual
    /// connect outcome is observed on `OP_RESPONSE` (success) or
    /// `OP_RST` (no listener / kernel refused) over the muxer's
    /// rxq. Callers that need a "did it connect" handshake on the
    /// stream should expect to see EOF immediately if the guest
    /// rejected the request.
    pub fn open_native_to_guest(
        &self,
        stream: MuxerStream,
        guest_port: u32,
    ) -> std::io::Result<()> {
        let cid = self.cid;
        let host_src_port = crate::devices::virtio::vsock::tsi_stream::alloc_host_src_port();
        self.pending_inbound.lock().unwrap().insert(
            (cid, host_src_port),
            PendingInbound {
                cid,
                host_src_port,
                vm_port: guest_port,
                stream: Some(stream),
            },
        );
        let req = RxPacket {
            hdr: Header {
                src_cid: VSOCK_HOST_CID,
                dst_cid: cid,
                src_port: host_src_port,
                dst_port: guest_port,
                len: 0,
                type_: VSOCK_TYPE_STREAM,
                op: VSOCK_OP_REQUEST,
                flags: 0,
                buf_alloc: OUR_BUF_ALLOC,
                fwd_cnt: 0,
            },
            data: Vec::new(),
        };
        let kick = self.kick_device.lock().unwrap().clone();
        push_rxq_and_kick(&self.rxq, &kick, req);
        Ok(())
    }

    pub fn open_tcp_to_guest(&self, tcp: TcpStream, vm_port: Option<u32>) -> std::io::Result<()> {
        self.open_stream_to_guest(MuxerStream::Tcp(tcp), vm_port)
    }

    /// Same as `open_tcp_to_guest` but injects `prefix` bytes into the
    /// guest as if they had been read from the TCP socket first. Used
    /// for SCM_RIGHTS handoff, where the router has already consumed
    /// some bytes from the client's TCP buffer to make a routing
    /// decision and now needs the guest to see them.
    pub fn open_tcp_to_guest_with_prefix(
        &self,
        tcp: TcpStream,
        prefix: Vec<u8>,
        vm_port: Option<u32>,
    ) -> std::io::Result<()> {
        self.open_stream_to_guest(MuxerStream::TcpWithPrefix(tcp, prefix), vm_port)
    }

    /// Snapshot just the TSI listeners (the only host-side state worth
    /// surviving restore). In-flight TCP/UDP conns and pending hand-
    /// shakes are intentionally dropped — peer TCPs see RST, retried
    /// requests come back.
    pub fn capture_tsi_listeners(&self) -> Vec<TsiListenerSnapshot> {
        let s = self.tsi.lock().unwrap();
        s.iter()
            .filter_map(|((cid, peer_port), st)| {
                // Only listeners — entries past TSI_LISTEN have vm_port set.
                st.vm_port.map(|vm_port| TsiListenerSnapshot {
                    cid: *cid,
                    peer_port: *peer_port,
                    vm_port,
                    family: st.family,
                    socktype: st.socktype,
                })
            })
            .collect()
    }

    pub fn listener_count(&self) -> usize {
        let s = self.tsi.lock().unwrap();
        s.values().filter(|st| st.vm_port.is_some()).count()
    }

    /// Re-bind every captured TSI listener: spawn a fresh accept
    /// thread on a new ephemeral host port and re-insert the
    /// TsiState. The guest only knows its `vm_port`; the host port
    /// changes across restore (transparent — guest never sees it).
    pub fn restore_tsi_listeners(self: &Arc<Self>, snaps: &[TsiListenerSnapshot]) {
        for s in snaps {
            // Build the on_accept closure exactly as TSI_LISTEN does.
            let pending = self.pending_inbound.clone();
            let rxq = self.rxq.clone();
            let kick = self.kick_device.lock().unwrap().clone();
            let on_accept: Arc<dyn Fn(u64, u32, u32, std::net::TcpStream) + Send + Sync> =
                Arc::new(move |cid, host_src_port, vm_port, tcp| {
                    pending.lock().unwrap().insert(
                        (cid, host_src_port),
                        PendingInbound {
                            cid,
                            host_src_port,
                            vm_port,
                            stream: Some(MuxerStream::Tcp(tcp)),
                        },
                    );
                    let req = RxPacket {
                        hdr: Header {
                            src_cid: VSOCK_HOST_CID,
                            dst_cid: cid,
                            src_port: host_src_port,
                            dst_port: vm_port,
                            len: 0,
                            type_: VSOCK_TYPE_STREAM,
                            op: VSOCK_OP_REQUEST,
                            flags: 0,
                            buf_alloc: OUR_BUF_ALLOC,
                            fwd_cnt: 0,
                        },
                        data: Vec::new(),
                    };
                    push_rxq_and_kick(&rxq, &kick, req);
                });
            match TsiListener::bind(s.cid, s.vm_port, on_accept) {
                Ok(listener) => {
                    if vsock_trace_enabled() {
                        eprintln!(
                            "[muxer] restored TSI listener cid={} vm_port={} -> host {}",
                            s.cid, s.vm_port, listener.host_addr
                        );
                    }
                    self.tsi.lock().unwrap().insert(
                        (s.cid, s.peer_port),
                        TsiState {
                            family: s.family,
                            socktype: s.socktype,
                            vm_port: Some(s.vm_port),
                            listener: Some(listener),
                        },
                    );
                }
                Err(e) => {
                    eprintln!(
                        "[muxer] restore listener cid={} vm_port={} ERR: {e}",
                        s.cid, s.vm_port
                    );
                }
            }
        }
    }

    fn pending_inbound_arc(&self) -> Arc<Mutex<HashMap<(u64, u32), PendingInbound>>> {
        self.pending_inbound.clone()
    }

    /// Push a packet to the muxer RX queue + wake the device.
    pub fn submit(&self, pkt: RxPacket) {
        let kick = self.kick_device.lock().unwrap().clone();
        push_rxq_and_kick(&self.rxq, &kick, pkt);
    }

    /// Process one TX packet from the guest. Returns RxPackets to
    /// be queued back. This is the muxer's main entry point;
    /// vsock::device calls it for every TX-queue head.
    pub fn handle_tx(&self, hdr: &Header, payload: &[u8]) -> Vec<RxPacket> {
        // TSI control DGRAMs: dst_port in [1024, 1031], type=DGRAM.
        if hdr.type_ == VSOCK_TYPE_DGRAM
            && hdr.dst_port >= TSI_PROXY_CREATE
            && hdr.dst_port <= TSI_PROXY_RELEASE
        {
            return self.handle_tsi_control(hdr, payload);
        }
        if hdr.type_ == VSOCK_TYPE_DGRAM && hdr.op == VSOCK_OP_RW {
            // Connected UDP sockets send payloads as regular vsock
            // DGRAMs to their per-socket peer port after TSI_CONNECT.
            // sendto(2) with an explicit address still uses the
            // TSI_SENDTO_DATA control port handled below.
            self.send_udp_payload(hdr.src_cid, hdr.dst_port, hdr.src_port, payload);
            return Vec::new();
        }

        // Stream connection ops.
        match hdr.op {
            VSOCK_OP_REQUEST => self.handle_request(hdr),
            VSOCK_OP_RESPONSE => self.handle_response(hdr),
            VSOCK_OP_RW => self.handle_rw(hdr, payload),
            VSOCK_OP_SHUTDOWN | VSOCK_OP_RST => self.handle_close(hdr),
            VSOCK_OP_CREDIT_UPDATE | VSOCK_OP_CREDIT_REQUEST => self.handle_credit(hdr),
            _ => Vec::new(),
        }
    }

    fn handle_request(&self, hdr: &Header) -> Vec<RxPacket> {
        // Env service: guest opens AF_VSOCK to (HOST, VSOCK_ENV_PORT)
        // and reads JSON. Reply RESPONSE + RW(json) + SHUTDOWN. No
        // muxer state lives past the request.
        if hdr.dst_port == VSOCK_ENV_PORT {
            let json = VSOCK_ENV_JSON
                .lock()
                .unwrap()
                .clone()
                .unwrap_or_else(|| r#"{"env":{},"secrets":{}}"#.into());
            let mk = |op: u16, flags: u32, data: Vec<u8>| RxPacket {
                hdr: Header {
                    src_cid: hdr.dst_cid,
                    dst_cid: hdr.src_cid,
                    src_port: hdr.dst_port,
                    dst_port: hdr.src_port,
                    len: data.len() as u32,
                    type_: VSOCK_TYPE_STREAM,
                    op,
                    flags,
                    buf_alloc: OUR_BUF_ALLOC,
                    fwd_cnt: 0,
                },
                data,
            };
            return vec![
                mk(VSOCK_OP_RESPONSE, 0, Vec::new()),
                mk(VSOCK_OP_RW, 0, json.into_bytes()),
                mk(VSOCK_OP_SHUTDOWN, 1u32 | 2u32, Vec::new()),
            ];
        }

        // Outbound CONNECT phase 2: guest opens vsock to
        // (HOST, TSI_PROXY_PORT=620) carrying src_port=peer_port
        // (the same peer_port the guest sent in the prior
        // TSI_CONNECT). Pair it with the parked TcpStream.
        if hdr.dst_port == TSI_PROXY_PORT {
            let key = (hdr.src_cid, hdr.src_port);
            let mut pending = self.pending_outbound.lock().unwrap();
            if let Some(mut po) = pending.remove(&key) {
                let tcp = match po.tcp.take() {
                    Some(t) => t,
                    None => return vec![RxPacket::rst_for(hdr)],
                };
                drop(pending);
                // For outbound, we use guest's peer_port as the
                // host-side port too (so RW packets we send back
                // have src=hdr.dst_port=TSI_PROXY_PORT and the
                // guest's RW to us has dst=our_host_port=peer_port).
                // Track via inbound_conns keyed the same way; the
                // Write/Close routing in handle_rw matches on
                // (cid, hdr.dst_port = our side port).
                let host_src_port = hdr.src_port; // = peer_port
                let guest_dst_port = hdr.src_port;
                let rxq = self.rxq.clone();
                let kick = self.kick_device.lock().unwrap().clone();
                let cid = po.cid;
                let state = Arc::new(Mutex::new(ProxyConnState::new(
                    guest_dst_port,
                    host_src_port,
                )));
                // Seed peer_buf_alloc/peer_fwd_cnt from the
                // REQUEST header — outbound conns start with
                // whatever the guest advertised on its REQUEST.
                {
                    let mut st = state.lock().unwrap();
                    st.peer_buf_alloc = hdr.buf_alloc;
                    st.peer_fwd_cnt = hdr.fwd_cnt;
                }
                let pause_flag = state.lock().unwrap().pause_flag.clone();
                let state_for_cb = state.clone();
                let on_data: Arc<dyn Fn(Vec<u8>) + Send + Sync> = Arc::new(move |data| {
                    let (op, flags) = if data.is_empty() {
                        (VSOCK_OP_SHUTDOWN, 1u32 | 2u32)
                    } else {
                        (VSOCK_OP_RW, 0u32)
                    };
                    let (our_fwd_cnt, _avail_before) = {
                        let mut st = state_for_cb.lock().unwrap();
                        let n = data.len() as u32;
                        st.tx_cnt = st.tx_cnt.wrapping_add(n);
                        let avail = st.available_credit();
                        if !data.is_empty() && avail <= CREDIT_LOW_WATERMARK {
                            st.paused_readable = true;
                            st.pause_flag.store(true, Ordering::Release);
                            if vsock_trace_enabled() {
                                eprintln!(
                                    "[muxer-ondata] pause host_src_port={} tx_cnt={} peer_fwd_cnt={} peer_buf_alloc={} avail={}",
                                    st.host_src_port,
                                    st.tx_cnt,
                                    st.peer_fwd_cnt,
                                    st.peer_buf_alloc,
                                    avail,
                                );
                            }
                        }
                        (st.our_fwd_cnt, avail)
                    };
                    let pkt = RxPacket {
                        hdr: Header {
                            src_cid: VSOCK_HOST_CID,
                            dst_cid: cid,
                            src_port: TSI_PROXY_PORT,
                            dst_port: guest_dst_port,
                            len: data.len() as u32,
                            type_: VSOCK_TYPE_STREAM,
                            op,
                            flags,
                            buf_alloc: OUR_BUF_ALLOC,
                            fwd_cnt: our_fwd_cnt,
                        },
                        data,
                    };
                    push_rxq_and_kick(&rxq, &kick, pkt);
                });
                let _ = self.io_tx.send(MuxerCmd::Register {
                    host_src_port,
                    stream: MuxerStream::Tcp(tcp),
                    on_data,
                    pause_flag,
                });
                let _ = self.io_waker.wake();
                self.inbound_conns
                    .lock()
                    .unwrap()
                    .insert((cid, host_src_port), state);
                if vsock_trace_enabled() {
                    eprintln!(
                        "[muxer] outbound conn established cid={cid} peer_port={host_src_port}"
                    );
                }
                // Send RESPONSE to acknowledge the vsock REQUEST.
                return vec![RxPacket {
                    hdr: Header {
                        src_cid: VSOCK_HOST_CID,
                        dst_cid: cid,
                        src_port: TSI_PROXY_PORT,
                        dst_port: hdr.src_port,
                        len: 0,
                        type_: VSOCK_TYPE_STREAM,
                        op: VSOCK_OP_RESPONSE,
                        flags: 0,
                        buf_alloc: OUR_BUF_ALLOC,
                        fwd_cnt: 0,
                    },
                    data: Vec::new(),
                }];
            }
            return vec![RxPacket::rst_for(hdr)];
        }
        // No matching listener for any other dst_port → RST.
        vec![RxPacket::rst_for(hdr)]
    }

    fn handle_response(&self, hdr: &Header) -> Vec<RxPacket> {
        // Guest accepted our host-initiated REQUEST. Match by
        // (cid, our_host_src_port = hdr.dst_port). Spawn the
        // TsiConn bridge: TCP→on_data closure pushes RW packets
        // to the muxer's rxq + kicks the device.
        let key = (hdr.src_cid, hdr.dst_port);
        let pending = self.pending_inbound.lock().unwrap().remove(&key);
        if let Some(mut pi) = pending {
            let stream = match pi.stream.take() {
                Some(s) => s,
                None => return vec![RxPacket::rst_for(hdr)],
            };
            let rxq = self.rxq.clone();
            let kick = self.kick_device.lock().unwrap().clone();
            let cid = pi.cid;
            let host_src_port = pi.host_src_port;
            // guest_dst_port = hdr.src_port — the guest's accepted
            // socket's source port, used as the dst for our RW.
            let guest_dst_port = hdr.src_port;
            let state = Arc::new(Mutex::new(ProxyConnState::new(
                guest_dst_port,
                host_src_port,
            )));
            // Seed peer credit from the RESPONSE header.
            {
                let mut st = state.lock().unwrap();
                st.peer_buf_alloc = hdr.buf_alloc;
                st.peer_fwd_cnt = hdr.fwd_cnt;
            }
            let pause_flag = state.lock().unwrap().pause_flag.clone();
            let state_for_cb = state.clone();
            // on_data: empty Vec means "host closed, send SHUTDOWN".
            let on_data: Arc<dyn Fn(Vec<u8>) + Send + Sync> = Arc::new(move |data| {
                let (op, flags) = if data.is_empty() {
                    // VSOCK_FLAGS_SHUTDOWN_RCV | VSOCK_FLAGS_SHUTDOWN_SEND
                    (VSOCK_OP_SHUTDOWN, 1u32 | 2u32)
                } else {
                    (VSOCK_OP_RW, 0u32)
                };
                let our_fwd_cnt = {
                    let mut st = state_for_cb.lock().unwrap();
                    let n = data.len() as u32;
                    st.tx_cnt = st.tx_cnt.wrapping_add(n);
                    let avail = st.available_credit();
                    if !data.is_empty() && avail <= CREDIT_LOW_WATERMARK {
                        st.paused_readable = true;
                        st.pause_flag.store(true, Ordering::Release);
                        if vsock_trace_enabled() {
                            eprintln!(
                                "[muxer-ondata] pause host_src_port={} tx_cnt={} peer_fwd_cnt={} peer_buf_alloc={} avail={}",
                                st.host_src_port,
                                st.tx_cnt,
                                st.peer_fwd_cnt,
                                st.peer_buf_alloc,
                                avail,
                            );
                        }
                    }
                    st.our_fwd_cnt
                };
                let pkt = RxPacket {
                    hdr: Header {
                        src_cid: VSOCK_HOST_CID,
                        dst_cid: cid,
                        src_port: host_src_port,
                        dst_port: guest_dst_port,
                        len: data.len() as u32,
                        type_: VSOCK_TYPE_STREAM,
                        op,
                        flags,
                        buf_alloc: OUR_BUF_ALLOC,
                        fwd_cnt: our_fwd_cnt,
                    },
                    data,
                };
                push_rxq_and_kick(&rxq, &kick, pkt);
            });
            // Hand the TCP stream to the single muxer-io thread.
            // No per-conn pthread_create — see the single-epoll-
            // thread design in muxer_thread.rs.
            let _ = self.io_tx.send(MuxerCmd::Register {
                host_src_port,
                stream,
                on_data,
                pause_flag,
            });
            let _ = self.io_waker.wake();
            self.inbound_conns.lock().unwrap().insert(key, state);
            if vsock_trace_enabled() {
                eprintln!("[muxer] inbound conn established cid={cid} host_src={host_src_port} guest_dst={guest_dst_port}");
            }
        }
        Vec::new()
    }

    fn handle_rw(&self, hdr: &Header, payload: &[u8]) -> Vec<RxPacket> {
        // Inbound vs outbound routing:
        //   inbound  RW: dst_port = our_host_src_port (we registered under it)
        //   outbound RW: dst_port = TSI_PROXY_PORT, src_port = peer_port
        //                (we registered under peer_port)
        let conn_key = if hdr.dst_port == TSI_PROXY_PORT {
            (hdr.src_cid, hdr.src_port)
        } else {
            (hdr.src_cid, hdr.dst_port)
        };
        // Look up the state Arc, then drop the outer map lock
        // before we lock the per-conn Mutex — the on_data path
        // also takes that inner lock, and holding both at once
        // (outer then inner) while another thread might take
        // them inner-only would be fine, but keeping locks
        // narrow is cheaper anyway.
        let state_arc = {
            let conns = self.inbound_conns.lock().unwrap();
            conns.get(&conn_key).cloned()
        };
        let Some(state_arc) = state_arc else {
            // Not an inbound conn — try fallback proxy.
            let key = proxy_key(hdr.dst_port, hdr.src_port);
            let mut proxies = self.proxies.lock().unwrap();
            if let Some(p) = proxies.get_mut(&key) {
                return p.handle_packet(hdr, payload);
            }
            return vec![RxPacket::rst_for(hdr)];
        };

        let (needs_credit_update, needs_resume, host_src_port) = {
            let mut st = state_arc.lock().unwrap();
            // EVERY vsock packet carries buf_alloc + fwd_cnt; we
            // refresh peer credit on every RW, not just on
            // dedicated CREDIT_UPDATE messages. This is the
            // critical flow-control fix — without it,
            // peer_fwd_cnt stays at 0 and we infinitely accrete
            // tx_cnt against a wall that never moves.
            let resume = st.observe_peer(hdr.buf_alloc, hdr.fwd_cnt);

            st.our_fwd_cnt = st.our_fwd_cnt.wrapping_add(payload.len() as u32);
            let consumed_since_last = st.our_fwd_cnt.wrapping_sub(st.last_credit_fwd_cnt);
            if vsock_trace_enabled() && st.our_fwd_cnt % 50000 < 1500 {
                eprintln!(
                    "[muxer-rw-fc] our_fwd_cnt={} last_credit={} consumed_since_last={} thr={} peer_buf_alloc={} peer_fwd_cnt={} tx_cnt={} avail={}",
                    st.our_fwd_cnt,
                    st.last_credit_fwd_cnt,
                    consumed_since_last,
                    CREDIT_UPDATE_THRESHOLD,
                    st.peer_buf_alloc,
                    st.peer_fwd_cnt,
                    st.tx_cnt,
                    st.available_credit(),
                );
            }
            let needs_update = if consumed_since_last >= CREDIT_UPDATE_THRESHOLD {
                let fwd_cnt = st.our_fwd_cnt;
                let guest_dst = st.guest_dst_port;
                st.last_credit_fwd_cnt = fwd_cnt;
                Some((fwd_cnt, guest_dst))
            } else {
                None
            };
            (needs_update, resume, st.host_src_port)
        };

        if !payload.is_empty() {
            let t0 = Instant::now();
            let _ = self.io_tx.send(MuxerCmd::Write {
                host_src_port: conn_key.1,
                bytes: payload.to_vec(),
            });
            let _ = self.io_waker.wake();
            mux_profile::record(
                Stage::GuestToTcpSend,
                payload.len(),
                t0.elapsed().as_micros() as u64,
            );
        }

        if needs_resume {
            let _ = self
                .io_tx
                .send(MuxerCmd::ResumeReadable { host_src_port });
            let _ = self.io_waker.wake();
            if vsock_trace_enabled() {
                eprintln!(
                    "[muxer-rw-fc] resume host_src_port={host_src_port} (credit returned)"
                );
            }
        }

        if let Some((fwd_cnt, guest_dst_port)) = needs_credit_update {
            return vec![RxPacket {
                hdr: Header {
                    src_cid: VSOCK_HOST_CID,
                    dst_cid: hdr.src_cid,
                    src_port: hdr.dst_port,
                    dst_port: guest_dst_port,
                    len: 0,
                    type_: VSOCK_TYPE_STREAM,
                    op: VSOCK_OP_CREDIT_UPDATE,
                    flags: 0,
                    buf_alloc: OUR_BUF_ALLOC,
                    fwd_cnt,
                },
                data: Vec::new(),
            }];
        }
        Vec::new()
    }

    fn handle_close(&self, hdr: &Header) -> Vec<RxPacket> {
        let key = proxy_key(hdr.dst_port, hdr.src_port);
        self.proxies.lock().unwrap().remove(&key);
        let conn_key = if hdr.dst_port == TSI_PROXY_PORT {
            (hdr.src_cid, hdr.src_port)
        } else {
            (hdr.src_cid, hdr.dst_port)
        };
        if self
            .inbound_conns
            .lock()
            .unwrap()
            .remove(&conn_key)
            .is_some()
        {
            let _ = self.io_tx.send(MuxerCmd::Close {
                host_src_port: conn_key.1,
            });
            let _ = self.io_waker.wake();
        }
        Vec::new()
    }

    fn handle_credit(&self, hdr: &Header) -> Vec<RxPacket> {
        // CREDIT_UPDATE / CREDIT_REQUEST carry the same
        // buf_alloc + fwd_cnt fields every header does. Refresh
        // peer credit on the matching conn and, if we were
        // paused, send ResumeReadable.
        let conn_key = if hdr.dst_port == TSI_PROXY_PORT {
            (hdr.src_cid, hdr.src_port)
        } else {
            (hdr.src_cid, hdr.dst_port)
        };
        let state_arc = {
            let conns = self.inbound_conns.lock().unwrap();
            conns.get(&conn_key).cloned()
        };
        let Some(state_arc) = state_arc else {
            return Vec::new();
        };
        let (resume, host_src_port, our_fwd_cnt, guest_dst_port) = {
            let mut st = state_arc.lock().unwrap();
            let resume = st.observe_peer(hdr.buf_alloc, hdr.fwd_cnt);
            (resume, st.host_src_port, st.our_fwd_cnt, st.guest_dst_port)
        };
        if resume {
            let _ = self
                .io_tx
                .send(MuxerCmd::ResumeReadable { host_src_port });
            let _ = self.io_waker.wake();
            if vsock_trace_enabled() {
                eprintln!(
                    "[muxer-credit] resume host_src_port={host_src_port} (credit returned)"
                );
            }
        }
        // On CREDIT_REQUEST the guest is explicitly asking us to
        // tell it our fwd_cnt — reply with a CREDIT_UPDATE so it
        // can compute its own send credit. (We always have room
        // since we forward straight to host TCP.)
        if hdr.op == VSOCK_OP_CREDIT_REQUEST {
            return vec![RxPacket {
                hdr: Header {
                    src_cid: VSOCK_HOST_CID,
                    dst_cid: hdr.src_cid,
                    src_port: hdr.dst_port,
                    dst_port: guest_dst_port,
                    len: 0,
                    type_: VSOCK_TYPE_STREAM,
                    op: VSOCK_OP_CREDIT_UPDATE,
                    flags: 0,
                    buf_alloc: OUR_BUF_ALLOC,
                    fwd_cnt: our_fwd_cnt,
                },
                data: Vec::new(),
            }];
        }
        Vec::new()
    }

    /// TSI control DGRAM dispatcher. Each control op is a separate
    /// DGRAM packet whose dst_port IS the opcode (1024..1031).
    /// Replies are DGRAM RW packets pushed back via the same
    /// (local_port, peer_port) — guest reads the result from a
    /// recvmsg on the AF_TSI control socket.
    fn handle_tsi_control(&self, hdr: &Header, body: &[u8]) -> Vec<RxPacket> {
        // Authentication: if this muxer was constructed with an
        // expected token (i.e. a snapshot baked under 0.6.0+),
        // every control DGRAM must be prefixed with that token.
        // Anything else — short body, wrong prefix — gets silently
        // dropped. We do NOT reply with an error code: an error
        // reply would itself confirm to a probing attacker that
        // they're talking to a real TSI muxer. The legitimate
        // kernel sender never sees mismatches, so the silent drop
        // is invisible to the only consumer that matters.
        //
        // Constant-time compare (see `ct_eq`) closes the obvious
        // timing oracle a naïve byte-by-byte `==` would open.
        let body: &[u8] = if let Some(expected) = self.tsi_token.as_ref() {
            if body.len() < TSI_TOKEN_LEN || !ct_eq(&body[..TSI_TOKEN_LEN], &expected[..]) {
                return Vec::new();
            }
            &body[TSI_TOKEN_LEN..]
        } else {
            body
        };
        let mk_resp = |result: i32| -> RxPacket {
            // TSI replies ride on RW DGRAMs.
            let bytes = result.to_le_bytes().to_vec();
            RxPacket {
                hdr: Header {
                    src_cid: hdr.dst_cid,
                    dst_cid: hdr.src_cid,
                    src_port: hdr.dst_port,
                    dst_port: hdr.src_port,
                    len: bytes.len() as u32,
                    type_: VSOCK_TYPE_DGRAM,
                    op: VSOCK_OP_RW,
                    flags: 0,
                    buf_alloc: OUR_BUF_ALLOC,
                    fwd_cnt: 0,
                },
                data: bytes,
            }
        };

        match hdr.dst_port {
            TSI_PROXY_CREATE if body.len() >= 8 => {
                // struct { u32 peer_port; u16 family; u16 type; }
                let (Some(peer_port), Some(family), Some(socktype)) = (
                    body_le_u32(body, 0),
                    body_le_u16(body, 4),
                    body_le_u16(body, 6),
                ) else {
                    return Vec::new();
                };
                eprintln!(
                    "[muxer] PROXY_CREATE peer_port={peer_port} family={family} type={socktype}"
                );
                self.tsi.lock().unwrap().insert(
                    (hdr.src_cid, peer_port),
                    TsiState {
                        family,
                        socktype,
                        vm_port: None,
                        listener: None,
                    },
                );
                // No reply: the AF_TSI control op doesn't emit a
                // recvmsg on the guest side for this op.
                Vec::new()
            }
            TSI_LISTEN if body.len() >= 8 => {
                // struct { u32 peer_port; u32 vm_port; ... }
                let (Some(peer_port), Some(vm_port)) = (body_le_u32(body, 0), body_le_u32(body, 4))
                else {
                    return vec![mk_resp(-22)];
                };
                // Accept-thread closure must hold weak refs to our
                // own state. We can't get Arc<Self> from inside an
                // &self method, so we capture the inner Arc-able
                // pieces directly.
                let pending = self.pending_inbound_arc();
                let rxq = self.rxq.clone();
                let kick = self.kick_device.lock().unwrap().clone();
                let on_accept: Arc<dyn Fn(u64, u32, u32, TcpStream) + Send + Sync> =
                    Arc::new(move |cid, host_src_port, vm_port, tcp| {
                        // Stash the TcpStream until guest sends RESPONSE.
                        pending.lock().unwrap().insert(
                            (cid, host_src_port),
                            PendingInbound {
                                cid,
                                host_src_port,
                                vm_port,
                                stream: Some(MuxerStream::Tcp(tcp)),
                            },
                        );
                        // Push vsock REQUEST to the guest.
                        let req = RxPacket {
                            hdr: Header {
                                src_cid: VSOCK_HOST_CID,
                                dst_cid: cid,
                                src_port: host_src_port,
                                dst_port: vm_port,
                                len: 0,
                                type_: VSOCK_TYPE_STREAM,
                                op: VSOCK_OP_REQUEST,
                                flags: 0,
                                buf_alloc: OUR_BUF_ALLOC,
                                fwd_cnt: 0,
                            },
                            data: Vec::new(),
                        };
                        push_rxq_and_kick(&rxq, &kick, req);
                    });
                let listener = TsiListener::bind(hdr.src_cid, vm_port, on_accept).ok();
                let host = listener.as_ref().map(|l| l.host_addr);
                eprintln!(
                    "[muxer] LISTEN peer_port={peer_port} vm_port={vm_port} -> host {host:?}"
                );
                if let Some(s) = self.tsi.lock().unwrap().get_mut(&(hdr.src_cid, peer_port)) {
                    s.vm_port = Some(vm_port);
                    s.listener = listener;
                }
                vec![mk_resp(if host.is_some() { 0 } else { -22 })]
            }
            TSI_ACCEPT if body.len() >= 4 => {
                let Some(peer_port) = body_le_u32(body, 0) else {
                    return vec![mk_resp(-22)];
                };
                let s = self.tsi.lock().unwrap();
                let r = if s.contains_key(&(hdr.src_cid, peer_port)) {
                    0
                } else {
                    -22
                };
                vec![mk_resp(r)]
            }
            TSI_GETNAME if body.len() >= 12 => {
                // Build a fake AF_INET sockaddr response — addr_len=16,
                // family=AF_INET, port=0, ip=127.0.0.1.
                let mut buf = vec![0u8; 4 + 4 + 128];
                buf[0..4].copy_from_slice(&0i32.to_le_bytes()); // result=0
                buf[4..8].copy_from_slice(&16u32.to_le_bytes()); // addr_len=16
                buf[8..10].copy_from_slice(&2u16.to_le_bytes()); // AF_INET
                buf[12..16].copy_from_slice(&[127, 0, 0, 1]); // 127.0.0.1
                vec![RxPacket {
                    hdr: Header {
                        src_cid: hdr.dst_cid,
                        dst_cid: hdr.src_cid,
                        src_port: hdr.dst_port,
                        dst_port: hdr.src_port,
                        len: buf.len() as u32,
                        type_: VSOCK_TYPE_DGRAM,
                        op: VSOCK_OP_RW,
                        flags: 0,
                        buf_alloc: OUR_BUF_ALLOC,
                        fwd_cnt: 0,
                    },
                    data: buf,
                }]
            }
            TSI_CONNECT if body.len() >= 8 => {
                // struct tsi_connect_req {
                //   u32 peer_port; u32 addr_len; sockaddr addr;
                // }
                // sockaddr_in: u16 family; u16 port_be; u32 addr_be4;
                let (Some(peer_port), Some(addr_len)) =
                    (body_le_u32(body, 0), body_le_u32(body, 4))
                else {
                    return vec![mk_resp(-22)];
                };
                let addr_len = addr_len as usize;
                if body.len() < 8 + addr_len.min(128) || addr_len < 16 {
                    return vec![mk_resp(-22)]; // -EINVAL
                }
                let Some(family) = body_le_u16(body, 8) else {
                    return vec![mk_resp(-22)];
                };
                if family != 2 {
                    // AF_INET only for now (no IPv6).
                    return vec![mk_resp(-97)]; // -EAFNOSUPPORT
                }
                let Some(port) = body_be_u16(body, 10) else {
                    return vec![mk_resp(-22)];
                };
                let ip = std::net::Ipv4Addr::new(body[12], body[13], body[14], body[15]);
                let target = std::net::SocketAddr::from((ip, port));
                let socktype = self
                    .tsi
                    .lock()
                    .unwrap()
                    .get(&(hdr.src_cid, peer_port))
                    .map(|s| s.socktype);
                if socktype == Some(libc::SOCK_DGRAM as u16) {
                    if let Err(why) = crate::vmm::egress_policy::check_addr(target) {
                        eprintln!("[muxer] UDP CONNECT cid={} peer_port={peer_port} -> {target} BLOCKED: {why}",
                            hdr.src_cid);
                        return vec![mk_resp(-13)]; // -EACCES
                    }
                    self.udp_dst
                        .lock()
                        .unwrap()
                        .insert((hdr.src_cid, peer_port), target);
                    eprintln!(
                        "[muxer] UDP CONNECT cid={} peer_port={peer_port} -> {target} OK",
                        hdr.src_cid
                    );
                    return vec![mk_resp(0)];
                }
                if let Err(why) = crate::vmm::egress_policy::check_addr(target) {
                    eprintln!(
                        "[muxer] CONNECT cid={} peer_port={peer_port} -> {target} BLOCKED: {why}",
                        hdr.src_cid
                    );
                    return vec![mk_resp(-13)]; // -EACCES
                }
                let res = std::net::TcpStream::connect_timeout(
                    &target,
                    std::time::Duration::from_secs(2),
                );
                match res {
                    Ok(tcp) => {
                        let _ = tcp.set_nodelay(true);
                        let _ = tcp.set_nonblocking(true);
                        eprintln!(
                            "[muxer] CONNECT cid={} peer_port={peer_port} -> {target} OK",
                            hdr.src_cid
                        );
                        self.pending_outbound.lock().unwrap().insert(
                            (hdr.src_cid, peer_port),
                            PendingOutbound {
                                cid: hdr.src_cid,
                                peer_port,
                                tcp: Some(tcp),
                            },
                        );
                        vec![mk_resp(0)]
                    }
                    Err(e) => {
                        eprintln!(
                            "[muxer] CONNECT cid={} peer_port={peer_port} -> {target} ERR: {e}",
                            hdr.src_cid
                        );
                        vec![mk_resp(-111)] // -ECONNREFUSED
                    }
                }
            }
            TSI_PROXY_RELEASE => {
                // No-op: vsock accepts inherit the listener's port,
                // so releasing on conn close kills the listener
                // permanently. Listener stays alive for the VM.
                Vec::new()
            }
            TSI_SENDTO_ADDR if body.len() >= 8 => {
                // struct tsi_sendto_addr { u32 peer_port; u32 addr_len;
                //                          sockaddr_in addr; }
                let (Some(peer_port), Some(addr_len)) =
                    (body_le_u32(body, 0), body_le_u32(body, 4))
                else {
                    return Vec::new();
                };
                let addr_len = addr_len as usize;
                if body.len() < 8 + addr_len.min(128) || addr_len < 16 {
                    return Vec::new();
                }
                let Some(family) = body_le_u16(body, 8) else {
                    return Vec::new();
                };
                if family == 2 {
                    let Some(port) = body_be_u16(body, 10) else {
                        return Vec::new();
                    };
                    let ip = std::net::Ipv4Addr::new(body[12], body[13], body[14], body[15]);
                    let addr = SocketAddr::from((ip, port));
                    if let Err(why) = crate::vmm::egress_policy::check_addr(addr) {
                        eprintln!(
                            "[muxer] SENDTO_ADDR cid={} -> {addr} BLOCKED: {why}",
                            hdr.src_cid
                        );
                    } else {
                        self.udp_dst
                            .lock()
                            .unwrap()
                            .insert((hdr.src_cid, peer_port), addr);
                    }
                }
                Vec::new()
            }
            TSI_SENDTO_DATA => {
                // The DGRAM socket sends data here directly; peer_port
                // is in hdr.src_port (the AF_INET socket's local TSI
                // port). Body is just the datagram payload.
                let peer_port = hdr.src_port;
                self.send_udp_payload(hdr.src_cid, peer_port, hdr.src_port, body);
                Vec::new()
            }
            _ => Vec::new(),
        }
    }

    fn send_udp_payload(&self, cid: u64, peer_port: u32, guest_dst_port: u32, data: &[u8]) {
        if data.is_empty() {
            return;
        }
        let key = (cid, peer_port);
        let dst = match self.udp_dst.lock().unwrap().get(&key).cloned() {
            Some(a) => a,
            None => {
                if vsock_trace_enabled() {
                    eprintln!("[muxer] UDP DATA cid={cid} peer_port={peer_port} with no dst");
                }
                return;
            }
        };
        let sock_arc = {
            let mut s = self.udp_sockets.lock().unwrap();
            if let Some(s) = s.get(&key) {
                s.clone()
            } else {
                let udp = match UdpSocket::bind("0.0.0.0:0") {
                    Ok(u) => u,
                    Err(e) => {
                        eprintln!("[muxer] UDP bind cid={cid} peer_port={peer_port}: {e}");
                        return;
                    }
                };
                let recv_fd = match udp.try_clone() {
                    Ok(fd) => fd,
                    Err(e) => {
                        eprintln!("[muxer] UDP clone cid={cid} peer_port={peer_port}: {e}");
                        return;
                    }
                };
                let arc = Arc::new(udp);
                s.insert(key, arc.clone());

                let rxq = self.rxq.clone();
                let kick = self.kick_device.lock().unwrap().clone();
                let on_data: Arc<dyn Fn(Vec<u8>) + Send + Sync> = Arc::new(move |data| {
                    let pkt = RxPacket {
                        hdr: Header {
                            src_cid: VSOCK_HOST_CID,
                            dst_cid: cid,
                            src_port: peer_port,
                            dst_port: guest_dst_port,
                            len: data.len() as u32,
                            type_: VSOCK_TYPE_DGRAM,
                            op: VSOCK_OP_RW,
                            flags: 0,
                            buf_alloc: OUR_BUF_ALLOC,
                            fwd_cnt: 0,
                        },
                        data,
                    };
                    push_rxq_and_kick(&rxq, &kick, pkt);
                });
                let _ = self.io_tx.send(MuxerCmd::RegisterUdp {
                    key: peer_port,
                    udp: recv_fd,
                    on_data,
                });
                let _ = self.io_waker.wake();
                arc
            }
        };
        if vsock_trace_enabled() {
            eprintln!(
                "[muxer] UDP DATA cid={cid} peer_port={peer_port} guest_dst={guest_dst_port} -> {dst} {}B",
                data.len()
            );
        }
        let _ = sock_arc.send_to(data, dst);
    }
}

fn push_rxq_and_kick(
    rxq: &Arc<Mutex<MuxerRxQ>>,
    kick: &Option<Arc<dyn Fn() + Send + Sync>>,
    pkt: RxPacket,
) {
    let bytes = pkt.data.len();
    let t0 = Instant::now();
    // Coalesce kicks: only raise the SPI when we transition the
    // queue from empty → non-empty. Subsequent pushes ride the
    // in-flight drain. Saves one `hv_gic_set_spi` syscall
    // (~3.6% of muxer-io non-syscall CPU under load) per
    // coalesced packet during bursts. See
    // docs/design/concurrency-floor-2026-05-04.md.
    //
    // Note: `push_was_empty` returns None when the queue is
    // saturated (back-pressure). In that case we don't kick
    // either — the existing in-flight drain has more than
    // enough work to do; one more SPI won't help.
    let was_empty = rxq.lock().unwrap().push_was_empty(pkt);
    if was_empty == Some(true) {
        if let Some(k) = kick.as_ref() {
            k();
        }
    }
    mux_profile::record(Stage::RxqKick, bytes, t0.elapsed().as_micros() as u64);
}

// Use HOST_CID so the unused-import warning goes away.
const _: u64 = VSOCK_HOST_CID;

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

    fn control_header(dst_port: u32) -> Header {
        Header {
            src_cid: 3,
            dst_cid: VSOCK_HOST_CID,
            src_port: 40_000,
            dst_port,
            len: 0,
            type_: VSOCK_TYPE_DGRAM,
            op: VSOCK_OP_RW,
            flags: 0,
            buf_alloc: OUR_BUF_ALLOC,
            fwd_cnt: 0,
        }
    }

    fn response_i32(resp: &[RxPacket]) -> Option<i32> {
        let bytes = resp.first()?.data.get(0..4)?.try_into().ok()?;
        Some(i32::from_le_bytes(bytes))
    }

    #[test]
    fn malformed_tsi_connect_returns_einval() -> Result<(), Box<dyn std::error::Error>> {
        let muxer = VsockMuxer::new(3)?;
        let hdr = control_header(TSI_CONNECT);

        let mut body = Vec::new();
        body.extend_from_slice(&1234u32.to_le_bytes());
        body.extend_from_slice(&16u32.to_le_bytes());

        let resp = muxer.handle_tsi_control(&hdr, &body);
        assert_eq!(response_i32(&resp), Some(-22));
        Ok(())
    }

    #[test]
    fn malformed_tsi_sendto_addr_is_ignored() -> Result<(), Box<dyn std::error::Error>> {
        let muxer = VsockMuxer::new(3)?;
        let hdr = control_header(TSI_SENDTO_ADDR);

        let mut body = Vec::new();
        body.extend_from_slice(&1234u32.to_le_bytes());
        body.extend_from_slice(&16u32.to_le_bytes());

        let resp = muxer.handle_tsi_control(&hdr, &body);
        assert!(resp.is_empty());
        Ok(())
    }

    /// Build a body for a `TSI_PROXY_CREATE` op: u32 peer_port +
    /// u16 family + u16 socktype = 8 bytes minimum. Used by the
    /// auth tests below — these don't care about the meaning,
    /// only that a valid body shape gets either dispatched or
    /// dropped depending on the auth prefix.
    fn proxy_create_body(peer_port: u32) -> Vec<u8> {
        let mut b = Vec::new();
        b.extend_from_slice(&peer_port.to_le_bytes());
        // family=AF_INET=2, socktype=SOCK_STREAM=1
        b.extend_from_slice(&2u16.to_le_bytes());
        b.extend_from_slice(&1u16.to_le_bytes());
        b
    }

    #[test]
    fn tsi_control_rejects_missing_token_prefix() -> Result<(), Box<dyn std::error::Error>> {
        // Muxer armed with auth; body has NO prefix — must be
        // silently dropped (no response, no state change).
        let token = [0x11u8; TSI_TOKEN_LEN];
        let muxer = VsockMuxer::with_tsi_token(3, Some(token))?;
        let hdr = control_header(TSI_PROXY_CREATE);
        let body = proxy_create_body(1234);

        let resp = muxer.handle_tsi_control(&hdr, &body);
        assert!(
            resp.is_empty(),
            "unprefixed TSI control DGRAM must be dropped, got {} pkts",
            resp.len()
        );
        // Bookkeeping side-effect: PROXY_CREATE inserts into the
        // `tsi` map. The drop must prevent that insertion.
        assert!(
            muxer.tsi.lock().unwrap().is_empty(),
            "dropped DGRAM still mutated tsi state"
        );
        Ok(())
    }

    #[test]
    fn tsi_control_rejects_wrong_token_prefix() -> Result<(), Box<dyn std::error::Error>> {
        let token = [0x11u8; TSI_TOKEN_LEN];
        let muxer = VsockMuxer::with_tsi_token(3, Some(token))?;
        let hdr = control_header(TSI_PROXY_CREATE);

        // Wrong prefix (all 0x22) followed by the real body.
        let mut body = vec![0x22u8; TSI_TOKEN_LEN];
        body.extend_from_slice(&proxy_create_body(1234));

        let resp = muxer.handle_tsi_control(&hdr, &body);
        assert!(resp.is_empty(), "wrong-token DGRAM must be dropped");
        assert!(muxer.tsi.lock().unwrap().is_empty());
        Ok(())
    }

    #[test]
    fn tsi_control_accepts_correct_token_prefix() -> Result<(), Box<dyn std::error::Error>> {
        let token = [0x33u8; TSI_TOKEN_LEN];
        let muxer = VsockMuxer::with_tsi_token(3, Some(token))?;
        let hdr = control_header(TSI_PROXY_CREATE);

        // Real token prefix + valid body. PROXY_CREATE doesn't
        // emit a response, but it does insert into the `tsi`
        // map — we assert on the side-effect, not the (empty)
        // response.
        let mut body = token.to_vec();
        body.extend_from_slice(&proxy_create_body(1234));

        let resp = muxer.handle_tsi_control(&hdr, &body);
        // PROXY_CREATE legitimately emits no reply (see the
        // existing dispatcher), so an empty resp here is the
        // CORRECT signal that the message was dispatched.
        assert!(resp.is_empty(), "PROXY_CREATE should not reply");
        assert_eq!(
            muxer.tsi.lock().unwrap().len(),
            1,
            "authenticated PROXY_CREATE should have inserted exactly one state"
        );
        Ok(())
    }

    #[test]
    fn tsi_control_disabled_token_accepts_all() -> Result<(), Box<dyn std::error::Error>> {
        // No token armed → legacy mode. Existing kernels that
        // boot without `supermachine.tsi_token=` send bare bodies;
        // the muxer must still dispatch them.
        let muxer = VsockMuxer::new(3)?;
        let hdr = control_header(TSI_PROXY_CREATE);

        let body = proxy_create_body(1234);
        let resp = muxer.handle_tsi_control(&hdr, &body);
        assert!(resp.is_empty(), "PROXY_CREATE doesn't reply");
        assert_eq!(muxer.tsi.lock().unwrap().len(), 1);
        Ok(())
    }

    #[test]
    fn ct_eq_matches_pairs_and_short_inputs() {
        let a = [0x55u8; TSI_TOKEN_LEN];
        let b = [0x55u8; TSI_TOKEN_LEN];
        assert!(ct_eq(&a, &b));

        // One byte different anywhere — equal-length case.
        let mut c = a;
        c[0] ^= 1;
        assert!(!ct_eq(&a, &c));
        c = a;
        c[TSI_TOKEN_LEN - 1] ^= 1;
        assert!(!ct_eq(&a, &c));

        // Length mismatch: short input must NEVER compare equal,
        // even if its bytes are a prefix of `a`. This is the
        // short-input case from the test plan.
        assert!(!ct_eq(&a[..16], &a[..]));
        assert!(!ct_eq(&[], &a[..]));
        assert!(!ct_eq(&a[..], &[]));
    }

    // -----------------------------------------------------------
    // Credit / flow-control regression suite.
    //
    // Pre-fix bug: handle_credit was a no-op, peer_buf_alloc and
    // peer_fwd_cnt were never tracked, on_data closures pushed
    // RW packets with fwd_cnt:0 unconditionally. Result: large
    // host→guest streams (e.g. 4 MiB HTTP GET via TSI host proxy)
    // overflowed the guest's vsock RX socket buffer and bytes
    // were silently dropped — the integrator's npm install died
    // with `ossl_gcm_stream_update: cipher operation failed`
    // because the TLS stream truncated mid-block.
    //
    // These tests pin the per-conn ProxyConnState invariants:
    // credit observed on every inbound packet, tx_cnt advances
    // on every outbound, pause flag trips when available credit
    // drops to the low watermark, ResumeReadable issued when
    // credit returns.
    // -----------------------------------------------------------

    fn stream_header(src_cid: u64, src_port: u32, dst_port: u32, op: u16) -> Header {
        Header {
            src_cid,
            dst_cid: VSOCK_HOST_CID,
            src_port,
            dst_port,
            len: 0,
            type_: VSOCK_TYPE_STREAM,
            op,
            flags: 0,
            buf_alloc: 0,
            fwd_cnt: 0,
        }
    }

    /// Insert a synthetic ProxyConnState into the muxer's
    /// `inbound_conns` map so handle_rw / handle_credit can run
    /// without a real TCP stream. Returns the Arc so the test
    /// can inspect state after dispatch.
    fn inject_state(
        muxer: &VsockMuxer,
        cid: u64,
        host_src_port: u32,
        guest_dst_port: u32,
    ) -> Arc<Mutex<ProxyConnState>> {
        let st = Arc::new(Mutex::new(ProxyConnState::new(
            guest_dst_port,
            host_src_port,
        )));
        muxer
            .inbound_conns
            .lock()
            .unwrap()
            .insert((cid, host_src_port), st.clone());
        st
    }

    #[test]
    fn credit_initial_peer_buf_alloc_is_zero() -> Result<(), Box<dyn std::error::Error>> {
        // T1: a freshly-constructed conn state advertises no
        // peer credit yet — we haven't observed any inbound
        // header. available_credit() must therefore be zero,
        // not OUR_BUF_ALLOC (the wrong default would silently
        // grant unbounded send credit and reproduce the bug).
        let st = ProxyConnState::new(/*guest_dst*/ 12_345, /*host_src*/ 50_000);
        assert_eq!(st.peer_buf_alloc, 0);
        assert_eq!(st.peer_fwd_cnt, 0);
        assert_eq!(st.tx_cnt, 0);
        assert_eq!(st.available_credit(), 0);
        assert!(!st.paused_readable);
        Ok(())
    }

    #[test]
    fn credit_update_handler_updates_state() -> Result<(), Box<dyn std::error::Error>> {
        // T2: dispatch a CREDIT_UPDATE packet and verify the
        // tracked state reflects the buf_alloc + fwd_cnt the
        // header carried. Pre-fix, handle_credit was a no-op
        // and this test would observe peer_buf_alloc still 0.
        let muxer = VsockMuxer::new(3)?;
        let state = inject_state(&muxer, 3, 50_000, 12_345);
        let mut hdr = stream_header(3, 12_345, 50_000, VSOCK_OP_CREDIT_UPDATE);
        hdr.buf_alloc = 262_144;
        hdr.fwd_cnt = 1_000;
        let _ = muxer.handle_credit(&hdr);
        let st = state.lock().unwrap();
        assert_eq!(st.peer_buf_alloc, 262_144);
        assert_eq!(st.peer_fwd_cnt, 1_000);
        Ok(())
    }

    #[test]
    fn rw_packet_updates_credit() -> Result<(), Box<dyn std::error::Error>> {
        // T3: an RW packet ALSO carries credit info in its
        // header; the muxer must observe it on every RW, not
        // just CREDIT_UPDATE. Without this, the guest never
        // gets credit-acked for what it pushed; with the inverse
        // miss (host-side), peer_fwd_cnt would never advance.
        let muxer = VsockMuxer::new(3)?;
        let state = inject_state(&muxer, 3, 50_000, 12_345);
        let mut hdr = stream_header(3, 12_345, 50_000, VSOCK_OP_RW);
        hdr.buf_alloc = 131_072;
        hdr.fwd_cnt = 4_096;
        let _ = muxer.handle_rw(&hdr, b"hello");
        let st = state.lock().unwrap();
        assert_eq!(st.peer_buf_alloc, 131_072);
        assert_eq!(st.peer_fwd_cnt, 4_096);
        // our_fwd_cnt advances by payload size — the guest→host
        // direction's credit accounting (unchanged by this fix
        // but exercised here as a sanity check).
        assert_eq!(st.our_fwd_cnt, b"hello".len() as u32);
        Ok(())
    }

    #[test]
    fn tx_cnt_advances_on_outbound() -> Result<(), Box<dyn std::error::Error>> {
        // T4: when on_data pushes N bytes toward the guest,
        // tx_cnt must advance by exactly N. We don't have a
        // real on_data closure here, but ProxyConnState exposes
        // the same accounting the closure performs.
        let st = Arc::new(Mutex::new(ProxyConnState::new(12_345, 50_000)));
        // Seed peer credit so available is positive.
        st.lock().unwrap().peer_buf_alloc = 65_536;
        // Simulate three on_data invocations.
        for n in [1024u32, 2048u32, 512u32] {
            let mut s = st.lock().unwrap();
            s.tx_cnt = s.tx_cnt.wrapping_add(n);
        }
        let s = st.lock().unwrap();
        assert_eq!(s.tx_cnt, 1024 + 2048 + 512);
        // Available credit dropped by exactly the total pushed.
        assert_eq!(s.available_credit(), 65_536 - (1024 + 2048 + 512));
        Ok(())
    }

    #[test]
    fn pause_when_credit_exhausted() -> Result<(), Box<dyn std::error::Error>> {
        // T5: peer advertises 4096 bytes. After pushing 4096
        // bytes, available_credit drops to 0 ≤ low_watermark,
        // so the on_data path must trip the pause flag. The
        // muxer-io thread checks this after each on_data and
        // drops READABLE interest — exactly the backpressure
        // the original bug was missing.
        let st = Arc::new(Mutex::new(ProxyConnState::new(12_345, 50_000)));
        {
            let mut s = st.lock().unwrap();
            s.peer_buf_alloc = 4096;
        }
        // Simulate the on_data body: decrement available, set
        // pause when at-or-below watermark.
        {
            let mut s = st.lock().unwrap();
            s.tx_cnt = s.tx_cnt.wrapping_add(4096);
            let avail = s.available_credit();
            if avail <= CREDIT_LOW_WATERMARK {
                s.paused_readable = true;
                s.pause_flag.store(true, Ordering::Release);
            }
        }
        let s = st.lock().unwrap();
        assert_eq!(s.available_credit(), 0);
        assert!(s.paused_readable, "pause flag must be set when credit exhausted");
        assert!(s.pause_flag.load(Ordering::Acquire));
        Ok(())
    }

    #[test]
    fn resume_after_credit_update() -> Result<(), Box<dyn std::error::Error>> {
        // T6: a paused conn receives a CREDIT_UPDATE that
        // advances peer_fwd_cnt past the prior in-flight bytes.
        // observe_peer must clear pause and return true so the
        // caller (handle_credit) sends ResumeReadable.
        let muxer = VsockMuxer::new(3)?;
        let state = inject_state(&muxer, 3, 50_000, 12_345);
        // Set up: peer advertised 65_536, we sent 65_536, paused.
        // Picking a larger window than CREDIT_LOW_WATERMARK so
        // an fwd_cnt advance creates strictly-positive headroom
        // above the watermark.
        {
            let mut s = state.lock().unwrap();
            s.peer_buf_alloc = 65_536;
            s.tx_cnt = 65_536;
            s.paused_readable = true;
            s.pause_flag.store(true, Ordering::Release);
            assert_eq!(s.available_credit(), 0);
        }
        // Now CREDIT_UPDATE with fwd_cnt=65_536 (guest consumed
        // everything). buf_alloc unchanged.
        let mut hdr = stream_header(3, 12_345, 50_000, VSOCK_OP_CREDIT_UPDATE);
        hdr.buf_alloc = 65_536;
        hdr.fwd_cnt = 65_536;
        let _ = muxer.handle_credit(&hdr);
        let s = state.lock().unwrap();
        assert_eq!(s.peer_fwd_cnt, 65_536);
        assert_eq!(s.available_credit(), 65_536);
        assert!(!s.paused_readable, "pause must clear once credit > watermark");
        assert!(!s.pause_flag.load(Ordering::Acquire));
        Ok(())
    }

    #[test]
    fn outgoing_fwd_cnt_reflects_our_consumption() -> Result<(), Box<dyn std::error::Error>> {
        // Regression test for the `fwd_cnt: 0` bug — outbound
        // packets must advertise our_fwd_cnt so the guest can
        // compute its own send credit. We verify by routing
        // through handle_rw and checking the emitted
        // CREDIT_UPDATE packet's fwd_cnt.
        let muxer = VsockMuxer::new(3)?;
        let _state = inject_state(&muxer, 3, 50_000, 12_345);
        let mut hdr = stream_header(3, 12_345, 50_000, VSOCK_OP_RW);
        hdr.buf_alloc = 262_144;
        hdr.fwd_cnt = 0;
        // Push enough bytes (in one shot) to cross the
        // CREDIT_UPDATE threshold (OUR_BUF_ALLOC / 2).
        let big = vec![0u8; (CREDIT_UPDATE_THRESHOLD + 1) as usize];
        let out = muxer.handle_rw(&hdr, &big);
        let pkt = out
            .into_iter()
            .find(|p| p.hdr.op == VSOCK_OP_CREDIT_UPDATE)
            .expect("expected CREDIT_UPDATE emission");
        assert_eq!(pkt.hdr.fwd_cnt, big.len() as u32);
        Ok(())
    }
}