punktfunk/crates/punktfunk-core/src/session.rs

//! Session lifecycle and the two hot-path state machines.
//!
//! - **Host** ([`Session::submit_frame`]): encoded access unit → FEC + packetize →
//!   optional AES-GCM seal → transport send.
//! - **Client** ([`Session::poll_frame`]): transport recv → optional open → reorder +
//!   FEC recover + reassemble → whole access unit.
//!
//! Both directions also carry input: a client [`Session::send_input`]s events; the host
//! drains them with [`Session::poll_input`].

use crate::config::{Config, Role};
use crate::crypto::SessionCrypto;
use crate::error::{PunktfunkError, Result};
use crate::fec::{coder_for, ErasureCoder};
use crate::input::InputEvent;
use crate::packet::{Packetizer, Reassembler, ReassemblerLimits, MAX_DATAGRAM_BYTES};
use crate::stats::{Stats, StatsCounters};
use crate::transport::Transport;

/// A reassembled, FEC-recovered access unit, ready to hand to the platform decoder.
pub struct Frame {
    pub data: Vec<u8>,
    pub frame_index: u32,
    pub pts_ns: u64,
    pub flags: u32,
}

/// One end of a stream. Constructed for a single [`Role`]; calling the other role's
/// methods returns [`PunktfunkError::InvalidArg`].
///
/// Anti-replay: the receive path runs each opened datagram's AEAD-authenticated sequence through a
/// sliding-window filter ([`ReplayWindow`]), so a captured, validly-sealed datagram can't be replayed
/// by an on-path attacker — closing the input-replay gap that previously rested solely on the
/// LAN/VPN transport assumption (plan §1). Genuine reordering within the window is still accepted;
/// video additionally benefits from the reassembler's per-frame dedup.
pub struct Session {
    config: Config,
    coder: Box<dyn ErasureCoder>,
    crypto: Option<SessionCrypto>,
    /// Anti-replay window over the peer's authenticated sequence (receive side). `Some` exactly when
    /// `crypto` is — the plaintext probe path carries no sequence to filter on.
    replay: Option<ReplayWindow>,
    transport: Box<dyn Transport>,
    packetizer: Packetizer,
    reassembler: Reassembler,
    stats: StatsCounters,
    /// Monotonic wire sequence, also the AES-GCM nonce counter.
    next_seq: u64,
    /// Client recv ring (reused across [`poll_frame`](Self::poll_frame)): `recvmmsg` drains a batch
    /// of datagrams into `recv_scratch` in one syscall, and poll_frame consumes them one at a time
    /// across calls (`recv_idx`..`recv_count`), refilling when drained. Allocated lazily on the
    /// first client poll so host sessions don't carry it. No per-packet recv alloc at line rate.
    recv_scratch: Vec<Vec<u8>>,
    recv_lens: Vec<usize>,
    recv_count: usize,
    recv_idx: usize,
    /// Host send pool: reused wire buffers (`seal_frame` seals in place into these, the caller sends
    /// then returns them via [`reclaim_wires`](Self::reclaim_wires)). After warmup each buffer keeps
    /// its capacity, so the per-packet ciphertext + wire `Vec` allocations vanish from the hot path.
    wire_pool: Vec<Vec<u8>>,
}

/// Datagrams drained per `recvmmsg` syscall on the client (the reused ring's size). At ~125k
/// pkt/s this is ~4k syscalls/s instead of 125k; the buffers cost `RECV_BATCH × RECV_BUF` (~64 KB).
const RECV_BATCH: usize = 32;

impl Session {
    pub fn new(config: Config, transport: Box<dyn Transport>) -> Result<Session> {
        config.validate()?;
        let coder = coder_for(config.fec.scheme);
        let crypto = config
            .encrypt
            .then(|| SessionCrypto::new(&config.key, config.salt, config.role));
        // A receive-side replay window exists exactly when the datagrams are sealed (they carry the
        // authenticated sequence the window keys on). Both roles receive from their peer.
        let replay = config.encrypt.then(ReplayWindow::new);
        let packetizer = Packetizer::new(&config);
        let reassembler = Reassembler::new(ReassemblerLimits::from_config(&config));
        Ok(Session {
            coder,
            crypto,
            replay,
            transport,
            packetizer,
            reassembler,
            stats: StatsCounters::default(),
            next_seq: 0,
            recv_scratch: Vec::new(),
            recv_lens: Vec::new(),
            recv_count: 0,
            recv_idx: 0,
            wire_pool: Vec::new(),
            config,
        })
    }

    pub fn role(&self) -> Role {
        self.config.role
    }

    pub fn stats(&self) -> Stats {
        self.stats.snapshot()
    }

    /// Wrap a packet for the wire: when encrypting, prepend the 8-byte big-endian
    /// sequence (the receiver derives the GCM nonce from it) then the ciphertext.
    /// Seal one plaintext packet into the reused `wire` buffer in place (no allocation): the wire is
    /// `seq(8) || ciphertext || tag` with crypto on, or just the packet with crypto off (probe).
    /// Byte-identical to the previous `seal` + concat path; `clear()` keeps the buffer's capacity.
    fn seal_into(&mut self, packet: &[u8], wire: &mut Vec<u8>) -> Result<()> {
        let seq = self.next_seq;
        self.next_seq = self.next_seq.wrapping_add(1);
        wire.clear();
        match &self.crypto {
            Some(c) => {
                wire.extend_from_slice(&seq.to_be_bytes()); // [0..8] plaintext seq prefix
                wire.extend_from_slice(packet); // [8..8+n] plaintext to encrypt
                wire.resize(wire.len() + crate::crypto::TAG_LEN, 0); // tag scratch
                c.seal_in_place(seq, &mut wire[8..])?; // encrypt [8..] in place, tag written at the end
            }
            None => wire.extend_from_slice(packet),
        }
        Ok(())
    }

    /// Unwrap a wire datagram back into a plaintext packet.
    fn open_from_wire(&self, wire: &[u8]) -> Result<Vec<u8>> {
        match &self.crypto {
            Some(c) => {
                if wire.len() < 8 {
                    return Err(PunktfunkError::BadPacket);
                }
                let seq = u64::from_be_bytes(wire[..8].try_into().unwrap());
                c.open(seq, &wire[8..])
            }
            None => Ok(wire.to_vec()),
        }
    }

    /// Feed an opened datagram's authenticated sequence to the anti-replay window: `true` = fresh
    /// (accept), `false` = a replay or older than the window (drop). Returns `true` when the session
    /// isn't encrypting (no window, and no sequence on the wire to key on).
    fn accept_seq(&mut self, seq: u64) -> bool {
        match self.replay.as_mut() {
            Some(w) => w.accept(seq),
            None => true,
        }
    }

    // -- Host path --------------------------------------------------------

    /// Host: FEC-protect, packetize, and seal one encoded access unit into wire packets WITHOUT
    /// sending them. Counts the frame + its packets/bytes as submitted; the caller transmits the
    /// returned packets via [`send_sealed`](Self::send_sealed) — in one call, or in chunks paced
    /// over the frame interval so a real NIC doesn't drop the whole frame as a line-rate burst (the
    /// 1 Gbps+ freeze fix). The nonce counter advances per packet, in order, so seal once and send
    /// the result intact. (Holding the `Vec<Vec<u8>>` also keeps the buffers alive for the batch.)
    pub fn seal_frame(
        &mut self,
        data: &[u8],
        pts_ns: u64,
        user_flags: u32,
    ) -> Result<Vec<Vec<u8>>> {
        if self.config.role != Role::Host {
            return Err(PunktfunkError::InvalidArg(
                "seal_frame called on a client session",
            ));
        }
        let packets = self
            .packetizer
            .packetize(data, pts_ns, user_flags, self.coder.as_ref())?;
        StatsCounters::add(&self.stats.frames_submitted, 1);
        // Reuse the wire-buffer pool the caller returns via `reclaim_wires`: one buffer per packet,
        // sealed in place — after warmup there is no per-packet ciphertext/wire allocation. (`wires`
        // is a local, so `seal_into`'s `&mut self` doesn't alias the `&mut` iteration over it.)
        let mut wires = std::mem::take(&mut self.wire_pool);
        wires.resize_with(packets.len(), Vec::new);
        for (wire, pkt) in wires.iter_mut().zip(packets.iter()) {
            self.seal_into(pkt, wire)?;
        }
        let bytes: u64 = wires.iter().map(|w| w.len() as u64).sum();
        StatsCounters::add(&self.stats.packets_sent, wires.len() as u64);
        StatsCounters::add(&self.stats.bytes_sent, bytes);
        Ok(wires)
    }

    /// Return the wire buffers from [`seal_frame`](Self::seal_frame) to the reuse pool once the caller
    /// has finished sending them, so the next frame reseals in place with no allocation. Optional —
    /// dropping the buffers instead just forfeits the reuse (correctness is unaffected).
    pub fn reclaim_wires(&mut self, wires: Vec<Vec<u8>>) {
        self.wire_pool = wires;
    }

    /// Host: transmit one chunk of already-[`seal_frame`](Self::seal_frame)ed packets in a single
    /// batched `sendmmsg`, returning how many the kernel accepted. The rest (`packets.len() - n`)
    /// are counted as send-buffer drops. Call once for the whole frame, or per paced chunk.
    pub fn send_sealed(&self, packets: &[&[u8]]) -> Result<usize> {
        // GSO when enabled (UdpTransport/Linux), else sendmmsg — same short-count drop contract.
        let sent = self.transport.send_gso(packets)?;
        if sent < packets.len() {
            StatsCounters::add(
                &self.stats.packets_send_dropped,
                (packets.len() - sent) as u64,
            );
        }
        Ok(sent)
    }

    /// Host: FEC-protect, packetize, seal, and send one encoded access unit (the whole frame in one
    /// batched send). Convenience composition of [`seal_frame`](Self::seal_frame) +
    /// [`send_sealed`](Self::send_sealed) for callers that don't pace (synthetic source, probe).
    pub fn submit_frame(&mut self, data: &[u8], pts_ns: u64, user_flags: u32) -> Result<()> {
        let wires = self.seal_frame(data, pts_ns, user_flags)?;
        let refs: Vec<&[u8]> = wires.iter().map(|w| w.as_slice()).collect();
        let r = self.send_sealed(&refs);
        drop(refs); // release the borrow of `wires` before returning the buffers to the pool
        self.reclaim_wires(wires);
        r.map(|_| ())
    }

    /// Host: live-adjust the FEC recovery percentage (adaptive FEC). Affects the next
    /// [`submit_frame`](Self::submit_frame)/[`seal_frame`](Self::seal_frame); the receiver needs no
    /// notification (each packet's header carries its block's data/recovery shard counts).
    pub fn set_fec_percent(&mut self, pct: u8) {
        self.packetizer.set_fec_percent(pct);
    }

    /// The current FEC recovery percentage (host side).
    pub fn fec_percent(&self) -> u8 {
        self.packetizer.fec_percent()
    }

    /// Host: drain one pending input event from the client, if any.
    pub fn poll_input(&mut self) -> Result<Option<InputEvent>> {
        if self.config.role != Role::Host {
            return Err(PunktfunkError::InvalidArg(
                "poll_input called on a client session",
            ));
        }
        while let Some(wire) = self.transport.recv()? {
            let pkt = match self.open_from_wire(&wire) {
                Ok(p) => p,
                Err(_) => continue, // drop undecryptable noise
            };
            // Anti-replay: a captured input datagram replayed by an on-path attacker opens cleanly
            // (its sequence + tag are still valid) — the window is what rejects the second copy.
            // `len >= 8` is guaranteed because the sealed-path open above succeeded.
            if self.replay.is_some() && !self.accept_seq(seq_of(&wire)) {
                StatsCounters::add(&self.stats.packets_dropped, 1);
                continue;
            }
            StatsCounters::add(&self.stats.packets_received, 1);
            if let Some(ev) = InputEvent::decode(&pkt) {
                return Ok(Some(ev));
            }
            // Not an input datagram (e.g. stray video) — ignore and keep draining.
        }
        Ok(None)
    }

    // -- Client path ------------------------------------------------------

    /// Client: drain the transport until a whole access unit is recovered, or no more
    /// packets are pending ([`PunktfunkError::NoFrame`]).
    pub fn poll_frame(&mut self) -> Result<Frame> {
        if self.config.role != Role::Client {
            return Err(PunktfunkError::InvalidArg(
                "poll_frame called on a host session",
            ));
        }
        // Lazily allocate the recv ring on first client poll (host sessions never get here).
        if self.recv_scratch.is_empty() {
            // Each buffer holds a max datagram + 1 (an oversized read fills it → reassembler rejects).
            self.recv_scratch = (0..RECV_BATCH)
                .map(|_| vec![0u8; MAX_DATAGRAM_BYTES + 1])
                .collect();
            self.recv_lens = vec![0usize; RECV_BATCH];
        }
        loop {
            // Refill the ring with one `recvmmsg` batch when the current one is drained.
            if self.recv_idx >= self.recv_count {
                self.recv_count = self
                    .transport
                    .recv_batch(&mut self.recv_scratch, &mut self.recv_lens)?;
                self.recv_idx = 0;
                if self.recv_count == 0 {
                    return Err(PunktfunkError::NoFrame);
                }
            }
            let i = self.recv_idx;
            self.recv_idx += 1;
            let len = self.recv_lens[i];
            // An oversized datagram fills the whole buffer (recvmmsg truncates + caps msg_len at the
            // buffer size) — drop it rather than hand up a truncated, corrupt packet, mirroring the
            // scalar `recv`'s `n >= RECV_BUF` check.
            if len > MAX_DATAGRAM_BYTES {
                continue;
            }
            // Open in place inside the ring buffer — no per-datagram allocation at line rate
            // (~125k pkt/s at 1 Gbps; the recv ring killed the recv alloc, this kills the decrypt
            // one). The plaintext lands at [8..8+n] of the sealed wire (behind the seq prefix); an
            // unencrypted (probe) datagram IS the packet. Field-precise borrows keep the slice into
            // `recv_scratch` alive across the replay/reassembler calls below.
            let (pkt_range, seq) = match &self.crypto {
                Some(c) => {
                    // A sealed datagram is at least seq prefix + tag; anything shorter is noise.
                    if len < 8 + crate::crypto::TAG_LEN {
                        continue;
                    }
                    let seq = u64::from_be_bytes(self.recv_scratch[i][..8].try_into().unwrap());
                    match c.open_in_place(seq, &mut self.recv_scratch[i][8..len]) {
                        Ok(n) => (8..8 + n, Some(seq)),
                        Err(_) => continue, // undecryptable noise — drop, keep draining
                    }
                }
                None => (0..len, None),
            };
            // Anti-replay (same rationale as poll_input): reject a datagram whose authenticated
            // sequence was already seen. Video also dedups per-frame downstream, but filtering here
            // is uniform and cheap.
            if let (Some(w), Some(seq)) = (self.replay.as_mut(), seq) {
                if !w.accept(seq) {
                    StatsCounters::add(&self.stats.packets_dropped, 1);
                    continue;
                }
            }
            let pkt = &self.recv_scratch[i][pkt_range];
            StatsCounters::add(&self.stats.packets_received, 1);
            StatsCounters::add(&self.stats.bytes_received, pkt.len() as u64);
            // The reassembler validates the packet via its parsed header (`magic`),
            // ignoring anything that isn't a well-formed video packet.
            if let Some(frame) = self
                .reassembler
                .push(pkt, self.coder.as_ref(), &self.stats)?
            {
                StatsCounters::add(&self.stats.frames_completed, 1);
                return Ok(frame);
            }
        }
    }

    /// Client: discard the ENTIRE pending receive backlog — the current recv ring plus everything
    /// queued in the kernel socket buffer — and reset the reassembler. Returns how many datagrams
    /// were thrown away (counted into `packets_dropped`).
    ///
    /// This is the latency-bound escape hatch: the receive path has no other way to skip ahead.
    /// Packets arrive strictly in order, so once a standing queue forms (the pump transiently
    /// slower than the wire, a Wi-Fi stall, power-save delivery clumping), the client plays that
    /// far behind FOREVER — it consumes at exactly the arrival rate, so the backlog never shrinks
    /// (observed live: a stream stuck 6–7 s behind, socket buffers full end to end). Discarding
    /// is memcpy-speed (no decrypt/reassembly/allocation), so this empties even a 32 MB buffer in
    /// milliseconds; the caller then requests a keyframe and the stream resumes live. The iteration
    /// cap (4096 batches ≈ 128k datagrams ≈ 190 MB) only guards against a line-rate sender
    /// outpacing the discard loop indefinitely.
    pub fn flush_backlog(&mut self) -> Result<u64> {
        if self.config.role != Role::Client {
            return Err(PunktfunkError::InvalidArg(
                "flush_backlog called on a host session",
            ));
        }
        // The undelivered tail of the current ring is backlog too.
        let mut flushed = self.recv_count.saturating_sub(self.recv_idx) as u64;
        self.recv_count = 0;
        self.recv_idx = 0;
        if !self.recv_scratch.is_empty() {
            for _ in 0..4096 {
                let n = self
                    .transport
                    .recv_batch(&mut self.recv_scratch, &mut self.recv_lens)?;
                if n == 0 {
                    break;
                }
                flushed += n as u64;
            }
        }
        self.reassembler.reset();
        StatsCounters::add(&self.stats.packets_dropped, flushed);
        Ok(flushed)
    }

    /// Client: serialize and send one input event to the host.
    pub fn send_input(&mut self, event: &InputEvent) -> Result<()> {
        if self.config.role != Role::Client {
            return Err(PunktfunkError::InvalidArg(
                "send_input called on a host session",
            ));
        }
        let pkt = event.encode();
        let mut wire = Vec::new(); // input is rare + per-event; no pool needed
        self.seal_into(&pkt, &mut wire)?;
        StatsCounters::add(&self.stats.packets_sent, 1);
        StatsCounters::add(&self.stats.bytes_sent, wire.len() as u64);
        if !self.transport.send(&wire)? {
            StatsCounters::add(&self.stats.packets_send_dropped, 1);
        }
        Ok(())
    }
}

/// Extract the AEAD-authenticated 8-byte big-endian sequence prefix from a sealed wire datagram.
/// Only called on the encrypted receive path, where a preceding successful open has already
/// established `wire.len() >= 8`.
fn seq_of(wire: &[u8]) -> u64 {
    u64::from_be_bytes(wire[..8].try_into().unwrap())
}

/// Depth of the anti-replay window, in sequences. The sender advances its sequence once per
/// datagram, so this must cover the reassembler's 120 ms loss window
/// ([`LOSS_WINDOW_NS`](crate::packet)) at line-rate packet rates — otherwise the replay filter
/// silently re-tightens the "late ≠ lost" fix: a Wi-Fi-retry-delayed shard the reassembler would
/// still use gets dropped here as "older than the window" first (4096 was only ~33 ms at the
/// ~125k pkt/s of a 1 Gbps stream). 32768 covers 120 ms up to ~270k pkt/s (≈2 Gbps+) and is
/// effectively unbounded for the sparse input stream, while still bounding how far back a replay
/// could hide; the bitmap costs 4 KiB per session.
const REPLAY_WINDOW: u64 = 32768;
const REPLAY_WORDS: usize = (REPLAY_WINDOW / 64) as usize;

/// Sliding-window anti-replay filter over the AEAD-authenticated wire sequence. The sender counts
/// its datagrams from 0, and the protocol never legitimately re-sends a sequence (FEC recovery
/// shards get fresh ones), so a sequence seen twice is a replay. The AEAD tag already authenticates
/// the sequence — a forged one can't open — so this only has to reject *duplicates* of validly
/// sealed datagrams (and anything older than the window, which we can no longer prove is fresh).
/// Genuine reordering within the window is accepted. Bitmap-per-sequence, indexed `seq % WINDOW`.
struct ReplayWindow {
    /// Highest sequence accepted so far; `seen` stays false until the first datagram.
    highest: u64,
    seen: bool,
    /// One bit per in-window sequence in `(highest - WINDOW, highest]`.
    bits: [u64; REPLAY_WORDS],
}

impl ReplayWindow {
    fn new() -> ReplayWindow {
        ReplayWindow {
            highest: 0,
            seen: false,
            bits: [0; REPLAY_WORDS],
        }
    }

    #[inline]
    fn word_bit(seq: u64) -> (usize, u64) {
        let idx = (seq % REPLAY_WINDOW) as usize;
        (idx / 64, 1u64 << (idx % 64))
    }
    fn is_set(&self, seq: u64) -> bool {
        let (w, b) = Self::word_bit(seq);
        self.bits[w] & b != 0
    }
    fn set(&mut self, seq: u64) {
        let (w, b) = Self::word_bit(seq);
        self.bits[w] |= b;
    }
    fn unset(&mut self, seq: u64) {
        let (w, b) = Self::word_bit(seq);
        self.bits[w] &= !b;
    }

    /// Record `seq`, returning `true` if it's fresh (accept) or `false` if it's a replay / too old.
    fn accept(&mut self, seq: u64) -> bool {
        if !self.seen {
            self.seen = true;
            self.highest = seq;
            self.set(seq);
            return true;
        }
        if seq > self.highest {
            // Advance the window. Sequences between the old and new high slide in unseen, so clear
            // their (possibly stale, from a full window ago) slots — unless we jumped an entire
            // window, in which case wipe the bitmap wholesale.
            if seq - self.highest >= REPLAY_WINDOW {
                self.bits = [0; REPLAY_WORDS];
            } else {
                let mut s = self.highest + 1;
                while s < seq {
                    self.unset(s);
                    s += 1;
                }
            }
            self.highest = seq;
            self.set(seq);
            true
        } else if self.highest - seq >= REPLAY_WINDOW || self.is_set(seq) {
            // Older than the window (can't prove it isn't a replay) or already seen (a duplicate) —
            // either way, drop it.
            false
        } else {
            self.set(seq); // in-window and not yet seen — a genuine reorder
            true
        }
    }
}

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

    #[test]
    fn accepts_in_order_and_rejects_duplicates() {
        let mut w = ReplayWindow::new();
        for seq in 0..1000 {
            assert!(w.accept(seq), "fresh in-order seq {seq} must be accepted");
        }
        // Every one of those is now a replay.
        for seq in 0..1000 {
            assert!(!w.accept(seq), "replayed seq {seq} must be rejected");
        }
    }

    #[test]
    fn accepts_reorder_within_window_once() {
        let mut w = ReplayWindow::new();
        assert!(w.accept(100));
        // Earlier-but-in-window sequences (a genuine reorder) are accepted exactly once.
        assert!(w.accept(80));
        assert!(!w.accept(80), "second copy of a reordered seq is a replay");
        assert!(w.accept(99));
        assert!(
            !w.accept(100),
            "the high-water seq itself can't be replayed"
        );
    }

    #[test]
    fn rejects_older_than_window() {
        let mut w = ReplayWindow::new();
        assert!(w.accept(REPLAY_WINDOW * 2));
        // Anything a full window or more behind the high-water mark is dropped (can't prove fresh).
        assert!(!w.accept(REPLAY_WINDOW * 2 - REPLAY_WINDOW));
        assert!(!w.accept(0));
        // But just inside the window is still accepted.
        assert!(w.accept(REPLAY_WINDOW * 2 - (REPLAY_WINDOW - 1)));
    }

    #[test]
    fn large_forward_jump_wipes_stale_bits() {
        let mut w = ReplayWindow::new();
        assert!(w.accept(5));
        // Jump far forward (more than a window). The slot for an old seq that aliases 5 mod WINDOW
        // must read as unseen afterward, i.e. the jump cleared it — so a NEW seq there is accepted.
        let far = 10 * REPLAY_WINDOW + 5;
        assert!(w.accept(far));
        assert!(
            !w.accept(5),
            "the pre-jump seq is now far older than the window"
        );
        // A fresh seq aliasing 5 (mod WINDOW) but inside the new window is accepted, proving the
        // stale bit was cleared rather than mistaken for a replay.
        assert!(w.accept(far - REPLAY_WINDOW + 1));
    }

    #[test]
    fn first_seq_need_not_be_zero() {
        // Startup loss can mean the first datagram we ever open isn't seq 0.
        let mut w = ReplayWindow::new();
        assert!(w.accept(42));
        assert!(!w.accept(42));
        assert!(w.accept(43));
    }

    #[test]
    fn seq_of_reads_the_big_endian_prefix() {
        let mut wire = 0x0102_0304_0506_0708u64.to_be_bytes().to_vec();
        wire.extend_from_slice(b"ciphertext-and-tag");
        assert_eq!(seq_of(&wire), 0x0102_0304_0506_0708);
    }
}