use crate::error::{Error, Result};
use crate::pipeline::{CodecConfig, Sample, TrackSpec};
use crate::rtcp::SenderReport;
use crate::rtp::{RtpMediaKind, parse_rtp_header, reassemble_audio, reassemble_video};
use alloc::vec::Vec;
use broadcast_common::Parse;
const NTP_FRACTION_SCALE: f64 = 4_294_967_296.0;
const MAX_AU_BUFFER_BYTES: usize = 4 * 1024 * 1024;
#[non_exhaustive]
pub struct RtpStreamTrack {
pub track_id: u32,
pub kind: RtpMediaKind,
pub config: CodecConfig,
pub clock_rate: u32,
}
impl RtpStreamTrack {
pub fn new(track_id: u32, kind: RtpMediaKind, config: CodecConfig, clock_rate: u32) -> Self {
Self {
track_id,
kind,
config,
clock_rate,
}
}
}
struct TrackState {
kind: RtpMediaKind,
config: CodecConfig,
clock_rate: u32,
cur_ts: Option<u32>,
cur_pkts: Vec<Vec<u8>>,
cur_bytes: usize,
last_unwrapped: Option<u64>,
first_unwrapped: Option<u64>,
pending: Option<PendingAu>,
last_duration: u32,
sr_anchor: Option<SrAnchor>,
}
struct SrAnchor {
ntp_seconds: f64,
raw_rtp_ts: u32,
}
fn wall_seconds(anchor: &SrAnchor, clock_rate: u32, au_unwrapped_ts: u64) -> f64 {
let anchor_unwrapped = unwrap_ts(Some(au_unwrapped_ts), anchor.raw_rtp_ts);
let delta_ticks = au_unwrapped_ts as i128 - i128::from(anchor_unwrapped);
anchor.ntp_seconds + (delta_ticks as f64) / f64::from(clock_rate)
}
struct PendingAu {
unwrapped_ts: u64,
is_sync: bool,
data: Vec<u8>,
}
pub struct RtpStreamDepacketiser {
tracks: Vec<(u32, TrackState)>,
}
fn unwrap_ts(prev: Option<u64>, ts: u32) -> u64 {
let Some(prev) = prev else {
return u64::from(ts);
};
let prev_low = prev as u32;
let delta = ts.wrapping_sub(prev_low) as i32;
if delta >= 0 {
prev + u64::from(delta as u32)
} else {
prev.saturating_sub(u64::from(delta.unsigned_abs()))
}
}
impl RtpStreamDepacketiser {
pub fn new(tracks: Vec<RtpStreamTrack>) -> Self {
let tracks = tracks
.into_iter()
.map(|t| {
(
t.track_id,
TrackState {
kind: t.kind,
config: t.config,
clock_rate: t.clock_rate,
cur_ts: None,
cur_pkts: Vec::new(),
cur_bytes: 0,
last_unwrapped: None,
first_unwrapped: None,
pending: None,
last_duration: 0,
sr_anchor: None,
},
)
})
.collect();
Self { tracks }
}
pub fn track_specs(&self) -> Vec<TrackSpec> {
self.tracks
.iter()
.map(|(id, st)| TrackSpec::new(*id, st.clock_rate, st.config.clone()))
.collect()
}
fn state(&mut self, track_id: u32) -> Option<&mut TrackState> {
self.tracks
.iter_mut()
.find(|(id, _)| *id == track_id)
.map(|(_, st)| st)
}
pub fn push_sender_report(&mut self, track_id: u32, sr: SenderReport) {
if let Some(st) = self.state(track_id) {
let ntp_seconds = f64::from(sr.ntp_msw) + f64::from(sr.ntp_lsw) / NTP_FRACTION_SCALE;
st.sr_anchor = Some(SrAnchor {
ntp_seconds,
raw_rtp_ts: sr.rtp_timestamp,
});
}
}
pub fn push_rtcp(&mut self, track_id: u32, bytes: &[u8]) {
if let Ok(sr) = SenderReport::parse(bytes) {
self.push_sender_report(track_id, sr);
}
}
pub fn sync_start_decode_times(&self) -> Vec<(u32, u64)> {
let anchored: Vec<(u32, f64, u32)> = self
.tracks
.iter()
.filter_map(|(id, st)| {
let anchor = st.sr_anchor.as_ref()?;
let first_ts = st.first_unwrapped?;
Some((
*id,
wall_seconds(anchor, st.clock_rate, first_ts),
st.clock_rate,
))
})
.collect();
if anchored.len() < 2 {
return Vec::new();
}
let origin = anchored
.iter()
.map(|(_, wall, _)| *wall)
.fold(f64::INFINITY, f64::min);
anchored
.into_iter()
.map(|(id, wall, clock_rate)| {
let raw_ticks = ((wall - origin) * f64::from(clock_rate)).max(0.0);
let ticks = (raw_ticks + 0.5) as u64;
(id, ticks)
})
.collect()
}
pub fn push(&mut self, track_id: u32, rtp_packet: &[u8]) -> Result<Vec<Sample>> {
let Some(st) = self.state(track_id) else {
return Ok(Vec::new());
};
let hdr = parse_rtp_header(rtp_packet)?;
let ts = hdr.timestamp;
let mut out = Vec::new();
if let Some(cur) = st.cur_ts {
if cur != ts && !st.cur_pkts.is_empty() {
Self::drain_complete(st, &mut out)?;
}
}
st.cur_ts = Some(ts);
st.cur_bytes += rtp_packet.len();
st.cur_pkts.push(rtp_packet.to_vec());
if st.cur_bytes > MAX_AU_BUFFER_BYTES {
st.cur_pkts.clear();
st.cur_bytes = 0;
st.cur_ts = None;
return Err(Error::BufferCapExceeded {
what: "RTP access-unit reassembly",
cap: MAX_AU_BUFFER_BYTES,
});
}
if matches!(st.kind, RtpMediaKind::H264) && hdr.marker {
Self::drain_complete(st, &mut out)?;
st.cur_ts = None;
}
Ok(out)
}
pub fn flush(&mut self, track_id: u32) -> Result<Vec<Sample>> {
let Some(st) = self.state(track_id) else {
return Ok(Vec::new());
};
let mut out = Vec::new();
if !st.cur_pkts.is_empty() {
Self::drain_complete(st, &mut out)?;
st.cur_ts = None;
}
if let Some(p) = st.pending.take() {
out.push(Sample::new(p.data, st.last_duration, p.is_sync, 0));
}
Ok(out)
}
fn drain_complete(st: &mut TrackState, out: &mut Vec<Sample>) -> Result<()> {
let pkts = core::mem::take(&mut st.cur_pkts);
st.cur_bytes = 0;
let aus = match st.kind {
RtpMediaKind::H264 => reassemble_video(&pkts)?,
RtpMediaKind::Aac => reassemble_audio(&pkts)?,
};
for au in aus {
let unwrapped = unwrap_ts(st.last_unwrapped, au.timestamp);
st.last_unwrapped = Some(unwrapped);
if st.first_unwrapped.is_none() {
st.first_unwrapped = Some(unwrapped);
}
if let Some(prev) = st.pending.take() {
let delta = unwrapped.saturating_sub(prev.unwrapped_ts);
let duration = u32::try_from(delta).unwrap_or(u32::MAX);
st.last_duration = duration;
out.push(Sample::new(prev.data, duration, prev.is_sync, 0));
}
st.pending = Some(PendingAu {
unwrapped_ts: unwrapped,
is_sync: au.is_sync,
data: au.data,
});
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::avc_config::{AVCConfigurationBox, AVCDecoderConfigurationRecord};
fn dummy_avc() -> CodecConfig {
CodecConfig::Avc {
config: AVCConfigurationBox::new(AVCDecoderConfigurationRecord {
configuration_version: 1,
profile_indication: 0x42,
profile_compatibility: 0,
level_indication: 0x1E,
length_size_minus_one: 3,
sps: alloc::vec![],
pps: alloc::vec![],
chroma_format: None,
bit_depth_luma_minus8: None,
bit_depth_chroma_minus8: None,
sps_ext: alloc::vec![],
}),
width: 0,
height: 0,
}
}
fn vpkt(seq: u16, ts: u32, marker: bool, nal: &[u8]) -> Vec<u8> {
let mut p = alloc::vec![0x80u8, if marker { 0x80 | 96 } else { 96 }];
p.extend_from_slice(&seq.to_be_bytes());
p.extend_from_slice(&ts.to_be_bytes());
p.extend_from_slice(&[0, 0, 0, 0]);
p.extend_from_slice(nal);
p
}
#[test]
fn video_stream_recovers_durations_and_sync() {
let mut d = RtpStreamDepacketiser::new(alloc::vec![RtpStreamTrack::new(
1,
RtpMediaKind::H264,
dummy_avc(),
90_000,
)]);
let idr = [0x65u8, 0xAA];
let non = [0x41u8, 0xBB];
assert!(d.push(1, &vpkt(1, 1000, true, &idr)).unwrap().is_empty());
let s0 = d.push(1, &vpkt(2, 4000, true, &non)).unwrap();
assert_eq!(s0.len(), 1);
assert_eq!(s0[0].duration, 3000);
assert!(s0[0].is_sync);
assert_eq!(s0[0].composition_offset, 0);
let s1 = d.push(1, &vpkt(3, 7000, true, &non)).unwrap();
assert_eq!(s1.len(), 1);
assert_eq!(s1[0].duration, 3000);
assert!(!s1[0].is_sync);
let s2 = d.flush(1).unwrap();
assert_eq!(s2.len(), 1);
assert_eq!(s2[0].duration, 3000);
}
fn fu_a_fragment(start: bool, extra_len: usize) -> Vec<u8> {
const NAL_TYPE_FU_A: u8 = 28;
const FU_START: u8 = 0x80;
const ORIG_TYPE_IDR: u8 = 5;
let fu_header = if start {
FU_START | ORIG_TYPE_IDR
} else {
ORIG_TYPE_IDR
};
let mut payload = alloc::vec![NAL_TYPE_FU_A, fu_header];
payload.extend(core::iter::repeat_n(0xABu8, extra_len));
payload
}
#[test]
fn runaway_fu_a_without_end_bit_is_bounded_not_unbounded() {
let mut d = RtpStreamDepacketiser::new(alloc::vec![RtpStreamTrack::new(
1,
RtpMediaKind::H264,
dummy_avc(),
90_000,
)]);
const FRAGMENT_FILLER: usize = 2048;
let mut hit_cap = false;
for i in 0..4096u16 {
match d.push(
1,
&vpkt(i, 1000, false, &fu_a_fragment(i == 0, FRAGMENT_FILLER)),
) {
Ok(samples) => assert!(
samples.is_empty(),
"a never-completing AU must not emit a sample"
),
Err(e) => {
assert!(
matches!(e, crate::error::Error::BufferCapExceeded { .. }),
"unexpected error variant: {e:?}"
);
hit_cap = true;
break;
}
}
}
assert!(
hit_cap,
"expected MAX_AU_BUFFER_BYTES to trip well within {} fragments \
(never grow unbounded)",
4096
);
let idr = [0x65u8, 0xAA];
let non = [0x41u8, 0xBB];
assert!(d.push(1, &vpkt(1, 4000, true, &idr)).unwrap().is_empty());
let s = d.push(1, &vpkt(2, 7000, true, &non)).unwrap();
assert_eq!(s.len(), 1);
assert_eq!(s[0].duration, 3000);
assert!(s[0].is_sync);
}
#[test]
fn track_specs_use_clock_rate_as_timescale() {
let d = RtpStreamDepacketiser::new(alloc::vec![RtpStreamTrack::new(
7,
RtpMediaKind::H264,
dummy_avc(),
90_000,
)]);
let specs = d.track_specs();
assert_eq!(specs.len(), 1);
assert_eq!(specs[0].track_id, 7);
assert_eq!(specs[0].timescale, 90_000);
}
}