Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

aes: rework backends #442

Merged
merged 9 commits into from
Aug 7, 2024
Merged

aes: rework backends #442

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
b286c09
aes: rework backends
newpavlov Jul 31, 2024
42728f6
temporarily allow unused macros
newpavlov Jul 31, 2024
f5cec24
Fix zeroize impl
newpavlov Jul 31, 2024
69ead99
Remove TODO comment
newpavlov Jul 31, 2024
9c7bcc4
Add ignore attribute to armv8 key expansion tests
newpavlov Jul 31, 2024
0dd075a
Change `par_size` for armv8 backends
newpavlov Jul 31, 2024
898fd85
fix
newpavlov Jul 31, 2024
5901c8a
Tweak par_size for AES-NI backends
newpavlov Aug 1, 2024
e42ac9e
Rename state to blocks
newpavlov Aug 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
178 changes: 90 additions & 88 deletions aes/src/armv8.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,20 +17,55 @@ mod expand;
#[cfg(test)]
mod test_expand;

use self::{
encdec::{decrypt1, decrypt8, encrypt1, encrypt8},
expand::{expand_key, inv_expanded_keys},
};
use crate::{Block, Block8};
use cipher::{
consts::{U16, U24, U32, U8},
inout::InOut,
AlgorithmName, BlockBackend, BlockCipher, BlockCipherDecrypt, BlockCipherEncrypt, BlockClosure,
BlockSizeUser, Key, KeyInit, KeySizeUser, ParBlocksSizeUser,
consts::{self, U16, U24, U32},
AlgorithmName, BlockCipher, BlockCipherDecrypt, BlockCipherEncrypt, BlockClosure,
BlockSizeUser, Key, KeyInit, KeySizeUser,
};
use core::arch::aarch64::*;
use core::fmt;

impl_backends!(
enc_name = Aes128BackEnc,
dec_name = Aes128BackDec,
key_size = consts::U16,
keys_ty = expand::Aes128RoundKeys,
par_size = consts::U15,
Copy link
Member Author

@newpavlov newpavlov Jul 31, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ARMv8 NEON has 32 SIMD registers, so techincally we can process 21, 19, and 17 blocks in parallel for AES-128, 192, and 256 respectively while keeping round keys in registers. But since the code forces inlining, it also balloons binary size, so additional benchmarks are needed.

Generated assembly for AES-128 looks approximately like this: https://rust.godbolt.org/z/EWzPe47c6

expand_keys = expand::expand_key,
inv_keys = expand::inv_expanded_keys,
encrypt = encdec::encrypt,
encrypt_par = encdec::encrypt_par,
decrypt = encdec::decrypt,
decrypt_par = encdec::decrypt_par,
);

impl_backends!(
enc_name = Aes192BackEnc,
dec_name = Aes192BackDec,
key_size = consts::U24,
keys_ty = expand::Aes192RoundKeys,
par_size = consts::U15,
expand_keys = expand::expand_key,
inv_keys = expand::inv_expanded_keys,
encrypt = encdec::encrypt,
encrypt_par = encdec::encrypt_par,
decrypt = encdec::decrypt,
decrypt_par = encdec::decrypt_par,
);

impl_backends!(
enc_name = Aes256BackEnc,
dec_name = Aes256BackDec,
key_size = consts::U32,
keys_ty = expand::Aes256RoundKeys,
par_size = consts::U15,
expand_keys = expand::expand_key,
inv_keys = expand::inv_expanded_keys,
encrypt = encdec::encrypt,
encrypt_par = encdec::encrypt_par,
decrypt = encdec::decrypt,
decrypt_par = encdec::decrypt_par,
);

macro_rules! define_aes_impl {
(
$name:ident,
Expand All @@ -46,19 +81,19 @@ macro_rules! define_aes_impl {
#[doc = "block cipher"]
#[derive(Clone)]
pub struct $name {
encrypt: $name_enc,
decrypt: $name_dec,
encrypt: $name_back_enc,
decrypt: $name_back_dec,
}

impl $name {
#[inline(always)]
pub(crate) fn get_enc_backend(&self) -> $name_back_enc<'_> {
self.encrypt.get_enc_backend()
pub(crate) fn get_enc_backend(&self) -> &$name_back_enc {
&self.encrypt
}

#[inline(always)]
pub(crate) fn get_dec_backend(&self) -> $name_back_dec<'_> {
self.decrypt.get_dec_backend()
pub(crate) fn get_dec_backend(&self) -> &$name_back_dec {
&self.decrypt
}
}

Expand All @@ -71,15 +106,16 @@ macro_rules! define_aes_impl {
impl KeyInit for $name {
#[inline]
fn new(key: &Key<Self>) -> Self {
let encrypt = $name_enc::new(key);
let decrypt = $name_dec::from(&encrypt);
let encrypt = $name_back_enc::new(key);
let decrypt = $name_back_dec::from(&encrypt);
Self { encrypt, decrypt }
}
}

impl From<$name_enc> for $name {
#[inline]
fn from(encrypt: $name_enc) -> $name {
let encrypt = encrypt.backend.clone();
let decrypt = (&encrypt).into();
Self { encrypt, decrypt }
}
Expand All @@ -88,8 +124,8 @@ macro_rules! define_aes_impl {
impl From<&$name_enc> for $name {
#[inline]
fn from(encrypt: &$name_enc) -> $name {
let decrypt = encrypt.into();
let encrypt = encrypt.clone();
let encrypt = encrypt.backend.clone();
let decrypt = (&encrypt).into();
Self { encrypt, decrypt }
}
}
Expand All @@ -100,13 +136,13 @@ macro_rules! define_aes_impl {

impl BlockCipherEncrypt for $name {
fn encrypt_with_backend(&self, f: impl BlockClosure<BlockSize = U16>) {
self.encrypt.encrypt_with_backend(f)
f.call(&mut &self.encrypt)
}
}

impl BlockCipherDecrypt for $name {
fn decrypt_with_backend(&self, f: impl BlockClosure<BlockSize = U16>) {
self.decrypt.decrypt_with_backend(f)
f.call(&mut &self.decrypt)
}
}

Expand All @@ -122,20 +158,30 @@ macro_rules! define_aes_impl {
}
}

impl Drop for $name {
#[inline]
fn drop(&mut self) {
#[cfg(feature = "zeroize")]
unsafe {
zeroize::zeroize_flat_type(self);
}
}
}

#[cfg(feature = "zeroize")]
impl zeroize::ZeroizeOnDrop for $name {}

#[doc=$doc]
#[doc = "block cipher (encrypt-only)"]
#[derive(Clone)]
pub struct $name_enc {
round_keys: [uint8x16_t; $rounds],
backend: $name_back_enc,
}

impl $name_enc {
#[inline(always)]
pub(crate) fn get_enc_backend(&self) -> $name_back_enc<'_> {
$name_back_enc(self)
pub(crate) fn get_enc_backend(&self) -> &$name_back_enc {
&self.backend
}
}

Expand All @@ -148,9 +194,8 @@ macro_rules! define_aes_impl {
impl KeyInit for $name_enc {
#[inline]
fn new(key: &Key<Self>) -> Self {
Self {
round_keys: unsafe { expand_key(key.as_ref()) },
}
let backend = $name_back_enc::new(key);
Self { backend }
}
}

Expand All @@ -160,7 +205,7 @@ macro_rules! define_aes_impl {

impl BlockCipherEncrypt for $name_enc {
fn encrypt_with_backend(&self, f: impl BlockClosure<BlockSize = U16>) {
f.call(&mut self.get_enc_backend())
f.call(&mut &self.backend)
}
}

Expand All @@ -180,7 +225,9 @@ macro_rules! define_aes_impl {
#[inline]
fn drop(&mut self) {
#[cfg(feature = "zeroize")]
zeroize::Zeroize::zeroize(&mut self.round_keys);
unsafe {
zeroize::zeroize_flat_type(self);
}
}
}

Expand All @@ -191,13 +238,13 @@ macro_rules! define_aes_impl {
#[doc = "block cipher (decrypt-only)"]
#[derive(Clone)]
pub struct $name_dec {
round_keys: [uint8x16_t; $rounds],
backend: $name_back_dec,
}

impl $name_dec {
#[inline(always)]
pub(crate) fn get_dec_backend(&self) -> $name_back_dec<'_> {
$name_back_dec(self)
pub(crate) fn get_dec_backend(&self) -> &$name_back_dec {
&self.backend
}
}

Expand All @@ -210,7 +257,9 @@ macro_rules! define_aes_impl {
impl KeyInit for $name_dec {
#[inline]
fn new(key: &Key<Self>) -> Self {
$name_enc::new(key).into()
let encrypt = $name_back_enc::new(key);
let backend = (&encrypt).into();
Self { backend }
}
}

Expand All @@ -222,10 +271,9 @@ macro_rules! define_aes_impl {
}

impl From<&$name_enc> for $name_dec {
fn from(enc: &$name_enc) -> $name_dec {
let mut round_keys = enc.round_keys;
unsafe { inv_expanded_keys(&mut round_keys) };
Self { round_keys }
fn from(encrypt: &$name_enc) -> $name_dec {
let backend = (&encrypt.backend).into();
Self { backend }
}
}

Expand All @@ -235,7 +283,7 @@ macro_rules! define_aes_impl {

impl BlockCipherDecrypt for $name_dec {
fn decrypt_with_backend(&self, f: impl BlockClosure<BlockSize = U16>) {
f.call(&mut self.get_dec_backend());
f.call(&mut &self.backend);
}
}

Expand All @@ -255,60 +303,14 @@ macro_rules! define_aes_impl {
#[inline]
fn drop(&mut self) {
#[cfg(feature = "zeroize")]
zeroize::Zeroize::zeroize(&mut self.round_keys);
}
}

#[cfg(feature = "zeroize")]
impl zeroize::ZeroizeOnDrop for $name_dec {}

pub(crate) struct $name_back_enc<'a>(&'a $name_enc);

impl<'a> BlockSizeUser for $name_back_enc<'a> {
type BlockSize = U16;
}

impl<'a> ParBlocksSizeUser for $name_back_enc<'a> {
type ParBlocksSize = U8;
}

impl<'a> BlockBackend for $name_back_enc<'a> {
#[inline(always)]
fn proc_block(&mut self, block: InOut<'_, '_, Block>) {
unsafe {
encrypt1(&self.0.round_keys, block);
zeroize::zeroize_flat_type(self);
}
}

#[inline(always)]
fn proc_par_blocks(&mut self, blocks: InOut<'_, '_, Block8>) {
unsafe { encrypt8(&self.0.round_keys, blocks) }
}
}

pub(crate) struct $name_back_dec<'a>(&'a $name_dec);

impl<'a> BlockSizeUser for $name_back_dec<'a> {
type BlockSize = U16;
}

impl<'a> ParBlocksSizeUser for $name_back_dec<'a> {
type ParBlocksSize = U8;
}

impl<'a> BlockBackend for $name_back_dec<'a> {
#[inline(always)]
fn proc_block(&mut self, block: InOut<'_, '_, Block>) {
unsafe {
decrypt1(&self.0.round_keys, block);
}
}

#[inline(always)]
fn proc_par_blocks(&mut self, blocks: InOut<'_, '_, Block8>) {
unsafe { decrypt8(&self.0.round_keys, blocks) }
}
}
#[cfg(feature = "zeroize")]
impl zeroize::ZeroizeOnDrop for $name_dec {}
};
}

Expand Down
Loading
Loading