Compiler Explorer

Source code

use core::{
    arch::{aarch64::*, asm},
    mem, slice
};

pub type Block = [u8; 16];
pub type Block8 = [Block; 8];

/// There are 4 AES words in a block.
const BLOCK_WORDS: usize = 4;

/// The AES (nee Rijndael) notion of a word is always 32-bits, or 4-bytes.
const WORD_SIZE: usize = 4;

/// AES round constants.
const ROUND_CONSTS: [u32; 10] = [0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36];

#[inline(always)]
pub unsafe fn vaeseq_u8(mut data: uint8x16_t, key: uint8x16_t) -> uint8x16_t {
    asm!(
        "AESE {d:v}.16B, {k:v}.16B",
        d = inout(vreg) data,
        k = in(vreg) key,
        options(pure, nomem, nostack, preserves_flags)
    );
    data
}

#[inline(always)]
pub unsafe fn vaesdq_u8(mut data: uint8x16_t, key: uint8x16_t) -> uint8x16_t {
    asm!(
        "AESD {d:v}.16B, {k:v}.16B",
        d = inout(vreg) data,
        k = in(vreg) key,
        options(pure, nomem, nostack, preserves_flags)
    );
    data
}

#[inline(always)]
pub unsafe fn vaesmcq_u8(mut data: uint8x16_t) -> uint8x16_t {
    asm!(
        "AESMC {d:v}.16B, {d:v}.16B",
        d = inout(vreg) data,
        options(pure, nomem, nostack, preserves_flags)
    );
    data
}

#[inline(always)]
pub unsafe fn vaesimcq_u8(mut data: uint8x16_t) -> uint8x16_t {
    asm!(
        "AESIMC {d:v}.16B, {d:v}.16B",
        d = inout(vreg) data,
        options(pure, nomem, nostack, preserves_flags)
    );
    data
}

#[inline(always)]
pub unsafe fn vaeseq_u8_and_vaesmcq_u8(mut data: uint8x16_t, key: uint8x16_t) -> uint8x16_t {
    asm!(
        "AESE {d:v}.16B, {k:v}.16B",
        "AESMC {d:v}.16B, {d:v}.16B",
        d = inout(vreg) data,
        k = in(vreg) key,
        options(pure, nomem, nostack, preserves_flags)
    );
    data
}

#[inline(always)]
pub unsafe fn vaesdq_u8_and_vaesimcq_u8(
    mut data: uint8x16_t,
    key: uint8x16_t,
) -> uint8x16_t {
    asm!(
        "AESD {d:v}.16B, {k:v}.16B",
        "AESIMC {d:v}.16B, {d:v}.16B",
        d = inout(vreg) data,
        k = in(vreg) key,
        options(pure, nomem, nostack, preserves_flags)
    );
    data
}

#[target_feature(enable = "aes")]
pub unsafe fn expand_key<const L: usize, const N: usize>(key: &[u8; L]) -> [uint8x16_t; N] {
    assert!((L == 16 && N == 11) || (L == 24 && N == 13) || (L == 32 && N == 15));

let mut expanded_keys: [uint8x16_t; N] = mem::zeroed();

// TODO(tarcieri): construct expanded keys using `vreinterpretq_u8_u32`
    let ek_words = slice::from_raw_parts_mut(expanded_keys.as_mut_ptr() as *mut u32, N * BLOCK_WORDS);

for (i, chunk) in key.chunks_exact(WORD_SIZE).enumerate() {
        ek_words[i] = u32::from_ne_bytes(chunk.try_into().unwrap());
    }

// From "The Rijndael Block Cipher" Section 4.1:
    // > The number of columns of the Cipher Key is denoted by `Nk` and is
    // > equal to the key length divided by 32 [bits].
    let nk = L / WORD_SIZE;

for i in nk..(N * BLOCK_WORDS) {
        let mut word = ek_words[i - 1];

if i % nk == 0 {
            word = sub_word(word).rotate_right(8) ^ ROUND_CONSTS[i / nk - 1];
        } else if nk > 6 && i % nk == 4 {
            word = sub_word(word)
        }

ek_words[i] = ek_words[i - nk] ^ word;
    }

expanded_keys
}

#[inline(always)]
pub fn inv_expanded_keys<const N: usize>(expanded_keys: &mut [uint8x16_t; N]) {
    assert!(N == 11 || N == 13 || N == 15);

for ek in expanded_keys.iter_mut().take(N - 1).skip(1) {
        unsafe { *ek = vaesimcq_u8(*ek) }
    }

expanded_keys.reverse();
}

/// Sub bytes for a single AES word: used for key expansion.
#[inline(always)]
pub fn sub_word(input: u32) -> u32 {
    unsafe {
        let input = vreinterpretq_u8_u32(vdupq_n_u32(input));

// AES single round encryption (with a "round" key of all zeros)
        let sub_input = vaeseq_u8(input, vdupq_n_u8(0));

vgetq_lane_u32(vreinterpretq_u32_u8(sub_input), 0)
    }
}

#[target_feature(enable = "aes")]
#[target_feature(enable = "neon")]
pub unsafe fn encrypt1<const N: usize>(
    expanded_keys: &[uint8x16_t; N],
    block: &mut Block,
) {
    let rounds = N - 1;
    assert!(rounds == 10 || rounds == 12 || rounds == 14);

let mut state = vld1q_u8(block.as_ptr());

for k in expanded_keys.iter().take(rounds - 1) {
        // AES single round encryption and mix columns
        state = vaeseq_u8_and_vaesmcq_u8(state, *k);
    }

// AES single round encryption
    state = vaeseq_u8(state, expanded_keys[rounds - 1]);

// Final add (bitwise XOR)
    state = veorq_u8(state, expanded_keys[rounds]);

vst1q_u8(block.as_mut_ptr(), state);
}

#[target_feature(enable = "aes")]
#[target_feature(enable = "neon")]
pub unsafe fn encrypt8<const N: usize>(
    expanded_keys: &[uint8x16_t; N],
    blocks: &mut Block8,
) {
    let rounds = N - 1;
    assert!(rounds == 10 || rounds == 12 || rounds == 14);

let mut state = [
        vld1q_u8(blocks[0].as_ptr()),
        vld1q_u8(blocks[1].as_ptr()),
        vld1q_u8(blocks[2].as_ptr()),
        vld1q_u8(blocks[3].as_ptr()),
        vld1q_u8(blocks[4].as_ptr()),
        vld1q_u8(blocks[5].as_ptr()),
        vld1q_u8(blocks[6].as_ptr()),
        vld1q_u8(blocks[7].as_ptr()),
    ];

for k in expanded_keys.iter().take(rounds - 1) {
        for i in 0..8 {
            // AES single round encryption and mix columns
            state[i] = vaeseq_u8_and_vaesmcq_u8(state[i], *k);
        }
    }

for i in 0..8 {
        // AES single round encryption
        state[i] = vaeseq_u8(state[i], expanded_keys[rounds - 1]);

// Final add (bitwise XOR)
        state[i] = veorq_u8(state[i], expanded_keys[rounds]);

vst1q_u8(blocks[i].as_mut_ptr(), state[i]);
    }
}

#[target_feature(enable = "aes")]
#[target_feature(enable = "neon")]
pub unsafe fn decrypt1<const N: usize>(
    expanded_keys: &[uint8x16_t; N],
    block: &mut Block,
) {
    let rounds = N - 1;
    assert!(rounds == 10 || rounds == 12 || rounds == 14);

let mut state = vld1q_u8(block.as_ptr());

for k in expanded_keys.iter().take(rounds - 1) {
        // AES single round decryption and inverse mix columns
        state = vaesdq_u8_and_vaesimcq_u8(state, *k);
    }

// AES single round decryption
    state = vaesdq_u8(state, expanded_keys[rounds - 1]);

// Final add (bitwise XOR)
    state = veorq_u8(state, expanded_keys[rounds]);

vst1q_u8(block.as_mut_ptr(), state);
}