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
15use super::{
16 block::{Block, BLOCK_LEN},
17 nonce::Nonce,
18 quic::Sample,
19};
20use crate::{
21 bits::{BitLength, FromUsizeBytes},
22 c, cpu,
23 endian::BigEndian,
24 error,
25 polyfill::{self, ArraySplitMap},
26};
27use core::ops::RangeFrom;
28
29#[derive(Clone)]
30pub(super) struct Key {
31 inner: AES_KEY,
32}
33
34macro_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]
44fn 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
59macro_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]
69fn 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
81macro_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]
97fn 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
127impl 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)]
304pub(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.
310const MAX_ROUNDS: usize = 14;
311
312pub enum Variant {
313 AES_128,
314 AES_256,
315}
316
317/// Nonce || Counter, all big-endian.
318#[repr(transparent)]
319pub(super) struct Counter([BigEndian<u32>; 4]);
320
321impl 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.
343pub struct Iv(Block);
344
345impl 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
352impl 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)]
363pub 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
384fn 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)]
428mod 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