Skip to content

Commit

Permalink
aes: Split FFI helpers into their own submodule.
Browse files Browse the repository at this point in the history 
Make it clearer which types are used in the FFI and start
separating the (unsafe) code for dealing with the C/assembly
implementations from the (mostly safe) rest of the code.

`git difftool HEAD^1:src/aead/aes.rs src/aead/aes/ffi.rs`.
  • Loading branch information
@briansmith
briansmith committed Jun 15, 2024
1 parent 20a7c48 commit b454dfd
Show file tree
Hide file tree
Showing 2 changed files with 238 additions and 179 deletions.
206 changes: 27 additions & 179 deletions src/aead/aes.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright 2018 Brian Smith.
// Copyright 2018-2024 Brian Smith.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
Expand All @@ -13,155 +13,25 @@
// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

use super::{nonce::Nonce, quic::Sample, NONCE_LEN};
use crate::{bits::BitLength, c, constant_time, cpu, error, polyfill::slice};
use core::{num::NonZeroUsize, ops::RangeFrom};
use crate::{constant_time, cpu, error};
use cfg_if::cfg_if;
use core::ops::RangeFrom;

#[derive(Clone)]
pub(super) struct Key {
inner: AES_KEY,
}
pub(super) use ffi::Counter;
#[macro_use]
mod ffi;

// SAFETY:
// * The function `$name` must read `bits` bits from `user_key`; `bits` will
// always be a valid AES key length, i.e. a whole number of bytes.
// * `$name` must set `key.rounds` to the value expected by the corresponding
// encryption/decryption functions and return 0, or otherwise must return
// non-zero to indicate failure.
// * `$name` may inspect CPU features.
//
// In BoringSSL, the C prototypes for these are in
// crypto/fipsmodule/aes/internal.h.
macro_rules! set_encrypt_key {
( $name:ident, $key_bytes:expr, $key:expr, $cpu_features:expr ) => {{
prefixed_extern! {
fn $name(user_key: *const u8, bits: BitLength<c::int>, key: *mut AES_KEY) -> c::int;
}
set_encrypt_key($name, $key_bytes, $key, $cpu_features)
}};
}

#[inline]
unsafe fn set_encrypt_key(
f: unsafe extern "C" fn(*const u8, BitLength<c::int>, *mut AES_KEY) -> c::int,
bytes: KeyBytes<'_>,
key: &mut AES_KEY,
_cpu_features: cpu::Features,
) -> Result<(), error::Unspecified> {
let (bytes, key_bits) = match bytes {
KeyBytes::AES_128(bytes) => (&bytes[..], BitLength::from_bits(128)),
KeyBytes::AES_256(bytes) => (&bytes[..], BitLength::from_bits(256)),
};

// Unusually, in this case zero means success and non-zero means failure.
if 0 == unsafe { f(bytes.as_ptr(), key_bits, key) } {
debug_assert_ne!(key.rounds, 0); // Sanity check initialization.
Ok(())
cfg_if! {
if #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] {
pub(super) use ffi::AES_KEY;
} else {
Err(error::Unspecified)
use ffi::AES_KEY;
}
}

macro_rules! encrypt_block {
($name:ident, $block:expr, $key:expr) => {{
prefixed_extern! {
fn $name(a: &Block, r: *mut Block, key: &AES_KEY);
}
encrypt_block_($name, $block, $key)
}};
}

#[inline]
fn encrypt_block_(
f: unsafe extern "C" fn(&Block, *mut Block, &AES_KEY),
a: Block,
key: &Key,
) -> Block {
let mut result = core::mem::MaybeUninit::uninit();
unsafe {
f(&a, result.as_mut_ptr(), &key.inner);
result.assume_init()
}
}

/// SAFETY:
/// * The caller must ensure that `$key` was initialized with the
/// `set_encrypt_key!` invocation that `$name` requires.
/// * The caller must ensure that fhe function `$name` satisfies the conditions
/// for the `f` parameter to `ctr32_encrypt_blocks`.
macro_rules! ctr32_encrypt_blocks {
($name:ident, $in_out:expr, $src:expr, $key:expr, $ctr:expr, $cpu_features:expr ) => {{
prefixed_extern! {
fn $name(
input: *const [u8; BLOCK_LEN],
output: *mut [u8; BLOCK_LEN],
blocks: c::NonZero_size_t,
key: &AES_KEY,
ivec: &Counter,
);
}
ctr32_encrypt_blocks($name, $in_out, $src, $key, $ctr, $cpu_features)
}};
}

/// SAFETY:
/// * `f` must not read more than `blocks` blocks from `input`.
/// * `f` must write exactly `block` blocks to `output`.
/// * In particular, `f` must handle blocks == 0 without reading from `input`
/// or writing to `output`.
/// * `f` must support the input overlapping with the output exactly or
/// with any nonnegative offset `n` (i.e. `input == output.add(n)`);
/// `f` does NOT need to support the cases where input < output.
/// * `key` must have been initialized with the `set_encrypt_key!` invocation
/// that corresponds to `f`.
/// * `f` may inspect CPU features.
#[inline]
unsafe fn ctr32_encrypt_blocks(
f: unsafe extern "C" fn(
input: *const [u8; BLOCK_LEN],
output: *mut [u8; BLOCK_LEN],
blocks: c::NonZero_size_t,
key: &AES_KEY,
ivec: &Counter,
),
in_out: &mut [u8],
src: RangeFrom<usize>,
key: &AES_KEY,
ctr: &mut Counter,
cpu_features: cpu::Features,
) {
let (input, leftover) = slice::as_chunks(&in_out[src]);
debug_assert_eq!(leftover.len(), 0);

let blocks = match NonZeroUsize::new(input.len()) {
Some(blocks) => blocks,
None => {
return;
}
};

let blocks_u32: u32 = blocks.get().try_into().unwrap();

let input = input.as_ptr();
let output: *mut [u8; BLOCK_LEN] = in_out.as_mut_ptr().cast();

let _: cpu::Features = cpu_features;

// SAFETY:
// * `input` points to `blocks` blocks.
// * `output` points to space for `blocks` blocks to be written.
// * input == output.add(n), where n == src.start, and the caller is
// responsible for ensuing this sufficient for `f` to work correctly.
// * The caller is responsible for ensuring `f` can handle any value of
// `blocks` including zero.
// * The caller is responsible for ensuring `key` was initialized by the
// `set_encrypt_key!` invocation required by `f`.
// * CPU feature detection has been done so `f` can inspect
// CPU features.
unsafe {
f(input, output, blocks, key, ctr);
}

ctr.increment_by_less_safe(blocks_u32);
#[derive(Clone)]
pub(super) struct Key {
inner: AES_KEY,
}

impl Key {
Expand All @@ -170,17 +40,12 @@ impl Key {
bytes: KeyBytes<'_>,
cpu_features: cpu::Features,
) -> Result<Self, error::Unspecified> {
let mut key = AES_KEY {
rd_key: [0u32; 4 * (MAX_ROUNDS + 1)],
rounds: 0,
};

match detect_implementation(cpu_features) {
let key = match detect_implementation(cpu_features) {
#[cfg(any(target_arch = "aarch64", target_arch = "x86_64", target_arch = "x86"))]
// SAFETY: `aes_hw_set_encrypt_key` satisfies the `set_encrypt_key!`
// contract for these target architectures.
Implementation::HWAES => unsafe {
set_encrypt_key!(aes_hw_set_encrypt_key, bytes, &mut key, cpu_features)?;
set_encrypt_key!(aes_hw_set_encrypt_key, bytes, cpu_features)
},

#[cfg(any(
Expand All @@ -192,15 +57,15 @@ impl Key {
// SAFETY: `vpaes_set_encrypt_key` satisfies the `set_encrypt_key!`
// contract for these target architectures.
Implementation::VPAES_BSAES => unsafe {
set_encrypt_key!(vpaes_set_encrypt_key, bytes, &mut key, cpu_features)?;
set_encrypt_key!(vpaes_set_encrypt_key, bytes, cpu_features)
},

// SAFETY: `aes_nohw_set_encrypt_key` satisfies the `set_encrypt_key!`
// contract.
Implementation::NOHW => unsafe {
set_encrypt_key!(aes_nohw_set_encrypt_key, bytes, &mut key, cpu_features)?;
set_encrypt_key!(aes_nohw_set_encrypt_key, bytes, cpu_features)
},
};
}?;

Ok(Self { inner: key })
}
Expand All @@ -218,9 +83,9 @@ impl Key {

// `encrypt_iv_xor_block` calls `encrypt_block` on `target_arch = "x86"`.
#[cfg(target_arch = "x86")]
Implementation::VPAES_BSAES => encrypt_block!(vpaes_encrypt, a, self),
Implementation::VPAES_BSAES => unsafe { encrypt_block!(vpaes_encrypt, a, &self.inner) },

Implementation::NOHW => encrypt_block!(aes_nohw_encrypt, a, self),
Implementation::NOHW => unsafe { encrypt_block!(aes_nohw_encrypt, a, &self.inner) },
}
}

Expand Down Expand Up @@ -372,17 +237,6 @@ impl Key {
}
}

// Keep this in sync with AES_KEY in aes.h.
#[repr(C)]
#[derive(Clone)]
pub(super) struct AES_KEY {
pub rd_key: [u32; 4 * (MAX_ROUNDS + 1)],
pub rounds: c::uint,
}

// Keep this in sync with `AES_MAXNR` in aes.h.
const MAX_ROUNDS: usize = 14;

pub const AES_128_KEY_LEN: usize = 128 / 8;
pub const AES_256_KEY_LEN: usize = 256 / 8;

Expand All @@ -391,10 +245,8 @@ pub enum KeyBytes<'a> {
AES_256(&'a [u8; AES_256_KEY_LEN]),
}

/// nonce || big-endian counter.
#[repr(transparent)]
pub(super) struct Counter([u8; BLOCK_LEN]);

// `Counter` is `ffi::Counter` as its representation is dictated by its use in
// the FFI.
impl Counter {
pub fn one(nonce: Nonce) -> Self {
let mut value = [0u8; BLOCK_LEN];
Expand Down Expand Up @@ -525,22 +377,18 @@ unsafe fn bsaes_ctr32_encrypt_blocks_with_vpaes_key(
fn vpaes_encrypt_key_to_bsaes(bsaes_key: *mut AES_KEY, vpaes_key: &AES_KEY);
}

let mut bsaes_key = AES_KEY {
rd_key: [0u32; 4 * (MAX_ROUNDS + 1)],
rounds: 0,
};
// SAFETY:
// * The caller ensures `vpaes_key` was initialized by
// `vpaes_set_encrypt_key`.
// * `bsaes_key was zeroed above, and `vpaes_encrypt_key_to_bsaes`
// is assumed to initialize `bsaes_key`.
unsafe {
vpaes_encrypt_key_to_bsaes(&mut bsaes_key, vpaes_key);
}
let bsaes_key =
unsafe { AES_KEY::derive(vpaes_encrypt_key_to_bsaes, &vpaes_key, cpu_features) };

// The code for `vpaes_encrypt_key_to_bsaes` notes "vpaes stores one
// fewer round count than bsaes, but the number of keys is the same,"
// so use this as a sanity check.
debug_assert_eq!(bsaes_key.rounds, vpaes_key.rounds + 1);
debug_assert_eq!(bsaes_key.rounds(), vpaes_key.rounds() + 1);

// SAFETY:
// * `bsaes_key` is in bsaes format after calling
Expand Down
Loading

0 comments on commit b454dfd

Please sign in to comment.