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::{counter, iv::Iv, quic::Sample, Block, Direction, BLOCK_LEN};
16 use crate::{bits::BitLength, c, cpu, endian::*, error, polyfill};
17
18 pub(crate) struct Key {
19 inner: AES_KEY,
20 cpu_features: cpu::Features,
21 }
22
23 macro_rules! set_encrypt_key {
24 ( $name:ident, $bytes:expr, $key_bits:expr, $key:expr ) => {{
25 extern "C" {
26 fn $name(user_key: *const u8, bits: c::uint, key: &mut AES_KEY) -> c::int;
27 }
28 set_encrypt_key($name, $bytes, $key_bits, $key)
29 }};
30 }
31
32 #[inline]
set_encrypt_key( f: unsafe extern "C" fn(*const u8, c::uint, &mut AES_KEY) -> c::int, bytes: &[u8], key_bits: BitLength, key: &mut AES_KEY, ) -> Result<(), error::Unspecified>33 fn set_encrypt_key(
34 f: unsafe extern "C" fn(*const u8, c::uint, &mut AES_KEY) -> c::int,
35 bytes: &[u8],
36 key_bits: BitLength,
37 key: &mut AES_KEY,
38 ) -> Result<(), error::Unspecified> {
39 // Unusually, in this case zero means success and non-zero means failure.
40 if 0 == unsafe { f(bytes.as_ptr(), key_bits.as_usize_bits() as c::uint, key) } {
41 Ok(())
42 } else {
43 Err(error::Unspecified)
44 }
45 }
46
47 macro_rules! encrypt_block {
48 ($name:ident, $block:expr, $key:expr) => {{
49 extern "C" {
50 fn $name(a: &Block, r: *mut Block, key: &AES_KEY);
51 }
52 encrypt_block_($name, $block, $key)
53 }};
54 }
55
56 #[inline]
encrypt_block_( f: unsafe extern "C" fn(&Block, *mut Block, &AES_KEY), a: Block, key: &Key, ) -> Block57 fn encrypt_block_(
58 f: unsafe extern "C" fn(&Block, *mut Block, &AES_KEY),
59 a: Block,
60 key: &Key,
61 ) -> Block {
62 let mut result = core::mem::MaybeUninit::uninit();
63 unsafe {
64 f(&a, result.as_mut_ptr(), &key.inner);
65 result.assume_init()
66 }
67 }
68
69 macro_rules! ctr32_encrypt_blocks {
70 ($name:ident, $in_out:expr, $in_prefix_len:expr, $key:expr, $ivec:expr ) => {{
71 extern "C" {
72 fn $name(
73 input: *const u8,
74 output: *mut u8,
75 blocks: c::size_t,
76 key: &AES_KEY,
77 ivec: &Counter,
78 );
79 }
80 ctr32_encrypt_blocks_($name, $in_out, $in_prefix_len, $key, $ivec)
81 }};
82 }
83
84 #[inline]
ctr32_encrypt_blocks_( f: unsafe extern "C" fn( input: *const u8, output: *mut u8, blocks: c::size_t, key: &AES_KEY, ivec: &Counter, ), in_out: &mut [u8], in_prefix_len: usize, key: &AES_KEY, ctr: &mut Counter, )85 fn ctr32_encrypt_blocks_(
86 f: unsafe extern "C" fn(
87 input: *const u8,
88 output: *mut u8,
89 blocks: c::size_t,
90 key: &AES_KEY,
91 ivec: &Counter,
92 ),
93 in_out: &mut [u8],
94 in_prefix_len: usize,
95 key: &AES_KEY,
96 ctr: &mut Counter,
97 ) {
98 let in_out_len = in_out.len().checked_sub(in_prefix_len).unwrap();
99 assert_eq!(in_out_len % BLOCK_LEN, 0);
100
101 let blocks = in_out_len / BLOCK_LEN;
102 let blocks_u32 = blocks as u32;
103 assert_eq!(blocks, polyfill::usize_from_u32(blocks_u32));
104
105 let input = in_out[in_prefix_len..].as_ptr();
106 let output = in_out.as_mut_ptr();
107
108 unsafe {
109 f(input, output, blocks, &key, ctr);
110 }
111 ctr.increment_by_less_safe(blocks_u32);
112 }
113
114 impl Key {
115 #[inline]
new( bytes: &[u8], variant: Variant, cpu_features: cpu::Features, ) -> Result<Self, error::Unspecified>116 pub fn new(
117 bytes: &[u8],
118 variant: Variant,
119 cpu_features: cpu::Features,
120 ) -> Result<Self, error::Unspecified> {
121 let key_bits = match variant {
122 Variant::AES_128 => BitLength::from_usize_bits(128),
123 Variant::AES_256 => BitLength::from_usize_bits(256),
124 };
125 if BitLength::from_usize_bytes(bytes.len())? != key_bits {
126 return Err(error::Unspecified);
127 }
128
129 let mut key = AES_KEY {
130 rd_key: [0u32; 4 * (MAX_ROUNDS + 1)],
131 rounds: 0,
132 };
133
134 match detect_implementation(cpu_features) {
135 #[cfg(any(
136 target_arch = "aarch64",
137 target_arch = "arm",
138 target_arch = "x86_64",
139 target_arch = "x86"
140 ))]
141 Implementation::HWAES => {
142 set_encrypt_key!(GFp_aes_hw_set_encrypt_key, bytes, key_bits, &mut key)?
143 }
144
145 #[cfg(any(
146 target_arch = "aarch64",
147 target_arch = "arm",
148 target_arch = "x86_64",
149 target_arch = "x86"
150 ))]
151 Implementation::VPAES_BSAES => {
152 set_encrypt_key!(GFp_vpaes_set_encrypt_key, bytes, key_bits, &mut key)?
153 }
154
155 #[cfg(not(target_arch = "aarch64"))]
156 Implementation::NOHW => {
157 set_encrypt_key!(GFp_aes_nohw_set_encrypt_key, bytes, key_bits, &mut key)?
158 }
159 };
160
161 Ok(Self {
162 inner: key,
163 cpu_features,
164 })
165 }
166
167 #[inline]
encrypt_block(&self, a: Block) -> Block168 pub fn encrypt_block(&self, a: Block) -> Block {
169 match detect_implementation(self.cpu_features) {
170 #[cfg(any(
171 target_arch = "aarch64",
172 target_arch = "arm",
173 target_arch = "x86_64",
174 target_arch = "x86"
175 ))]
176 Implementation::HWAES => encrypt_block!(GFp_aes_hw_encrypt, a, self),
177
178 #[cfg(any(
179 target_arch = "aarch64",
180 target_arch = "arm",
181 target_arch = "x86_64",
182 target_arch = "x86"
183 ))]
184 Implementation::VPAES_BSAES => encrypt_block!(GFp_vpaes_encrypt, a, self),
185
186 #[cfg(not(target_arch = "aarch64"))]
187 Implementation::NOHW => encrypt_block!(GFp_aes_nohw_encrypt, a, self),
188 }
189 }
190
191 #[inline]
encrypt_iv_xor_block(&self, iv: Iv, input: Block) -> Block192 pub fn encrypt_iv_xor_block(&self, iv: Iv, input: Block) -> Block {
193 let mut output = self.encrypt_block(Block::from(&iv.into_bytes_less_safe()));
194 output.bitxor_assign(input);
195 output
196 }
197
198 #[inline]
ctr32_encrypt_blocks( &self, in_out: &mut [u8], direction: Direction, ctr: &mut Counter, )199 pub(super) fn ctr32_encrypt_blocks(
200 &self,
201 in_out: &mut [u8],
202 direction: Direction,
203 ctr: &mut Counter,
204 ) {
205 let in_prefix_len = match direction {
206 Direction::Opening { in_prefix_len } => in_prefix_len,
207 Direction::Sealing => 0,
208 };
209
210 let in_out_len = in_out.len().checked_sub(in_prefix_len).unwrap();
211
212 assert_eq!(in_out_len % BLOCK_LEN, 0);
213
214 match detect_implementation(self.cpu_features) {
215 #[cfg(any(
216 target_arch = "aarch64",
217 target_arch = "arm",
218 target_arch = "x86_64",
219 target_arch = "x86"
220 ))]
221 Implementation::HWAES => ctr32_encrypt_blocks!(
222 GFp_aes_hw_ctr32_encrypt_blocks,
223 in_out,
224 in_prefix_len,
225 &self.inner,
226 ctr
227 ),
228
229 #[cfg(any(target_arch = "aarch64", target_arch = "arm", target_arch = "x86_64"))]
230 Implementation::VPAES_BSAES => {
231 // 8 blocks is the cut-off point where it's faster to use BSAES.
232 #[cfg(target_arch = "arm")]
233 let in_out = if in_out_len >= 8 * BLOCK_LEN {
234 let remainder = in_out_len % (8 * BLOCK_LEN);
235 let bsaes_in_out_len = if remainder < (4 * BLOCK_LEN) {
236 in_out_len - remainder
237 } else {
238 in_out_len
239 };
240
241 let mut bsaes_key = AES_KEY {
242 rd_key: [0u32; 4 * (MAX_ROUNDS + 1)],
243 rounds: 0,
244 };
245 extern "C" {
246 fn GFp_vpaes_encrypt_key_to_bsaes(
247 bsaes_key: &mut AES_KEY,
248 vpaes_key: &AES_KEY,
249 );
250 }
251 unsafe {
252 GFp_vpaes_encrypt_key_to_bsaes(&mut bsaes_key, &self.inner);
253 }
254 ctr32_encrypt_blocks!(
255 GFp_bsaes_ctr32_encrypt_blocks,
256 &mut in_out[..(bsaes_in_out_len + in_prefix_len)],
257 in_prefix_len,
258 &bsaes_key,
259 ctr
260 );
261
262 &mut in_out[bsaes_in_out_len..]
263 } else {
264 in_out
265 };
266
267 ctr32_encrypt_blocks!(
268 GFp_vpaes_ctr32_encrypt_blocks,
269 in_out,
270 in_prefix_len,
271 &self.inner,
272 ctr
273 )
274 }
275
276 #[cfg(any(target_arch = "x86"))]
277 Implementation::VPAES_BSAES => {
278 super::shift::shift_full_blocks(in_out, in_prefix_len, |input| {
279 self.encrypt_iv_xor_block(ctr.increment(), Block::from(input))
280 });
281 }
282
283 #[cfg(not(target_arch = "aarch64"))]
284 Implementation::NOHW => ctr32_encrypt_blocks!(
285 GFp_aes_nohw_ctr32_encrypt_blocks,
286 in_out,
287 in_prefix_len,
288 &self.inner,
289 ctr
290 ),
291 }
292 }
293
new_mask(&self, sample: Sample) -> [u8; 5]294 pub fn new_mask(&self, sample: Sample) -> [u8; 5] {
295 let block = self.encrypt_block(Block::from(&sample));
296
297 let mut out: [u8; 5] = [0; 5];
298 out.copy_from_slice(&block.as_ref()[..5]);
299
300 out
301 }
302
303 // TODO: use `matches!` when MSRV increases to 1.42.0 and remove this
304 // `#[allow(...)]`
305 #[allow(clippy::unknown_clippy_lints)]
306 #[allow(clippy::match_like_matches_macro)]
307 #[cfg(target_arch = "x86_64")]
308 #[must_use]
is_aes_hw(&self) -> bool309 pub fn is_aes_hw(&self) -> bool {
310 match detect_implementation(self.cpu_features) {
311 Implementation::HWAES => true,
312 _ => false,
313 }
314 }
315
316 #[cfg(target_arch = "x86_64")]
317 #[must_use]
inner_less_safe(&self) -> &AES_KEY318 pub(super) fn inner_less_safe(&self) -> &AES_KEY {
319 &self.inner
320 }
321 }
322
323 // Keep this in sync with AES_KEY in aes.h.
324 #[repr(C)]
325 pub(super) struct AES_KEY {
326 pub rd_key: [u32; 4 * (MAX_ROUNDS + 1)],
327 pub rounds: c::uint,
328 }
329
330 // Keep this in sync with `AES_MAXNR` in aes.h.
331 const MAX_ROUNDS: usize = 14;
332
333 pub enum Variant {
334 AES_128,
335 AES_256,
336 }
337
338 pub type Counter = counter::Counter<BigEndian<u32>>;
339
340 #[repr(C)] // Only so `Key` can be `#[repr(C)]`
341 #[derive(Clone, Copy)]
342 pub enum Implementation {
343 #[cfg(any(
344 target_arch = "aarch64",
345 target_arch = "arm",
346 target_arch = "x86_64",
347 target_arch = "x86"
348 ))]
349 HWAES = 1,
350
351 // On "arm" only, this indicates that the bsaes implementation may be used.
352 #[cfg(any(
353 target_arch = "aarch64",
354 target_arch = "arm",
355 target_arch = "x86_64",
356 target_arch = "x86"
357 ))]
358 VPAES_BSAES = 2,
359
360 #[cfg(not(target_arch = "aarch64"))]
361 NOHW = 3,
362 }
363
detect_implementation(cpu_features: cpu::Features) -> Implementation364 fn detect_implementation(cpu_features: cpu::Features) -> Implementation {
365 // `cpu_features` is only used for specific platforms.
366 #[cfg(not(any(
367 target_arch = "aarch64",
368 target_arch = "arm",
369 target_arch = "x86_64",
370 target_arch = "x86"
371 )))]
372 let _cpu_features = cpu_features;
373
374 #[cfg(any(
375 target_arch = "aarch64",
376 target_arch = "arm",
377 target_arch = "x86_64",
378 target_arch = "x86"
379 ))]
380 {
381 if cpu::intel::AES.available(cpu_features) || cpu::arm::AES.available(cpu_features) {
382 return Implementation::HWAES;
383 }
384 }
385
386 #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
387 {
388 if cpu::intel::SSSE3.available(cpu_features) {
389 return Implementation::VPAES_BSAES;
390 }
391 }
392
393 #[cfg(target_arch = "arm")]
394 {
395 if cpu::arm::NEON.available(cpu_features) {
396 return Implementation::VPAES_BSAES;
397 }
398 }
399
400 #[cfg(target_arch = "aarch64")]
401 {
402 Implementation::VPAES_BSAES
403 }
404
405 #[cfg(not(target_arch = "aarch64"))]
406 {
407 Implementation::NOHW
408 }
409 }
410
411 #[cfg(test)]
412 mod tests {
413 use super::{super::BLOCK_LEN, *};
414 use crate::test;
415 use core::convert::TryInto;
416
417 #[test]
test_aes()418 pub fn test_aes() {
419 test::run(test_file!("aes_tests.txt"), |section, test_case| {
420 assert_eq!(section, "");
421 let key = consume_key(test_case, "Key");
422 let input = test_case.consume_bytes("Input");
423 let input: &[u8; BLOCK_LEN] = input.as_slice().try_into()?;
424 let expected_output = test_case.consume_bytes("Output");
425
426 let block = Block::from(input);
427 let output = key.encrypt_block(block);
428 assert_eq!(output.as_ref(), &expected_output[..]);
429
430 Ok(())
431 })
432 }
433
consume_key(test_case: &mut test::TestCase, name: &str) -> Key434 fn consume_key(test_case: &mut test::TestCase, name: &str) -> Key {
435 let key = test_case.consume_bytes(name);
436 let variant = match key.len() {
437 16 => Variant::AES_128,
438 32 => Variant::AES_256,
439 _ => unreachable!(),
440 };
441 Key::new(&key[..], variant, cpu::features()).unwrap()
442 }
443 }
444