use alloc::vec::Vec;
use broadcast_common::{Package, Parse, crc32_mpeg2};
use mpeg_pes::{Pts as PesPts, StreamId};
use mpeg_ts::mux::SectionPacketizer;
use mpeg_ts::ts::{Pcr, TS_PACKET_SIZE, TsHeader};
use crate::aac_asc::AudioSpecificConfig;
use crate::annexb::{iter_length_prefixed_nals, length_prefixed_to_annexb};
use crate::error::{Error, Result};
use crate::media::{Media, Track};
use crate::mp4esds::EsdsBox;
use crate::nal::{NalCodec, nal_unit_type};
use crate::pipeline::{CodecConfig, DataCarriage, Sample};
const PAT_PID: u16 = 0x0000;
const PMT_PID: u16 = 0x1000;
const ES_PID_BASE: u16 = 0x0100;
const PROGRAM_NUMBER: u16 = 1;
const TABLE_ID_PAT: u8 = 0x00;
const TABLE_ID_PMT: u8 = 0x02;
const CRC32_LEN: usize = 4;
const SECTION_SYNTAX_FLAGS_HI: u8 = 0xB0;
const SECTION_LENGTH_HI_MASK: u8 = 0x0F;
const VERSION_CURRENT_NEXT: u8 = 0xC1;
const PID_RESERVED_HI: u8 = 0xE0;
const INFO_RESERVED_HI: u8 = 0xF0;
const STREAM_TYPE_AVC: u8 = 0x1B;
const STREAM_TYPE_HEVC: u8 = 0x24;
const STREAM_TYPE_MPEG2_VIDEO: u8 = 0x02;
const STREAM_TYPE_AAC_ADTS: u8 = 0x0F;
const STREAM_TYPE_MPEG1_AUDIO: u8 = 0x03;
const STREAM_TYPE_MPEG2_AUDIO: u8 = 0x04;
const STREAM_TYPE_AC3: u8 = 0x81;
const STREAM_TYPE_EAC3: u8 = 0x87;
const STREAM_TYPE_DTS: u8 = 0x82;
const STREAM_TYPE_MPEGH: u8 = 0x2D;
const OTI_MPEG1_AUDIO: u8 = 0x6B;
const OTI_MPEG2_AUDIO: u8 = 0x69;
const MAX_ES_INFO_LENGTH: usize = 0x0FFF;
const DESCRIPTOR_TAG_EXTENSION: u8 = 0x3F;
const MPEGH_3DAUDIO_EXTENSION_TAG: u8 = 0x08;
const MPEGH_3DAUDIO_DESCRIPTOR_BODY_LEN: u8 = 2;
fn mpegh_3daudio_descriptor(profile_level_indication: u8) -> Vec<u8> {
alloc::vec![
DESCRIPTOR_TAG_EXTENSION,
MPEGH_3DAUDIO_DESCRIPTOR_BODY_LEN,
MPEGH_3DAUDIO_EXTENSION_TAG,
profile_level_indication,
]
}
const STREAM_ID_VIDEO_BASE: u8 = 0xE0;
const STREAM_ID_AUDIO_BASE: u8 = 0xC0;
const STREAM_ID_PRIVATE_1: u8 = 0xBD;
const PES_START_CODE: [u8; 3] = [0x00, 0x00, 0x01];
const HEADER_FIXED: usize = 3;
const MIN_LEN: usize = 6;
const PES_OPTIONAL_MARKER: u8 = 0x80;
const PTS_DTS_FLAGS_PTS_ONLY: u8 = 0x80;
const PTS_DTS_FLAGS_BOTH: u8 = 0xC0;
const TS_PREFIX_PTS_WITH_DTS: u8 = 0b0011;
const TS_PREFIX_DTS: u8 = 0b0001;
const TS_VALUE_MASK: u64 = TS_TIMESTAMP_MOD - 1;
const H264_NAL_TYPE_MASK: u8 = 0x1F;
const H264_NAL_AUD: u8 = 9;
const H264_NAL_SPS: u8 = 7;
const HEVC_NAL_VPS: u8 = 32;
const HEVC_NAL_SPS: u8 = 33;
const HEVC_NAL_PPS: u8 = 34;
const HEVC_NAL_AUD: u8 = 35;
const AF_CTRL_ADAPTATION: u8 = 0x20;
const AF_CTRL_PAYLOAD: u8 = 0x10;
const AF_PCR_FLAG: u8 = 0x10;
const PCR_FIELD_LEN: usize = 6;
const STUFFING_BYTE: u8 = 0xFF;
const TS_CLOCK_HZ: u64 = 90_000;
const TS_TIMESTAMP_MOD: u64 = 1 << 33;
const PCR_LEAD_TICKS: u64 = 9_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum EsKind {
Avc,
Hevc,
Mpeg2Video,
Aac,
MpegAudio {
is_mpeg2: bool,
},
Ac3,
Eac3,
Dts,
MpegH,
Data {
stream_type: u8,
carriage: DataCarriage,
},
}
impl EsKind {
fn stream_type(self) -> u8 {
match self {
EsKind::Avc => STREAM_TYPE_AVC,
EsKind::Hevc => STREAM_TYPE_HEVC,
EsKind::Mpeg2Video => STREAM_TYPE_MPEG2_VIDEO,
EsKind::Aac => STREAM_TYPE_AAC_ADTS,
EsKind::MpegAudio { is_mpeg2: true } => STREAM_TYPE_MPEG2_AUDIO,
EsKind::MpegAudio { is_mpeg2: false } => STREAM_TYPE_MPEG1_AUDIO,
EsKind::Ac3 => STREAM_TYPE_AC3,
EsKind::Eac3 => STREAM_TYPE_EAC3,
EsKind::Dts => STREAM_TYPE_DTS,
EsKind::MpegH => STREAM_TYPE_MPEGH,
EsKind::Data { stream_type, .. } => stream_type,
}
}
fn is_video(self) -> bool {
matches!(self, EsKind::Avc | EsKind::Hevc | EsKind::Mpeg2Video)
}
fn is_section_carried(self) -> bool {
matches!(
self,
EsKind::Data {
carriage: DataCarriage::Sections,
..
}
)
}
fn from_config(config: &CodecConfig) -> Option<Self> {
match config {
CodecConfig::Avc { .. } => Some(EsKind::Avc),
CodecConfig::Hevc { .. } => Some(EsKind::Hevc),
CodecConfig::Mpeg2Video { .. } => Some(EsKind::Mpeg2Video),
CodecConfig::Aac { .. } => Some(EsKind::Aac),
CodecConfig::MpegAudio { esds, .. } => Some(EsKind::MpegAudio {
is_mpeg2: mpeg_audio_is_mpeg2(esds),
}),
CodecConfig::Ac3 { .. } => Some(EsKind::Ac3),
CodecConfig::Eac3 { .. } => Some(EsKind::Eac3),
CodecConfig::Dts { .. } => Some(EsKind::Dts),
CodecConfig::MpegH { .. } => Some(EsKind::MpegH),
CodecConfig::Data {
stream_type,
carriage,
..
} => Some(EsKind::Data {
stream_type: *stream_type,
carriage: *carriage,
}),
_ => None,
}
}
}
fn mpeg_audio_is_mpeg2(esds: &EsdsBox) -> bool {
let oti = esds
.es_descriptor
.decoder_config
.as_ref()
.map(|dc| dc.object_type_indication.0)
.unwrap_or(OTI_MPEG1_AUDIO);
oti == OTI_MPEG2_AUDIO
}
pub(crate) struct EsPlan {
pid: u16,
stream_id: StreamId,
kind: EsKind,
asc: Option<AudioSpecificConfig>,
avc_sps_pps: Vec<Vec<u8>>,
hevc_vps_sps_pps: Vec<Vec<u8>>,
descriptors: Vec<u8>,
}
struct TaggedPacket {
sort_key: u64,
packet: [u8; TS_PACKET_SIZE],
}
#[derive(Debug, Default, Clone)]
pub struct TsMux {
_private: (),
}
impl TsMux {
pub fn new() -> Self {
Self { _private: () }
}
}
impl Package for TsMux {
type Media = Media;
type Output = Vec<u8>;
type Error = Error;
fn package(&mut self, media: &Media) -> Result<Vec<u8>> {
if media.tracks.is_empty() {
return Err(Error::InvalidInput("cannot package a Media with no tracks"));
}
let samples: Vec<&[Sample]> = media.tracks.iter().map(|t| t.samples.as_slice()).collect();
mux_tracks(&media.tracks, &samples)
}
}
pub(crate) fn plan_elementary_streams(tracks: &[Track]) -> Result<(Vec<EsPlan>, Vec<usize>)> {
let mut plans: Vec<EsPlan> = Vec::new();
let mut planned_idx: Vec<usize> = Vec::new();
let mut next_pid = ES_PID_BASE;
let (mut n_video, mut n_audio) = (0u8, 0u8);
for (idx, track) in tracks.iter().enumerate() {
let Some(kind) = EsKind::from_config(&track.spec.config) else {
continue; };
let stream_id = if kind.is_video() {
let id = StreamId(STREAM_ID_VIDEO_BASE + n_video);
n_video += 1;
id
} else {
match kind {
EsKind::Data {
carriage: DataCarriage::Pes,
..
} => StreamId(STREAM_ID_PRIVATE_1),
EsKind::Data {
carriage: DataCarriage::Sections,
..
} => StreamId(0),
_ => {
let id = StreamId(STREAM_ID_AUDIO_BASE + n_audio);
n_audio += 1;
id
}
}
};
let asc = match &track.spec.config {
CodecConfig::Aac { esds, .. } => Some(asc_from_esds(esds)?),
_ => None,
};
let avc_sps_pps = match &track.spec.config {
CodecConfig::Avc { config, .. } => {
let r = &config.config;
let mut sets = Vec::new();
for sps in &r.sps {
sets.push(sps.0.clone());
}
for pps in &r.pps {
sets.push(pps.0.clone());
}
sets
}
_ => Vec::new(),
};
let hevc_vps_sps_pps = match &track.spec.config {
CodecConfig::Hevc { config, .. } => hevc_parameter_sets(&config.config),
_ => Vec::new(),
};
let descriptors = match &track.spec.config {
CodecConfig::Data { descriptors, .. } => descriptors.clone(),
CodecConfig::MpegH { config, .. } => {
mpegh_3daudio_descriptor(config.mpegh3da_profile_level_indication)
}
_ => Vec::new(),
};
plans.push(EsPlan {
pid: next_pid,
stream_id,
kind,
asc,
avc_sps_pps,
hevc_vps_sps_pps,
descriptors,
});
planned_idx.push(idx);
next_pid += 1;
}
if plans.is_empty() {
return Err(Error::InvalidInput(
"no track carries a TS-representable codec (AVC/HEVC/MPEG-2-video/AAC/\
MPEG-audio/AC-3/E-AC-3/DTS/MPEG-H/Data)",
));
}
Ok((plans, planned_idx))
}
fn hevc_parameter_sets(
record: &crate::hevc_config::HEVCDecoderConfigurationRecord,
) -> Vec<Vec<u8>> {
let mut vps = Vec::new();
let mut sps = Vec::new();
let mut pps = Vec::new();
for array in &record.arrays {
let bucket = match array.nal_unit_type {
HEVC_NAL_VPS => &mut vps,
HEVC_NAL_SPS => &mut sps,
HEVC_NAL_PPS => &mut pps,
_ => continue,
};
for nalu in &array.nalus {
bucket.push(nalu.0.clone());
}
}
vps.into_iter().chain(sps).chain(pps).collect()
}
pub(crate) fn mux_tracks(tracks: &[Track], samples: &[&[Sample]]) -> Result<Vec<u8>> {
let zero = alloc::vec![0u64; tracks.len()];
mux_tracks_at(tracks, samples, &zero)
}
pub(crate) fn mux_tracks_at(
tracks: &[Track],
samples: &[&[Sample]],
base_dts_ticks: &[u64],
) -> Result<Vec<u8>> {
debug_assert_eq!(tracks.len(), samples.len());
debug_assert_eq!(tracks.len(), base_dts_ticks.len());
let (plans, planned_idx) = plan_elementary_streams(tracks)?;
let pcr_pid = plans
.iter()
.find(|p| p.kind.is_video())
.or_else(|| plans.iter().find(|p| !p.kind.is_section_carried()))
.map(|p| p.pid)
.unwrap_or(plans[0].pid);
let mut out: Vec<u8> = Vec::new();
let pat = build_pat_section(PMT_PID);
for pkt in packetize_section(PAT_PID, &pat) {
out.extend_from_slice(&pkt);
}
let pmt = build_pmt_section(pcr_pid, &plans);
for pkt in packetize_section(PMT_PID, &pmt) {
out.extend_from_slice(&pkt);
}
let mut tagged: Vec<TaggedPacket> = Vec::new();
for (plan, &track_idx) in plans.iter().zip(&planned_idx) {
let track = &tracks[track_idx];
let ts_scale = track.spec.timescale.max(1) as u64;
let mut dts_ticks_local: u64 = base_dts_ticks[track_idx];
let mut cc: u8 = 0;
let mut section_packetizer = SectionPacketizer::new(plan.pid);
for sample in samples[track_idx] {
let dts90 = rescale(dts_ticks_local, ts_scale) + PCR_LEAD_TICKS;
let sort_key = rescale_for_ordering(dts_ticks_local, ts_scale);
if plan.kind.is_section_carried() {
for pkt in section_packetizer.packetize(&[sample.data.as_slice()]) {
tagged.push(TaggedPacket {
sort_key,
packet: pkt,
});
}
} else {
let pts_local = dts_ticks_local as i64 + sample.composition_offset as i64;
let pts90 = rescale_signed(pts_local, ts_scale) + PCR_LEAD_TICKS;
let es_payload = build_es_payload(plan, sample)?;
let carry_pcr = plan.pid == pcr_pid;
packetize_pes(
plan,
&es_payload,
pts90,
dts90,
carry_pcr,
&mut cc,
sort_key,
&mut tagged,
);
}
dts_ticks_local += sample.duration as u64;
}
}
tagged.sort_by_key(|t| t.sort_key);
for t in &tagged {
out.extend_from_slice(&t.packet);
}
debug_assert_eq!(out.len() % TS_PACKET_SIZE, 0);
Ok(out)
}
fn rescale(ticks: u64, timescale: u64) -> u64 {
let scaled = (ticks * TS_CLOCK_HZ + timescale / 2) / timescale;
scaled % TS_TIMESTAMP_MOD
}
fn rescale_signed(ticks: i64, timescale: u64) -> u64 {
if ticks <= 0 {
return 0;
}
rescale(ticks as u64, timescale)
}
fn rescale_for_ordering(ticks: u64, timescale: u64) -> u64 {
let scaled = (ticks as u128 * TS_CLOCK_HZ as u128 + timescale as u128 / 2) / timescale as u128;
scaled.min(u64::MAX as u128) as u64
}
fn asc_from_esds(esds: &crate::mp4esds::EsdsBox) -> Result<AudioSpecificConfig> {
let dsi = esds
.es_descriptor
.decoder_config
.as_ref()
.and_then(|dc| dc.decoder_specific_info.as_ref())
.ok_or(Error::InvalidInput(
"AAC esds carries no DecoderSpecificInfo (AudioSpecificConfig)",
))?;
AudioSpecificConfig::parse(&dsi.data)
}
fn build_es_payload(plan: &EsPlan, sample: &Sample) -> Result<Vec<u8>> {
match plan.kind {
EsKind::Avc => build_annexb_au(&sample.data, sample.is_sync, &plan.avc_sps_pps),
EsKind::Hevc => build_hevc_annexb_au(&sample.data, sample.is_sync, &plan.hevc_vps_sps_pps),
EsKind::Aac => {
let asc = plan
.asc
.as_ref()
.ok_or(Error::InvalidInput("AAC ES has no AudioSpecificConfig"))?;
let frame_len = (sample.data.len() + 7) as u16; let header = asc.to_adts_header(frame_len)?;
let mut out = Vec::with_capacity(header.len() + sample.data.len());
out.extend_from_slice(&header);
out.extend_from_slice(&sample.data);
Ok(out)
}
EsKind::Mpeg2Video
| EsKind::MpegAudio { .. }
| EsKind::Ac3
| EsKind::Eac3
| EsKind::Dts
| EsKind::MpegH
| EsKind::Data { .. } => Ok(sample.data.clone()),
}
}
fn build_annexb_au(length_prefixed: &[u8], is_sync: bool, sps_pps: &[Vec<u8>]) -> Result<Vec<u8>> {
let nals = iter_length_prefixed_nals(length_prefixed)?;
let needs_params = is_sync
&& !sps_pps.is_empty()
&& !nals
.iter()
.any(|n| !n.is_empty() && (n[0] & H264_NAL_TYPE_MASK) == H264_NAL_SPS);
if !needs_params {
return length_prefixed_to_annexb(length_prefixed);
}
let mut out = Vec::with_capacity(length_prefixed.len() + total_param_len(sps_pps));
let mut inserted = false;
for nal in &nals {
let nal_type = nal.first().map(|b| b & H264_NAL_TYPE_MASK);
if !inserted && nal_type != Some(H264_NAL_AUD) {
append_param_sets(&mut out, sps_pps);
inserted = true;
}
out.extend_from_slice(&[0, 0, 0, 1]);
out.extend_from_slice(nal);
}
if !inserted {
append_param_sets(&mut out, sps_pps);
}
Ok(out)
}
fn build_hevc_annexb_au(
length_prefixed: &[u8],
is_sync: bool,
vps_sps_pps: &[Vec<u8>],
) -> Result<Vec<u8>> {
let nals = iter_length_prefixed_nals(length_prefixed)?;
let needs_params = is_sync
&& !vps_sps_pps.is_empty()
&& !nals
.iter()
.any(|n| nal_unit_type(NalCodec::Hevc, n) == Some(HEVC_NAL_SPS));
if !needs_params {
return length_prefixed_to_annexb(length_prefixed);
}
let mut out = Vec::with_capacity(length_prefixed.len() + total_param_len(vps_sps_pps));
let mut inserted = false;
for nal in &nals {
let nal_type = nal_unit_type(NalCodec::Hevc, nal);
if !inserted && nal_type != Some(HEVC_NAL_AUD) {
append_param_sets(&mut out, vps_sps_pps);
inserted = true;
}
out.extend_from_slice(&[0, 0, 0, 1]);
out.extend_from_slice(nal);
}
if !inserted {
append_param_sets(&mut out, vps_sps_pps);
}
Ok(out)
}
fn total_param_len(sps_pps: &[Vec<u8>]) -> usize {
sps_pps.iter().map(|p| 4 + p.len()).sum()
}
fn append_param_sets(out: &mut Vec<u8>, sps_pps: &[Vec<u8>]) {
for p in sps_pps {
out.extend_from_slice(&[0, 0, 0, 1]);
out.extend_from_slice(p);
}
}
fn build_pat_section(pmt_pid: u16) -> Vec<u8> {
let mut body = Vec::new();
body.extend_from_slice(&1u16.to_be_bytes()); body.push(VERSION_CURRENT_NEXT);
body.push(0); body.push(0); body.extend_from_slice(&PROGRAM_NUMBER.to_be_bytes());
body.push(PID_RESERVED_HI | ((pmt_pid >> 8) as u8 & !PID_RESERVED_HI));
body.push((pmt_pid & 0xFF) as u8);
finish_section(TABLE_ID_PAT, body)
}
fn build_pmt_section(pcr_pid: u16, plans: &[EsPlan]) -> Vec<u8> {
let mut body = Vec::new();
body.extend_from_slice(&PROGRAM_NUMBER.to_be_bytes());
body.push(VERSION_CURRENT_NEXT);
body.push(0); body.push(0); body.push(PID_RESERVED_HI | ((pcr_pid >> 8) as u8 & !PID_RESERVED_HI));
body.push((pcr_pid & 0xFF) as u8);
body.push(INFO_RESERVED_HI);
body.push(0);
for p in plans {
body.push(p.kind.stream_type());
body.push(PID_RESERVED_HI | ((p.pid >> 8) as u8 & !PID_RESERVED_HI));
body.push((p.pid & 0xFF) as u8);
let es_info_length = p.descriptors.len().min(MAX_ES_INFO_LENGTH);
body.push(INFO_RESERVED_HI | ((es_info_length >> 8) as u8 & !INFO_RESERVED_HI));
body.push((es_info_length & 0xFF) as u8);
body.extend_from_slice(&p.descriptors[..es_info_length]);
}
finish_section(TABLE_ID_PMT, body)
}
fn finish_section(table_id: u8, body: Vec<u8>) -> Vec<u8> {
let section_length = body.len() + CRC32_LEN;
let mut section = Vec::with_capacity(3 + section_length);
section.push(table_id);
section.push(SECTION_SYNTAX_FLAGS_HI | ((section_length >> 8) as u8 & SECTION_LENGTH_HI_MASK));
section.push((section_length & 0xFF) as u8);
section.extend_from_slice(&body);
let crc = crc32_mpeg2::compute(§ion);
section.extend_from_slice(&crc.to_be_bytes());
section
}
fn packetize_section(pid: u16, section: &[u8]) -> Vec<[u8; TS_PACKET_SIZE]> {
let mut packets = Vec::new();
let mut cc: u8 = 0;
let mut pos = 0usize;
let mut first = true;
while pos < section.len() || first {
let mut pkt = [STUFFING_BYTE; TS_PACKET_SIZE];
let hdr = TsHeader {
tei: false,
pusi: first,
pid,
scrambling: 0,
has_adaptation: false,
has_payload: true,
continuity_counter: cc,
};
hdr.serialize_into(&mut pkt[..4]).expect("4-byte TS header");
cc = (cc + 1) & 0x0F;
let mut w = 4usize;
let cap = if first {
pkt[w] = 0; w += 1;
TS_PACKET_SIZE - w
} else {
TS_PACKET_SIZE - w
};
let take = (section.len() - pos).min(cap);
pkt[w..w + take].copy_from_slice(§ion[pos..pos + take]);
pos += take;
packets.push(pkt);
first = false;
}
packets
}
#[allow(clippy::too_many_arguments)]
fn packetize_pes(
plan: &EsPlan,
es_payload: &[u8],
pts90: u64,
dts90: u64,
carry_pcr: bool,
cc: &mut u8,
sort_key: u64,
tagged: &mut Vec<TaggedPacket>,
) {
let pes = build_pes_bytes(plan, es_payload, pts90, dts90);
let mut pos = 0usize;
let mut first = true;
while pos < pes.len() {
let mut pkt = [STUFFING_BYTE; TS_PACKET_SIZE];
let remaining = pes.len() - pos;
let want_pcr = first && carry_pcr;
let pcr_af_content = if want_pcr { 1 + PCR_FIELD_LEN } else { 0 };
let forced_header = 4 + if want_pcr { 1 + pcr_af_content } else { 0 };
let cap = TS_PACKET_SIZE - forced_header;
let is_last = remaining <= cap;
let to_copy = remaining.min(cap);
let stuff = cap - to_copy;
if want_pcr {
let af_len = pcr_af_content + stuff;
write_af_packet(
&mut pkt,
plan.pid,
first,
*cc,
af_len,
true,
Some(pcr_for(dts90)),
&pes[pos..pos + to_copy],
);
pos += to_copy;
} else if is_last && stuff > 0 {
let af_len = TS_PACKET_SIZE - 4 - 1 - to_copy;
write_af_packet(
&mut pkt,
plan.pid,
first,
*cc,
af_len,
false,
None,
&pes[pos..pos + to_copy],
);
pos += to_copy;
} else {
let hdr = TsHeader {
tei: false,
pusi: first,
pid: plan.pid,
scrambling: 0,
has_adaptation: false,
has_payload: true,
continuity_counter: *cc,
};
hdr.serialize_into(&mut pkt[..4]).expect("4-byte TS header");
pkt[4..4 + to_copy].copy_from_slice(&pes[pos..pos + to_copy]);
pos += to_copy;
}
*cc = (*cc + 1) & 0x0F;
tagged.push(TaggedPacket {
sort_key,
packet: pkt,
});
first = false;
}
}
fn pcr_for(dts90: u64) -> Pcr {
let base = dts90.saturating_sub(PCR_LEAD_TICKS);
Pcr::from_27mhz(base * 300)
}
#[allow(clippy::too_many_arguments)]
fn write_af_packet(
pkt: &mut [u8; TS_PACKET_SIZE],
pid: u16,
pusi: bool,
cc: u8,
af_len: usize,
has_pcr: bool,
pcr: Option<Pcr>,
payload: &[u8],
) {
let hdr = TsHeader {
tei: false,
pusi,
pid,
scrambling: 0,
has_adaptation: true,
has_payload: true,
continuity_counter: cc,
};
hdr.serialize_into(&mut pkt[..4]).expect("4-byte TS header");
debug_assert_eq!(pkt[3] & (AF_CTRL_ADAPTATION | AF_CTRL_PAYLOAD), 0x30);
pkt[4] = af_len as u8;
if af_len >= 1 {
pkt[5] = if has_pcr { AF_PCR_FLAG } else { 0 };
if has_pcr {
if let Some(p) = pcr {
pkt[6..6 + PCR_FIELD_LEN].copy_from_slice(&p.to_field_bytes());
}
}
}
let payload_start = 5 + af_len;
pkt[payload_start..payload_start + payload.len()].copy_from_slice(payload);
}
fn build_pes_bytes(plan: &EsPlan, es_payload: &[u8], pts90: u64, dts90: u64) -> Vec<u8> {
let include_dts = dts90 != pts90;
let opt_content = if include_dts { 10 } else { 5 };
let after_len = HEADER_FIXED + opt_content + es_payload.len();
let pes_packet_length = if plan.kind.is_video() {
0u16
} else {
after_len.min(u16::MAX as usize) as u16
};
let mut out = Vec::with_capacity(MIN_LEN + HEADER_FIXED + opt_content + es_payload.len());
out.extend_from_slice(&PES_START_CODE);
out.push(plan.stream_id.0);
out.extend_from_slice(&pes_packet_length.to_be_bytes());
out.push(PES_OPTIONAL_MARKER); out.push(if include_dts {
PTS_DTS_FLAGS_BOTH
} else {
PTS_DTS_FLAGS_PTS_ONLY
});
out.push(opt_content as u8); if include_dts {
out.extend_from_slice(&encode_timestamp(pts90, TS_PREFIX_PTS_WITH_DTS));
out.extend_from_slice(&encode_timestamp(dts90, TS_PREFIX_DTS));
} else {
out.extend_from_slice(&PesPts(pts90).to_field_bytes());
}
out.extend_from_slice(es_payload);
out
}
fn encode_timestamp(ts: u64, prefix: u8) -> [u8; 5] {
let ts = ts & TS_VALUE_MASK;
[
(prefix << 4) | ((((ts >> 30) & 0x07) as u8) << 1) | 0x01,
((ts >> 22) & 0xFF) as u8,
((((ts >> 15) & 0x7F) as u8) << 1) | 0x01,
((ts >> 7) & 0xFF) as u8,
(((ts & 0x7F) as u8) << 1) | 0x01,
]
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn es_kind_stream_types_mirror_demux() {
assert_eq!(EsKind::Avc.stream_type(), 0x1B);
assert_eq!(EsKind::Hevc.stream_type(), 0x24);
assert_eq!(EsKind::Mpeg2Video.stream_type(), 0x02);
assert_eq!(EsKind::Aac.stream_type(), 0x0F);
assert_eq!(
EsKind::MpegAudio { is_mpeg2: false }.stream_type(),
0x03,
"MPEG-1 audio"
);
assert_eq!(
EsKind::MpegAudio { is_mpeg2: true }.stream_type(),
0x04,
"MPEG-2 audio"
);
assert_eq!(EsKind::Ac3.stream_type(), 0x81);
assert_eq!(EsKind::Eac3.stream_type(), 0x87);
assert_eq!(EsKind::Dts.stream_type(), 0x82);
assert_eq!(EsKind::MpegH.stream_type(), 0x2D);
assert!(EsKind::Avc.is_video());
assert!(EsKind::Hevc.is_video());
assert!(EsKind::Mpeg2Video.is_video());
assert!(!EsKind::Aac.is_video());
assert!(!EsKind::MpegAudio { is_mpeg2: false }.is_video());
assert_eq!(
EsKind::Data {
stream_type: 0x06,
carriage: DataCarriage::Pes,
}
.stream_type(),
0x06
);
assert_eq!(
EsKind::Data {
stream_type: 0x86,
carriage: DataCarriage::Sections,
}
.stream_type(),
0x86
);
assert!(
EsKind::Data {
stream_type: 0x86,
carriage: DataCarriage::Sections,
}
.is_section_carried()
);
assert!(
!EsKind::Data {
stream_type: 0x06,
carriage: DataCarriage::Pes,
}
.is_section_carried()
);
}
#[test]
fn pat_section_crc_is_valid() {
let pat = build_pat_section(PMT_PID);
assert_eq!(crc32_mpeg2::compute(&pat), 0);
assert_eq!(pat[0], TABLE_ID_PAT);
}
#[test]
fn pmt_section_crc_is_valid() {
let plans = alloc::vec![EsPlan {
pid: ES_PID_BASE,
stream_id: StreamId(STREAM_ID_VIDEO_BASE),
kind: EsKind::Avc,
asc: None,
avc_sps_pps: Vec::new(),
hevc_vps_sps_pps: Vec::new(),
descriptors: Vec::new(),
}];
let pmt = build_pmt_section(ES_PID_BASE, &plans);
assert_eq!(crc32_mpeg2::compute(&pmt), 0);
assert_eq!(pmt[0], TABLE_ID_PMT);
}
#[test]
fn pmt_section_carries_es_info_descriptors() {
let descriptors = alloc::vec![0x59, 0x02, 0xAA, 0xBB]; let plans = alloc::vec![EsPlan {
pid: ES_PID_BASE,
stream_id: StreamId(0),
kind: EsKind::Data {
stream_type: 0x06,
carriage: DataCarriage::Pes,
},
asc: None,
avc_sps_pps: Vec::new(),
hevc_vps_sps_pps: Vec::new(),
descriptors: descriptors.clone(),
}];
let pmt = build_pmt_section(ES_PID_BASE, &plans);
assert_eq!(crc32_mpeg2::compute(&pmt), 0);
let es_loop_start = 8 + 4;
assert_eq!(pmt[es_loop_start], 0x06, "stream_type");
let es_info_length =
(((pmt[es_loop_start + 3] & 0x0F) as usize) << 8) | pmt[es_loop_start + 4] as usize;
assert_eq!(es_info_length, descriptors.len());
let desc_start = es_loop_start + 5;
assert_eq!(
&pmt[desc_start..desc_start + es_info_length],
&descriptors[..]
);
}
#[test]
fn section_packets_are_whole_and_pusi() {
let pat = build_pat_section(PMT_PID);
let pkts = packetize_section(PAT_PID, &pat);
assert_eq!(pkts.len(), 1);
assert_eq!(pkts[0][0], 0x47);
assert_ne!(pkts[0][1] & 0x40, 0, "PUSI must be set on the first packet");
}
#[test]
fn mpegh_3daudio_descriptor_carries_profile_level() {
let bytes = mpegh_3daudio_descriptor(0x0B);
assert_eq!(
bytes,
alloc::vec![
DESCRIPTOR_TAG_EXTENSION,
MPEGH_3DAUDIO_DESCRIPTOR_BODY_LEN,
MPEGH_3DAUDIO_EXTENSION_TAG,
0x0B,
]
);
assert_ne!(bytes, mpegh_3daudio_descriptor(0x10));
}
}