feat(moq-mux): add Avcc/Hvcc param-set parsers and annexb prefix helpers

Expose the out-of-band parameter sets a length-prefixed (avc1/hvc1) stream
needs to re-inject at each keyframe when converting back to Annex-B:

- `h264::Avcc` gains `sps`/`pps` fields (and `#[non_exhaustive]`); `Avcc::parse`
  now collects the NAL lists, deriving resolution from the first SPS.
- `h265::Hvcc` + `Hvcc::parse`: the HEVC analogue, sorting VPS/SPS/PPS by type.
- `annexb::from_length_prefixed` (with an optional keyframe prefix) and
  `annexb::build_prefix`.

The internal `avcc_params`/`hvcc_params` flatteners now delegate to the typed
parsers, so there is a single byte-parser per codec. This also fixes
`hvcc_params` to keep only VPS/SPS/PPS as its doc already claimed (it previously
flattened every NAL array).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: kixelated

rs/moq-mux/src/codec/annexb.rs

@@ -1,4 +1,4 @@
-use anyhow::{self};
+use anyhow::Context;
 use bytes::{Buf, Bytes, BytesMut};
 
 pub const START_CODE: Bytes = Bytes::from_static(&[0, 0, 0, 1]);
@@ -148,6 +148,55 @@ pub fn find_start_code(b: &[u8]) -> Option<(usize, usize)> {
 	}
 }
 
+/// Convert a length-prefixed NALU payload (avc1 / hvc1 wire shape) to Annex-B,
+/// optionally prepending `prefix` bytes (typically VPS/SPS/PPS NAL units already
+/// in Annex-B form, for keyframe parameter-set injection).
+pub fn from_length_prefixed(payload: &[u8], length_size: usize, prefix: Option<&[u8]>) -> anyhow::Result<Bytes> {
+	anyhow::ensure!(
+		(1..=4).contains(&length_size),
+		"invalid avc1/hvc1 length size {length_size}"
+	);
+
+	let mut out = BytesMut::with_capacity(payload.len() + prefix.map(|p| p.len()).unwrap_or(0) + 16);
+	if let Some(p) = prefix {
+		out.extend_from_slice(p);
+	}
+
+	let mut pos = 0;
+	while pos < payload.len() {
+		let after_prefix = pos
+			.checked_add(length_size)
+			.context("truncated length-prefixed NAL unit")?;
+		anyhow::ensure!(payload.len() >= after_prefix, "truncated length-prefixed NAL unit");
+		let mut len = 0usize;
+		for byte in &payload[pos..after_prefix] {
+			len = (len << 8) | (*byte as usize);
+		}
+		let after_nal = after_prefix
+			.checked_add(len)
+			.context("truncated length-prefixed NAL unit")?;
+		anyhow::ensure!(payload.len() >= after_nal, "truncated length-prefixed NAL unit");
+		out.extend_from_slice(&START_CODE);
+		out.extend_from_slice(&payload[after_prefix..after_nal]);
+		pos = after_nal;
+	}
+
+	Ok(out.freeze())
+}
+
+/// Concatenate `start_code | nal` for every NAL in `nals` and freeze the result.
+/// Used to build a keyframe parameter-set prefix for an Annex-B elementary stream.
+pub fn build_prefix<'a, I: IntoIterator<Item = &'a Bytes>>(nals: I) -> Bytes {
+	let nals: Vec<&Bytes> = nals.into_iter().collect();
+	let total: usize = nals.iter().map(|n| n.len() + START_CODE.len()).sum();
+	let mut out = BytesMut::with_capacity(total);
+	for nal in nals {
+		out.extend_from_slice(&START_CODE);
+		out.extend_from_slice(nal);
+	}
+	out.freeze()
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;

rs/moq-mux/src/codec/h264/mod.rs

@@ -53,12 +53,17 @@ impl Sps {
 /// avcC bytes are still what gets stored as the catalog `description`; this
 /// struct is for the field extraction.
 #[derive(Debug, Clone)]
+#[non_exhaustive]
 pub struct Avcc {
 	pub profile: u8,
 	pub constraints: u8,
 	pub level: u8,
 	/// NALU length size in bytes (typically 4).
 	pub length_size: usize,
+	/// SPS NAL units carried out-of-band in the record.
+	pub sps: Vec<Bytes>,
+	/// PPS NAL units carried out-of-band in the record.
+	pub pps: Vec<Bytes>,
 	/// Resolution from the embedded SPS, if one was present and parseable.
 	pub coded_width: Option<u32>,
 	pub coded_height: Option<u32>,
@@ -67,33 +72,40 @@ pub struct Avcc {
 impl Avcc {
 	/// Parse an AVCDecoderConfigurationRecord buffer.
 	pub fn parse(avcc: &[u8]) -> anyhow::Result<Self> {
-		anyhow::ensure!(avcc.len() >= 6, "AVCDecoderConfigurationRecord too short");
+		anyhow::ensure!(avcc.len() >= 7, "AVCDecoderConfigurationRecord too short");
 
 		let profile = avcc[1];
 		let constraints = avcc[2];
 		let level = avcc[3];
 		let length_size = (avcc[4] & 0x03) as usize + 1;
-		let num_sps = avcc[5] & 0x1f;
+		let num_sps = (avcc[5] & 0x1f) as usize;
 
+		let mut sps = Vec::with_capacity(num_sps);
+		let mut pos = read_param_set_array(avcc, 6, num_sps, &mut sps)?;
+
+		anyhow::ensure!(avcc.len() > pos, "AVCDecoderConfigurationRecord truncated");
+		let num_pps = avcc[pos] as usize;
+		pos += 1;
+		let mut pps = Vec::with_capacity(num_pps);
+		read_param_set_array(avcc, pos, num_pps, &mut pps)?;
+
+		// Resolution from the first parseable SPS.
 		let (mut coded_width, mut coded_height) = (None, None);
-		if num_sps > 0 && avcc.len() >= 8 {
-			let sps_len = u16::from_be_bytes([avcc[6], avcc[7]]) as usize;
-			let sps_start = 8;
-			let sps_end = sps_start + sps_len;
-			if sps_end <= avcc.len()
-				&& sps_len > 1
-				&& let Ok(sps) = Sps::parse(&avcc[sps_start..sps_end])
-			{
-				coded_width = Some(sps.coded_width);
-				coded_height = Some(sps.coded_height);
-			}
+		if let Some(first) = sps.first()
+			&& first.len() > 1
+			&& let Ok(parsed) = Sps::parse(first)
+		{
+			coded_width = Some(parsed.coded_width);
+			coded_height = Some(parsed.coded_height);
 		}
 
 		Ok(Self {
 			profile,
 			constraints,
 			level,
 			length_size,
+			sps,
+			pps,
 			coded_width,
 			coded_height,
 		})
@@ -168,20 +180,13 @@ pub(crate) fn build_avcc(sps_nals: &[Bytes], pps_nals: &[Bytes]) -> anyhow::Resu
 /// Extract the parameter-set NALs (SPS then PPS) and the NALU length size from
 /// an AVCDecoderConfigurationRecord. The inverse of [`build_avcc`]; used to
 /// re-emit out-of-band avc1 parameter sets as inline Annex-B (e.g. for MPEG-TS).
+/// The SPS+PPS NALs (in that order) and NALU length size, flattened for callers
+/// that re-emit every parameter set as one Annex-B prefix (e.g. MPEG-TS export).
 pub(crate) fn avcc_params(avcc: &[u8]) -> anyhow::Result<(usize, Vec<Bytes>)> {
-	anyhow::ensure!(avcc.len() >= 6, "AVCDecoderConfigurationRecord too short");
-	let length_size = (avcc[4] & 0x03) as usize + 1;
-
-	let mut params = Vec::new();
-	let num_sps = avcc[5] & 0x1f;
-	let mut pos = read_param_set_array(avcc, 6, num_sps as usize, &mut params)?;
-
-	anyhow::ensure!(avcc.len() > pos, "avcC missing PPS count");
-	let num_pps = avcc[pos];
-	pos += 1;
-	read_param_set_array(avcc, pos, num_pps as usize, &mut params)?;
-
-	Ok((length_size, params))
+	let avcc = Avcc::parse(avcc)?;
+	let mut params = avcc.sps;
+	params.extend(avcc.pps);
+	Ok((avcc.length_size, params))
 }
 
 /// Read `count` u16-length-prefixed NALs starting at `pos`, appending each to
@@ -389,6 +394,20 @@ mod tests {
 		assert_eq!(params[1], pps);
 	}
 
+	#[test]
+	fn avcc_parse_separates_sps_and_pps() {
+		let sps = Bytes::from_static(&[0x67, 0x42, 0xc0, 0x1f, 0xde]);
+		let pps0 = Bytes::from_static(&[0x68, 0xce, 0x3c, 0x80]);
+		let pps1 = Bytes::from_static(&[0x68, 0xce, 0x3c, 0x81]);
+
+		let avcc = build_avcc(std::slice::from_ref(&sps), &[pps0.clone(), pps1.clone()]).unwrap();
+		let parsed = Avcc::parse(&avcc).unwrap();
+
+		assert_eq!(parsed.length_size, 4);
+		assert_eq!(parsed.sps, vec![sps]);
+		assert_eq!(parsed.pps, vec![pps0, pps1]);
+	}
+
 	#[test]
 	fn build_avcc_carries_multiple_pps() {
 		// A source with one SPS and two PPS (ids 0 and 1): the avcC must keep both,

rs/moq-mux/src/codec/h265/mod.rs

@@ -282,34 +282,69 @@ pub(crate) fn build_hvcc(vps_nals: &[Bytes], sps_nals: &[Bytes], pps_nals: &[Byt
 	Ok(out.freeze())
 }
 
-/// Extract the parameter-set NALs (VPS, SPS, PPS in array order) and the NALU
-/// length size from an HEVCDecoderConfigurationRecord. The inverse of
-/// [`build_hvcc`]; used to re-emit out-of-band hvc1 parameter sets as inline
-/// Annex-B (e.g. for MPEG-TS).
+/// The VPS+SPS+PPS NALs (in that order) and NALU length size, flattened for
+/// callers that re-emit every parameter set as one Annex-B prefix (e.g. MPEG-TS
+/// export).
 pub(crate) fn hvcc_params(hvcc: &[u8]) -> anyhow::Result<(usize, Vec<Bytes>)> {
-	anyhow::ensure!(hvcc.len() >= 23, "HEVCDecoderConfigurationRecord too short");
-	let length_size = (hvcc[21] & 0x03) as usize + 1;
-	let num_arrays = hvcc[22];
+	let hvcc = Hvcc::parse(hvcc)?;
+	let mut params = hvcc.vps;
+	params.extend(hvcc.sps);
+	params.extend(hvcc.pps);
+	Ok((hvcc.length_size, params))
+}
 
-	let mut params = Vec::new();
-	let mut pos = 23;
-	for _ in 0..num_arrays {
-		// Skip the array_completeness | NAL_unit_type byte.
-		anyhow::ensure!(hvcc.len() >= pos + 3, "truncated hvcC NAL array header");
-		pos += 1;
-		let num_nalus = u16::from_be_bytes([hvcc[pos], hvcc[pos + 1]]);
-		pos += 2;
-		for _ in 0..num_nalus {
-			anyhow::ensure!(hvcc.len() >= pos + 2, "truncated hvcC NAL length");
-			let len = u16::from_be_bytes([hvcc[pos], hvcc[pos + 1]]) as usize;
-			pos += 2;
-			anyhow::ensure!(hvcc.len() >= pos + len, "hvcC NAL exceeds buffer");
-			params.push(Bytes::copy_from_slice(&hvcc[pos..pos + len]));
-			pos += len;
+/// The parameter sets carried out-of-band in an HEVCDecoderConfigurationRecord,
+/// split by NAL type.
+#[derive(Debug, Clone)]
+#[non_exhaustive]
+pub struct Hvcc {
+	/// NALU length size in bytes (typically 4).
+	pub length_size: usize,
+	pub vps: Vec<Bytes>,
+	pub sps: Vec<Bytes>,
+	pub pps: Vec<Bytes>,
+}
+
+impl Hvcc {
+	/// Parse an HEVCDecoderConfigurationRecord, sorting the VPS/SPS/PPS NAL units
+	/// by type. The HEVC analogue of [`super::h264::Avcc::parse`]: used to recover
+	/// the Annex-B parameter sets a length-prefixed (hvc1) stream needs at each
+	/// keyframe for in-band injection (VPS→SPS→PPS order).
+	pub fn parse(hvcc: &[u8]) -> anyhow::Result<Self> {
+		anyhow::ensure!(hvcc.len() >= 23, "HEVCDecoderConfigurationRecord too short");
+		let length_size = (hvcc[21] & 0x03) as usize + 1;
+		let num_arrays = hvcc[22];
+
+		let mut out = Self {
+			length_size,
+			vps: Vec::new(),
+			sps: Vec::new(),
+			pps: Vec::new(),
+		};
+		let mut pos = 23;
+		for _ in 0..num_arrays {
+			anyhow::ensure!(hvcc.len() >= pos + 3, "truncated hvcC NAL array header");
+			let nal_type = hvcc[pos] & 0x3f;
+			let num_nalus = u16::from_be_bytes([hvcc[pos + 1], hvcc[pos + 2]]);
+			pos += 3;
+			for _ in 0..num_nalus {
+				anyhow::ensure!(hvcc.len() >= pos + 2, "truncated hvcC NAL length");
+				let len = u16::from_be_bytes([hvcc[pos], hvcc[pos + 1]]) as usize;
+				pos += 2;
+				anyhow::ensure!(hvcc.len() >= pos + len, "hvcC NAL exceeds buffer");
+				let nal = Bytes::copy_from_slice(&hvcc[pos..pos + len]);
+				pos += len;
+				match NALUnitType::from(nal_type) {
+					NALUnitType::VpsNut => out.vps.push(nal),
+					NALUnitType::SpsNut => out.sps.push(nal),
+					NALUnitType::PpsNut => out.pps.push(nal),
+					_ => {}
+				}
+			}
 		}
-	}
 
-	Ok((length_size, params))
+		Ok(out)
+	}
 }
 
 /// Pack the constraint flags from ITU H.265 V10 §7.3.3 Profile, tier and level syntax.
@@ -356,5 +391,12 @@ mod tests {
 		assert_eq!(params[0].as_ref(), vps);
 		assert_eq!(params[1].as_ref(), sps);
 		assert_eq!(params[2].as_ref(), pps);
+
+		// The typed parser keys the same NALs by type.
+		let parsed = Hvcc::parse(&hvcc).unwrap();
+		assert_eq!(parsed.length_size, 4);
+		assert_eq!(parsed.vps, vec![Bytes::copy_from_slice(vps)]);
+		assert_eq!(parsed.sps, vec![Bytes::copy_from_slice(sps)]);
+		assert_eq!(parsed.pps, vec![Bytes::copy_from_slice(pps)]);
 	}
 }