Skip to main content

sheathe_crypto/
lib.rs

1//! Common Encryption (ISO/IEC 23001-7) for **sheathe**.
2//!
3//! Mirrors Shaka Packager's `media/crypto`: the four CENC protection schemes
4//! plus a sample [`Encryptor`] that applies them, all on top of a pure-Rust
5//! AES-128 block cipher.
6//!
7//! - **`cenc`** — AES-128 **CTR**, full-region. The keystream runs continuously
8//!   across a sample's protected byte ranges (clear ranges do not advance the
9//!   counter).
10//! - **`cens`** — AES-128 **CTR** with [pattern](Pattern) encryption: only the
11//!   crypt-phase blocks consume keystream; skipped blocks pass through clear,
12//!   the counter continuing across them.
13//! - **`cbc1`** — AES-128 **CBC**, full-region. CBC chaining runs continuously
14//!   across the sample (clear ranges are skipped); a trailing partial block
15//!   (< 16 bytes) of each protected range is left in the clear.
16//! - **`cbcs`** — AES-128 **CBC** with pattern encryption, chaining reset to the
17//!   constant IV at the start of each subsample; trailing partial blocks clear.
18//!
19//! Pattern encryption (`cens`/`cbcs`) is, by convention, applied to video only;
20//! audio uses [`Pattern::NONE`] (full-region) even under those schemes. The
21//! caller decides the per-track pattern and the NAL-aware clear/protected split
22//! via the [`Subsample`] list, so this crate stays format-agnostic.
23
24use aes::Aes128;
25use aes::cipher::{BlockCipherEncrypt, KeyInit};
26use sheathe_core::{Error, Result};
27
28mod pssh;
29pub use pssh::ProtectionSystem;
30
31/// A CENC protection scheme (the `schm` `scheme_type`).
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum Scheme {
34    /// `cenc` — AES-128 CTR, (sub)sample encryption.
35    Cenc,
36    /// `cens` — AES-128 CTR with pattern encryption.
37    Cens,
38    /// `cbc1` — AES-128 CBC, full-region (sub)sample encryption.
39    Cbc1,
40    /// `cbcs` — AES-128 CBC, pattern encryption (Apple FairPlay friendly).
41    Cbcs,
42}
43
44impl Scheme {
45    /// The four-character scheme type written into the `schm` box.
46    pub fn scheme_type(self) -> [u8; 4] {
47        match self {
48            Scheme::Cenc => *b"cenc",
49            Scheme::Cens => *b"cens",
50            Scheme::Cbc1 => *b"cbc1",
51            Scheme::Cbcs => *b"cbcs",
52        }
53    }
54
55    /// CBC-based schemes (`cbc1`, `cbcs`); the others are CTR (`cenc`, `cens`).
56    /// Non-pattern CBC requires 16-byte-aligned protected subsample ranges.
57    pub fn is_cbc(self) -> bool {
58        matches!(self, Scheme::Cbc1 | Scheme::Cbcs)
59    }
60
61    /// Pattern-capable schemes (`cens`, `cbcs`) — written with a version-1
62    /// `tenc` carrying the crypt/skip block counts.
63    pub fn is_pattern(self) -> bool {
64        matches!(self, Scheme::Cens | Scheme::Cbcs)
65    }
66
67    /// `cbcs` reuses one constant IV for every sample; the others derive a
68    /// unique per-sample IV.
69    pub fn uses_constant_iv(self) -> bool {
70        matches!(self, Scheme::Cbcs)
71    }
72}
73
74/// A crypt/skip block pattern (ISO/IEC 23001-7 §9.6): encrypt `crypt_blocks`
75/// 16-byte blocks, then leave `skip_blocks` blocks clear, repeating across a
76/// protected range. [`Pattern::NONE`] (`crypt_blocks == 0`) means full-region
77/// encryption — used by `cenc`/`cbc1` and for audio under `cens`/`cbcs`.
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub struct Pattern {
80    /// Number of 16-byte blocks encrypted per pattern cycle.
81    pub crypt_blocks: u8,
82    /// Number of 16-byte blocks skipped (left clear) per cycle.
83    pub skip_blocks: u8,
84}
85
86impl Pattern {
87    /// No pattern: encrypt the whole protected range.
88    pub const NONE: Pattern = Pattern { crypt_blocks: 0, skip_blocks: 0 };
89    /// The standard CMAF video pattern: encrypt 1 block, skip 9.
90    pub const VIDEO: Pattern = Pattern { crypt_blocks: 1, skip_blocks: 9 };
91
92    /// Whether this pattern leaves any blocks clear (i.e. is a real pattern).
93    fn is_patterned(self) -> bool {
94        self.crypt_blocks != 0
95    }
96}
97
98/// A content key plus its 16-byte Key ID (`KID`).
99#[derive(Debug, Clone)]
100pub struct ContentKey {
101    /// The 16-byte key identifier referenced by `tenc`/`pssh`.
102    pub kid: [u8; 16],
103    /// The 16-byte AES content key.
104    pub key: [u8; 16],
105}
106
107impl ContentKey {
108    /// The key and KID for crypto period `period`, derived by left-rotating both
109    /// by `period % 16` bytes — the naive scheme Shaka Packager uses for raw-key
110    /// rotation. Period 0 returns the key unchanged.
111    pub fn rotated(&self, period: u32) -> ContentKey {
112        let n = (period % 16) as usize;
113        let mut kid = [0u8; 16];
114        let mut key = [0u8; 16];
115        for i in 0..16 {
116            kid[i] = self.kid[(i + n) % 16];
117            key[i] = self.key[(i + n) % 16];
118        }
119        ContentKey { kid, key }
120    }
121}
122
123/// A contiguous run within a sample: `clear` plaintext bytes followed by
124/// `protected` bytes to encrypt.
125#[derive(Debug, Clone, Copy, PartialEq, Eq)]
126pub struct Subsample {
127    /// Number of leading clear (unencrypted) bytes.
128    pub clear: u32,
129    /// Number of following protected (encrypted) bytes.
130    pub protected: u32,
131}
132
133/// An AES-128 sample encryptor bound to one content key.
134pub struct Encryptor {
135    cipher: Aes128,
136}
137
138impl Encryptor {
139    /// Build an encryptor for a 16-byte AES-128 key.
140    pub fn new(key: &[u8; 16]) -> Self {
141        Self { cipher: Aes128::new_from_slice(key).expect("AES-128 key is 16 bytes") }
142    }
143
144    /// Encrypt `data` in place under `scheme` with the given `pattern`, treating
145    /// it as the given subsample layout. `iv` is the 16-byte initialization
146    /// vector (per-sample, or the constant IV for `cbcs`). `pattern` must be
147    /// [`Pattern::NONE`] for the non-pattern schemes `cenc`/`cbc1`.
148    pub fn encrypt(
149        &self,
150        scheme: Scheme,
151        pattern: Pattern,
152        iv: &[u8; 16],
153        data: &mut [u8],
154        subsamples: &[Subsample],
155    ) -> Result<()> {
156        // Validate the layout covers exactly `data`.
157        let total: u64 =
158            subsamples.iter().map(|s| u64::from(s.clear) + u64::from(s.protected)).sum();
159        if total != data.len() as u64 {
160            return Err(Error::malformed("subsample layout does not cover sample"));
161        }
162        if pattern.is_patterned() && !scheme.is_pattern() {
163            return Err(Error::malformed("pattern set on a non-pattern scheme"));
164        }
165        if scheme.is_cbc() {
166            self.cbc(pattern, iv, data, subsamples);
167        } else {
168            self.ctr(pattern, iv, data, subsamples);
169        }
170        Ok(())
171    }
172
173    /// AES-128-CTR (`cenc`/`cens`). The counter runs continuously over the
174    /// encrypted byte ranges of the whole sample; clear bytes and skipped
175    /// pattern blocks do not advance it. For [`Pattern::NONE`] the encrypted
176    /// ranges are the full protected ranges (`cenc`); otherwise they are the
177    /// crypt-phase blocks of each range (`cens`).
178    fn ctr(&self, pattern: Pattern, iv: &[u8; 16], data: &mut [u8], subsamples: &[Subsample]) {
179        let mut counter = *iv;
180        let mut keystream = [0u8; 16];
181        let mut ks_pos = 16usize; // force a fresh block on first use
182
183        for_each_crypt_range(pattern, subsamples, |start, len| {
184            for byte in &mut data[start..start + len] {
185                if ks_pos == 16 {
186                    keystream = counter;
187                    self.encrypt_block(&mut keystream);
188                    incr_be(&mut counter);
189                    ks_pos = 0;
190                }
191                *byte ^= keystream[ks_pos];
192                ks_pos += 1;
193            }
194        });
195    }
196
197    /// AES-128-CBC (`cbc1`/`cbcs`). For [`Pattern::NONE`] (`cbc1`) the CBC chain
198    /// runs continuously across the sample's protected ranges, seeded from `iv`;
199    /// each range's trailing partial block (< 16 bytes) is left clear. For a
200    /// pattern (`cbcs`) the chain resets to the constant `iv` at the start of
201    /// each subsample and advances only over crypt-phase blocks.
202    fn cbc(&self, pattern: Pattern, iv: &[u8; 16], data: &mut [u8], subsamples: &[Subsample]) {
203        let mut off = 0usize;
204        // `cbc1` chains across subsamples; `cbcs` reuses the constant IV per
205        // subsample. A single running chain handles both: it is only reset per
206        // subsample when a pattern is in effect.
207        let mut chain = *iv;
208        for s in subsamples {
209            off += s.clear as usize;
210            if pattern.is_patterned() {
211                chain = *iv;
212            }
213            let mut remaining = s.protected as usize;
214            let mut block_index = 0usize;
215            let cycle = pattern.crypt_blocks as usize + pattern.skip_blocks as usize;
216            while remaining >= 16 {
217                let encrypt =
218                    !pattern.is_patterned() || block_index % cycle < pattern.crypt_blocks as usize;
219                if encrypt {
220                    let mut block = [0u8; 16];
221                    block.copy_from_slice(&data[off..off + 16]);
222                    for (b, c) in block.iter_mut().zip(chain.iter()) {
223                        *b ^= *c;
224                    }
225                    self.encrypt_block(&mut block);
226                    data[off..off + 16].copy_from_slice(&block);
227                    chain = block;
228                }
229                off += 16;
230                remaining -= 16;
231                block_index += 1;
232            }
233            off += remaining; // trailing partial block stays clear
234        }
235    }
236
237    /// Encrypt one 16-byte block in place (AES-128-ECB primitive).
238    fn encrypt_block(&self, block: &mut [u8; 16]) {
239        let mut ga = (*block).into();
240        self.cipher.encrypt_block(&mut ga);
241        block.copy_from_slice(&ga);
242    }
243}
244
245/// Invoke `f(start, len)` for each contiguous byte range that gets encrypted,
246/// in order, given a subsample layout and pattern. With [`Pattern::NONE`] this
247/// is each protected range whole; with a crypt/skip pattern it is the
248/// crypt-phase blocks of each protected range. Per ISO/IEC 23001-7 §9.6, a
249/// trailing partial block (< 16 bytes) of a protected range is left in the
250/// clear under pattern encryption, so the crypt phase only covers whole blocks.
251fn for_each_crypt_range(
252    pattern: Pattern,
253    subsamples: &[Subsample],
254    mut f: impl FnMut(usize, usize),
255) {
256    let mut off = 0usize;
257    for s in subsamples {
258        off += s.clear as usize;
259        let protected = s.protected as usize;
260        if !pattern.is_patterned() {
261            if protected > 0 {
262                f(off, protected);
263            }
264            off += protected;
265            continue;
266        }
267        let crypt = pattern.crypt_blocks as usize * 16;
268        let skip = pattern.skip_blocks as usize * 16;
269        let mut pos = 0usize;
270        while pos < protected {
271            let phase = crypt.min(protected - pos);
272            // Encrypt only whole 16-byte blocks; any partial block at the very
273            // end of the range stays clear.
274            let whole = phase - phase % 16;
275            if whole > 0 {
276                f(off + pos, whole);
277            }
278            pos += phase;
279            pos += skip.min(protected - pos);
280        }
281        off += protected;
282    }
283}
284
285/// Increment a 16-byte big-endian counter by one (wrapping).
286fn incr_be(counter: &mut [u8; 16]) {
287    for byte in counter.iter_mut().rev() {
288        let (v, carry) = byte.overflowing_add(1);
289        *byte = v;
290        if !carry {
291            break;
292        }
293    }
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299
300    const KEY: [u8; 16] = [
301        0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f,
302        0x3c,
303    ];
304
305    fn hex(s: &str) -> Vec<u8> {
306        (0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap()).collect()
307    }
308
309    #[test]
310    fn cenc_matches_nist_ctr_vector() {
311        // NIST SP800-38A, F.5.1 (CTR-AES128.Encrypt), first block.
312        let iv = hex("f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff");
313        let mut data = hex("6bc1bee22e409f96e93d7e117393172a");
314        let enc = Encryptor::new(&KEY);
315        let subs = [Subsample { clear: 0, protected: 16 }];
316        enc.encrypt(Scheme::Cenc, Pattern::NONE, iv[..].try_into().unwrap(), &mut data, &subs)
317            .unwrap();
318        assert_eq!(data, hex("874d6191b620e3261bef6864990db6ce"));
319    }
320
321    #[test]
322    fn cbc_schemes_match_nist_cbc_vector() {
323        // NIST SP800-38A, F.2.1 (CBC-AES128.Encrypt), first block. Both `cbc1`
324        // (no pattern) and `cbcs` (1:9, first block is in the crypt phase)
325        // encrypt the first block identically.
326        let iv = hex("000102030405060708090a0b0c0d0e0f");
327        let subs = [Subsample { clear: 0, protected: 16 }];
328        for (scheme, pattern) in [(Scheme::Cbc1, Pattern::NONE), (Scheme::Cbcs, Pattern::VIDEO)] {
329            let mut data = hex("6bc1bee22e409f96e93d7e117393172a");
330            let enc = Encryptor::new(&KEY);
331            enc.encrypt(scheme, pattern, iv[..].try_into().unwrap(), &mut data, &subs).unwrap();
332            assert_eq!(data, hex("7649abac8119b246cee98e9b12e9197d"), "{scheme:?}");
333        }
334    }
335
336    #[test]
337    fn cenc_leaves_clear_bytes_untouched() {
338        let iv = [0u8; 16];
339        let mut data = vec![0xAAu8; 32];
340        let enc = Encryptor::new(&KEY);
341        // 8 clear, 24 protected (8+24=32).
342        enc.encrypt(
343            Scheme::Cenc,
344            Pattern::NONE,
345            &iv,
346            &mut data,
347            &[Subsample { clear: 8, protected: 24 }],
348        )
349        .unwrap();
350        assert!(data[..8].iter().all(|&b| b == 0xAA), "clear prefix must be untouched");
351        assert!(data[8..].iter().any(|&b| b != 0xAA), "protected region must change");
352    }
353
354    #[test]
355    fn pattern_schemes_skip_blocks() {
356        // 10 blocks under a 1:9 pattern: only block 0 is encrypted; blocks 1..9
357        // are skipped (left clear). Holds for both `cbcs` (CBC) and `cens` (CTR).
358        let iv = [0u8; 16];
359        for scheme in [Scheme::Cbcs, Scheme::Cens] {
360            let mut data = vec![0x11u8; 160];
361            let original = data.clone();
362            let enc = Encryptor::new(&KEY);
363            enc.encrypt(
364                scheme,
365                Pattern::VIDEO,
366                &iv,
367                &mut data,
368                &[Subsample { clear: 0, protected: 160 }],
369            )
370            .unwrap();
371            assert_ne!(data[..16], original[..16], "{scheme:?}: first block encrypted");
372            assert_eq!(data[16..], original[16..], "{scheme:?}: blocks 1..9 skipped");
373        }
374    }
375
376    #[test]
377    fn rejects_mismatched_layout() {
378        let enc = Encryptor::new(&KEY);
379        let mut data = vec![0u8; 10];
380        let err = enc.encrypt(
381            Scheme::Cenc,
382            Pattern::NONE,
383            &[0u8; 16],
384            &mut data,
385            &[Subsample { clear: 0, protected: 9 }],
386        );
387        assert!(err.is_err());
388    }
389
390    #[test]
391    fn content_key_rotation_left_rotates_by_period() {
392        let base = ContentKey { kid: KEY, key: KEY };
393        assert_eq!(base.rotated(0).kid, KEY, "period 0 is unchanged");
394        // Period 1 left-rotates by one byte.
395        let mut expect = KEY;
396        expect.rotate_left(1);
397        assert_eq!(base.rotated(1).kid, expect);
398        assert_eq!(base.rotated(1).key, expect);
399        // Period 16 wraps back to the original.
400        assert_eq!(base.rotated(16).kid, KEY);
401    }
402
403    #[test]
404    fn rejects_pattern_on_non_pattern_scheme() {
405        let enc = Encryptor::new(&KEY);
406        let mut data = vec![0u8; 16];
407        let err = enc.encrypt(
408            Scheme::Cenc,
409            Pattern::VIDEO,
410            &[0u8; 16],
411            &mut data,
412            &[Subsample { clear: 0, protected: 16 }],
413        );
414        assert!(err.is_err());
415    }
416
417    /// Every scheme is symmetric here: CTR is self-inverse, and applying the CBC
418    /// path twice with these helpers is not — so we only round-trip the CTR
419    /// schemes via re-encryption, and check CBC via known structure elsewhere.
420    #[test]
421    fn ctr_schemes_round_trip_across_subsamples() {
422        let enc = Encryptor::new(&KEY);
423        let iv = [3u8; 16];
424        let subs = [Subsample { clear: 5, protected: 40 }, Subsample { clear: 10, protected: 65 }];
425        for (scheme, pattern) in [(Scheme::Cenc, Pattern::NONE), (Scheme::Cens, Pattern::VIDEO)] {
426            let original: Vec<u8> = (0..120u8).collect();
427            let mut data = original.clone();
428            enc.encrypt(scheme, pattern, &iv, &mut data, &subs).unwrap();
429            assert_ne!(data, original, "{scheme:?}: ciphertext must differ");
430            assert_eq!(&data[..5], &original[..5], "{scheme:?}: leading clear bytes preserved");
431            enc.encrypt(scheme, pattern, &iv, &mut data, &subs).unwrap();
432            assert_eq!(data, original, "{scheme:?}: CTR round-trip restores plaintext");
433        }
434    }
435
436    /// `cbc1` decrypts back to plaintext: decrypt the full-block portion of one
437    /// protected range and confirm the trailing partial block was left clear.
438    #[test]
439    fn cbc1_leaves_trailing_partial_clear() {
440        let enc = Encryptor::new(&KEY);
441        let iv = [7u8; 16];
442        // 37 protected bytes = 2 full blocks (32) + 5 trailing clear.
443        let original: Vec<u8> = (0..40u8).collect();
444        let mut data = original.clone();
445        enc.encrypt(
446            Scheme::Cbc1,
447            Pattern::NONE,
448            &iv,
449            &mut data,
450            &[Subsample { clear: 3, protected: 37 }],
451        )
452        .unwrap();
453        assert_eq!(&data[..3], &original[..3], "leading clear preserved");
454        assert_ne!(&data[3..35], &original[3..35], "full blocks encrypted");
455        assert_eq!(&data[35..], &original[35..], "trailing partial block left clear");
456    }
457}