1use aes::Aes128;
25use aes::cipher::generic_array::GenericArray;
26use aes::cipher::{BlockEncrypt, KeyInit};
27use sheathe_core::{Error, Result};
28
29mod pssh;
30pub use pssh::ProtectionSystem;
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum Scheme {
35 Cenc,
37 Cens,
39 Cbc1,
41 Cbcs,
43}
44
45impl Scheme {
46 pub fn scheme_type(self) -> [u8; 4] {
48 match self {
49 Scheme::Cenc => *b"cenc",
50 Scheme::Cens => *b"cens",
51 Scheme::Cbc1 => *b"cbc1",
52 Scheme::Cbcs => *b"cbcs",
53 }
54 }
55
56 pub fn is_cbc(self) -> bool {
59 matches!(self, Scheme::Cbc1 | Scheme::Cbcs)
60 }
61
62 pub fn is_pattern(self) -> bool {
65 matches!(self, Scheme::Cens | Scheme::Cbcs)
66 }
67
68 pub fn uses_constant_iv(self) -> bool {
71 matches!(self, Scheme::Cbcs)
72 }
73}
74
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80pub struct Pattern {
81 pub crypt_blocks: u8,
83 pub skip_blocks: u8,
85}
86
87impl Pattern {
88 pub const NONE: Pattern = Pattern { crypt_blocks: 0, skip_blocks: 0 };
90 pub const VIDEO: Pattern = Pattern { crypt_blocks: 1, skip_blocks: 9 };
92
93 fn is_patterned(self) -> bool {
95 self.crypt_blocks != 0
96 }
97}
98
99#[derive(Debug, Clone)]
101pub struct ContentKey {
102 pub kid: [u8; 16],
104 pub key: [u8; 16],
106}
107
108impl ContentKey {
109 pub fn rotated(&self, period: u32) -> ContentKey {
113 let n = (period % 16) as usize;
114 let mut kid = [0u8; 16];
115 let mut key = [0u8; 16];
116 for i in 0..16 {
117 kid[i] = self.kid[(i + n) % 16];
118 key[i] = self.key[(i + n) % 16];
119 }
120 ContentKey { kid, key }
121 }
122}
123
124#[derive(Debug, Clone, Copy, PartialEq, Eq)]
127pub struct Subsample {
128 pub clear: u32,
130 pub protected: u32,
132}
133
134pub struct Encryptor {
136 cipher: Aes128,
137}
138
139impl Encryptor {
140 pub fn new(key: &[u8; 16]) -> Self {
142 Self { cipher: Aes128::new(GenericArray::from_slice(key)) }
143 }
144
145 pub fn encrypt(
150 &self,
151 scheme: Scheme,
152 pattern: Pattern,
153 iv: &[u8; 16],
154 data: &mut [u8],
155 subsamples: &[Subsample],
156 ) -> Result<()> {
157 let total: u64 =
159 subsamples.iter().map(|s| u64::from(s.clear) + u64::from(s.protected)).sum();
160 if total != data.len() as u64 {
161 return Err(Error::malformed("subsample layout does not cover sample"));
162 }
163 if pattern.is_patterned() && !scheme.is_pattern() {
164 return Err(Error::malformed("pattern set on a non-pattern scheme"));
165 }
166 if scheme.is_cbc() {
167 self.cbc(pattern, iv, data, subsamples);
168 } else {
169 self.ctr(pattern, iv, data, subsamples);
170 }
171 Ok(())
172 }
173
174 fn ctr(&self, pattern: Pattern, iv: &[u8; 16], data: &mut [u8], subsamples: &[Subsample]) {
180 let mut counter = *iv;
181 let mut keystream = [0u8; 16];
182 let mut ks_pos = 16usize; for_each_crypt_range(pattern, subsamples, |start, len| {
185 for byte in &mut data[start..start + len] {
186 if ks_pos == 16 {
187 keystream = counter;
188 self.encrypt_block(&mut keystream);
189 incr_be(&mut counter);
190 ks_pos = 0;
191 }
192 *byte ^= keystream[ks_pos];
193 ks_pos += 1;
194 }
195 });
196 }
197
198 fn cbc(&self, pattern: Pattern, iv: &[u8; 16], data: &mut [u8], subsamples: &[Subsample]) {
204 let mut off = 0usize;
205 let mut chain = *iv;
209 for s in subsamples {
210 off += s.clear as usize;
211 if pattern.is_patterned() {
212 chain = *iv;
213 }
214 let mut remaining = s.protected as usize;
215 let mut block_index = 0usize;
216 let cycle = pattern.crypt_blocks as usize + pattern.skip_blocks as usize;
217 while remaining >= 16 {
218 let encrypt =
219 !pattern.is_patterned() || block_index % cycle < pattern.crypt_blocks as usize;
220 if encrypt {
221 let mut block = [0u8; 16];
222 block.copy_from_slice(&data[off..off + 16]);
223 for (b, c) in block.iter_mut().zip(chain.iter()) {
224 *b ^= *c;
225 }
226 self.encrypt_block(&mut block);
227 data[off..off + 16].copy_from_slice(&block);
228 chain = block;
229 }
230 off += 16;
231 remaining -= 16;
232 block_index += 1;
233 }
234 off += remaining; }
236 }
237
238 fn encrypt_block(&self, block: &mut [u8; 16]) {
240 let mut ga = GenericArray::clone_from_slice(block);
241 self.cipher.encrypt_block(&mut ga);
242 block.copy_from_slice(&ga);
243 }
244}
245
246fn for_each_crypt_range(
253 pattern: Pattern,
254 subsamples: &[Subsample],
255 mut f: impl FnMut(usize, usize),
256) {
257 let mut off = 0usize;
258 for s in subsamples {
259 off += s.clear as usize;
260 let protected = s.protected as usize;
261 if !pattern.is_patterned() {
262 if protected > 0 {
263 f(off, protected);
264 }
265 off += protected;
266 continue;
267 }
268 let crypt = pattern.crypt_blocks as usize * 16;
269 let skip = pattern.skip_blocks as usize * 16;
270 let mut pos = 0usize;
271 while pos < protected {
272 let phase = crypt.min(protected - pos);
273 let whole = phase - phase % 16;
276 if whole > 0 {
277 f(off + pos, whole);
278 }
279 pos += phase;
280 pos += skip.min(protected - pos);
281 }
282 off += protected;
283 }
284}
285
286fn incr_be(counter: &mut [u8; 16]) {
288 for byte in counter.iter_mut().rev() {
289 let (v, carry) = byte.overflowing_add(1);
290 *byte = v;
291 if !carry {
292 break;
293 }
294 }
295}
296
297#[cfg(test)]
298mod tests {
299 use super::*;
300
301 const KEY: [u8; 16] = [
302 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f,
303 0x3c,
304 ];
305
306 fn hex(s: &str) -> Vec<u8> {
307 (0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap()).collect()
308 }
309
310 #[test]
311 fn cenc_matches_nist_ctr_vector() {
312 let iv = hex("f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff");
314 let mut data = hex("6bc1bee22e409f96e93d7e117393172a");
315 let enc = Encryptor::new(&KEY);
316 let subs = [Subsample { clear: 0, protected: 16 }];
317 enc.encrypt(Scheme::Cenc, Pattern::NONE, iv[..].try_into().unwrap(), &mut data, &subs)
318 .unwrap();
319 assert_eq!(data, hex("874d6191b620e3261bef6864990db6ce"));
320 }
321
322 #[test]
323 fn cbc_schemes_match_nist_cbc_vector() {
324 let iv = hex("000102030405060708090a0b0c0d0e0f");
328 let subs = [Subsample { clear: 0, protected: 16 }];
329 for (scheme, pattern) in [(Scheme::Cbc1, Pattern::NONE), (Scheme::Cbcs, Pattern::VIDEO)] {
330 let mut data = hex("6bc1bee22e409f96e93d7e117393172a");
331 let enc = Encryptor::new(&KEY);
332 enc.encrypt(scheme, pattern, iv[..].try_into().unwrap(), &mut data, &subs).unwrap();
333 assert_eq!(data, hex("7649abac8119b246cee98e9b12e9197d"), "{scheme:?}");
334 }
335 }
336
337 #[test]
338 fn cenc_leaves_clear_bytes_untouched() {
339 let iv = [0u8; 16];
340 let mut data = vec![0xAAu8; 32];
341 let enc = Encryptor::new(&KEY);
342 enc.encrypt(
344 Scheme::Cenc,
345 Pattern::NONE,
346 &iv,
347 &mut data,
348 &[Subsample { clear: 8, protected: 24 }],
349 )
350 .unwrap();
351 assert!(data[..8].iter().all(|&b| b == 0xAA), "clear prefix must be untouched");
352 assert!(data[8..].iter().any(|&b| b != 0xAA), "protected region must change");
353 }
354
355 #[test]
356 fn pattern_schemes_skip_blocks() {
357 let iv = [0u8; 16];
360 for scheme in [Scheme::Cbcs, Scheme::Cens] {
361 let mut data = vec![0x11u8; 160];
362 let original = data.clone();
363 let enc = Encryptor::new(&KEY);
364 enc.encrypt(
365 scheme,
366 Pattern::VIDEO,
367 &iv,
368 &mut data,
369 &[Subsample { clear: 0, protected: 160 }],
370 )
371 .unwrap();
372 assert_ne!(data[..16], original[..16], "{scheme:?}: first block encrypted");
373 assert_eq!(data[16..], original[16..], "{scheme:?}: blocks 1..9 skipped");
374 }
375 }
376
377 #[test]
378 fn rejects_mismatched_layout() {
379 let enc = Encryptor::new(&KEY);
380 let mut data = vec![0u8; 10];
381 let err = enc.encrypt(
382 Scheme::Cenc,
383 Pattern::NONE,
384 &[0u8; 16],
385 &mut data,
386 &[Subsample { clear: 0, protected: 9 }],
387 );
388 assert!(err.is_err());
389 }
390
391 #[test]
392 fn content_key_rotation_left_rotates_by_period() {
393 let base = ContentKey { kid: KEY, key: KEY };
394 assert_eq!(base.rotated(0).kid, KEY, "period 0 is unchanged");
395 let mut expect = KEY;
397 expect.rotate_left(1);
398 assert_eq!(base.rotated(1).kid, expect);
399 assert_eq!(base.rotated(1).key, expect);
400 assert_eq!(base.rotated(16).kid, KEY);
402 }
403
404 #[test]
405 fn rejects_pattern_on_non_pattern_scheme() {
406 let enc = Encryptor::new(&KEY);
407 let mut data = vec![0u8; 16];
408 let err = enc.encrypt(
409 Scheme::Cenc,
410 Pattern::VIDEO,
411 &[0u8; 16],
412 &mut data,
413 &[Subsample { clear: 0, protected: 16 }],
414 );
415 assert!(err.is_err());
416 }
417
418 #[test]
422 fn ctr_schemes_round_trip_across_subsamples() {
423 let enc = Encryptor::new(&KEY);
424 let iv = [3u8; 16];
425 let subs = [Subsample { clear: 5, protected: 40 }, Subsample { clear: 10, protected: 65 }];
426 for (scheme, pattern) in [(Scheme::Cenc, Pattern::NONE), (Scheme::Cens, Pattern::VIDEO)] {
427 let original: Vec<u8> = (0..120u8).collect();
428 let mut data = original.clone();
429 enc.encrypt(scheme, pattern, &iv, &mut data, &subs).unwrap();
430 assert_ne!(data, original, "{scheme:?}: ciphertext must differ");
431 assert_eq!(&data[..5], &original[..5], "{scheme:?}: leading clear bytes preserved");
432 enc.encrypt(scheme, pattern, &iv, &mut data, &subs).unwrap();
433 assert_eq!(data, original, "{scheme:?}: CTR round-trip restores plaintext");
434 }
435 }
436
437 #[test]
440 fn cbc1_leaves_trailing_partial_clear() {
441 let enc = Encryptor::new(&KEY);
442 let iv = [7u8; 16];
443 let original: Vec<u8> = (0..40u8).collect();
445 let mut data = original.clone();
446 enc.encrypt(
447 Scheme::Cbc1,
448 Pattern::NONE,
449 &iv,
450 &mut data,
451 &[Subsample { clear: 3, protected: 37 }],
452 )
453 .unwrap();
454 assert_eq!(&data[..3], &original[..3], "leading clear preserved");
455 assert_ne!(&data[3..35], &original[3..35], "full blocks encrypted");
456 assert_eq!(&data[35..], &original[35..], "trailing partial block left clear");
457 }
458}