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 CENC protection schemes plus a
4//! sample [`Encryptor`] that applies them. Two schemes are implemented on top of
5//! a pure-Rust AES-128 block cipher:
6//!
7//! - **`cenc`** — AES-128 **CTR**. The keystream runs continuously across a
8//!   sample's protected byte ranges (clear ranges do not advance the counter).
9//! - **`cbcs`** — AES-128 **CBC** with pattern encryption (crypt 1 block, skip
10//!   9), CBC chaining reset to the constant IV at the start of each subsample;
11//!   a trailing partial block (< 16 bytes) is left in the clear.
12//!
13//! Encryption operates on a list of [`Subsample`] (clear/protected byte runs),
14//! so the caller (the MP4 muxer) decides the NAL-aware clear/protected split;
15//! this crate stays format-agnostic.
16
17use aes::Aes128;
18use aes::cipher::generic_array::GenericArray;
19use aes::cipher::{BlockEncrypt, KeyInit};
20use sheathe_core::{Error, Result};
21
22/// A CENC protection scheme (the `schm` `scheme_type`).
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum Scheme {
25    /// `cenc` — AES-128 CTR, (sub)sample encryption.
26    Cenc,
27    /// `cbcs` — AES-128 CBC, pattern encryption (Apple FairPlay friendly).
28    Cbcs,
29}
30
31impl Scheme {
32    /// The four-character scheme type written into the `schm` box.
33    pub fn scheme_type(self) -> [u8; 4] {
34        match self {
35            Scheme::Cenc => *b"cenc",
36            Scheme::Cbcs => *b"cbcs",
37        }
38    }
39}
40
41/// A content key plus its 16-byte Key ID (`KID`).
42#[derive(Debug, Clone)]
43pub struct ContentKey {
44    /// The 16-byte key identifier referenced by `tenc`/`pssh`.
45    pub kid: [u8; 16],
46    /// The 16-byte AES content key.
47    pub key: [u8; 16],
48}
49
50/// A contiguous run within a sample: `clear` plaintext bytes followed by
51/// `protected` bytes to encrypt.
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub struct Subsample {
54    /// Number of leading clear (unencrypted) bytes.
55    pub clear: u32,
56    /// Number of following protected (encrypted) bytes.
57    pub protected: u32,
58}
59
60/// `cbcs` pattern: encrypt 1 of every 10 sixteen-byte blocks.
61const CBCS_CRYPT_BLOCKS: usize = 1;
62const CBCS_PATTERN_BLOCKS: usize = 10;
63
64/// An AES-128 sample encryptor bound to one content key.
65pub struct Encryptor {
66    cipher: Aes128,
67}
68
69impl Encryptor {
70    /// Build an encryptor for a 16-byte AES-128 key.
71    pub fn new(key: &[u8; 16]) -> Self {
72        Self { cipher: Aes128::new(GenericArray::from_slice(key)) }
73    }
74
75    /// Encrypt `data` in place under `scheme`, treating it as the given
76    /// subsample layout. `iv` is the 16-byte per-sample initialization vector.
77    pub fn encrypt(
78        &self,
79        scheme: Scheme,
80        iv: &[u8; 16],
81        data: &mut [u8],
82        subsamples: &[Subsample],
83    ) -> Result<()> {
84        // Validate the layout covers exactly `data`.
85        let total: u64 =
86            subsamples.iter().map(|s| u64::from(s.clear) + u64::from(s.protected)).sum();
87        if total != data.len() as u64 {
88            return Err(Error::malformed("subsample layout does not cover sample"));
89        }
90        match scheme {
91            Scheme::Cenc => self.cenc(iv, data, subsamples),
92            Scheme::Cbcs => self.cbcs(iv, data, subsamples),
93        }
94        Ok(())
95    }
96
97    /// AES-128-CTR with a counter continuous across protected bytes.
98    fn cenc(&self, iv: &[u8; 16], data: &mut [u8], subsamples: &[Subsample]) {
99        let mut counter = *iv;
100        let mut keystream = [0u8; 16];
101        let mut ks_pos = 16usize; // force a fresh block on first use
102        let mut off = 0usize;
103
104        for s in subsamples {
105            off += s.clear as usize;
106            let end = off + s.protected as usize;
107            while off < end {
108                if ks_pos == 16 {
109                    keystream = counter;
110                    self.encrypt_block(&mut keystream);
111                    incr_be(&mut counter);
112                    ks_pos = 0;
113                }
114                data[off] ^= keystream[ks_pos];
115                ks_pos += 1;
116                off += 1;
117            }
118        }
119    }
120
121    /// AES-128-CBC with 1:9 pattern encryption per subsample.
122    fn cbcs(&self, iv: &[u8; 16], data: &mut [u8], subsamples: &[Subsample]) {
123        let mut off = 0usize;
124        for s in subsamples {
125            off += s.clear as usize;
126            let mut remaining = s.protected as usize;
127            let mut chain = *iv;
128            let mut block_index = 0usize;
129            while remaining >= 16 {
130                if block_index % CBCS_PATTERN_BLOCKS < CBCS_CRYPT_BLOCKS {
131                    let mut block = [0u8; 16];
132                    block.copy_from_slice(&data[off..off + 16]);
133                    for (b, c) in block.iter_mut().zip(chain.iter()) {
134                        *b ^= *c;
135                    }
136                    self.encrypt_block(&mut block);
137                    data[off..off + 16].copy_from_slice(&block);
138                    chain = block;
139                }
140                off += 16;
141                remaining -= 16;
142                block_index += 1;
143            }
144            off += remaining; // trailing partial block stays clear
145        }
146    }
147
148    /// Encrypt one 16-byte block in place (AES-128-ECB primitive).
149    fn encrypt_block(&self, block: &mut [u8; 16]) {
150        let mut ga = GenericArray::clone_from_slice(block);
151        self.cipher.encrypt_block(&mut ga);
152        block.copy_from_slice(&ga);
153    }
154}
155
156/// Increment a 16-byte big-endian counter by one (wrapping).
157fn incr_be(counter: &mut [u8; 16]) {
158    for byte in counter.iter_mut().rev() {
159        let (v, carry) = byte.overflowing_add(1);
160        *byte = v;
161        if !carry {
162            break;
163        }
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170
171    const KEY: [u8; 16] = [
172        0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f,
173        0x3c,
174    ];
175
176    fn hex(s: &str) -> Vec<u8> {
177        (0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap()).collect()
178    }
179
180    #[test]
181    fn cenc_matches_nist_ctr_vector() {
182        // NIST SP800-38A, F.5.1 (CTR-AES128.Encrypt), first block.
183        let iv = hex("f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff");
184        let mut data = hex("6bc1bee22e409f96e93d7e117393172a");
185        let enc = Encryptor::new(&KEY);
186        let subs = [Subsample { clear: 0, protected: 16 }];
187        enc.encrypt(Scheme::Cenc, iv[..].try_into().unwrap(), &mut data, &subs).unwrap();
188        assert_eq!(data, hex("874d6191b620e3261bef6864990db6ce"));
189    }
190
191    #[test]
192    fn cbcs_first_block_matches_nist_cbc_vector() {
193        // NIST SP800-38A, F.2.1 (CBC-AES128.Encrypt), first block.
194        let iv = hex("000102030405060708090a0b0c0d0e0f");
195        let mut data = hex("6bc1bee22e409f96e93d7e117393172a");
196        let enc = Encryptor::new(&KEY);
197        let subs = [Subsample { clear: 0, protected: 16 }];
198        enc.encrypt(Scheme::Cbcs, iv[..].try_into().unwrap(), &mut data, &subs).unwrap();
199        assert_eq!(data, hex("7649abac8119b246cee98e9b12e9197d"));
200    }
201
202    #[test]
203    fn cenc_leaves_clear_bytes_untouched() {
204        let iv = [0u8; 16];
205        let mut data = vec![0xAAu8; 32];
206        let enc = Encryptor::new(&KEY);
207        // 8 clear, 24 protected (8+24=32).
208        enc.encrypt(Scheme::Cenc, &iv, &mut data, &[Subsample { clear: 8, protected: 24 }])
209            .unwrap();
210        assert!(data[..8].iter().all(|&b| b == 0xAA), "clear prefix must be untouched");
211        assert!(data[8..].iter().any(|&b| b != 0xAA), "protected region must change");
212    }
213
214    #[test]
215    fn cbcs_pattern_skips_blocks() {
216        let iv = [0u8; 16];
217        // 10 blocks: only block 0 is encrypted, blocks 1..9 skipped (clear).
218        let mut data = vec![0x11u8; 160];
219        let original = data.clone();
220        let enc = Encryptor::new(&KEY);
221        enc.encrypt(Scheme::Cbcs, &iv, &mut data, &[Subsample { clear: 0, protected: 160 }])
222            .unwrap();
223        assert_ne!(data[..16], original[..16], "first block encrypted");
224        assert_eq!(data[16..], original[16..], "blocks 1..9 skipped");
225    }
226
227    #[test]
228    fn rejects_mismatched_layout() {
229        let enc = Encryptor::new(&KEY);
230        let mut data = vec![0u8; 10];
231        let err = enc.encrypt(
232            Scheme::Cenc,
233            &[0u8; 16],
234            &mut data,
235            &[Subsample { clear: 0, protected: 9 }],
236        );
237        assert!(err.is_err());
238    }
239
240    #[test]
241    fn cenc_round_trips_across_subsamples() {
242        // CTR is symmetric: encrypting the ciphertext again recovers the input,
243        // exercising the continuous counter across multiple subsamples and the
244        // preservation of clear regions.
245        let enc = Encryptor::new(&KEY);
246        let iv = [3u8; 16];
247        let subs = [Subsample { clear: 5, protected: 20 }, Subsample { clear: 10, protected: 65 }];
248        let original: Vec<u8> = (0..100u8).collect();
249        let mut data = original.clone();
250
251        enc.encrypt(Scheme::Cenc, &iv, &mut data, &subs).unwrap();
252        assert_ne!(data, original, "ciphertext must differ");
253        assert_eq!(&data[..5], &original[..5], "leading clear bytes preserved");
254
255        enc.encrypt(Scheme::Cenc, &iv, &mut data, &subs).unwrap();
256        assert_eq!(data, original, "CTR round-trip restores plaintext");
257    }
258}