arcly-stream 0.1.4

//! A focused MPEG-TS demuxer for SRT ingest.
//!
//! SRT carries an MPEG-TS bytestream in its data packets. This demuxer consumes
//! 188-byte TS packets, follows the PAT → PMT → elementary-PID chain, reassembles
//! PES packets on the video PID, and emits one [`TsPayload`] per access unit in
//! Annex-B form with a decoded PTS.
//!
//! It is deliberately small: single program, first video **and** first audio
//! elementary stream. Video: H.264 (`stream_type` 0x1B) and H.265 (0x24); audio:
//! AAC (0x0F ADTS, 0x11 LATM) and MP3 (0x03/0x04). Continuity-counter gaps are
//! tolerated — a lost packet simply truncates the in-progress PES, which is
//! dropped.

use crate::CodecId;
use bytes::Bytes;

/// Whether a demuxed [`TsPayload`] is a video or audio access unit.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TsTrackKind {
    /// A video access unit (Annex-B framed).
    Video,
    /// An audio access unit (e.g. ADTS AAC).
    Audio,
}

/// One reassembled access unit demuxed from the TS stream.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TsPayload {
    /// Elementary-stream access unit (Annex-B framed for H.264/H.265 video,
    /// raw elementary bytes for audio).
    pub data: Bytes,
    /// Codec identified from the PMT `stream_type`.
    pub codec: CodecId,
    /// Whether this is a video or audio access unit.
    pub kind: TsTrackKind,
    /// Presentation timestamp in milliseconds (PES PTS / 90).
    pub pts_ms: i64,
    /// Whether the access unit holds a keyframe (IDR). Always `false` for audio.
    pub keyframe: bool,
}

/// Per-elementary-stream PES reassembly state (one per video / audio PID).
#[derive(Debug)]
struct Track {
    /// Elementary PID this track reassembles.
    pid: u16,
    /// Codec from the PMT `stream_type`.
    codec: CodecId,
    /// Video vs audio.
    kind: TsTrackKind,
    /// PES reassembly buffer.
    pes: Vec<u8>,
    /// PTS (90 kHz) of the PES currently being reassembled.
    pts: i64,
    /// Whether a PES is open (between two PUSI markers).
    open: bool,
}

impl Track {
    fn new(pid: u16, codec: CodecId, kind: TsTrackKind) -> Self {
        Self {
            pid,
            codec,
            kind,
            pes: Vec::new(),
            pts: 0,
            open: false,
        }
    }

    /// Feed one TS payload for this PID, flushing the prior access unit on PUSI.
    fn feed(&mut self, payload: &[u8], pusi: bool, out: &mut Vec<TsPayload>) {
        if pusi {
            self.flush(out);
            if let Some((pts, es_offset)) = parse_pes_header(payload) {
                self.pts = pts;
                self.open = true;
                self.pes.extend_from_slice(&payload[es_offset..]);
            }
        } else if self.open {
            self.pes.extend_from_slice(payload);
        }
    }

    /// Emit the buffered PES as an access unit, if any.
    fn flush(&mut self, out: &mut Vec<TsPayload>) {
        if !self.open || self.pes.is_empty() {
            self.pes.clear();
            self.open = false;
            return;
        }
        let es = std::mem::take(&mut self.pes);
        let keyframe = matches!(self.kind, TsTrackKind::Video) && is_keyframe(&es, self.codec);
        out.push(TsPayload {
            data: Bytes::from(es),
            codec: self.codec,
            kind: self.kind,
            pts_ms: self.pts / 90,
            keyframe,
        });
        self.open = false;
    }
}

const TS_PACKET_LEN: usize = 188;
const TS_SYNC: u8 = 0x47;

/// Stateful MPEG-TS demuxer. Feed it TS bytes with [`push`](Self::push).
#[derive(Debug)]
pub struct TsDemuxer {
    /// PID carrying the PMT, learned from the PAT (PID 0).
    pmt_pid: Option<u16>,
    /// The first video elementary stream, learned from the PMT.
    video: Option<Track>,
    /// The first audio elementary stream, learned from the PMT.
    audio: Option<Track>,
    /// Carry for TS bytes that span `push` calls but don't fill a packet.
    carry: Vec<u8>,
}

impl Default for TsDemuxer {
    fn default() -> Self {
        Self::new()
    }
}

impl TsDemuxer {
    /// A fresh demuxer that has not yet seen a PAT.
    pub fn new() -> Self {
        Self {
            pmt_pid: None,
            video: None,
            audio: None,
            carry: Vec::new(),
        }
    }

    /// Push a chunk of the TS bytestream, returning any access units completed.
    pub fn push(&mut self, bytes: &[u8]) -> Vec<TsPayload> {
        let mut out = Vec::new();
        // Prepend any carried partial packet.
        let mut data = std::mem::take(&mut self.carry);
        data.extend_from_slice(bytes);

        let mut i = 0;
        while i + TS_PACKET_LEN <= data.len() {
            let pkt = &data[i..i + TS_PACKET_LEN];
            if pkt[0] == TS_SYNC {
                self.handle_packet(pkt, &mut out);
                i += TS_PACKET_LEN;
            } else {
                // Resync: skip a byte and look for the next sync.
                i += 1;
            }
        }
        // Carry the remainder (a partial packet) to the next call.
        self.carry = data[i..].to_vec();
        out
    }

    /// Process one 188-byte TS packet.
    fn handle_packet(&mut self, pkt: &[u8], out: &mut Vec<TsPayload>) {
        let pusi = pkt[1] & 0x40 != 0;
        let pid = (((pkt[1] & 0x1F) as u16) << 8) | pkt[2] as u16;
        let adaptation = (pkt[3] >> 4) & 0x03;
        let has_payload = adaptation == 1 || adaptation == 3;
        if !has_payload {
            return;
        }
        // Skip the adaptation field if present.
        let mut payload_start = 4;
        if adaptation == 3 {
            let af_len = pkt[4] as usize;
            payload_start = 5 + af_len;
        }
        if payload_start >= TS_PACKET_LEN {
            return;
        }
        let payload = &pkt[payload_start..];

        if pid == 0 {
            self.parse_pat(payload, pusi);
        } else if Some(pid) == self.pmt_pid {
            self.parse_pmt(payload, pusi);
        } else if let Some(track) = self
            .video
            .as_mut()
            .filter(|t| t.pid == pid)
            .or_else(|| self.audio.as_mut().filter(|t| t.pid == pid))
        {
            track.feed(payload, pusi, out);
        }
    }

    /// Parse the PAT to learn the PMT PID (first program).
    fn parse_pat(&mut self, payload: &[u8], pusi: bool) {
        let section = section_body(payload, pusi);
        let Some(section) = section else { return };
        // PAT entries start after the 8-byte section header, 4 bytes each.
        let mut i = 8;
        while i + 4 <= section.len().saturating_sub(4) {
            let program = u16::from_be_bytes([section[i], section[i + 1]]);
            let pid = (((section[i + 2] & 0x1F) as u16) << 8) | section[i + 3] as u16;
            if program != 0 {
                self.pmt_pid = Some(pid);
                return;
            }
            i += 4;
        }
    }

    /// Parse the PMT to learn the first video and first audio elementary PID.
    fn parse_pmt(&mut self, payload: &[u8], pusi: bool) {
        let Some(section) = section_body(payload, pusi) else {
            return;
        };
        if section.len() < 12 {
            return;
        }
        let program_info_len = (((section[10] & 0x0F) as usize) << 8) | section[11] as usize;
        let mut i = 12 + program_info_len;
        while i + 5 <= section.len().saturating_sub(4) {
            let stream_type = section[i];
            let pid = (((section[i + 1] & 0x1F) as u16) << 8) | section[i + 2] as u16;
            let es_info_len = (((section[i + 3] & 0x0F) as usize) << 8) | section[i + 4] as usize;
            match stream_type_to_track(stream_type) {
                Some((codec, TsTrackKind::Video)) if self.video.is_none() => {
                    self.video = Some(Track::new(pid, codec, TsTrackKind::Video));
                }
                Some((codec, TsTrackKind::Audio)) if self.audio.is_none() => {
                    self.audio = Some(Track::new(pid, codec, TsTrackKind::Audio));
                }
                _ => {}
            }
            i += 5 + es_info_len;
        }
    }
}

/// Map a PMT `stream_type` to a codec and track kind, or `None` if unsupported.
fn stream_type_to_track(stream_type: u8) -> Option<(CodecId, TsTrackKind)> {
    match stream_type {
        0x1B => Some((CodecId::H264, TsTrackKind::Video)),
        0x24 => Some((CodecId::H265, TsTrackKind::Video)),
        0x0F | 0x11 => Some((CodecId::AAC, TsTrackKind::Audio)), // ADTS / LATM AAC
        0x03 | 0x04 => Some((CodecId::MP3, TsTrackKind::Audio)),
        _ => None,
    }
}

/// Extract the PSI section body from a TS payload, honoring the `pointer_field`
/// that precedes a section in a PUSI packet.
fn section_body(payload: &[u8], pusi: bool) -> Option<&[u8]> {
    if pusi {
        let pointer = *payload.first()? as usize;
        payload.get(1 + pointer..)
    } else {
        Some(payload)
    }
}

/// Parse a PES header, returning `(pts_90khz, es_payload_offset)`.
fn parse_pes_header(p: &[u8]) -> Option<(i64, usize)> {
    // Start code 00 00 01, stream_id, 2-byte length, then the optional header.
    if p.len() < 9 || p[0] != 0 || p[1] != 0 || p[2] != 1 {
        return None;
    }
    let header_data_len = p[8] as usize;
    // Clamp to the payload: a corrupt/hostile PES header may declare a length
    // that runs past the packet, so `es_offset` could otherwise exceed `p.len()`
    // and panic when the caller slices `payload[es_offset..]`.
    let es_offset = (9 + header_data_len).min(p.len());
    let pts_dts_flags = p[7] >> 6;
    let pts = if pts_dts_flags & 0x02 != 0 && p.len() >= 14 {
        // 33-bit PTS spread across 5 bytes with marker bits.
        let b = &p[9..14];
        (((b[0] as i64 >> 1) & 0x07) << 30)
            | ((b[1] as i64) << 22)
            | (((b[2] as i64 >> 1) & 0x7F) << 15)
            | ((b[3] as i64) << 7)
            | ((b[4] as i64 >> 1) & 0x7F)
    } else {
        0
    };
    Some((pts, es_offset))
}

/// Detect a keyframe in an Annex-B elementary access unit.
fn is_keyframe(es: &[u8], codec: CodecId) -> bool {
    let mut i = 0;
    while i + 4 < es.len() {
        // Match 3- or 4-byte start codes.
        let sc3 = es[i] == 0 && es[i + 1] == 0 && es[i + 2] == 1;
        let sc4 = es[i] == 0 && es[i + 1] == 0 && es[i + 2] == 0 && es[i + 3] == 1;
        if sc3 || sc4 {
            let nal_off = if sc4 { i + 4 } else { i + 3 };
            if let Some(&hdr) = es.get(nal_off) {
                match codec {
                    CodecId::H264 if hdr & 0x1F == 5 => return true,
                    CodecId::H265 if (16..=21).contains(&((hdr >> 1) & 0x3F)) => return true,
                    _ => {}
                }
            }
            i = nal_off;
        } else {
            i += 1;
        }
    }
    false
}

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

    /// Build a 188-byte TS packet for `pid` with `pusi` and the given payload.
    fn ts_packet(pid: u16, pusi: bool, payload: &[u8]) -> Vec<u8> {
        let mut pkt = vec![0u8; TS_PACKET_LEN];
        pkt[0] = TS_SYNC;
        pkt[1] = if pusi { 0x40 } else { 0 } | ((pid >> 8) as u8 & 0x1F);
        pkt[2] = (pid & 0xFF) as u8;
        pkt[3] = 0x10; // payload only, cc 0
        let n = payload.len().min(TS_PACKET_LEN - 4);
        pkt[4..4 + n].copy_from_slice(&payload[..n]);
        pkt
    }

    /// A PAT pointing program 1 at PMT PID 0x1000.
    fn pat() -> Vec<u8> {
        let mut sec = vec![0u8]; // pointer_field = 0
                                 // table_id 0, section header (8 bytes from table_id), then one program.
        sec.extend_from_slice(&[0x00, 0xB0, 0x0D, 0, 0, 0xC1, 0, 0]);
        sec.extend_from_slice(&[0x00, 0x01]); // program_number 1
        sec.extend_from_slice(&[0xE0 | 0x10, 0x00]); // PMT PID 0x1000
        sec.extend_from_slice(&[0, 0, 0, 0]); // CRC placeholder
        ts_packet(0, true, &sec)
    }

    /// A PMT declaring an H.264 video stream on PID 0x0100, optionally followed
    /// by an AAC audio stream on PID 0x0101.
    fn pmt_with(audio: bool) -> Vec<u8> {
        let mut sec = vec![0u8]; // pointer_field
                                 // table_id 2, header to program_info_length.
        sec.extend_from_slice(&[0x02, 0xB0, 0x12, 0, 0x01, 0xC1, 0, 0]);
        sec.extend_from_slice(&[0xE1, 0x00]); // PCR PID
        sec.extend_from_slice(&[0xF0, 0x00]); // program_info_length 0
        sec.extend_from_slice(&[0x1B, 0xE1, 0x00, 0xF0, 0x00]); // H.264 on PID 0x100
        if audio {
            sec.extend_from_slice(&[0x0F, 0xE1, 0x01, 0xF0, 0x00]); // AAC on PID 0x101
        }
        sec.extend_from_slice(&[0, 0, 0, 0]); // CRC placeholder
        ts_packet(0x1000, true, &sec)
    }

    fn pmt() -> Vec<u8> {
        pmt_with(false)
    }

    /// A PES packet on `pid` with `stream_id` wrapping `es` with a PTS.
    fn pes_on(pid: u16, stream_id: u8, es: &[u8], pts: i64) -> Vec<u8> {
        let mut p = vec![0x00, 0x00, 0x01, stream_id, 0x00, 0x00, 0x80, 0x80, 0x05];
        // 5-byte PTS with marker bits.
        let pts = pts as u64;
        p.push((0x21 | (((pts >> 30) & 0x07) << 1)) as u8);
        p.push(((pts >> 22) & 0xFF) as u8);
        p.push((0x01 | (((pts >> 15) & 0x7F) << 1)) as u8);
        p.push(((pts >> 7) & 0xFF) as u8);
        p.push((0x01 | ((pts & 0x7F) << 1)) as u8);
        p.extend_from_slice(es);
        ts_packet(pid, true, &p)
    }

    /// A PES packet on the video PID 0x100 (stream_id 0xE0).
    fn video_pes(es: &[u8], pts: i64) -> Vec<u8> {
        pes_on(0x0100, 0xE0, es, pts)
    }

    /// A PES packet on the audio PID 0x101 (stream_id 0xC0).
    fn audio_pes(es: &[u8], pts: i64) -> Vec<u8> {
        pes_on(0x0101, 0xC0, es, pts)
    }

    #[test]
    fn pes_header_decodes_pts() {
        // Reuse the builder's PES bytes (strip the 4-byte TS header).
        let pes = video_pes(&[], 90_000);
        let (pts, _off) = parse_pes_header(&pes[4..]).unwrap();
        assert_eq!(pts, 90_000);
    }

    #[test]
    fn keyframe_detection_h264_idr() {
        let idr = [0, 0, 0, 1, 0x65, 0xAA];
        assert!(is_keyframe(&idr, CodecId::H264));
        let non_idr = [0, 0, 0, 1, 0x41, 0xAA];
        assert!(!is_keyframe(&non_idr, CodecId::H264));
    }

    #[test]
    fn full_chain_pat_pmt_pes_emits_access_unit() {
        let mut d = TsDemuxer::new();
        assert!(d.push(&pat()).is_empty());
        assert!(d.push(&pmt()).is_empty());
        assert_eq!(d.video.as_ref().unwrap().pid, 0x0100);
        assert_eq!(d.video.as_ref().unwrap().codec, CodecId::H264);

        // First PES opens; its AU is emitted when the next PES (PUSI) arrives.
        let idr = [0, 0, 0, 1, 0x65, 0x11, 0x22];
        assert!(d.push(&video_pes(&idr, 9000)).is_empty());
        let delta = [0, 0, 0, 1, 0x41, 0x33];
        let out = d.push(&video_pes(&delta, 12000));
        assert_eq!(out.len(), 1);
        assert_eq!(out[0].codec, CodecId::H264);
        assert_eq!(out[0].kind, TsTrackKind::Video);
        assert_eq!(out[0].pts_ms, 100); // 9000 / 90
        assert!(out[0].keyframe);
        // The access unit begins with the IDR NAL (the fixed-size test packet
        // zero-pads past the payload; valid TS uses adaptation-field stuffing).
        assert!(out[0].data.starts_with(&idr));
    }

    #[test]
    fn carries_partial_packet_across_pushes() {
        let mut d = TsDemuxer::new();
        let p = pat();
        // Feed the PAT split mid-packet; the demuxer must carry the remainder.
        assert!(d.push(&p[..100]).is_empty());
        assert!(d.push(&p[100..]).is_empty());
        // Then a PMT in one shot resolves the video PID.
        d.push(&pmt());
        assert_eq!(d.video.as_ref().unwrap().pid, 0x0100);
    }

    #[test]
    fn demuxes_audio_track_alongside_video() {
        let mut d = TsDemuxer::new();
        d.push(&pat());
        d.push(&pmt_with(true));
        // Both tracks are now known.
        assert_eq!(d.audio.as_ref().unwrap().pid, 0x0101);
        assert_eq!(d.audio.as_ref().unwrap().codec, CodecId::AAC);

        // Open an audio PES, then a second one to flush the first.
        let adts = [0xFF, 0xF1, 0x4C, 0x80, 0x01, 0x23];
        assert!(d.push(&audio_pes(&adts, 18000)).is_empty());
        let out = d.push(&audio_pes(&[0xFF, 0xF1, 0x00], 19000));
        assert_eq!(out.len(), 1);
        let au = &out[0];
        assert_eq!(au.kind, TsTrackKind::Audio);
        assert_eq!(au.codec, CodecId::AAC);
        assert!(!au.keyframe, "audio access units are never keyframes");
        assert_eq!(au.pts_ms, 200); // 18000 / 90
        assert!(au.data.starts_with(&adts));
    }

    #[test]
    fn pes_header_with_oversized_declared_length_is_clamped() {
        // Regression (found by the ts_demux fuzz target): a PES header declaring
        // header_data_len far past the payload must not panic when sliced.
        let p = [0x00, 0x00, 0x01, 0xE0, 0x00, 0x00, 0x80, 0x00, 0xFF, 0xAA];
        let (_pts, es_offset) = parse_pes_header(&p).unwrap();
        assert_eq!(es_offset, p.len(), "offset clamped to payload length");
        // Slicing at the returned offset is always valid.
        let _ = &p[es_offset..];
    }

    #[test]
    fn demuxer_survives_oversized_pes_header() {
        // The full path the fuzzer hit: a video PES whose declared header length
        // overruns the TS packet.
        let mut d = TsDemuxer::new();
        d.push(&pat());
        d.push(&pmt());
        let mut pes = vec![0x00, 0x00, 0x01, 0xE0, 0x00, 0x00, 0x80, 0x00, 0xFF];
        pes.extend_from_slice(&[0x11, 0x22]); // far less data than declared
                                              // Must not panic.
        let _ = d.push(&ts_packet(0x0100, true, &pes));
    }

    #[test]
    fn audio_only_stream_type_maps_to_aac() {
        assert_eq!(
            stream_type_to_track(0x0F),
            Some((CodecId::AAC, TsTrackKind::Audio))
        );
        assert_eq!(
            stream_type_to_track(0x03),
            Some((CodecId::MP3, TsTrackKind::Audio))
        );
        assert!(stream_type_to_track(0x99).is_none());
    }
}