use aes::cipher::generic_array::GenericArray;
use aes::cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit, StreamCipher};
use bytes::{Bytes, BytesMut};
use crate::cenc::{SampleEncryptionEntry, SubSampleEntry, TrackEncryptionBox};
use crate::error::{Error, Result};
pub(crate) fn rewrite_in_place(
data: &mut Bytes,
f: impl FnOnce(&mut [u8]) -> Result<()>,
) -> Result<bool> {
let owned = core::mem::take(data);
let (mut buf, fast_path) = match owned.try_into_mut() {
Ok(buf) => (buf, true),
Err(shared) => (BytesMut::from(&shared[..]), false),
};
let result = f(&mut buf);
*data = buf.freeze();
result.map(|()| fast_path)
}
fn validate_subsample_map(subsamples: &[SubSampleEntry], data_len: usize) -> Result<()> {
let mut offset = 0usize;
for sub in subsamples {
offset = offset
.checked_add(sub.bytes_of_clear_data as usize)
.ok_or(Error::InvalidInput("CENC subsample clear length overflow"))?;
let end = offset
.checked_add(sub.bytes_of_protected_data as usize)
.ok_or(Error::InvalidInput(
"CENC subsample protected length overflow",
))?;
if end > data_len {
return Err(Error::BufferTooShort {
need: end,
have: data_len,
what: "CENC subsample range exceeds sample",
});
}
offset = end;
}
if offset != data_len {
return Err(Error::InvalidInput(
"CENC subsample map does not cover the whole sample (ISO/IEC 23001-7 §9.3): the \
uncovered bytes would be passed through unprotected",
));
}
Ok(())
}
const VALID_IV_LENS: [usize; 2] = [8, 16];
fn check_iv_len(len: usize) -> Result<()> {
if !VALID_IV_LENS.contains(&len) {
return Err(Error::InvalidValue {
field: "CENC IV length",
value: len as u64,
reason: "ISO/IEC 23001-7 §9.2/§12.2 permit only 8 or 16 bytes",
});
}
Ok(())
}
type Aes128Ctr = ctr::Ctr128BE<aes::Aes128>;
type Aes128CbcDec = cbc::Decryptor<aes::Aes128>;
type Aes128CbcEnc = cbc::Encryptor<aes::Aes128>;
const KEY_LEN: usize = 16;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum CbcsOp {
Encrypt,
Decrypt,
}
pub(crate) fn apply_ctr(
iv: &[u8],
key: &[u8; KEY_LEN],
subsamples: &[SubSampleEntry],
data: &mut [u8],
) -> Result<()> {
if iv.is_empty() {
return Err(Error::InvalidInput(
"cenc (AES-CTR) sample has no per-sample IV: an all-zero counter block would reuse \
one keystream for every sample. cenc requires a per-sample IV in senc — a \
tenc.default_constant_IV (default_per_sample_iv_size == 0) is cbcs-only",
));
}
check_iv_len(iv.len())?;
if !subsamples.is_empty() {
validate_subsample_map(subsamples, data.len())?;
}
let mut counter = [0u8; KEY_LEN];
counter[..iv.len()].copy_from_slice(iv);
let mut cipher = Aes128Ctr::new(key.into(), (&counter).into());
if subsamples.is_empty() {
cipher.apply_keystream(data);
return Ok(());
}
let mut offset = 0usize;
for sub in subsamples {
offset += sub.bytes_of_clear_data as usize;
let end = offset + sub.bytes_of_protected_data as usize;
cipher.apply_keystream(&mut data[offset..end]);
offset = end;
}
Ok(())
}
fn resolve_cbcs_iv(
entry: &SampleEncryptionEntry,
tenc: &TrackEncryptionBox,
) -> Result<[u8; KEY_LEN]> {
let src: &[u8] = if !entry.initialization_vector.is_empty() {
&entry.initialization_vector
} else if let Some(civ) = tenc.default_constant_iv.as_deref() {
civ
} else {
return Err(Error::InvalidInput(
"cbcs sample has no per-sample IV and tenc carries no default_constant_IV",
));
};
check_iv_len(src.len())?;
let mut iv = [0u8; KEY_LEN];
iv[..src.len()].copy_from_slice(src);
Ok(iv)
}
pub(crate) fn cbcs_pattern(
key: &[u8; KEY_LEN],
chain_iv: &mut [u8; KEY_LEN],
crypt_byte_block: u8,
skip_byte_block: u8,
range: &mut [u8],
op: CbcsOp,
) {
let (crypt_blocks, skip_blocks) = if crypt_byte_block == 0 && skip_byte_block == 0 {
(1usize, 0usize)
} else {
(crypt_byte_block as usize, skip_byte_block as usize)
};
let mut offset = 0usize;
while offset < range.len() {
let remaining = range.len() - offset;
let want = crypt_blocks * KEY_LEN;
let run_len = (want.min(remaining) / KEY_LEN) * KEY_LEN;
if run_len == 0 {
break;
}
match op {
CbcsOp::Decrypt => {
let mut next_chain = [0u8; KEY_LEN];
next_chain.copy_from_slice(&range[offset + run_len - KEY_LEN..offset + run_len]);
let mut dec = Aes128CbcDec::new(key.into(), (&*chain_iv).into());
for chunk in range[offset..offset + run_len].chunks_exact_mut(KEY_LEN) {
let block = GenericArray::from_mut_slice(chunk);
dec.decrypt_block_mut(block);
}
*chain_iv = next_chain;
}
CbcsOp::Encrypt => {
let mut enc = Aes128CbcEnc::new(key.into(), (&*chain_iv).into());
for chunk in range[offset..offset + run_len].chunks_exact_mut(KEY_LEN) {
let block = GenericArray::from_mut_slice(chunk);
enc.encrypt_block_mut(block);
}
chain_iv.copy_from_slice(&range[offset + run_len - KEY_LEN..offset + run_len]);
}
}
offset += run_len;
if run_len < want {
break;
}
offset += (skip_blocks * KEY_LEN).min(range.len() - offset);
}
}
pub(crate) fn cbcs_sample(
tenc: &TrackEncryptionBox,
entry: &SampleEncryptionEntry,
key: &[u8; KEY_LEN],
data: &mut [u8],
op: CbcsOp,
) -> Result<()> {
let crypt_blocks = tenc.default_crypt_byte_block;
let skip_blocks = tenc.default_skip_byte_block;
if crypt_blocks == 0 && skip_blocks != 0 {
return Err(Error::InvalidInput(
"cbcs pattern crypt_byte_block=0 with nonzero skip leaves data unprotected",
));
}
let seed_iv = resolve_cbcs_iv(entry, tenc)?;
if !entry.subsamples.is_empty() {
validate_subsample_map(&entry.subsamples, data.len())?;
}
if entry.subsamples.is_empty() {
let mut chain_iv = seed_iv;
cbcs_pattern(key, &mut chain_iv, crypt_blocks, skip_blocks, data, op);
return Ok(());
}
let mut offset = 0usize;
for sub in &entry.subsamples {
offset += sub.bytes_of_clear_data as usize;
let end = offset + sub.bytes_of_protected_data as usize;
let mut chain_iv = seed_iv;
cbcs_pattern(
key,
&mut chain_iv,
crypt_blocks,
skip_blocks,
&mut data[offset..end],
op,
);
offset = end;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
const KEY: [u8; KEY_LEN] = [
0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F,
0x10,
];
const IV8: [u8; 8] = [0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88];
#[test]
fn ctr_encrypt_then_decrypt_round_trips() {
let plaintext: Vec<u8> = (0u8..97).collect(); let subsamples = alloc::vec![
SubSampleEntry {
bytes_of_clear_data: 5,
bytes_of_protected_data: 32,
},
SubSampleEntry {
bytes_of_clear_data: 3,
bytes_of_protected_data: 57,
},
];
let mut buf = plaintext.clone();
apply_ctr(&IV8, &KEY, &subsamples, &mut buf).unwrap();
assert_ne!(
buf, plaintext,
"encryption should change the protected bytes"
);
apply_ctr(&IV8, &KEY, &subsamples, &mut buf).unwrap();
assert_eq!(buf, plaintext);
}
#[test]
fn cbcs_encrypt_then_decrypt_round_trips_with_pattern_and_trailing_partial() {
const CRYPT_BLOCKS: u8 = 1;
const SKIP_BLOCKS: u8 = 9;
let plaintext: Vec<u8> = (0u8..=255).cycle().take(320 + 10).collect();
let tenc = TrackEncryptionBox {
version: 1,
default_crypt_byte_block: CRYPT_BLOCKS,
default_skip_byte_block: SKIP_BLOCKS,
default_is_protected: 1,
default_per_sample_iv_size: 16,
default_kid: [0u8; KEY_LEN],
default_constant_iv: None,
};
let entry = SampleEncryptionEntry {
initialization_vector: IV8.to_vec(),
subsamples: Vec::new(),
};
let mut buf = plaintext.clone();
cbcs_sample(&tenc, &entry, &KEY, &mut buf, CbcsOp::Encrypt).unwrap();
assert_ne!(
buf, plaintext,
"encryption should change the protected blocks"
);
cbcs_sample(&tenc, &entry, &KEY, &mut buf, CbcsOp::Decrypt).unwrap();
assert_eq!(buf, plaintext);
}
#[test]
fn cbcs_sample_decrypt_rejects_zero_crypt_nonzero_skip() {
let tenc = TrackEncryptionBox {
version: 1,
default_crypt_byte_block: 0,
default_skip_byte_block: 9,
default_is_protected: 1,
default_per_sample_iv_size: 8,
default_kid: [0u8; KEY_LEN],
default_constant_iv: None,
};
let entry = SampleEncryptionEntry {
initialization_vector: IV8.to_vec(),
subsamples: Vec::new(),
};
let mut data: Vec<u8> = (0u8..64).collect();
let err = cbcs_sample(&tenc, &entry, &KEY, &mut data, CbcsOp::Decrypt).unwrap_err();
assert!(matches!(err, Error::InvalidInput(_)));
}
#[test]
fn ctr_rejects_empty_iv() {
let mut data: Vec<u8> = (0u8..64).collect();
let untouched = data.clone();
let err = apply_ctr(&[], &KEY, &[], &mut data).unwrap_err();
assert!(matches!(err, Error::InvalidInput(_)));
assert_eq!(data, untouched, "a rejected call must not cipher anything");
}
#[test]
fn ctr_rejects_non_8_or_16_byte_iv() {
for len in [1usize, 7, 9, 12, 15, 17, 20] {
let iv = alloc::vec![0x42u8; len];
let mut data: Vec<u8> = (0u8..64).collect();
let untouched = data.clone();
let err = apply_ctr(&iv, &KEY, &[], &mut data).unwrap_err();
assert!(
matches!(err, Error::InvalidValue { .. }),
"len {len}: expected InvalidValue, got {err:?}"
);
assert_eq!(
data, untouched,
"len {len}: a rejected call must not cipher anything"
);
}
}
#[test]
fn cbcs_rejects_non_8_or_16_byte_iv() {
let tenc = TrackEncryptionBox {
version: 1,
default_crypt_byte_block: 1,
default_skip_byte_block: 9,
default_is_protected: 1,
default_per_sample_iv_size: 12,
default_kid: [0u8; KEY_LEN],
default_constant_iv: None,
};
let entry = SampleEncryptionEntry {
initialization_vector: alloc::vec![0x42u8; 12],
subsamples: Vec::new(),
};
let mut data: Vec<u8> = (0u8..64).collect();
let err = cbcs_sample(&tenc, &entry, &KEY, &mut data, CbcsOp::Decrypt).unwrap_err();
assert!(matches!(err, Error::InvalidValue { .. }), "got {err:?}");
let tenc_constant = TrackEncryptionBox {
default_per_sample_iv_size: 0,
default_constant_iv: Some(alloc::vec![0x42u8; 12]),
..tenc
};
let entry_empty = SampleEncryptionEntry {
initialization_vector: Vec::new(),
subsamples: Vec::new(),
};
let mut data2: Vec<u8> = (0u8..64).collect();
let err = cbcs_sample(
&tenc_constant,
&entry_empty,
&KEY,
&mut data2,
CbcsOp::Decrypt,
)
.unwrap_err();
assert!(matches!(err, Error::InvalidValue { .. }), "got {err:?}");
}
#[test]
fn ctr_rejects_partial_subsample_coverage() {
let mut data: Vec<u8> = (0u8..=255).cycle().take(1000).collect();
let untouched = data.clone();
let subsamples = alloc::vec![SubSampleEntry {
bytes_of_clear_data: 4,
bytes_of_protected_data: 96,
}];
let err = apply_ctr(&IV8, &KEY, &subsamples, &mut data).unwrap_err();
assert!(matches!(err, Error::InvalidInput(_)), "got {err:?}");
assert_eq!(
data, untouched,
"the sample must be left byte-identical, not partially keystreamed"
);
}
#[test]
fn cbcs_rejects_partial_subsample_coverage() {
let tenc = TrackEncryptionBox {
version: 1,
default_crypt_byte_block: 1,
default_skip_byte_block: 9,
default_is_protected: 1,
default_per_sample_iv_size: 8,
default_kid: [0u8; KEY_LEN],
default_constant_iv: None,
};
let entry = SampleEncryptionEntry {
initialization_vector: IV8.to_vec(),
subsamples: alloc::vec![SubSampleEntry {
bytes_of_clear_data: 4,
bytes_of_protected_data: 96,
}],
};
let mut data: Vec<u8> = (0u8..=255).cycle().take(1000).collect();
let untouched = data.clone();
let err = cbcs_sample(&tenc, &entry, &KEY, &mut data, CbcsOp::Decrypt).unwrap_err();
assert!(matches!(err, Error::InvalidInput(_)), "got {err:?}");
assert_eq!(data, untouched, "the sample must be left byte-identical");
}
#[test]
fn second_subsample_overrun_leaves_the_sample_unchanged() {
let plaintext: Vec<u8> = (0u8..=255).cycle().take(200).collect();
let overrunning = alloc::vec![
SubSampleEntry {
bytes_of_clear_data: 4,
bytes_of_protected_data: 60, },
SubSampleEntry {
bytes_of_clear_data: 4,
bytes_of_protected_data: 4096, },
];
let tenc = TrackEncryptionBox {
version: 1,
default_crypt_byte_block: 1,
default_skip_byte_block: 9,
default_is_protected: 1,
default_per_sample_iv_size: 8,
default_kid: [0u8; KEY_LEN],
default_constant_iv: None,
};
let entry = SampleEncryptionEntry {
initialization_vector: IV8.to_vec(),
subsamples: overrunning.clone(),
};
for label in ["cenc", "cbcs"] {
let mut data = Bytes::from(plaintext.clone());
let err = rewrite_in_place(&mut data, |buf| {
if label == "cenc" {
apply_ctr(&IV8, &KEY, &overrunning, buf)
} else {
cbcs_sample(&tenc, &entry, &KEY, buf, CbcsOp::Encrypt)
}
})
.expect_err("an overrunning subsample map must be rejected");
assert!(
matches!(err, Error::BufferTooShort { .. }),
"{label}: {err:?}"
);
assert_eq!(
&data[..],
&plaintext[..],
"{label}: sample.data must be unchanged — not half-encrypted, not empty"
);
}
}
#[test]
fn rewrite_in_place_restores_the_buffer_on_err() {
let original: Vec<u8> = (0u8..32).collect();
let mut data = Bytes::from(original.clone());
let err = rewrite_in_place(&mut data, |_buf| {
Err(Error::InvalidInput("closure failed before mutating"))
})
.expect_err("the closure's error must propagate");
assert!(matches!(err, Error::InvalidInput(_)));
assert_eq!(
&data[..],
&original[..],
"a validate-then-mutate closure's failure must leave the sample intact"
);
}
#[test]
fn rewrite_in_place_takes_fast_path_when_unique() {
let mut data = Bytes::from(alloc::vec![1u8, 2, 3, 4]);
let took_fast_path = rewrite_in_place(&mut data, |buf| {
buf[0] = 0xFF;
Ok(())
})
.expect("rewrite ok");
assert!(
took_fast_path,
"uniquely-owned Bytes must take the zero-copy try_into_mut fast path"
);
assert_eq!(&data[..], &[0xFF, 2, 3, 4]);
}
#[test]
fn rewrite_in_place_copies_and_leaves_other_holder_untouched_when_shared() {
let original = Bytes::from(alloc::vec![9u8, 9, 9, 9]);
let mut shared_handle = original.clone(); let took_fast_path = rewrite_in_place(&mut shared_handle, |buf| {
buf[0] = 0x00;
Ok(())
})
.expect("rewrite ok");
assert!(
!took_fast_path,
"shared Bytes must not take the fast path — that would alias/corrupt the other holder"
);
assert_eq!(
&original[..],
&[9, 9, 9, 9],
"the other holder's bytes must be untouched by shared_handle's rewrite"
);
assert_eq!(
&shared_handle[..],
&[0x00, 9, 9, 9],
"the rewriting handle's own view must reflect the mutation"
);
}
}