openrtc 3.0.1

//! Bounded RTP-to-portable-media conversion for the loopback WHIP demo.
//!
//! This source edge does not own publication or connection recovery. The
//! `BroadcastSession` remains the sole owner of admission and media actions.

use anyhow::{ensure, Result};
use bytes::Bytes;
use webrtc::rtp::codecs::h264::H264Packet;
use webrtc::rtp::packetizer::Depacketizer;

use crate::media::{EncodedMediaSample, MAX_MEDIA_CHUNK_BYTES};

const START_CODE: [u8; 4] = [0, 0, 0, 1];
const DEFAULT_VIDEO_DURATION_US: u32 = 33_333;
const DEFAULT_AUDIO_DURATION_US: u32 = 20_000;

fn timestamp_us(timestamp: u32, base: u32, clock_rate: u64) -> u64 {
    u64::from(timestamp.wrapping_sub(base)) * 1_000_000 / clock_rate
}

fn annexb_nals(payload: &[u8]) -> impl Iterator<Item = &[u8]> {
    let starts: Vec<usize> = payload
        .windows(4)
        .enumerate()
        .filter_map(|(index, bytes)| (bytes == START_CODE).then_some(index))
        .collect();
    starts.into_iter().filter_map(move |start| {
        let end = payload[start + 4..]
            .windows(4)
            .position(|bytes| bytes == START_CODE)
            .map_or(payload.len(), |next| start + 4 + next);
        (start + 4 < end).then_some(&payload[start + 4..end])
    })
}

/// An H.264 access-unit assembler. It accepts RTP packet payloads after the
/// receiver's jitter/NACK path and emits one Annex B chunk per marked frame.
pub struct H264RtpIngest {
    depacketizer: H264Packet,
    first_timestamp: Option<u32>,
    access_timestamp: Option<u32>,
    last_sequence: Option<u16>,
    access_unit: Vec<u8>,
    sps: Option<Vec<u8>>,
    pps: Option<Vec<u8>>,
    waiting_for_keyframe: bool,
}

impl Default for H264RtpIngest {
    fn default() -> Self {
        Self {
            depacketizer: H264Packet::default(),
            first_timestamp: None,
            access_timestamp: None,
            last_sequence: None,
            access_unit: Vec::new(),
            sps: None,
            pps: None,
            waiting_for_keyframe: true,
        }
    }
}

impl H264RtpIngest {
    fn discard_incomplete(&mut self) {
        self.depacketizer = H264Packet::default();
        self.access_unit.clear();
        self.access_timestamp = None;
        self.waiting_for_keyframe = true;
    }

    /// Returns `None` until an RTP marker completes a decodable access unit.
    pub fn push(
        &mut self,
        sequence: u16,
        timestamp: u32,
        marker: bool,
        payload: &[u8],
    ) -> Result<Option<EncodedMediaSample>> {
        if let Some(last) = self.last_sequence {
            let distance = sequence.wrapping_sub(last);
            if distance == 0 || distance > 0x8000 {
                return Ok(None);
            }
            if distance != 1 {
                self.discard_incomplete();
            }
        }
        self.last_sequence = Some(sequence);
        // WebRTC senders interleave zero-payload RTP padding probes. They
        // consume sequence numbers but are not parts of an H.264 access unit.
        if payload.len() < 3 && payload.first().is_none_or(|header| header & 0x1f != 24) {
            return Ok(None);
        }
        if self
            .access_timestamp
            .is_some_and(|current| current != timestamp)
        {
            self.discard_incomplete();
        }
        self.access_timestamp = Some(timestamp);

        // The webrtc-rs depacketizer assumes well-formed STAP-A length fields.
        // Validate before passing untrusted RTP bytes to it.
        if payload.first().is_some_and(|header| header & 0x1f == 24) {
            let mut offset = 1;
            while offset < payload.len() {
                ensure!(offset + 2 <= payload.len(), "truncated H.264 STAP-A size");
                let size = u16::from_be_bytes([payload[offset], payload[offset + 1]]) as usize;
                offset += 2;
                ensure!(
                    size > 0 && offset + size <= payload.len(),
                    "invalid H.264 STAP-A size"
                );
                offset += size;
            }
        }
        let fragment = match self
            .depacketizer
            .depacketize(&Bytes::copy_from_slice(payload))
        {
            Ok(fragment) => fragment,
            Err(error) => {
                self.discard_incomplete();
                return Err(error.into());
            }
        };
        ensure!(
            self.access_unit.len().saturating_add(fragment.len()) <= MAX_MEDIA_CHUNK_BYTES,
            "H.264 access unit exceeds OpenRTC media limit"
        );
        self.access_unit.extend_from_slice(&fragment);
        if !marker {
            return Ok(None);
        }

        let mut has_sps = false;
        let mut has_pps = false;
        let mut is_idr = false;
        for nalu in annexb_nals(&self.access_unit) {
            match nalu[0] & 0x1f {
                7 => {
                    self.sps = Some(nalu.to_vec());
                    has_sps = true;
                }
                8 => {
                    self.pps = Some(nalu.to_vec());
                    has_pps = true;
                }
                5 => is_idr = true,
                _ => {}
            }
        }
        let mut assembled = std::mem::take(&mut self.access_unit);
        self.access_timestamp = None;
        if assembled.is_empty() {
            return Ok(None);
        }
        if self.waiting_for_keyframe && !is_idr {
            return Ok(None);
        }
        if is_idr {
            let (Some(sps), Some(pps)) = (&self.sps, &self.pps) else {
                self.waiting_for_keyframe = true;
                return Ok(None);
            };
            let mut prefix = Vec::new();
            if !has_sps {
                prefix.extend_from_slice(&START_CODE);
                prefix.extend_from_slice(sps);
            }
            if !has_pps {
                prefix.extend_from_slice(&START_CODE);
                prefix.extend_from_slice(pps);
            }
            ensure!(
                prefix.len().saturating_add(assembled.len()) <= MAX_MEDIA_CHUNK_BYTES,
                "H.264 keyframe exceeds OpenRTC media limit"
            );
            prefix.append(&mut assembled);
            assembled = prefix;
            self.waiting_for_keyframe = false;
        }
        let base = *self.first_timestamp.get_or_insert(timestamp);
        Ok(Some(EncodedMediaSample {
            timestamp_us: timestamp_us(timestamp, base, 90_000),
            duration_us: DEFAULT_VIDEO_DURATION_US,
            keyframe: is_idr,
            discardable: !is_idr,
            payload: assembled,
        }))
    }
}

/// Opus RTP payloads already contain one encoded packet; the edge normalizes
/// the 48 kHz RTP clock while leaving decoding to the browser.
#[derive(Default)]
pub struct OpusRtpIngest {
    first_timestamp: Option<u32>,
    last_timestamp: Option<u32>,
    last_sequence: Option<u16>,
}

impl OpusRtpIngest {
    pub fn push(
        &mut self,
        sequence: u16,
        timestamp: u32,
        payload: &[u8],
    ) -> Result<Option<EncodedMediaSample>> {
        if payload.is_empty() {
            return Ok(None);
        }
        ensure!(
            payload.len() <= MAX_MEDIA_CHUNK_BYTES,
            "Opus packet exceeds media limit"
        );
        if let Some(last) = self.last_sequence {
            let distance = sequence.wrapping_sub(last);
            if distance == 0 || distance > 0x8000 {
                return Ok(None);
            }
        }
        let duration_us = self
            .last_timestamp
            .map(|last| timestamp.wrapping_sub(last))
            .filter(|delta| (120..=5_760).contains(delta))
            .map_or(DEFAULT_AUDIO_DURATION_US, |delta| {
                (u64::from(delta) * 1_000_000 / 48_000) as u32
            });
        self.last_sequence = Some(sequence);
        self.last_timestamp = Some(timestamp);
        let base = *self.first_timestamp.get_or_insert(timestamp);
        Ok(Some(EncodedMediaSample {
            timestamp_us: timestamp_us(timestamp, base, 48_000),
            duration_us,
            // WebCodecs registers every Opus EncodedAudioChunk as a key chunk.
            keyframe: true,
            discardable: false,
            payload: payload.to_vec(),
        }))
    }
}

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

    #[test]
    fn h264_stap_a_idr_repeats_parameter_sets_after_loss() {
        let mut ingest = H264RtpIngest::default();
        let stap = [
            24, 0, 3, 0x67, 0x42, 0x1f, 0, 3, 0x68, 0xce, 0x06, 0, 3, 0x65, 0x88, 0x84,
        ];
        let first = ingest.push(1, 100, true, &stap).unwrap().unwrap();
        assert!(first.keyframe);
        assert_eq!(first.payload.iter().filter(|byte| **byte == 1).count(), 3);
        assert!(ingest
            .push(3, 3_100, true, &[0x41, 0x11, 0x22])
            .unwrap()
            .is_none());
        let second = ingest
            .push(4, 6_100, true, &[0x65, 0x99, 0x77])
            .unwrap()
            .unwrap();
        assert!(second.keyframe);
        assert!(second.payload.starts_with(&[0, 0, 0, 1, 0x67, 0x42]));
        assert_eq!(second.timestamp_us, 66_666);
    }

    #[test]
    fn h264_rejects_truncated_stap_and_waits_for_keyframe() {
        let mut ingest = H264RtpIngest::default();
        assert!(ingest.push(1, 0, true, &[24, 0]).is_err());
        assert!(ingest
            .push(2, 3_000, true, &[0x41, 0x11, 0x22])
            .unwrap()
            .is_none());
    }

    #[test]
    fn opus_clock_normalizes_and_skips_duplicates() {
        let mut ingest = OpusRtpIngest::default();
        let first = ingest
            .push(9, u32::MAX - 959, &[0xf8, 0xff])
            .unwrap()
            .unwrap();
        assert_eq!(first.timestamp_us, 0);
        assert!(first.keyframe);
        assert!(ingest.push(9, u32::MAX - 959, &[0xf8]).unwrap().is_none());
        let second = ingest.push(10, 0, &[0xf8, 0xff]).unwrap().unwrap();
        assert_eq!(second.timestamp_us, 20_000);
        assert_eq!(second.duration_us, 20_000);
    }
}