1use aes::Aes128;
18use aes::cipher::generic_array::GenericArray;
19use aes::cipher::{BlockEncrypt, KeyInit};
20use sheathe_core::{Error, Result};
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum Scheme {
25 Cenc,
27 Cbcs,
29}
30
31impl Scheme {
32 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#[derive(Debug, Clone)]
43pub struct ContentKey {
44 pub kid: [u8; 16],
46 pub key: [u8; 16],
48}
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub struct Subsample {
54 pub clear: u32,
56 pub protected: u32,
58}
59
60const CBCS_CRYPT_BLOCKS: usize = 1;
62const CBCS_PATTERN_BLOCKS: usize = 10;
63
64pub struct Encryptor {
66 cipher: Aes128,
67}
68
69impl Encryptor {
70 pub fn new(key: &[u8; 16]) -> Self {
72 Self { cipher: Aes128::new(GenericArray::from_slice(key)) }
73 }
74
75 pub fn encrypt(
78 &self,
79 scheme: Scheme,
80 iv: &[u8; 16],
81 data: &mut [u8],
82 subsamples: &[Subsample],
83 ) -> Result<()> {
84 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 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; 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 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; }
146 }
147
148 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
156fn 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 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 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 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 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 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}