1 | // Copyright 2018 Brian Smith. |
2 | // |
3 | // Permission to use, copy, modify, and/or distribute this software for any |
4 | // purpose with or without fee is hereby granted, provided that the above |
5 | // copyright notice and this permission notice appear in all copies. |
6 | // |
7 | // THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHORS DISCLAIM ALL WARRANTIES |
8 | // WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF |
9 | // MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY |
10 | // SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES |
11 | // WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION |
12 | // OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN |
13 | // CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. |
14 | |
15 | use super::{ |
16 | block::{Block, BLOCK_LEN}, |
17 | nonce::Nonce, |
18 | quic::Sample, |
19 | }; |
20 | use crate::{ |
21 | bits::{BitLength, FromUsizeBytes}, |
22 | c, cpu, |
23 | endian::BigEndian, |
24 | error, |
25 | polyfill::{self, ArraySplitMap}, |
26 | }; |
27 | use core::ops::RangeFrom; |
28 | |
29 | #[derive (Clone)] |
30 | pub(super) struct Key { |
31 | inner: AES_KEY, |
32 | } |
33 | |
34 | macro_rules! set_encrypt_key { |
35 | ( $name:ident, $bytes:expr, $key_bits:expr, $key:expr ) => {{ |
36 | prefixed_extern! { |
37 | fn $name(user_key: *const u8, bits: c::uint, key: &mut AES_KEY) -> c::int; |
38 | } |
39 | set_encrypt_key($name, $bytes, $key_bits, $key) |
40 | }}; |
41 | } |
42 | |
43 | #[inline ] |
44 | fn set_encrypt_key( |
45 | f: unsafe extern "C" fn(*const u8, c::uint, &mut AES_KEY) -> c::int, |
46 | bytes: &[u8], |
47 | key_bits: BitLength, |
48 | key: &mut AES_KEY, |
49 | ) -> Result<(), error::Unspecified> { |
50 | // Unusually, in this case zero means success and non-zero means failure. |
51 | #[allow (clippy::cast_possible_truncation)] |
52 | if 0 == unsafe { f(bytes.as_ptr(), key_bits.as_bits() as c::uint, key) } { |
53 | Ok(()) |
54 | } else { |
55 | Err(error::Unspecified) |
56 | } |
57 | } |
58 | |
59 | macro_rules! encrypt_block { |
60 | ($name:ident, $block:expr, $key:expr) => {{ |
61 | prefixed_extern! { |
62 | fn $name(a: &Block, r: *mut Block, key: &AES_KEY); |
63 | } |
64 | encrypt_block_($name, $block, $key) |
65 | }}; |
66 | } |
67 | |
68 | #[inline ] |
69 | fn encrypt_block_( |
70 | f: unsafe extern "C" fn(&Block, *mut Block, &AES_KEY), |
71 | a: Block, |
72 | key: &Key, |
73 | ) -> Block { |
74 | let mut result: MaybeUninit = core::mem::MaybeUninit::uninit(); |
75 | unsafe { |
76 | f(&a, result.as_mut_ptr(), &key.inner); |
77 | result.assume_init() |
78 | } |
79 | } |
80 | |
81 | macro_rules! ctr32_encrypt_blocks { |
82 | ($name:ident, $in_out:expr, $src:expr, $key:expr, $ivec:expr ) => {{ |
83 | prefixed_extern! { |
84 | fn $name( |
85 | input: *const [u8; BLOCK_LEN], |
86 | output: *mut [u8; BLOCK_LEN], |
87 | blocks: c::size_t, |
88 | key: &AES_KEY, |
89 | ivec: &Counter, |
90 | ); |
91 | } |
92 | ctr32_encrypt_blocks_($name, $in_out, $src, $key, $ivec) |
93 | }}; |
94 | } |
95 | |
96 | #[inline ] |
97 | fn ctr32_encrypt_blocks_( |
98 | f: unsafe extern "C" fn( |
99 | input: *const [u8; BLOCK_LEN], |
100 | output: *mut [u8; BLOCK_LEN], |
101 | blocks: c::size_t, |
102 | key: &AES_KEY, |
103 | ivec: &Counter, |
104 | ), |
105 | in_out: &mut [u8], |
106 | src: RangeFrom<usize>, |
107 | key: &AES_KEY, |
108 | ctr: &mut Counter, |
109 | ) { |
110 | let in_out_len: usize = in_out[src.clone()].len(); |
111 | assert_eq!(in_out_len % BLOCK_LEN, 0); |
112 | |
113 | let blocks: usize = in_out_len / BLOCK_LEN; |
114 | #[allow (clippy::cast_possible_truncation)] |
115 | let blocks_u32: u32 = blocks as u32; |
116 | assert_eq!(blocks, polyfill::usize_from_u32(blocks_u32)); |
117 | |
118 | let input: *const [u8; 16] = in_out[src].as_ptr().cast::<[u8; BLOCK_LEN]>(); |
119 | let output: *mut [u8; 16] = in_out.as_mut_ptr().cast::<[u8; BLOCK_LEN]>(); |
120 | |
121 | unsafe { |
122 | f(input, output, blocks, key, ctr); |
123 | } |
124 | ctr.increment_by_less_safe(increment_by:blocks_u32); |
125 | } |
126 | |
127 | impl Key { |
128 | #[inline ] |
129 | pub fn new( |
130 | bytes: &[u8], |
131 | variant: Variant, |
132 | cpu_features: cpu::Features, |
133 | ) -> Result<Self, error::Unspecified> { |
134 | let key_bits = match variant { |
135 | Variant::AES_128 => BitLength::from_usize_bits(128), |
136 | Variant::AES_256 => BitLength::from_usize_bits(256), |
137 | }; |
138 | if BitLength::from_usize_bytes(bytes.len())? != key_bits { |
139 | return Err(error::Unspecified); |
140 | } |
141 | |
142 | let mut key = AES_KEY { |
143 | rd_key: [0u32; 4 * (MAX_ROUNDS + 1)], |
144 | rounds: 0, |
145 | }; |
146 | |
147 | match detect_implementation(cpu_features) { |
148 | #[cfg (any( |
149 | target_arch = "aarch64" , |
150 | target_arch = "arm" , |
151 | target_arch = "x86_64" , |
152 | target_arch = "x86" |
153 | ))] |
154 | Implementation::HWAES => { |
155 | set_encrypt_key!(aes_hw_set_encrypt_key, bytes, key_bits, &mut key)? |
156 | } |
157 | |
158 | #[cfg (any( |
159 | target_arch = "aarch64" , |
160 | target_arch = "arm" , |
161 | target_arch = "x86_64" , |
162 | target_arch = "x86" |
163 | ))] |
164 | Implementation::VPAES_BSAES => { |
165 | set_encrypt_key!(vpaes_set_encrypt_key, bytes, key_bits, &mut key)? |
166 | } |
167 | |
168 | Implementation::NOHW => { |
169 | set_encrypt_key!(aes_nohw_set_encrypt_key, bytes, key_bits, &mut key)? |
170 | } |
171 | }; |
172 | |
173 | Ok(Self { inner: key }) |
174 | } |
175 | |
176 | #[inline ] |
177 | pub fn encrypt_block(&self, a: Block, cpu_features: cpu::Features) -> Block { |
178 | match detect_implementation(cpu_features) { |
179 | #[cfg (any( |
180 | target_arch = "aarch64" , |
181 | target_arch = "arm" , |
182 | target_arch = "x86_64" , |
183 | target_arch = "x86" |
184 | ))] |
185 | Implementation::HWAES => encrypt_block!(aes_hw_encrypt, a, self), |
186 | |
187 | #[cfg (any( |
188 | target_arch = "aarch64" , |
189 | target_arch = "arm" , |
190 | target_arch = "x86_64" , |
191 | target_arch = "x86" |
192 | ))] |
193 | Implementation::VPAES_BSAES => encrypt_block!(vpaes_encrypt, a, self), |
194 | |
195 | Implementation::NOHW => encrypt_block!(aes_nohw_encrypt, a, self), |
196 | } |
197 | } |
198 | |
199 | #[inline ] |
200 | pub fn encrypt_iv_xor_block(&self, iv: Iv, input: Block, cpu_features: cpu::Features) -> Block { |
201 | let encrypted_iv = self.encrypt_block(iv.into_block_less_safe(), cpu_features); |
202 | encrypted_iv ^ input |
203 | } |
204 | |
205 | #[inline ] |
206 | pub(super) fn ctr32_encrypt_within( |
207 | &self, |
208 | in_out: &mut [u8], |
209 | src: RangeFrom<usize>, |
210 | ctr: &mut Counter, |
211 | cpu_features: cpu::Features, |
212 | ) { |
213 | let in_out_len = in_out[src.clone()].len(); |
214 | |
215 | assert_eq!(in_out_len % BLOCK_LEN, 0); |
216 | |
217 | match detect_implementation(cpu_features) { |
218 | #[cfg (any( |
219 | target_arch = "aarch64" , |
220 | target_arch = "arm" , |
221 | target_arch = "x86_64" , |
222 | target_arch = "x86" |
223 | ))] |
224 | Implementation::HWAES => { |
225 | ctr32_encrypt_blocks!(aes_hw_ctr32_encrypt_blocks, in_out, src, &self.inner, ctr) |
226 | } |
227 | |
228 | #[cfg (any(target_arch = "aarch64" , target_arch = "arm" , target_arch = "x86_64" ))] |
229 | Implementation::VPAES_BSAES => { |
230 | // 8 blocks is the cut-off point where it's faster to use BSAES. |
231 | #[cfg (target_arch = "arm" )] |
232 | let in_out = if in_out_len >= 8 * BLOCK_LEN { |
233 | let remainder = in_out_len % (8 * BLOCK_LEN); |
234 | let bsaes_in_out_len = if remainder < (4 * BLOCK_LEN) { |
235 | in_out_len - remainder |
236 | } else { |
237 | in_out_len |
238 | }; |
239 | |
240 | let mut bsaes_key = AES_KEY { |
241 | rd_key: [0u32; 4 * (MAX_ROUNDS + 1)], |
242 | rounds: 0, |
243 | }; |
244 | prefixed_extern! { |
245 | fn vpaes_encrypt_key_to_bsaes(bsaes_key: &mut AES_KEY, vpaes_key: &AES_KEY); |
246 | } |
247 | unsafe { |
248 | vpaes_encrypt_key_to_bsaes(&mut bsaes_key, &self.inner); |
249 | } |
250 | ctr32_encrypt_blocks!( |
251 | bsaes_ctr32_encrypt_blocks, |
252 | &mut in_out[..(src.start + bsaes_in_out_len)], |
253 | src.clone(), |
254 | &bsaes_key, |
255 | ctr |
256 | ); |
257 | |
258 | &mut in_out[bsaes_in_out_len..] |
259 | } else { |
260 | in_out |
261 | }; |
262 | |
263 | ctr32_encrypt_blocks!(vpaes_ctr32_encrypt_blocks, in_out, src, &self.inner, ctr) |
264 | } |
265 | |
266 | #[cfg (target_arch = "x86" )] |
267 | Implementation::VPAES_BSAES => { |
268 | super::shift::shift_full_blocks(in_out, src, |input| { |
269 | self.encrypt_iv_xor_block(ctr.increment(), Block::from(input), cpu_features) |
270 | }); |
271 | } |
272 | |
273 | Implementation::NOHW => { |
274 | ctr32_encrypt_blocks!(aes_nohw_ctr32_encrypt_blocks, in_out, src, &self.inner, ctr) |
275 | } |
276 | } |
277 | } |
278 | |
279 | pub fn new_mask(&self, sample: Sample) -> [u8; 5] { |
280 | let block = self.encrypt_block(Block::from(&sample), cpu::features()); |
281 | |
282 | let mut out: [u8; 5] = [0; 5]; |
283 | out.copy_from_slice(&block.as_ref()[..5]); |
284 | |
285 | out |
286 | } |
287 | |
288 | #[cfg (any(target_arch = "x86_64" , target_arch = "aarch64" ))] |
289 | #[must_use ] |
290 | pub fn is_aes_hw(&self, cpu_features: cpu::Features) -> bool { |
291 | matches!(detect_implementation(cpu_features), Implementation::HWAES) |
292 | } |
293 | |
294 | #[cfg (any(target_arch = "x86_64" , target_arch = "aarch64" ))] |
295 | #[must_use ] |
296 | pub(super) fn inner_less_safe(&self) -> &AES_KEY { |
297 | &self.inner |
298 | } |
299 | } |
300 | |
301 | // Keep this in sync with AES_KEY in aes.h. |
302 | #[repr (C)] |
303 | #[derive (Clone)] |
304 | pub(super) struct AES_KEY { |
305 | pub rd_key: [u32; 4 * (MAX_ROUNDS + 1)], |
306 | pub rounds: c::uint, |
307 | } |
308 | |
309 | // Keep this in sync with `AES_MAXNR` in aes.h. |
310 | const MAX_ROUNDS: usize = 14; |
311 | |
312 | pub enum Variant { |
313 | AES_128, |
314 | AES_256, |
315 | } |
316 | |
317 | /// Nonce || Counter, all big-endian. |
318 | #[repr (transparent)] |
319 | pub(super) struct Counter([BigEndian<u32>; 4]); |
320 | |
321 | impl Counter { |
322 | pub fn one(nonce: Nonce) -> Self { |
323 | let [n0: BigEndian, n1: BigEndian, n2: BigEndian] = nonce.as_ref().array_split_map(BigEndian::<u32>::from); |
324 | Self([n0, n1, n2, 1.into()]) |
325 | } |
326 | |
327 | pub fn increment(&mut self) -> Iv { |
328 | let iv: [[u8; 4]; 4] = self.0.map(Into::into); |
329 | let iv: Iv = Iv(Block::from(iv)); |
330 | self.increment_by_less_safe(increment_by:1); |
331 | iv |
332 | } |
333 | |
334 | fn increment_by_less_safe(&mut self, increment_by: u32) { |
335 | let old_value: u32 = self.0[3].into(); |
336 | self.0[3] = (old_value + increment_by).into(); |
337 | } |
338 | } |
339 | |
340 | /// The IV for a single block encryption. |
341 | /// |
342 | /// Intentionally not `Clone` to ensure each is used only once. |
343 | pub struct Iv(Block); |
344 | |
345 | impl From<Counter> for Iv { |
346 | fn from(counter: Counter) -> Self { |
347 | let iv: [[u8; 4]; 4] = counter.0.map(Into::into); |
348 | Self(Block::from(iv)) |
349 | } |
350 | } |
351 | |
352 | impl Iv { |
353 | /// "Less safe" because it defeats attempts to use the type system to prevent reuse of the IV. |
354 | #[inline ] |
355 | pub(super) fn into_block_less_safe(self) -> Block { |
356 | self.0 |
357 | } |
358 | } |
359 | |
360 | #[repr (C)] // Only so `Key` can be `#[repr(C)]` |
361 | #[derive (Clone, Copy)] |
362 | #[allow (clippy::upper_case_acronyms)] |
363 | pub enum Implementation { |
364 | #[cfg (any( |
365 | target_arch = "aarch64" , |
366 | target_arch = "arm" , |
367 | target_arch = "x86_64" , |
368 | target_arch = "x86" |
369 | ))] |
370 | HWAES = 1, |
371 | |
372 | // On "arm" only, this indicates that the bsaes implementation may be used. |
373 | #[cfg (any( |
374 | target_arch = "aarch64" , |
375 | target_arch = "arm" , |
376 | target_arch = "x86_64" , |
377 | target_arch = "x86" |
378 | ))] |
379 | VPAES_BSAES = 2, |
380 | |
381 | NOHW = 3, |
382 | } |
383 | |
384 | fn detect_implementation(cpu_features: cpu::Features) -> Implementation { |
385 | // `cpu_features` is only used for specific platforms. |
386 | #[cfg (not(any( |
387 | target_arch = "aarch64" , |
388 | target_arch = "arm" , |
389 | target_arch = "x86_64" , |
390 | target_arch = "x86" |
391 | )))] |
392 | let _cpu_features = cpu_features; |
393 | |
394 | #[cfg (any(target_arch = "aarch64" , target_arch = "arm" ))] |
395 | { |
396 | if cpu::arm::AES.available(cpu_features) { |
397 | return Implementation::HWAES; |
398 | } |
399 | } |
400 | |
401 | #[cfg (any(target_arch = "x86_64" , target_arch = "x86" ))] |
402 | { |
403 | if cpu::intel::AES.available(cpu_features) { |
404 | return Implementation::HWAES; |
405 | } |
406 | } |
407 | |
408 | #[cfg (any(target_arch = "x86_64" , target_arch = "x86" ))] |
409 | { |
410 | if cpu::intel::SSSE3.available(cpu_features) { |
411 | return Implementation::VPAES_BSAES; |
412 | } |
413 | } |
414 | |
415 | #[cfg (any(target_arch = "aarch64" , target_arch = "arm" ))] |
416 | { |
417 | if cpu::arm::NEON.available(cpu_features) { |
418 | return Implementation::VPAES_BSAES; |
419 | } |
420 | } |
421 | |
422 | { |
423 | Implementation::NOHW |
424 | } |
425 | } |
426 | |
427 | #[cfg (test)] |
428 | mod tests { |
429 | use super::*; |
430 | use crate::test; |
431 | |
432 | #[test ] |
433 | pub fn test_aes() { |
434 | let cpu_features = cpu::features(); |
435 | test::run(test_file!("aes_tests.txt" ), |section, test_case| { |
436 | assert_eq!(section, "" ); |
437 | let key = consume_key(test_case , "Key" ); |
438 | let input = test_case .consume_bytes("Input" ); |
439 | let input: &[u8; BLOCK_LEN] = input.as_slice().try_into()?; |
440 | let expected_output = test_case .consume_bytes("Output" ); |
441 | |
442 | let block = Block::from(input); |
443 | let output = key.encrypt_block(block, cpu_features); |
444 | assert_eq!(output.as_ref(), &expected_output[..]); |
445 | |
446 | Ok(()) |
447 | }) |
448 | } |
449 | |
450 | fn consume_key(test_case: &mut test::TestCase, name: &str) -> Key { |
451 | let key = test_case .consume_bytes(name); |
452 | let variant = match key.len() { |
453 | 16 => Variant::AES_128, |
454 | 32 => Variant::AES_256, |
455 | _ => unreachable!(), |
456 | }; |
457 | Key::new(&key[..], variant, cpu::features()).unwrap() |
458 | } |
459 | } |
460 | |