use alloc::vec::Vec;
use broadcast_common::Encrypt;
use crate::annexb::{NAL_LENGTH_SIZE, iter_length_prefixed_nals};
use crate::cenc::{CencScheme, SampleEncryptionEntry, SubSampleEntry, TrackEncryptionBox};
use crate::cenc_crypto::{self, CbcsOp};
use crate::error::{Error, Result};
use crate::media::{Media, TrackEncryption};
use crate::nal::NalCodec;
use crate::pipeline::CodecConfig;
const KEY_LEN: usize = 16;
const PER_SAMPLE_IV_SIZE: u8 = 8;
const DEFAULT_CBCS_PATTERN: (u8, u8) = (1, 9);
const CBCS_PATTERN_MAX: u8 = 0x0F;
const VALID_EXPLICIT_IV_LENS: [usize; 2] = [8, 16];
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum IvGen {
Counter {
base: u64,
},
Explicit(Vec<Vec<u8>>),
Constant([u8; KEY_LEN]),
}
impl Default for IvGen {
fn default() -> Self {
IvGen::Counter { base: 0 }
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum SubsamplePolicy {
Video,
WholeSample,
}
#[derive(Debug, Clone)]
pub struct EncryptConfig {
pub scheme: CencScheme,
pub kid: [u8; KEY_LEN],
pub key: [u8; KEY_LEN],
pub iv: IvGen,
pub pattern: Option<(u8, u8)>,
pub subsample: SubsamplePolicy,
}
#[derive(Debug, Clone, Copy, Default)]
pub struct CencEncryptor;
impl Encrypt for CencEncryptor {
type Media = Media;
type Config = EncryptConfig;
type Error = Error;
fn encrypt(&self, media: &mut Media, cfg: &EncryptConfig) -> Result<()> {
let pattern = match cfg.scheme {
CencScheme::Cbcs => {
let p = cfg.pattern.unwrap_or(DEFAULT_CBCS_PATTERN);
if p.0 > CBCS_PATTERN_MAX || p.1 > CBCS_PATTERN_MAX {
return Err(Error::InvalidInput(
"cbcs pattern block counts must each be 0..=15",
));
}
p
}
CencScheme::Cenc => (0, 0),
};
let (per_sample_iv_size, default_constant_iv) = tenc_iv_fields(&cfg.iv)?;
let tenc = TrackEncryptionBox {
version: if cfg.scheme == CencScheme::Cbcs { 1 } else { 0 },
default_crypt_byte_block: pattern.0,
default_skip_byte_block: pattern.1,
default_is_protected: 1,
default_per_sample_iv_size: per_sample_iv_size,
default_kid: cfg.kid,
default_constant_iv,
};
for track in &mut media.tracks {
let nal_codec = nal_codec_for(&track.spec.config);
let sample_count = track.samples.len();
let mut entries = Vec::with_capacity(sample_count);
for (idx, sample) in track.samples.iter_mut().enumerate() {
let iv = resolve_iv(&cfg.iv, idx, sample_count)?;
let subsamples = match (cfg.subsample, nal_codec) {
(SubsamplePolicy::Video, Some(codec)) => nal_subsamples(codec, &sample.data)?,
_ => Vec::new(),
};
let entry = SampleEncryptionEntry {
initialization_vector: iv,
subsamples,
};
match cfg.scheme {
CencScheme::Cenc => cenc_crypto::apply_ctr(
&entry.initialization_vector,
&cfg.key,
&entry.subsamples,
&mut sample.data,
)?,
CencScheme::Cbcs => cenc_crypto::cbcs_sample(
&tenc,
&entry,
&cfg.key,
&mut sample.data,
CbcsOp::Encrypt,
)?,
}
entries.push(entry);
}
track.encryption = Some(TrackEncryption {
scheme: cfg.scheme,
tenc: tenc.clone(),
samples: entries,
});
}
Ok(())
}
}
fn nal_codec_for(config: &CodecConfig) -> Option<NalCodec> {
match config {
CodecConfig::Avc { .. } => Some(NalCodec::Avc),
CodecConfig::Hevc { .. } => Some(NalCodec::Hevc),
CodecConfig::Vvc { .. } => Some(NalCodec::Vvc),
_ => None,
}
}
fn nal_subsamples(codec: NalCodec, data: &[u8]) -> Result<Vec<SubSampleEntry>> {
let header_len: usize = match codec {
NalCodec::Avc => 1,
NalCodec::Hevc | NalCodec::Vvc => 2,
};
let nals = iter_length_prefixed_nals(data)?;
let mut out = Vec::with_capacity(nals.len());
for nal in nals {
let clear_header = header_len.min(nal.len());
out.push(SubSampleEntry {
bytes_of_clear_data: (NAL_LENGTH_SIZE + clear_header) as u16,
bytes_of_protected_data: (nal.len() - clear_header) as u32,
});
}
let total: usize = out
.iter()
.map(|s| s.bytes_of_clear_data as usize + s.bytes_of_protected_data as usize)
.sum();
if total != data.len() {
return Err(Error::InvalidInput(
"NAL subsample map does not cover the whole sample",
));
}
Ok(out)
}
fn resolve_iv(iv_gen: &IvGen, idx: usize, sample_count: usize) -> Result<Vec<u8>> {
match iv_gen {
IvGen::Counter { base } => {
let v = base.checked_add(idx as u64).ok_or(Error::InvalidInput(
"CENC IV counter overflow (base + sample_index)",
))?;
Ok(v.to_be_bytes().to_vec())
}
IvGen::Explicit(ivs) => {
if ivs.len() != sample_count {
return Err(Error::InvalidInput(
"IvGen::Explicit IV count does not match the track's sample count",
));
}
let iv = &ivs[idx];
if !VALID_EXPLICIT_IV_LENS.contains(&iv.len()) {
return Err(Error::InvalidInput(
"CENC per-sample IV must be 8 or 16 bytes",
));
}
Ok(iv.clone())
}
IvGen::Constant(_) => Ok(Vec::new()),
}
}
fn tenc_iv_fields(iv_gen: &IvGen) -> Result<(u8, Option<Vec<u8>>)> {
match iv_gen {
IvGen::Constant(iv) => Ok((0, Some(iv.to_vec()))),
IvGen::Counter { .. } => Ok((PER_SAMPLE_IV_SIZE, None)),
IvGen::Explicit(ivs) => {
let len = match ivs.first() {
Some(first) => {
if ivs.iter().any(|iv| iv.len() != first.len()) {
return Err(Error::InvalidInput(
"IvGen::Explicit IVs must all share one length (tenc.default_per_sample_iv_size is one value for the whole track)",
));
}
first.len()
}
None => PER_SAMPLE_IV_SIZE as usize,
};
if !VALID_EXPLICIT_IV_LENS.contains(&len) {
return Err(Error::InvalidInput(
"CENC per-sample IV must be 8 or 16 bytes",
));
}
Ok((len as u8, None))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use broadcast_common::Unpackage;
use crate::ts_demux::TsDemux;
const KID: [u8; 16] = [
0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF,
0x00,
];
const KEY: [u8; 16] = [
0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F,
0x10,
];
fn clear_media() -> Media {
let mut path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"));
path.push("..");
path.push("fixtures");
path.push("ts");
path.push("h264");
path.push("main.ts");
let bytes = std::fs::read(&path).unwrap_or_else(|e| panic!("read {path:?}: {e}"));
let mut demux = TsDemux::new();
let media = demux
.unpackage(bytes.as_slice())
.expect("demux fixtures/ts/h264/main.ts");
media
.select_tracks_by(|t| matches!(t.spec.config, CodecConfig::Avc { .. }))
.expect("AVC video track present")
}
fn snapshot(media: &Media) -> Vec<Vec<u8>> {
media.tracks[0]
.samples
.iter()
.map(|s| s.data.clone())
.collect()
}
#[test]
fn cenc_round_trip_reverses_byte_identical() {
let mut media = clear_media();
let original = snapshot(&media);
let cfg = EncryptConfig {
scheme: CencScheme::Cenc,
kid: KID,
key: KEY,
iv: IvGen::Counter { base: 7 },
pattern: None,
subsample: SubsamplePolicy::Video,
};
CencEncryptor.encrypt(&mut media, &cfg).expect("encrypt");
let track = &mut media.tracks[0];
let enc = track.encryption.clone().expect("track.encryption Some");
assert_eq!(enc.scheme, CencScheme::Cenc);
assert_eq!(enc.tenc.default_kid, KID);
assert_eq!(enc.samples.len(), track.samples.len());
assert!(
track
.samples
.iter()
.zip(original.iter())
.any(|(s, o)| &s.data != o),
"encrypt must change protected bytes"
);
for (sample, entry) in track.samples.iter_mut().zip(enc.samples.iter()) {
cenc_crypto::apply_ctr(
&entry.initialization_vector,
&KEY,
&entry.subsamples,
&mut sample.data,
)
.expect("reverse apply_ctr");
}
let reversed: Vec<Vec<u8>> = track.samples.iter().map(|s| s.data.clone()).collect();
assert_eq!(reversed, original, "cenc round trip must be byte-identical");
}
#[test]
fn cbcs_round_trip_reverses_byte_identical() {
let mut media = clear_media();
let original = snapshot(&media);
let cfg = EncryptConfig {
scheme: CencScheme::Cbcs,
kid: KID,
key: KEY,
iv: IvGen::Counter { base: 0 },
pattern: Some((1, 9)),
subsample: SubsamplePolicy::Video,
};
CencEncryptor.encrypt(&mut media, &cfg).expect("encrypt");
let track = &mut media.tracks[0];
let enc = track.encryption.clone().expect("track.encryption Some");
assert_eq!(enc.scheme, CencScheme::Cbcs);
assert_eq!(enc.tenc.default_crypt_byte_block, 1);
assert_eq!(enc.tenc.default_skip_byte_block, 9);
assert_eq!(enc.samples.len(), track.samples.len());
assert!(
track
.samples
.iter()
.zip(original.iter())
.any(|(s, o)| &s.data != o),
"encrypt must change protected bytes"
);
for (sample, entry) in track.samples.iter_mut().zip(enc.samples.iter()) {
cenc_crypto::cbcs_sample(&enc.tenc, entry, &KEY, &mut sample.data, CbcsOp::Decrypt)
.expect("reverse cbcs_sample");
}
let reversed: Vec<Vec<u8>> = track.samples.iter().map(|s| s.data.clone()).collect();
assert_eq!(reversed, original, "cbcs round trip must be byte-identical");
}
#[test]
fn whole_sample_policy_yields_empty_subsample_map() {
let mut media = clear_media();
let cfg = EncryptConfig {
scheme: CencScheme::Cenc,
kid: KID,
key: KEY,
iv: IvGen::default(),
pattern: None,
subsample: SubsamplePolicy::WholeSample,
};
CencEncryptor.encrypt(&mut media, &cfg).expect("encrypt");
let enc = media.tracks[0].encryption.as_ref().expect("Some");
assert!(
enc.samples.iter().all(|e| e.subsamples.is_empty()),
"WholeSample policy must record an empty subsample map"
);
}
#[test]
fn explicit_iv_count_mismatch_errors() {
let mut media = clear_media();
let n = media.tracks[0].samples.len();
assert!(n > 1, "fixture must have more than one sample to bite");
let cfg = EncryptConfig {
scheme: CencScheme::Cenc,
kid: KID,
key: KEY,
iv: IvGen::Explicit(alloc::vec![alloc::vec![0u8; 8]; n - 1]),
pattern: None,
subsample: SubsamplePolicy::WholeSample,
};
let err = CencEncryptor.encrypt(&mut media, &cfg).unwrap_err();
assert!(matches!(err, Error::InvalidInput(_)));
}
#[test]
fn explicit_iv_too_long_errors() {
let mut media = clear_media();
let n = media.tracks[0].samples.len();
let cfg = EncryptConfig {
scheme: CencScheme::Cenc,
kid: KID,
key: KEY,
iv: IvGen::Explicit(alloc::vec![alloc::vec![0u8; 17]; n]),
pattern: None,
subsample: SubsamplePolicy::WholeSample,
};
let err = CencEncryptor.encrypt(&mut media, &cfg).unwrap_err();
assert!(matches!(err, Error::InvalidInput(_)));
}
#[test]
fn explicit_iv_empty_errors() {
let mut media = clear_media();
let n = media.tracks[0].samples.len();
let cfg = EncryptConfig {
scheme: CencScheme::Cenc,
kid: KID,
key: KEY,
iv: IvGen::Explicit(alloc::vec![alloc::vec![]; n]),
pattern: None,
subsample: SubsamplePolicy::WholeSample,
};
let err = CencEncryptor.encrypt(&mut media, &cfg).unwrap_err();
assert!(matches!(err, Error::InvalidInput(_)));
}
#[test]
fn explicit_iv_wrong_uniform_length_errors() {
let mut media = clear_media();
let n = media.tracks[0].samples.len();
let cfg = EncryptConfig {
scheme: CencScheme::Cenc,
kid: KID,
key: KEY,
iv: IvGen::Explicit(alloc::vec![alloc::vec![0u8; 12]; n]),
pattern: None,
subsample: SubsamplePolicy::WholeSample,
};
let err = CencEncryptor.encrypt(&mut media, &cfg).unwrap_err();
assert!(matches!(err, Error::InvalidInput(_)));
}
#[test]
fn explicit_iv_valid_lengths_are_ok() {
for len in [8usize, 16] {
let mut media = clear_media();
let n = media.tracks[0].samples.len();
let cfg = EncryptConfig {
scheme: CencScheme::Cenc,
kid: KID,
key: KEY,
iv: IvGen::Explicit(alloc::vec![alloc::vec![0xABu8; len]; n]),
pattern: None,
subsample: SubsamplePolicy::WholeSample,
};
CencEncryptor
.encrypt(&mut media, &cfg)
.unwrap_or_else(|e| panic!("{len}-byte explicit IV must be accepted: {e:?}"));
let enc = media.tracks[0].encryption.as_ref().expect("Some");
assert_eq!(
enc.tenc.default_per_sample_iv_size, len as u8,
"tenc.default_per_sample_iv_size must match the actual IV length used"
);
}
}
#[test]
fn cbcs_pattern_zero_crypt_nonzero_skip_errors() {
let mut media = clear_media();
let cfg = EncryptConfig {
scheme: CencScheme::Cbcs,
kid: KID,
key: KEY,
iv: IvGen::Counter { base: 0 },
pattern: Some((0, 9)),
subsample: SubsamplePolicy::Video,
};
let err = CencEncryptor.encrypt(&mut media, &cfg).unwrap_err();
assert!(matches!(err, Error::InvalidInput(_)));
}
#[test]
fn cbcs_pattern_component_too_large_errors() {
let mut media = clear_media();
let cfg = EncryptConfig {
scheme: CencScheme::Cbcs,
kid: KID,
key: KEY,
iv: IvGen::Counter { base: 0 },
pattern: Some((17, 9)),
subsample: SubsamplePolicy::Video,
};
let err = CencEncryptor.encrypt(&mut media, &cfg).unwrap_err();
assert!(matches!(err, Error::InvalidInput(_)));
}
}