koh 0.9.1

//! The generic SSP [`Transport`]: one per peer, carrying the local state out and the
//! remote state in. A direct port of mosh's `TransportSender` + `Transport::recv`,
//! restructured as a pure, clock-injected state machine (no I/O, no async).

use std::collections::VecDeque;

use crate::wire::{Fragment, FragmentAssembly, Fragmenter, Instruction, PROTOCOL_VERSION};
use serde::de::DeserializeOwned;
use serde::Serialize;
use tracing::trace;

use crate::ssp::{
    RttEstimator, SyncState, ACK_DELAY, ACK_INTERVAL, ACTIVE_RETRY_TIMEOUT, NEVER,
    RECEIVED_STATES_CAP, SEND_MINDELAY, SENT_STATES_CAP, SHUTDOWN_RETRIES, SHUTDOWN_SENTINEL,
};

/// A state snapshot tagged with its sequence number and the wall-clock ms it was created.
///
/// Internal SSP machinery: `mod transport` is private and the `ssp` re-export was dropped, so this
/// is crate-only and not part of the public API (the never-empty-deque invariant is upheld inside
/// [`Transport`]).
#[derive(Debug, Clone)]
struct TimestampedState<S> {
    timestamp: u64,
    num: u64,
    state: S,
}

/// Outcome of feeding one datagram to [`Transport::recv`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RecvOutcome {
    /// Fragment buffered; the instruction it belongs to is not yet complete.
    Incomplete,
    /// A new, newest-in-order remote state was applied. The app should react.
    NewState,
    /// An older (out-of-order) state was inserted; the newest state is unchanged.
    OutOfOrder,
    /// Already had this `new_num`; nothing applied (still processed the ack).
    Duplicate,
    /// The diff's base (`old_num`) is not in our `received_states`; dropped (replay guard).
    MissingBase,
    /// Refused by the anti-accumulation bound: the received-state count ceiling
    /// ([`RECEIVED_STATES_CAP`]) or the per-direction resource budget
    /// ([`SyncState::RECEIVE_BUDGET_UNITS`]) would be exceeded.
    Quenched,
}

/// One synchronization channel to a single peer.
///
/// `Local` is the state this side authors and sends (`UserInput` on the client,
/// `TerminalScreen` on the server). `Remote` is the state it receives.
pub struct Transport<Local: SyncState, Remote: SyncState> {
    // ---- sender side (our authoritative local state) ----
    current_state: Local,
    /// Front = most-recent state known-acked by the peer (the diff base). Back = last
    /// transmitted. Never empty.
    sent_states: VecDeque<TimestampedState<Local>>,
    /// `num` of the newest sent state we believe the peer already has.
    assumed_receiver_num: u64,
    fragmenter: Fragmenter,
    next_ack_time: u64,
    next_send_time: u64,
    /// Newest remote `num` we have received in order — what we advertise as our ack.
    ack_num: u64,
    pending_data_ack: bool,
    last_heard: u64,
    /// Start of the current input-coalescing window, or [`NEVER`] when none is pending.
    mindelay_clock: u64,
    // ---- shutdown ----
    shutdown_in_progress: bool,
    shutdown_tries: u32,
    shutdown_start: u64,
    // ---- receiver side (peer's remote state) ----
    received_states: Vec<TimestampedState<Remote>>,
    assembly: FragmentAssembly,
    /// Snapshot of the remote state the app last consumed via [`get_remote_diff`](Transport::get_remote_diff).
    last_delivered_remote: Remote,
    // ---- shared ----
    rtt: RttEstimator,
    connected: bool,
    /// Datagram payload budget (bytes). Updated from `Connection::max_datagram_size()`.
    mtu: usize,
}

impl<Local: SyncState, Remote: SyncState> Transport<Local, Remote> {
    /// Create a transport at time `now` (ms). `mtu` is the datagram payload budget.
    pub fn new(now: u64, mtu: usize) -> Self {
        let mut sent_states = VecDeque::new();
        sent_states.push_back(TimestampedState {
            timestamp: now,
            num: 0,
            state: Local::default(),
        });
        let received_states = vec![TimestampedState {
            timestamp: now,
            num: 0,
            state: Remote::default(),
        }];
        Self {
            current_state: Local::default(),
            sent_states,
            assumed_receiver_num: 0,
            fragmenter: Fragmenter::new(),
            next_ack_time: now,
            next_send_time: now,
            ack_num: 0,
            pending_data_ack: false,
            last_heard: 0,
            mindelay_clock: NEVER,
            shutdown_in_progress: false,
            shutdown_tries: 0,
            shutdown_start: NEVER,
            received_states,
            // Reassemble inbound instructions under THIS direction's decompressed ceiling: the
            // server (Remote = UserInput) gets the tight keystroke cap, the client (Remote =
            // TerminalScreen) the screen cap (KOH-02).
            assembly: FragmentAssembly::with_limit(Remote::RECV_DECODE_LIMIT),
            last_delivered_remote: Remote::default(),
            rtt: RttEstimator::new(),
            connected: false,
            mtu,
        }
    }

    // ----- never-empty-deque invariant -----
    //
    // `sent_states` and `received_states` are seeded with one element in `new()` and every
    // mutation keeps at least one (the front/last is never the element removed). These helpers are
    // the ONLY place that invariant is unwrapped, so the panic surface is a handful of audited
    // one-liners and any *new* stray `.unwrap()` elsewhere in the impl is still caught by clippy.

    #[expect(clippy::unwrap_used, reason = "sent_states is never empty (invariant)")]
    fn sent_front(&self) -> &TimestampedState<Local> {
        self.sent_states.front().unwrap()
    }
    #[expect(clippy::unwrap_used, reason = "sent_states is never empty (invariant)")]
    fn sent_back(&self) -> &TimestampedState<Local> {
        self.sent_states.back().unwrap()
    }
    #[expect(clippy::unwrap_used, reason = "sent_states is never empty (invariant)")]
    fn sent_back_mut(&mut self) -> &mut TimestampedState<Local> {
        self.sent_states.back_mut().unwrap()
    }
    #[expect(
        clippy::unwrap_used,
        reason = "received_states is never empty (invariant)"
    )]
    fn received_first(&self) -> &TimestampedState<Remote> {
        self.received_states.first().unwrap()
    }
    #[expect(
        clippy::unwrap_used,
        reason = "received_states is never empty (invariant)"
    )]
    fn received_last(&self) -> &TimestampedState<Remote> {
        self.received_states.last().unwrap()
    }

    // ----- accessors / driver hooks -----

    /// Mutable access to the live local state (append input, update the screen, …).
    pub const fn current_mut(&mut self) -> &mut Local {
        &mut self.current_state
    }

    /// Read the live local state. Test-only: production mutates via [`current_mut`](Self::current_mut)
    /// and reads the peer's stream via [`remote_state`](Self::remote_state); it never reads the live
    /// local state back.
    #[cfg(test)]
    pub const fn current(&self) -> &Local {
        &self.current_state
    }

    /// The newest in-order remote state we hold (what the app should render/process).
    pub fn remote_state(&self) -> &Remote {
        &self.received_last().state
    }

    /// `num` of the newest in-order remote state (what we ack to the peer).
    pub fn remote_num(&self) -> u64 {
        self.received_last().num
    }

    /// Consume the change since the app last looked: the diff from the previously-delivered
    /// remote state to the newest one, then collapse stored received states (mosh
    /// `get_remote_diff`). The server uses this to drain newly-typed input for the PTY.
    ///
    /// Intended to be called after a [`recv`](Self::recv) returning [`RecvOutcome::NewState`]. It is
    /// always safe to call regardless: with no new state the returned diff is empty (`diff_from`
    /// against the already-delivered state), so an unconditional caller just gets a no-op.
    pub fn get_remote_diff(&mut self) -> Remote::Diff {
        let newest = self.received_last().state.clone();
        let diff = newest.diff_from(&self.last_delivered_remote);
        // Rationalize the received list against its oldest element (mirror of the send side).
        let oldest = self.received_first().state.clone();
        for s in &mut self.received_states {
            s.state.subtract_prefix(&oldest);
        }
        self.last_delivered_remote = self.received_last().state.clone();
        diff
    }

    /// `num` of the newest local state we have transmitted.
    pub fn newest_sent_num(&self) -> u64 {
        self.sent_back().num
    }

    /// Mark the QUIC connection up/down. While down, [`tick`](Self::tick) sends nothing
    /// and [`wait_time`](Self::wait_time) returns [`NEVER`].
    pub fn set_connected(&mut self, connected: bool) {
        self.connected = connected;
    }

    /// Update the datagram payload budget (from `Connection::max_datagram_size()`).
    pub fn set_mtu(&mut self, mtu: usize) {
        self.mtu = mtu;
    }

    /// Feed a smoothed RTT sample (ms), typically `Connection::rtt()` each tick.
    pub fn observe_rtt(&mut self, rtt_ms: f64) {
        self.rtt.sample(rtt_ms);
    }

    /// The send interval (ms) = `clamp(ceil(SRTT/2), MIN, MAX)`. This — NOT raw SRTT — is what
    /// the adaptive predictor's engage/flag thresholds are tuned against (mosh feeds the same
    /// quantity to `PredictionEngine`; see `terminaloverlay.cc` SRTT handling).
    pub fn send_interval(&self) -> u64 {
        self.rtt.send_interval()
    }

    /// Wall-clock (ms) of the most recent decoded inbound datagram, or 0 if we've never heard
    /// from the peer. Updated on *every* inbound (incl. duplicates/keepalives), so the driver
    /// can drive its "link down / resuming" UI off real liveness rather than only new state.
    pub fn last_heard(&self) -> u64 {
        self.last_heard
    }

    /// Whether the peer has been heard from within the last `window` ms. Returns `false` until
    /// the first datagram is received (so the UI shows "connecting", not "link down", at start).
    pub fn link_up_within(&self, now: u64, window: u64) -> bool {
        self.last_heard > 0 && now.saturating_sub(self.last_heard) <= window
    }

    // ----- shutdown -----

    /// Begin a clean shutdown: outgoing instructions carry the [`SHUTDOWN_SENTINEL`]
    /// `new_num` so the peer flushes our final state, then acks the close.
    pub fn start_shutdown(&mut self, now: u64) {
        if !self.shutdown_in_progress {
            self.shutdown_in_progress = true;
            self.shutdown_start = now;
        }
    }

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

    /// The peer has acknowledged our shutdown (our acked base is the sentinel).
    pub fn shutdown_acknowledged(&self) -> bool {
        self.sent_front().num == SHUTDOWN_SENTINEL
    }

    /// We have given up waiting for the peer to ack our shutdown.
    pub fn shutdown_ack_timed_out(&self, now: u64) -> bool {
        if !self.shutdown_in_progress {
            return false;
        }
        self.shutdown_tries >= SHUTDOWN_RETRIES
            || now.saturating_sub(self.shutdown_start) >= ACTIVE_RETRY_TIMEOUT
    }

    // ----- timers -----

    /// Recompute `assumed_receiver_num`, collapse states, and recompute send/ack deadlines.
    /// Idempotent; run at the top of [`tick`](Self::tick) and [`wait_time`](Self::wait_time).
    fn calculate_timers(&mut self, now: u64) {
        self.update_assumed_receiver_state(now);
        self.rationalize_states();

        if self.pending_data_ack && self.next_ack_time > now + ACK_DELAY {
            self.next_ack_time = now + ACK_DELAY;
        }

        let back_ts = self.sent_back().timestamp;
        let interval = self.rtt.send_interval();
        let rto = self.rtt.timeout();
        let recently_heard = self.last_heard + ACTIVE_RETRY_TIMEOUT > now;

        let current_eq_back = self.current_state == self.sent_back().state;
        let current_eq_assumed = self.current_state == *self.assumed_state();
        let current_eq_front = self.current_state == self.sent_front().state;

        if !current_eq_back {
            // (A) new unsent input — coalesce ≥ SEND_MINDELAY, but respect the frame rate.
            if self.mindelay_clock == NEVER {
                self.mindelay_clock = now;
            }
            self.next_send_time = (self.mindelay_clock + SEND_MINDELAY).max(back_ts + interval);
        } else if !current_eq_assumed && recently_heard {
            // (B) nothing new, but the peer may lack our latest — retransmit at frame rate.
            self.next_send_time = back_ts + interval;
            if self.mindelay_clock != NEVER {
                self.next_send_time = self.next_send_time.max(self.mindelay_clock + SEND_MINDELAY);
            }
        } else if !current_eq_front && recently_heard {
            // (C) peer assumed-current but hasn't acked our base — slow retransmit.
            self.next_send_time = back_ts + rto + ACK_DELAY;
        } else {
            // (D) fully in sync (or peer silent > 10s).
            self.next_send_time = NEVER;
        }

        if self.shutdown_in_progress || self.ack_num == SHUTDOWN_SENTINEL {
            self.next_ack_time = back_ts + interval;
        }
    }

    /// `assumed_receiver_num` = newest state we believe the peer holds: the acked base plus
    /// any state sent within `RTO + ACK_DELAY` of now ("benefit of the doubt").
    fn update_assumed_receiver_state(&mut self, now: u64) {
        let horizon = self.rtt.timeout() + ACK_DELAY;
        let mut assumed = self.sent_front().num;
        for s in self.sent_states.iter().skip(1) {
            if now.saturating_sub(s.timestamp) < horizon {
                assumed = s.num;
            } else {
                break;
            }
        }
        self.assumed_receiver_num = assumed;
    }

    /// Express the live state and every stored state relative to the acked base, so diffs
    /// stay small and acked input is physically dropped (see [`SyncState::subtract_prefix`]).
    fn rationalize_states(&mut self) {
        let known = self.sent_front().state.clone();
        self.current_state.subtract_prefix(&known);
        for s in &mut self.sent_states {
            s.state.subtract_prefix(&known);
        }
    }

    fn assumed_idx(&self) -> usize {
        self.sent_states
            .iter()
            .position(|s| s.num == self.assumed_receiver_num)
            .unwrap_or(0)
    }

    #[expect(
        clippy::indexing_slicing,
        reason = "assumed_idx() returns a valid in-bounds position (or 0, and sent_states is non-empty)"
    )]
    fn assumed_state(&self) -> &Local {
        &self.sent_states[self.assumed_idx()].state
    }

    /// Milliseconds until the next send/ack is due, or [`NEVER`] when idle/disconnected.
    pub fn wait_time(&mut self, now: u64) -> u64 {
        self.calculate_timers(now);
        if !self.connected {
            return NEVER;
        }
        let next = self.next_ack_time.min(self.next_send_time);
        if next == NEVER {
            NEVER
        } else {
            next.saturating_sub(now)
        }
    }

    // ----- send -----

    /// Decide whether to send this tick and return the datagrams (encoded [`Fragment`]s) to
    /// transmit. Empty when nothing is due. Mirrors mosh `TransportSender::tick`.
    #[expect(
        clippy::indexing_slicing,
        reason = "assumed_idx() and chosen_idx (0 or assumed_idx) are valid in-bounds positions"
    )]
    pub fn tick(&mut self, now: u64) -> Vec<Vec<u8>> {
        self.calculate_timers(now);
        if !self.connected {
            return Vec::new();
        }
        if now < self.next_ack_time && now < self.next_send_time {
            return Vec::new();
        }

        // Compute the diff against the assumed receiver state, then maybe retarget to the
        // acked base if that is cheaper / self-healing (prospective resend optimization).
        let assumed_idx = self.assumed_idx();
        let mut chosen_idx = assumed_idx;
        // Only the serialized bytes are transmitted; don't hold the typed diff (a repaint can be a
        // few MiB) alive past serialization.
        let mut diff_bytes = encode_diff(
            &self
                .current_state
                .diff_from(&self.sent_states[assumed_idx].state),
        );

        if self.assumed_receiver_num != self.sent_front().num {
            let resend_bytes = encode_diff(&self.current_state.diff_from(&self.sent_front().state));
            let shorter = resend_bytes.len() <= diff_bytes.len();
            let modestly_longer = resend_bytes.len() < 1000
                && resend_bytes.len().saturating_sub(diff_bytes.len()) < 100;
            if shorter || modestly_longer {
                trace!(
                    from_num = self.sent_states[assumed_idx].num,
                    to_num = self.sent_front().num,
                    "retargeting diff to the acked base (prospective resend)"
                );
                chosen_idx = 0;
                diff_bytes = resend_bytes;
            }
        }

        let chosen_base_num = self.sent_states[chosen_idx].num;
        // The diff is empty exactly when the live state equals the chosen base state.
        let is_empty = self.current_state == self.sent_states[chosen_idx].state;

        if is_empty {
            let mut out = Vec::new();
            if now >= self.next_ack_time {
                out = self.send_empty_ack(now);
                self.mindelay_clock = NEVER;
            }
            if now >= self.next_send_time {
                self.next_send_time = NEVER;
                self.mindelay_clock = NEVER;
            }
            out
        } else if now >= self.next_send_time || now >= self.next_ack_time {
            let out = self.send_to_receiver(now, chosen_base_num, diff_bytes);
            self.mindelay_clock = NEVER;
            out
        } else {
            Vec::new()
        }
    }

    /// Assign `new_num`, store the state, build the instruction, and fragment it.
    fn send_to_receiver(&mut self, now: u64, old_num: u64, diff: Vec<u8>) -> Vec<Vec<u8>> {
        let back_num = self.sent_back().num;
        let current_eq_back = self.current_state == self.sent_back().state;
        // saturating_add: once a shutdown sentinel state (num == u64::MAX) is the back, a
        // `+ 1` would overflow (debug panic) before the sentinel override below.
        let new_num = if self.shutdown_in_progress {
            SHUTDOWN_SENTINEL
        } else if current_eq_back {
            back_num
        } else {
            back_num.saturating_add(1)
        };

        if new_num == back_num {
            self.sent_back_mut().timestamp = now; // retransmit: bump ts only
        } else {
            self.add_sent_state(now, new_num, self.current_state.clone());
        }

        let out = self.send_in_fragments(old_num, new_num, diff);
        self.assumed_receiver_num = self.sent_back().num;
        self.next_ack_time = now + ACK_INTERVAL;
        self.next_send_time = NEVER;
        self.pending_data_ack = false;
        out
    }

    /// Pure ack / keep-alive: advances `new_num`, stores the (unchanged) state, empty diff.
    fn send_empty_ack(&mut self, now: u64) -> Vec<Vec<u8>> {
        let back_num = self.sent_back().num;
        // saturating_add so an already-sentinel back never overflows; override for shutdown.
        let new_num = if self.shutdown_in_progress {
            SHUTDOWN_SENTINEL
        } else {
            back_num.saturating_add(1)
        };
        let old_num = self.assumed_receiver_num;
        if new_num == back_num {
            // Repeat of an existing num (e.g. the shutdown sentinel every tick): bump the
            // timestamp, don't push a duplicate-num state and churn `sent_states`.
            self.sent_back_mut().timestamp = now;
        } else {
            self.add_sent_state(now, new_num, self.current_state.clone());
        }
        let out = self.send_in_fragments(old_num, new_num, Vec::new());
        // This empty ack discharges any pending data-ack obligation, so clear the flag (mirroring
        // `send_to_receiver`). Without this, `calculate_timers` keeps yanking `next_ack_time` back to
        // `now + ACK_DELAY` (the guard below uses `pending_data_ack`), so an idle ack-only side would
        // re-emit an empty ack every ~ACK_DELAY ms forever instead of settling onto the much slower
        // ACK_INTERVAL idle cadence (~10x the idle datagram rate — mobile battery/radio cost).
        self.pending_data_ack = false;
        self.next_ack_time = now + ACK_INTERVAL;
        self.next_send_time = NEVER;
        out
    }

    fn send_in_fragments(&mut self, old_num: u64, new_num: u64, diff: Vec<u8>) -> Vec<Vec<u8>> {
        let instr = Instruction {
            protocol_version: PROTOCOL_VERSION,
            old_num,
            new_num,
            ack_num: self.ack_num,
            throwaway_num: self.sent_front().num,
            diff,
        };
        if new_num == SHUTDOWN_SENTINEL {
            self.shutdown_tries += 1;
        }
        trace!(
            old_num,
            new_num,
            ack_num = self.ack_num,
            "sending instruction"
        );
        let frags = match self.fragmenter.fragment(&instr, self.mtu) {
            Ok(f) => f,
            Err(e) => {
                tracing::error!(error=%e, "fragmentation failed");
                return Vec::new();
            }
        };
        frags.iter().filter_map(|f| f.encode().ok()).collect()
    }

    fn add_sent_state(&mut self, now: u64, num: u64, state: Local) {
        self.sent_states.push_back(TimestampedState {
            timestamp: now,
            num,
            state,
        });
        if self.sent_states.len() > SENT_STATES_CAP {
            // Drop the 16th-from-end: keeps the acked base (front) and the recent tail.
            let idx = self.sent_states.len() - 16;
            self.sent_states.remove(idx);
        }
    }

    /// Drop every sent state below `ack` (peer confirmed it holds `ack`). No-op for a stale
    /// ack naming a state we already culled.
    fn process_acknowledgment_through(&mut self, ack: u64) {
        if self.sent_states.iter().any(|s| s.num == ack) {
            self.sent_states.retain(|s| s.num >= ack);
            trace!(
                ack,
                sent_states = self.sent_states.len(),
                "processed peer ack"
            );
        }
    }

    // ----- receive -----

    /// Feed one inbound datagram (an encoded [`Fragment`]). Returns the outcome; on
    /// [`RecvOutcome::NewState`] the app should consume [`remote_state`](Self::remote_state).
    #[expect(
        clippy::unwrap_used,
        clippy::indexing_slicing,
        reason = "ref_idx comes from .position() (valid); last() follows a push (non-empty)"
    )]
    pub fn recv(&mut self, now: u64, datagram: &[u8]) -> RecvOutcome {
        let frag = match Fragment::decode(datagram) {
            Ok(f) => f,
            Err(e) => {
                tracing::warn!(error=%e, "dropping undecodable fragment");
                return RecvOutcome::Incomplete;
            }
        };
        // Any decoded datagram is a sign of life from the peer, so refresh the
        // active-retransmission liveness gate here — NOT only when a new newest-in-order state
        // lands. On a lossy link the peer's retransmits/dups/acks may be all that arrives; if
        // those didn't refresh `last_heard`, we'd stop retransmitting our own state to a peer
        // that is demonstrably still connected (mosh sets last_heard on every recv).
        self.last_heard = now;
        let instr = match self.assembly.add(frag) {
            Ok(Some(i)) => i,
            Ok(None) => return RecvOutcome::Incomplete,
            Err(e) => {
                tracing::warn!(error=%e, "dropping unreassemblable instruction");
                return RecvOutcome::Incomplete;
            }
        };

        // The peer's ack of OUR stream is processed even for dup/out-of-order packets.
        self.process_acknowledgment_through(instr.ack_num);

        // Idempotency: already have this state.
        if self.received_states.iter().any(|s| s.num == instr.new_num) {
            return RecvOutcome::Duplicate;
        }
        // Must hold the diff base, else drop (out-of-order / replay defense).
        let Some(ref_idx) = self
            .received_states
            .iter()
            .position(|s| s.num == instr.old_num)
        else {
            tracing::trace!(
                old_num = instr.old_num,
                "dropping instruction: diff base not held (out-of-order / replay)"
            );
            return RecvOutcome::MissingBase;
        };
        // Clone the base BEFORE the throwaway GC. A peer controls `throwaway_num`, and
        // `process_throwaway_until` legitimately drops every state below it — including this
        // base when `throwaway_num > old_num`. Re-resolving the base after the GC and
        // `.expect()`-ing it is a peer-triggerable panic (remote DoS of a pure state machine).
        // Owning the clone makes the GC harmless.
        let mut new_state = self.received_states[ref_idx].state.clone();

        self.process_throwaway_until(instr.throwaway_num);

        // Anti-accumulation count ceiling: refuse (don't merely throttle) once the received list is
        // at the cap, so a peer that pins `old_num`/`throwaway_num` to prevent collapse can't grow
        // it without bound. Honest peers collapse well below the cap, so this never trips for them.
        if self.received_states.len() >= RECEIVED_STATES_CAP {
            // debug, not warn: off by default (no spam in normal ops), but lets an operator
            // distinguish a throttled hostile peer (KOH-01/02) from a benign duplicate under
            // `RUST_LOG=koh=debug`. Same for the byte-budget quenches below.
            tracing::debug!(
                cap = RECEIVED_STATES_CAP,
                "quenched: received-states count at cap"
            );
            return RecvOutcome::Quenched;
        }

        // Anti-accumulation *byte/units* budget across all retained received states (KOH-01).
        // Summed BEFORE the (potentially multi-MB) decode+apply so a budget-saturated receiver
        // short-circuits without doing the work (K-14): every new state contributes >= 0 units, so
        // once the retained set alone meets the budget the outcome is already Quench — decoding and
        // applying a state only to discard it is wasted CPU/alloc a hostile peer could repeat on
        // every datagram. `received_states <= RECEIVED_STATES_CAP`, so the sum is cheap, and it is
        // recomputed each datagram so it stays correct as `subtract_prefix` shrinks states.
        let retained: usize = self
            .received_states
            .iter()
            .map(|s| s.state.resource_units())
            .sum();
        if retained >= Remote::RECEIVE_BUDGET_UNITS {
            tracing::debug!(
                retained_units = retained,
                budget = Remote::RECEIVE_BUDGET_UNITS,
                "quenched: receive byte budget saturated (pre-decode)"
            );
            return RecvOutcome::Quenched;
        }

        if !instr.diff.is_empty() {
            match decode_diff::<Remote::Diff>(&instr.diff) {
                Ok(d) => new_state.apply(&d),
                Err(e) => {
                    tracing::warn!(error=%e, "dropping instruction with undecodable diff");
                    return RecvOutcome::Incomplete;
                }
            }
        }

        // Precise check now that the new state is fully built: its exact size against the
        // already-summed retained budget (received_states is unchanged by decode/apply above).
        if retained.saturating_add(new_state.resource_units()) > Remote::RECEIVE_BUDGET_UNITS {
            tracing::debug!(
                budget = Remote::RECEIVE_BUDGET_UNITS,
                "quenched: receive byte budget exceeded by new state (post-decode)"
            );
            return RecvOutcome::Quenched;
        }

        let ts = TimestampedState {
            timestamp: now,
            num: instr.new_num,
            state: new_state,
        };

        // Insert sorted by num (handles reordering).
        if let Some(pos) = self.received_states.iter().position(|s| s.num > ts.num) {
            self.received_states.insert(pos, ts);
            RecvOutcome::OutOfOrder
        } else {
            self.received_states.push(ts);
            // Newest in-order state: advance our ack, owe a fast ack. (`last_heard` was
            // already refreshed for this datagram above, on any decoded inbound.)
            self.ack_num = self.received_states.last().unwrap().num;
            if !instr.diff.is_empty() {
                self.pending_data_ack = true;
            }
            RecvOutcome::NewState
        }
    }

    /// GC received states below `throwaway_num` (the peer's acked base). Always keeps ≥ 1.
    fn process_throwaway_until(&mut self, throwaway_num: u64) {
        if self.received_states.len() <= 1 {
            return;
        }
        let keep_from = self
            .received_states
            .iter()
            .position(|s| s.num >= throwaway_num)
            .unwrap_or(0);
        if keep_from > 0 {
            self.received_states.drain(0..keep_from);
        }
    }
}

/// Serialize a typed diff for the wire. A no-change diff still serializes to a few bytes,
/// which is why emptiness is decided by state equality, not by this length.
#[expect(
    clippy::expect_used,
    reason = "postcard serialization of our own Serialize types into a Vec is infallible"
)]
fn encode_diff<D: Serialize>(diff: &D) -> Vec<u8> {
    postcard::to_allocvec(diff).expect("diff serialization is infallible for our types")
}

fn decode_diff<D: DeserializeOwned>(bytes: &[u8]) -> Result<D, postcard::Error> {
    postcard::from_bytes(bytes)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::wire::{Fragmenter, Instruction};
    use serde::{Deserialize, Serialize};

    /// A trivial absolute-value state: each diff fully describes the target, so we can craft
    /// arbitrary instructions without worrying about diff bases.
    #[derive(Clone, Default, PartialEq, Debug)]
    struct Abs(u64);
    #[derive(Serialize, Deserialize, Clone)]
    struct AbsDiff(u64);
    impl SyncState for Abs {
        type Diff = AbsDiff;
        // Cost-free test stub: declare the (now-required) bounds explicitly as unbounded/zero.
        const RECV_DECODE_LIMIT: usize = crate::wire::MAX_DECOMPRESSED;
        const RECEIVE_BUDGET_UNITS: usize = usize::MAX;
        fn resource_units(&self) -> usize {
            0
        }
        fn diff_from(&self, _base: &Self) -> AbsDiff {
            AbsDiff(self.0)
        }
        fn apply(&mut self, d: &AbsDiff) {
            self.0 = d.0;
        }
    }

    fn instr(old: u64, new: u64, throwaway: u64, val: u64) -> Instruction {
        Instruction {
            protocol_version: PROTOCOL_VERSION,
            old_num: old,
            new_num: new,
            ack_num: 0,
            throwaway_num: throwaway,
            diff: postcard::to_allocvec(&AbsDiff(val)).unwrap(),
        }
    }

    /// Encode a (small) instruction as a single datagram, the way the wire layer ships it.
    ///
    /// Each call models a DISTINCT send from the peer's single, monotonic [`Fragmenter`] (one per
    /// connection in production — see `Transport::fragmenter`), so it stamps a fresh increasing
    /// fragment id. Reusing the RETURNED bytes (not calling this again) reproduces the same id and is
    /// therefore a genuine byte-identical replay — which is exactly what the fragment-layer replay
    /// gate (K-06) drops. A fresh `Fragmenter::new()` per call (the old helper) instead reused id 1
    /// for every datagram, which no real sender does and which the replay gate would treat as a
    /// replay of the first instruction.
    fn datagram(i: &Instruction) -> Vec<u8> {
        thread_local! {
            static NEXT_FRAG_ID: std::cell::Cell<u64> = const { std::cell::Cell::new(1) };
        }
        let id = NEXT_FRAG_ID.with(|c| {
            let v = c.get();
            c.set(v + 1);
            v
        });
        let mut frags = Fragmenter::new().fragment(i, 1200).unwrap();
        assert_eq!(frags.len(), 1, "test instruction must fit one fragment");
        frags[0].id = id;
        frags[0].encode().unwrap()
    }

    /// An absolute state whose "resource cost" is its value, with a tiny receive budget — so the
    /// anti-accumulation bound can be exercised without crafting megabytes of input (KOH-01).
    #[derive(Clone, Default, PartialEq, Debug)]
    struct Grow(u64);
    #[derive(Serialize, Deserialize, Clone)]
    struct GrowDiff(u64);
    impl SyncState for Grow {
        type Diff = GrowDiff;
        const RECV_DECODE_LIMIT: usize = crate::wire::MAX_DECOMPRESSED;
        const RECEIVE_BUDGET_UNITS: usize = 100;
        fn resource_units(&self) -> usize {
            self.0 as usize
        }
        fn diff_from(&self, _base: &Self) -> GrowDiff {
            GrowDiff(self.0)
        }
        fn apply(&mut self, d: &GrowDiff) {
            self.0 = d.0;
        }
    }

    fn grow_datagram(old: u64, new: u64, throwaway: u64, val: u64) -> Vec<u8> {
        datagram(&Instruction {
            protocol_version: PROTOCOL_VERSION,
            old_num: old,
            new_num: new,
            ack_num: 0,
            throwaway_num: throwaway,
            diff: postcard::to_allocvec(&GrowDiff(val)).unwrap(),
        })
    }

    #[test]
    fn received_state_budget_refuses_accumulation() {
        // KOH-01: a peer that pins old_num/throwaway to prevent collapse is refused once the
        // retained received-state resource budget would be exceeded — not allowed to OOM us.
        let mut t = Transport::<Grow, Grow>::new(0, 1200);
        // First large state fits (base num-0 is 0 units + 60 = 60 <= the 100 budget).
        assert_eq!(
            t.recv(10, &grow_datagram(0, 1, 0, 60)),
            RecvOutcome::NewState
        );
        // A second large state from the same (un-collapsed) base would push retained 60 + 60 over
        // the budget, so it is refused rather than accumulated.
        assert_eq!(
            t.recv(20, &grow_datagram(0, 2, 0, 60)),
            RecvOutcome::Quenched
        );
        assert_eq!(
            t.remote_state().0,
            60,
            "the refused state must not have been applied"
        );
    }

    #[test]
    fn recv_missing_base_is_reported_and_state_unchanged() {
        // The diff base (old_num) must be held or the instruction is dropped as MissingBase — the
        // out-of-order / replay defense. A fresh transport holds only the num-0 base, so an old_num
        // of 5 references no held base.
        let mut t = Transport::<Abs, Abs>::new(0, 1200);
        assert_eq!(
            t.recv(10, &datagram(&instr(5, 6, 0, 99))),
            RecvOutcome::MissingBase
        );
        assert_eq!(t.remote_num(), 0, "no state was applied");
    }

    #[test]
    fn idle_empty_acks_settle_to_ack_interval_not_a_flood() {
        // Regression: after receiving data (which owes a fast ack), an idle ack-only side must emit
        // ONE settling empty ack and then fall back to the slow ACK_INTERVAL keepalive cadence — it
        // must NOT re-emit an empty ack every ACK_DELAY ms forever. The bug was that `send_empty_ack`
        // reset `next_ack_time` to ACK_INTERVAL but left `pending_data_ack` set, so `calculate_timers`
        // kept yanking the deadline back to `now + ACK_DELAY`. ~10-30x the idle datagram rate on a
        // quiesced link — exactly the mobile battery/radio cost ACK_INTERVAL exists to avoid.
        let mut t = Transport::<Abs, Abs>::new(0, 1200);
        t.set_connected(true);
        // A non-empty inbound state owes a fast ack; our local state stays at its default, so every
        // subsequent transmission is a pure empty ack (no diff to send).
        assert_eq!(
            t.recv(0, &datagram(&instr(0, 1, 0, 42))),
            RecvOutcome::NewState
        );
        assert!(t.pending_data_ack, "a received data state owes an ack");

        // Walk time forward in ACK_DELAY steps across a full ACK_INTERVAL, recording every tick that
        // actually emits a datagram.
        let mut ack_times = Vec::new();
        let mut now = 0;
        while now <= ACK_INTERVAL {
            now += ACK_DELAY;
            if !t.tick(now).is_empty() {
                ack_times.push(now);
            }
        }
        assert!(
            !t.pending_data_ack,
            "the empty ack must discharge the pending-data-ack obligation"
        );
        // No empty acks between the single settling ack and the slow keepalive: a flood would land
        // repeatedly inside this window.
        let flooded = ack_times
            .iter()
            .any(|&at| at > ACK_DELAY && at < ACK_INTERVAL);
        assert!(
            !flooded,
            "idle side flooded empty acks every ~ACK_DELAY instead of settling onto ACK_INTERVAL; \
             emissions at {ack_times:?}"
        );
    }

    #[test]
    fn recv_out_of_order_inserts_without_regressing_the_newest() {
        // A state whose num lands BEFORE an already-held newer state is OutOfOrder: inserted in num
        // order, but the newest-in-order value (what we render/ack) must not regress.
        let mut t = Transport::<Abs, Abs>::new(0, 1200);
        assert_eq!(
            t.recv(10, &datagram(&instr(0, 2, 0, 22))),
            RecvOutcome::NewState
        );
        assert_eq!(
            t.recv(20, &datagram(&instr(2, 10, 0, 33))),
            RecvOutcome::NewState
        );
        assert_eq!(t.remote_num(), 10);
        // Deliver num 5 (between the held 2 and 10): inserted out of order.
        assert_eq!(
            t.recv(30, &datagram(&instr(2, 5, 0, 44))),
            RecvOutcome::OutOfOrder
        );
        assert_eq!(
            t.remote_num(),
            10,
            "the newest in-order state must not regress"
        );
        assert_eq!(
            t.remote_state().0,
            33,
            "render still reflects the newest (num 10) state"
        );
    }

    /// Regression for P1a: a peer-supplied `throwaway_num > old_num` makes the throwaway GC
    /// drop the diff base. Before the fix, recv() re-resolved the base after the GC with
    /// `.expect()` and panicked on this peer-controlled input. After the fix, the base is
    /// cloned before the GC and applied safely.
    #[test]
    fn throwaway_gc_dropping_base_does_not_panic() {
        let mut t = Transport::<Abs, Abs>::new(0, 1200);
        // received_states = [0, 2]
        assert_eq!(
            t.recv(10, &datagram(&instr(0, 2, 0, 22))),
            RecvOutcome::NewState
        );
        assert_eq!(t.remote_state().0, 22);
        // old=0 base, but throwaway_num=1 GCs num 0 (the base) before apply.
        assert_eq!(
            t.recv(20, &datagram(&instr(0, 5, 1, 55))),
            RecvOutcome::NewState
        );
        assert_eq!(
            t.remote_state().0,
            55,
            "diff must apply against the base cloned before the throwaway GC"
        );
    }

    /// Regression for P1c: `last_heard` (the active-retransmission liveness gate) must refresh
    /// on EVERY decoded datagram, including duplicate keepalives — not only on a new state.
    /// Otherwise a peer whose only-arriving traffic is dups/retransmits falsely times out.
    #[test]
    fn last_heard_updates_on_duplicate() {
        let mut t = Transport::<Abs, Abs>::new(0, 1200);
        let dg = datagram(&instr(0, 2, 0, 22));
        assert_eq!(t.recv(10, &dg), RecvOutcome::NewState);
        assert_eq!(t.last_heard(), 10);
        // (a) A byte-identical replay (same fragment id) is dropped at the fragment layer as a
        // replay — it is NOT re-inflated/re-applied (K-06) — but `recv` still stamps last_heard for
        // any decoded datagram BEFORE reassembly, so a re-sent keepalive of unchanged content still
        // proves liveness.
        assert_eq!(t.recv(5000, &dg), RecvOutcome::Incomplete);
        assert_eq!(
            t.last_heard(),
            5000,
            "a replayed keepalive still refreshes last_heard"
        );
        // (b) A DISTINCT datagram re-announcing an already-held new_num (a fresh fragment id — e.g. a
        // keepalive whose carried ack advanced) reassembles past the replay gate and is caught by
        // the SSP-level duplicate check (new_num 2 already held) — and likewise advances liveness.
        assert_eq!(
            t.recv(9000, &datagram(&instr(0, 2, 0, 22))),
            RecvOutcome::Duplicate
        );
        assert_eq!(
            t.last_heard(),
            9000,
            "a duplicate keepalive must refresh last_heard"
        );
        assert!(t.link_up_within(9100, 10_000));
        assert!(!t.link_up_within(20_000, 10_000));
    }

    /// Regression for P1d: a long shutdown must not overflow (`back_num + 1` on the sentinel)
    /// nor push a fresh u64::MAX state every tick. Exactly one sentinel state should be resident.
    #[test]
    fn shutdown_dedups_sentinel_and_never_overflows() {
        let mut t = Transport::<Abs, Abs>::new(0, 1200);
        t.set_connected(true);
        t.start_shutdown(0);
        // Many ticks at the frame rate; pre-fix this churned sent_states with sentinels (and
        // risked a `u64::MAX + 1` overflow). Reaching the end without panicking is half the test.
        for i in 0..200u64 {
            let _ = t.tick(i * 100);
        }
        assert_eq!(t.newest_sent_num(), SHUTDOWN_SENTINEL);
        // base (num 0) + a single deduped sentinel — not a queue churned toward the 32 cap.
        assert!(
            t.sent_states.len() <= 2,
            "shutdown sentinel must be deduped (bump ts), got {} sent_states",
            t.sent_states.len()
        );
    }

    /// Ported from mosh src/tests/network-no-diff.test: the sender must not generate a new diff/
    /// state while the application state is unchanged (mosh's regression was the server busy-
    /// looping / repainting when nothing changed). An unchanged state may be retransmitted for
    /// reliability but mints no new state number; a real change always gets a fresh one.
    #[test]
    fn unchanged_state_mints_no_new_content_state() {
        let mut t = Transport::<Abs, Abs>::new(0, 1200);
        t.set_connected(true);
        t.observe_rtt(20.0);

        // A real change is sent.
        t.current_mut().0 = 1;
        let mut now = 0u64;
        let mut sent = Vec::new();
        for _ in 0..100 {
            now += 20;
            sent = t.tick(now);
            if !sent.is_empty() {
                break;
            }
        }
        assert!(!sent.is_empty(), "a changed state must be sent");
        let after_first = t.newest_sent_num();

        // Unchanged across many ticks (well inside the ACK_INTERVAL): the transport may RETRANSMIT
        // the still-unacked state (reliability), but it must not mint a NEW state number — there
        // is no new screen content to diff. This is the heart of mosh's no-diff guarantee.
        for _ in 0..10 {
            now += 20;
            let _ = t.tick(now); // retransmits allowed; content is unchanged
            assert_eq!(
                t.newest_sent_num(),
                after_first,
                "unchanged state must not mint a new content state (retransmit reuses the number)"
            );
        }

        // A subsequent real change is sent again, with a fresh state number.
        t.current_mut().0 = 2;
        let mut sent_again = Vec::new();
        for _ in 0..100 {
            now += 20;
            sent_again = t.tick(now);
            if !sent_again.is_empty() {
                break;
            }
        }
        assert!(!sent_again.is_empty(), "a later change must be sent");
        assert!(
            t.newest_sent_num() > after_first,
            "the changed state gets a fresh number"
        );
    }

    proptest::proptest! {
        #![proptest_config(proptest::prelude::ProptestConfig::with_cases(128))]

        /// Feed ARBITRARY SSP envelopes (old/new/ack/throwaway including `u64::MAX` / the shutdown
        /// sentinel, inverted `old > new`, repeats) into the receive path. `recv` is the SSP
        /// reconciliation trust boundary where every peer-controlled num drives held-base lookup,
        /// throwaway GC (which once shipped a panic), the insertion-sort, and the cap/budget gates —
        /// yet it had no property coverage. With release `overflow-checks` on, an unguarded `+`/`-` on
        /// a wire num would panic here; instead it must never panic, and `received_states` must stay
        /// non-empty and bounded under any sequence. (KSSP-01)
        #[test]
        fn recv_survives_arbitrary_envelopes(
            ops in proptest::collection::vec(
                (
                    proptest::prelude::any::<u64>(),
                    proptest::prelude::any::<u64>(),
                    proptest::prelude::any::<u64>(),
                    proptest::prelude::any::<u64>(),
                    0u64..200,
                ),
                0..64,
            ),
        ) {
            let mut t = Transport::<Grow, Grow>::new(0, 1200);
            let mut now = 0u64;
            for (old, new, ack, throwaway, val) in ops {
                now = now.saturating_add(1);
                let dg = datagram(&Instruction {
                    protocol_version: PROTOCOL_VERSION,
                    old_num: old,
                    new_num: new,
                    ack_num: ack,
                    throwaway_num: throwaway,
                    diff: postcard::to_allocvec(&GrowDiff(val)).unwrap(),
                });
                let _ = t.recv(now, &dg); // must never panic on any envelope
                proptest::prop_assert!(
                    !t.received_states.is_empty(),
                    "received_states must never be empty (the never-empty invariant)"
                );
                proptest::prop_assert!(
                    t.received_states.len() <= RECEIVED_STATES_CAP,
                    "received_states {} exceeded the {} cap",
                    t.received_states.len(),
                    RECEIVED_STATES_CAP
                );
            }
        }
    }
}