AES-128 using ARMv8 Cryptography Extensions in Rust

aes128_aarch64.rs

#![feature(stdsimd)]

// Based on the following C intrinsics implementation:
// <https://github.com/noloader/AES-Intrinsics/blob/master/aes-arm.c>
//
// Original C written and placed in public domain by Jeffrey Walton.
// Based on code from ARM, and by Johannes Schneiders, Skip Hovsmith and
// Barry O'Rourke for the mbedTLS project.

use core::arch::aarch64::*;
use core::ptr;

unsafe fn aes_process_arm(
    key: &[u8; 16],
    subkeys: &[[u8; 16]; 10],
    rounds: usize,
    input: &[u8; 16],
    output: &mut [u8; 16],
) {
    let mut block = vld1q_u8(input.as_ptr());

    // AES single round encryption
    block = vaeseq_u8(block, vld1q_u8(key.as_ptr()));
    // AES mix columns
    block = vaesmcq_u8(block);

    // AES single round encryption
    block = vaeseq_u8(block, vld1q_u8(subkeys[0].as_ptr()));
    // AES mix columns
    block = vaesmcq_u8(block);

    for i in 1..(rounds - 2) {
        // AES single round encryption
        block = vaeseq_u8(block, vld1q_u8(subkeys[i].as_ptr()));
        // AES mix columns
        block = vaesmcq_u8(block);
    }

    // AES single round encryption
    block = vaeseq_u8(block, vld1q_u8(subkeys[rounds - 2].as_ptr()));
    // Final Add (bitwise Xor)
    block = veorq_u8(block, vld1q_u8(subkeys[rounds - 1].as_ptr()));

    // TODO(tarcieri): use `vst1q_u8` when available
    // error[E0425]: cannot find function `vst1q_u8` in this scope
    // vst1q_u8(output, block);
    ptr::copy_nonoverlapping(&block as *const _ as *const u8, output.as_mut_ptr(), 16);
}

fn main() {
    /* FIPS 197, Appendix B input */
    /* user input, unaligned buffer */
    let input = [
        0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37, 0x07,
        0x34,
    ];

    /* FIPS 197, Appendix B key */
    /* user input, unaligned buffer */
    let key = [
        0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x9, 0xcf, 0x4f,
        0x3c,
    ];

    /* FIPS 197, Appendix B expanded subkeys */
    /* library controlled, aligned buffer */
    #[rustfmt::skip]
    let subkeys = [
        [0xA0, 0xFA, 0xFE, 0x17, 0x88, 0x54, 0x2c, 0xb1, 0x23, 0xa3, 0x39, 0x39, 0x2a, 0x6c, 0x76, 0x05],
        [0xF2, 0xC2, 0x95, 0xF2, 0x7a, 0x96, 0xb9, 0x43, 0x59, 0x35, 0x80, 0x7a, 0x73, 0x59, 0xf6, 0x7f],
        [0x3D, 0x80, 0x47, 0x7D, 0x47, 0x16, 0xFE, 0x3E, 0x1E, 0x23, 0x7E, 0x44, 0x6D, 0x7A, 0x88, 0x3B],
        [0xEF, 0x44, 0xA5, 0x41, 0xA8, 0x52, 0x5B, 0x7F, 0xB6, 0x71, 0x25, 0x3B, 0xDB, 0x0B, 0xAD, 0x00],
        [0xD4, 0xD1, 0xC6, 0xF8, 0x7C, 0x83, 0x9D, 0x87, 0xCA, 0xF2, 0xB8, 0xBC, 0x11, 0xF9, 0x15, 0xBC],
        [0x6D, 0x88, 0xA3, 0x7A, 0x11, 0x0B, 0x3E, 0xFD, 0xDB, 0xF9, 0x86, 0x41, 0xCA, 0x00, 0x93, 0xFD],
        [0x4E, 0x54, 0xF7, 0x0E, 0x5F, 0x5F, 0xC9, 0xF3, 0x84, 0xA6, 0x4F, 0xB2, 0x4E, 0xA6, 0xDC, 0x4F],
        [0xEA, 0xD2, 0x73, 0x21, 0xB5, 0x8D, 0xBA, 0xD2, 0x31, 0x2B, 0xF5, 0x60, 0x7F, 0x8D, 0x29, 0x2F],
        [0xAC, 0x77, 0x66, 0xF3, 0x19, 0xFA, 0xDC, 0x21, 0x28, 0xD1, 0x29, 0x41, 0x57, 0x5c, 0x00, 0x6E],
        [0xD0, 0x14, 0xF9, 0xA8, 0xC9, 0xEE, 0x25, 0x89, 0xE1, 0x3F, 0x0c, 0xC8, 0xB6, 0x63, 0x0C, 0xA6]
    ];

    /* Result */
    let mut result = [0u8; 16];

    unsafe { aes_process_arm(&key, &subkeys, 10, &input, &mut result) };

    println!("Input: {:?}", input);
    println!("Key: {:?}", key);
    println!("Output: {:?}", result);

    /* FIPS 197, Appendix B output */
    let expected = [
        0x39, 0x25, 0x84, 0x1D, 0x02, 0xDC, 0x09, 0xFB, 0xDC, 0x11, 0x85, 0x97, 0x19, 0x6A, 0x0B,
        0x32,
    ];

    assert_eq!(result, expected);
}