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::generic_array::GenericArray;
26use aes::cipher::{BlockEncrypt, KeyInit};
27use sheathe_core::{Error, Result};
28
29mod pssh;
30pub use pssh::ProtectionSystem;
31
32/// A CENC protection scheme (the `schm` `scheme_type`).
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum Scheme {
35    /// `cenc` — AES-128 CTR, (sub)sample encryption.
36    Cenc,
37    /// `cens` — AES-128 CTR with pattern encryption.
38    Cens,
39    /// `cbc1` — AES-128 CBC, full-region (sub)sample encryption.
40    Cbc1,
41    /// `cbcs` — AES-128 CBC, pattern encryption (Apple FairPlay friendly).
42    Cbcs,
43}
44
45impl Scheme {
46    /// The four-character scheme type written into the `schm` box.
47    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    /// CBC-based schemes (`cbc1`, `cbcs`); the others are CTR (`cenc`, `cens`).
57    /// Non-pattern CBC requires 16-byte-aligned protected subsample ranges.
58    pub fn is_cbc(self) -> bool {
59        matches!(self, Scheme::Cbc1 | Scheme::Cbcs)
60    }
61
62    /// Pattern-capable schemes (`cens`, `cbcs`) — written with a version-1
63    /// `tenc` carrying the crypt/skip block counts.
64    pub fn is_pattern(self) -> bool {
65        matches!(self, Scheme::Cens | Scheme::Cbcs)
66    }
67
68    /// `cbcs` reuses one constant IV for every sample; the others derive a
69    /// unique per-sample IV.
70    pub fn uses_constant_iv(self) -> bool {
71        matches!(self, Scheme::Cbcs)
72    }
73}
74
75/// A crypt/skip block pattern (ISO/IEC 23001-7 §9.6): encrypt `crypt_blocks`
76/// 16-byte blocks, then leave `skip_blocks` blocks clear, repeating across a
77/// protected range. [`Pattern::NONE`] (`crypt_blocks == 0`) means full-region
78/// encryption — used by `cenc`/`cbc1` and for audio under `cens`/`cbcs`.
79#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80pub struct Pattern {
81    /// Number of 16-byte blocks encrypted per pattern cycle.
82    pub crypt_blocks: u8,
83    /// Number of 16-byte blocks skipped (left clear) per cycle.
84    pub skip_blocks: u8,
85}
86
87impl Pattern {
88    /// No pattern: encrypt the whole protected range.
89    pub const NONE: Pattern = Pattern { crypt_blocks: 0, skip_blocks: 0 };
90    /// The standard CMAF video pattern: encrypt 1 block, skip 9.
91    pub const VIDEO: Pattern = Pattern { crypt_blocks: 1, skip_blocks: 9 };
92
93    /// Whether this pattern leaves any blocks clear (i.e. is a real pattern).
94    fn is_patterned(self) -> bool {
95        self.crypt_blocks != 0
96    }
97}
98
99/// A content key plus its 16-byte Key ID (`KID`).
100#[derive(Debug, Clone)]
101pub struct ContentKey {
102    /// The 16-byte key identifier referenced by `tenc`/`pssh`.
103    pub kid: [u8; 16],
104    /// The 16-byte AES content key.
105    pub key: [u8; 16],
106}
107
108impl ContentKey {
109    /// The key and KID for crypto period `period`, derived by left-rotating both
110    /// by `period % 16` bytes — the naive scheme Shaka Packager uses for raw-key
111    /// rotation. Period 0 returns the key unchanged.
112    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/// A contiguous run within a sample: `clear` plaintext bytes followed by
125/// `protected` bytes to encrypt.
126#[derive(Debug, Clone, Copy, PartialEq, Eq)]
127pub struct Subsample {
128    /// Number of leading clear (unencrypted) bytes.
129    pub clear: u32,
130    /// Number of following protected (encrypted) bytes.
131    pub protected: u32,
132}
133
134/// An AES-128 sample encryptor bound to one content key.
135pub struct Encryptor {
136    cipher: Aes128,
137}
138
139impl Encryptor {
140    /// Build an encryptor for a 16-byte AES-128 key.
141    pub fn new(key: &[u8; 16]) -> Self {
142        Self { cipher: Aes128::new(GenericArray::from_slice(key)) }
143    }
144
145    /// Encrypt `data` in place under `scheme` with the given `pattern`, treating
146    /// it as the given subsample layout. `iv` is the 16-byte initialization
147    /// vector (per-sample, or the constant IV for `cbcs`). `pattern` must be
148    /// [`Pattern::NONE`] for the non-pattern schemes `cenc`/`cbc1`.
149    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        // Validate the layout covers exactly `data`.
158        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    /// AES-128-CTR (`cenc`/`cens`). The counter runs continuously over the
175    /// encrypted byte ranges of the whole sample; clear bytes and skipped
176    /// pattern blocks do not advance it. For [`Pattern::NONE`] the encrypted
177    /// ranges are the full protected ranges (`cenc`); otherwise they are the
178    /// crypt-phase blocks of each range (`cens`).
179    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; // force a fresh block on first use
183
184        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    /// AES-128-CBC (`cbc1`/`cbcs`). For [`Pattern::NONE`] (`cbc1`) the CBC chain
199    /// runs continuously across the sample's protected ranges, seeded from `iv`;
200    /// each range's trailing partial block (< 16 bytes) is left clear. For a
201    /// pattern (`cbcs`) the chain resets to the constant `iv` at the start of
202    /// each subsample and advances only over crypt-phase blocks.
203    fn cbc(&self, pattern: Pattern, iv: &[u8; 16], data: &mut [u8], subsamples: &[Subsample]) {
204        let mut off = 0usize;
205        // `cbc1` chains across subsamples; `cbcs` reuses the constant IV per
206        // subsample. A single running chain handles both: it is only reset per
207        // subsample when a pattern is in effect.
208        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; // trailing partial block stays clear
235        }
236    }
237
238    /// Encrypt one 16-byte block in place (AES-128-ECB primitive).
239    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
246/// Invoke `f(start, len)` for each contiguous byte range that gets encrypted,
247/// in order, given a subsample layout and pattern. With [`Pattern::NONE`] this
248/// is each protected range whole; with a crypt/skip pattern it is the
249/// crypt-phase blocks of each protected range. Per ISO/IEC 23001-7 §9.6, a
250/// trailing partial block (< 16 bytes) of a protected range is left in the
251/// clear under pattern encryption, so the crypt phase only covers whole blocks.
252fn 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            // Encrypt only whole 16-byte blocks; any partial block at the very
274            // end of the range stays clear.
275            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
286/// Increment a 16-byte big-endian counter by one (wrapping).
287fn 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        // NIST SP800-38A, F.5.1 (CTR-AES128.Encrypt), first block.
313        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        // NIST SP800-38A, F.2.1 (CBC-AES128.Encrypt), first block. Both `cbc1`
325        // (no pattern) and `cbcs` (1:9, first block is in the crypt phase)
326        // encrypt the first block identically.
327        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        // 8 clear, 24 protected (8+24=32).
343        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        // 10 blocks under a 1:9 pattern: only block 0 is encrypted; blocks 1..9
358        // are skipped (left clear). Holds for both `cbcs` (CBC) and `cens` (CTR).
359        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        // Period 1 left-rotates by one byte.
396        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        // Period 16 wraps back to the original.
401        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    /// Every scheme is symmetric here: CTR is self-inverse, and applying the CBC
419    /// path twice with these helpers is not — so we only round-trip the CTR
420    /// schemes via re-encryption, and check CBC via known structure elsewhere.
421    #[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    /// `cbc1` decrypts back to plaintext: decrypt the full-block portion of one
438    /// protected range and confirm the trailing partial block was left clear.
439    #[test]
440    fn cbc1_leaves_trailing_partial_clear() {
441        let enc = Encryptor::new(&KEY);
442        let iv = [7u8; 16];
443        // 37 protected bytes = 2 full blocks (32) + 5 trailing clear.
444        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}