Feat: Aes128 precompile

crates/core/executor/Cargo.toml

@@ -45,6 +45,7 @@ bytemuck = "1.16.3"
 tiny-keccak = { version = "2.0.2", features = ["keccak"] }
 vec_map = { version = "0.8.2", features = ["serde"] }
 enum-map = { version = "2.7.3", features = ["serde"] }
+aes = "0.8.4"
 sha2 = { workspace = true }
 anyhow = { workspace = true }
 tracing-subscriber = "0.3.19"

crates/core/executor/src/air.rs

@@ -41,6 +41,8 @@ pub enum MipsAirId {
     Secp256r1DoubleAssign = 11,
     /// The Poseidon2 Permute chip
     Poseidon2Permute = 46,
+    /// The AES-128 Encrypt chip
+    Aes128Encrypt = 49,
     /// The Keccak sponge chip.
     KeccakSponge = 48,
     /// The bn254 add assign chip.
@@ -153,6 +155,7 @@ impl MipsAirId {
             Self::Secp256r1AddAssign => "Secp256r1AddAssign",
             Self::Secp256r1DoubleAssign => "Secp256r1DoubleAssign",
             Self::Poseidon2Permute => "Poseidon2Permute",
+            Self::Aes128Encrypt => "Aes128Encrypt",
             Self::KeccakSponge => "KeccakSponge",
             Self::Bn254AddAssign => "Bn254AddAssign",
             Self::Bn254DoubleAssign => "Bn254DoubleAssign",

crates/core/executor/src/artifacts/mips_costs.json

@@ -9,6 +9,7 @@
   "Secp256r1Decompress": 2686,
   "Secp256k1Decompress": 2686,
   "KeccakSponge": 102216,
+  "Aes128Encrypt": 38764,
   "Bn254AddAssign": 4013,
   "Bitwise": 42,
   "ShiftLeft": 68,

crates/core/executor/src/events/precompiles/aes128.rs

@@ -0,0 +1,38 @@
+use serde::{Deserialize, Serialize};
+
+use crate::events::{
+    memory::{MemoryReadRecord, MemoryWriteRecord},
+    MemoryLocalEvent,
+};
+
+pub const AES_128_BLOCK_U32S: usize = 4;
+pub const AES_128_BLOCK_BYTES: usize = 16;
+
+/// AES128 Encrypt Event
+///
+/// This event is emitted when a AES128 encrypt operation is performed.
+#[derive(Default, Debug, Clone, Serialize, Deserialize)]
+pub struct AES128EncryptEvent {
+    /// The shard number.
+    pub shard: u32,
+    /// The clock cycle.
+    pub clk: u32,
+    /// The address of the block
+    pub block_addr: u32,
+    /// The address of the key
+    pub key_addr: u32,
+    /// The input block as a [u32; AES_128_BLOCK_U32S] words.
+    pub input: [u32; AES_128_BLOCK_U32S],
+    /// The key as a [u32; AES_128_BLOCK_U32S] words.
+    pub key: [u32; AES_128_BLOCK_U32S],
+    /// The output block as a [u32; AES_128_BLOCK_U32S] words.
+    pub output: [u32; AES_128_BLOCK_U32S],
+    /// The memory records for the input
+    pub input_read_records: [MemoryReadRecord; AES_128_BLOCK_U32S],
+    /// The memory records for the key
+    pub key_read_records: [MemoryReadRecord; AES_128_BLOCK_U32S],
+    /// The memory records for the output
+    pub output_write_records: [MemoryWriteRecord; AES_128_BLOCK_U32S],
+    /// The local memory access records.
+    pub local_mem_access: Vec<MemoryLocalEvent>,
+}

crates/core/executor/src/events/precompiles/mod.rs

@@ -1,3 +1,4 @@
+mod aes128;
 mod ec;
 mod edwards;
 mod fptower;
@@ -11,6 +12,7 @@ mod uint256;
 
 use super::{MemoryLocalEvent, SyscallEvent};
 use crate::syscalls::SyscallCode;
+pub use aes128::*;
 pub use ec::*;
 pub use edwards::*;
 pub use fptower::*;
@@ -34,6 +36,8 @@ pub enum PrecompileEvent {
     ShaCompress(ShaCompressEvent),
     /// Keccak sponge precompile event.
     KeccakSponge(KeccakSpongeEvent),
+    /// AES-128 encrypt precompile event.
+    Aes128Encrypt(AES128EncryptEvent),
     /// Edwards curve add precompile event.
     EdAdd(EllipticCurveAddEvent),
     /// Edwards curve decompress precompile event.
@@ -105,6 +109,9 @@ impl PrecompileLocalMemory for Vec<(SyscallEvent, PrecompileEvent)> {
                 PrecompileEvent::KeccakSponge(e) => {
                     iterators.push(e.local_mem_access.iter());
                 }
+                PrecompileEvent::Aes128Encrypt(e) => {
+                    iterators.push(e.local_mem_access.iter());
+                }
                 PrecompileEvent::EdDecompress(e) => {
                     iterators.push(e.local_mem_access.iter());
                 }

crates/core/executor/src/record.rs

@@ -124,6 +124,7 @@ impl ExecutionRecord {
                 SyscallCode::KECCAK_SPONGE => opts.keccak,
                 SyscallCode::SHA_EXTEND => opts.sha_extend,
                 SyscallCode::SHA_COMPRESS => opts.sha_compress,
+                SyscallCode::AES128_ENCRYPT => opts.aes128_encrypt,
                 _ => opts.deferred,
             };
 

crates/core/executor/src/syscalls/code.rs

@@ -169,6 +169,9 @@ pub enum SyscallCode {
     POSEIDON2_PERMUTE = 0x00_01_00_30,
     SYS_LINUX = 5000,
 
+    /// Executes the `AES128_ENCRYPT` precompile.
+    AES128_ENCRYPT = 0x01_01_00_31,
+
     UNIMPLEMENTED = 0xFF_FF_FF_FF,
 }
 
@@ -190,6 +193,7 @@ impl SyscallCode {
             0x01_01_00_07 => SyscallCode::ED_ADD,
             0x00_01_00_08 => SyscallCode::ED_DECOMPRESS,
             0x01_01_00_09 => SyscallCode::KECCAK_SPONGE,
+            0x01_01_00_31 => SyscallCode::AES128_ENCRYPT,
             0x01_01_00_0A => SyscallCode::SECP256K1_ADD,
             0x00_01_00_0B => SyscallCode::SECP256K1_DOUBLE,
             0x00_01_00_0C => SyscallCode::SECP256K1_DECOMPRESS,

crates/core/executor/src/syscalls/mod.rs

@@ -22,6 +22,7 @@ pub use code::*;
 pub use context::*;
 use hint::{HintLenSyscall, HintReadSyscall};
 use precompiles::{
+    aes128::encrypt::AES128EncryptSyscall,
     edwards::{add::EdwardsAddAssignSyscall, decompress::EdwardsDecompressSyscall},
     fptower::{Fp2AddSubSyscall, Fp2MulSyscall, FpOpSyscall},
     keccak::sponge::KeccakSpongeSyscall,
@@ -102,6 +103,8 @@ pub fn default_syscall_map() -> HashMap<SyscallCode, Arc<dyn Syscall>> {
 
     syscall_map.insert(SyscallCode::POSEIDON2_PERMUTE, Arc::new(Poseidon2PermuteSyscall));
 
+    syscall_map.insert(SyscallCode::AES128_ENCRYPT, Arc::new(AES128EncryptSyscall));
+
     syscall_map.insert(SyscallCode::KECCAK_SPONGE, Arc::new(KeccakSpongeSyscall));
 
     syscall_map.insert(

crates/core/executor/src/syscalls/precompiles/aes128/encrypt.rs

@@ -0,0 +1,196 @@
+use crate::events::{AES128EncryptEvent, PrecompileEvent, AES_128_BLOCK_U32S};
+use crate::syscalls::precompiles::aes128::utils::mul_md5;
+use crate::syscalls::{Syscall, SyscallCode, SyscallContext};
+
+pub(crate) struct AES128EncryptSyscall;
+
+pub const AES128_RCON: [[u8; 4]; 10] = [
+    [0x01, 0x00, 0x00, 0x00],
+    [0x02, 0x00, 0x00, 0x00],
+    [0x04, 0x00, 0x00, 0x00],
+    [0x08, 0x00, 0x00, 0x00],
+    [0x10, 0x00, 0x00, 0x00],
+    [0x20, 0x00, 0x00, 0x00],
+    [0x40, 0x00, 0x00, 0x00],
+    [0x80, 0x00, 0x00, 0x00],
+    [0x1B, 0x00, 0x00, 0x00],
+    [0x36, 0x00, 0x00, 0x00],
+];
+
+pub const AES_SBOX: [u8; 256] = [
+    0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76,
+    0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0,
+    0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15,
+    0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75,
+    0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84,
+    0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf,
+    0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8,
+    0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2,
+    0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
+    0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb,
+    0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
+    0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08,
+    0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a,
+    0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e,
+    0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf,
+    0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16,
+];
+
+impl Syscall for AES128EncryptSyscall {
+    fn num_extra_cycles(&self) -> u32 {
+        1
+    }
+    fn execute(
+        &self,
+        rt: &mut SyscallContext,
+        syscall_code: SyscallCode,
+        arg1: u32,
+        arg2: u32,
+    ) -> Option<u32> {
+        let start_clk = rt.clk;
+        let block_ptr = arg1;
+        let key_ptr = arg2;
+
+        let mut input_read_records = Vec::new();
+        let mut key_read_records = Vec::new();
+        let mut output_write_records = Vec::new();
+
+        let mut input = Vec::new();
+        let mut key_u32s = Vec::new();
+        let mut state = Vec::new();
+        let mut key = Vec::new();
+        let mut output = Vec::new();
+
+        // read block input
+        for i in 0..AES_128_BLOCK_U32S {
+            let (record, value) = rt.mr(block_ptr + i as u32 * 4);
+            input_read_records.push(record);
+            input.push(value);
+            state.extend(value.to_le_bytes());
+        }
+
+        // read key
+        for i in 0..AES_128_BLOCK_U32S {
+            let (record, value) = rt.mr(key_ptr + i as u32 * 4);
+            key_read_records.push(record);
+            key_u32s.push(value);
+            key.extend(value.to_le_bytes());
+        }
+
+        // // Add Roundkey, Round 0
+        for i in 0..state.len() {
+            state[i] ^= key[i];
+        }
+
+        // perform AES
+        let mut round_key = key;
+        for i in 1..11 {
+            // compute round key
+            Self::compute_round_key(&mut round_key, i - 1);
+
+            // Subs_bytes
+            for j in 0..state.len() {
+                let value = AES_SBOX[state[j] as usize];
+                state[j] = value;
+            }
+
+            // Shift row
+            let shift_row = [
+                state[0], state[5], state[10], state[15], state[4], state[9], state[14], state[3],
+                state[8], state[13], state[2], state[7], state[12], state[1], state[6], state[11],
+            ]
+            .to_vec();
+
+            // Mix columns
+            let mix_columns = if i != 10 {
+                let mut mixed = shift_row.clone();
+                for col in 0..4 {
+                    let col_start = col * 4;
+                    let s0 = shift_row[col_start];
+                    let s1 = shift_row[col_start + 1];
+                    let s2 = shift_row[col_start + 2];
+                    let s3 = shift_row[col_start + 3];
+                    mixed[col_start] =
+                        mul_md5(s0, 2) ^ mul_md5(s1, 3) ^ mul_md5(s2, 1) ^ mul_md5(s3, 1);
+                    mixed[col_start + 1] =
+                        mul_md5(s0, 1) ^ mul_md5(s1, 2) ^ mul_md5(s2, 3) ^ mul_md5(s3, 1);
+                    mixed[col_start + 2] =
+                        mul_md5(s0, 1) ^ mul_md5(s1, 1) ^ mul_md5(s2, 2) ^ mul_md5(s3, 3);
+                    mixed[col_start + 3] =
+                        mul_md5(s0, 3) ^ mul_md5(s1, 1) ^ mul_md5(s2, 1) ^ mul_md5(s3, 2);
+                }
+                mixed
+            } else {
+                shift_row
+            };
+
+            // Add round key
+            for j in 0..state.len() {
+                state[j] = mix_columns[j] ^ round_key[j];
+            }
+        }
+
+        // write output
+        // Increment the clk by 1 before writing because we read from memory at start_clk.
+        rt.clk += 1;
+        assert_eq!(state.len(), 16);
+        for chunk in state.chunks(4) {
+            let value = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
+            output.push(value);
+        }
+        let write_records = rt.mw_slice(block_ptr, output.as_slice());
+        output_write_records.extend_from_slice(&write_records);
+
+        // Push the AES128 encrypt event.
+        let shard = rt.current_shard();
+        let aes128_event = PrecompileEvent::Aes128Encrypt(AES128EncryptEvent {
+            shard,
+            clk: start_clk,
+            block_addr: block_ptr,
+            key_addr: key_ptr,
+            input: input.as_slice().try_into().unwrap(),
+            key: key_u32s.as_slice().try_into().unwrap(),
+            output: output.as_slice().try_into().unwrap(),
+            input_read_records: input_read_records.as_slice().try_into().unwrap(),
+            key_read_records: key_read_records.as_slice().try_into().unwrap(),
+            output_write_records: output_write_records.as_slice().try_into().unwrap(),
+            local_mem_access: rt.postprocess(),
+        });
+        let aes128_syscall_event =
+            rt.rt.syscall_event(start_clk, None, rt.next_pc, syscall_code.syscall_id(), arg1, arg2);
+        rt.add_precompile_event(syscall_code, aes128_syscall_event, aes128_event);
+        None
+    }
+}
+
+impl AES128EncryptSyscall {
+    fn compute_round_key(previous_key: &mut [u8], round: usize) {
+        if previous_key.len() != 16 {
+            panic!("AES128: wrong previous key length");
+        }
+        // First 4 bytes
+        let g_w3 = {
+            let mut result =
+                [previous_key[13], previous_key[14], previous_key[15], previous_key[12]];
+            for (i, rcon) in AES128_RCON[round].iter().enumerate() {
+                let value = AES_SBOX[result[i] as usize];
+                result[i] = value ^ rcon;
+            }
+            result
+        };
+
+        let prev = previous_key.to_vec().clone();
+        for i in 0..4 {
+            let w = if i == 0 {
+                prev[0..4].iter().zip(g_w3.iter()).map(|(&a, &b)| a ^ b).collect::<Vec<u8>>()
+            } else {
+                prev[i * 4..(i + 1) * 4]
+                    .iter()
+                    .zip(previous_key[(i - 1) * 4..i * 4].iter())
+                    .map(|(&a, &b)| a ^ b)
+                    .collect::<Vec<u8>>()
+            };
+            previous_key[i * 4..(i + 1) * 4].copy_from_slice(&w);
+        }
+    }
+}

crates/core/executor/src/syscalls/precompiles/aes128/mod.rs

@@ -0,0 +1,2 @@
+pub mod encrypt;
+pub mod utils;

crates/core/executor/src/syscalls/precompiles/aes128/utils.rs

@@ -0,0 +1,18 @@
+/// Multiply a byte by 2 in GF(2^8) with AES polynomial 0x9B
+fn xtime(x: u8) -> u8 {
+    if x & 0x80 != 0 {
+        (x << 1) ^ 0x1b
+    } else {
+        x << 1
+    }
+}
+
+/// Multiply a byte by 1, 2, or 3 in GF(2^8)
+pub fn mul_md5(x: u8, by: u8) -> u8 {
+    match by {
+        1 => x,
+        2 => xtime(x),
+        3 => xtime(x) ^ x, // 3*x = (2*x) ⊕ x
+        _ => panic!("Only supports multipliers 1, 2, or 3"),
+    }
+}

crates/core/executor/src/syscalls/precompiles/mod.rs

@@ -1,3 +1,4 @@
+pub mod aes128;
 pub mod edwards;
 pub mod fptower;
 pub mod keccak;