M Cargo.toml → Cargo.toml

@@ -8,4 +8,5 @@ plonky2 = { git = "https://github.com/0xPARC/plonky2.git", rev = "109d517d09c210ae4c2cee381d3e3fbc04aa3812" }
 anyhow = "1.0.56"
 
 [dev-dependencies]
-rand = "0.10.0"+rand = "0.10.0"
+aes = "0.8.4"

M src/circuit.rs → src/circuit.rs

@@ -61,10 +61,10 @@ s: StateTarget,
     ) -> StateTarget;
 
     /// KeyExpansion
-    fn key_expansion<const NK: usize, const NR: usize>(
+    fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
         &mut self,
         sbox_lut_idx: usize,
-        key: [ByteArrayTarget; 4 * NK],
+        key: [ByteArrayTarget; NK * NB],
     ) -> [[ByteArrayTarget; 4]; 4 * (NR + 1)];
 
     /// GF(2^8) addition
@@ -164,10 +164,10 @@ array::from_fn(|j| self.gf_2_8_add_bits(s.0[i][j], round_key[j][i]))
         }))
     }
 
-    fn key_expansion<const NK: usize, const NR: usize>(
+    fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
         &mut self,
         sbox_lut_idx: usize,
-        key: [ByteArrayTarget; 4 * NK],
+        key: [ByteArrayTarget; NK * NB],
     ) -> [[ByteArrayTarget; 4]; 4 * (NR + 1)] {
         let rcon: [ByteArrayTarget; 11] = array::from_fn(|i| {
             let rcon_bits = le_bits_from_byte(RCON[i]);
@@ -547,10 +547,10 @@ // Circuit declaration
         let config = CircuitConfig::standard_recursion_config();
         let mut builder = CircuitBuilder::<F, D>::new(config);
 
-        let key_target: [ByteArrayTarget; _] =
+        let key_target: [ByteArrayTarget; NK * NB] =
             array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
         let sbox_lut = sbox_lut(&mut builder);
-        let expanded_key_target = builder.key_expansion::<NK, NR>(sbox_lut, key_target);
+        let expanded_key_target = builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target);
 
         println!(
             "key_expansion(NK:{}, NB:{}, NR:{}) num_gates: {}",
@@ -623,17 +623,18 @@ ) -> Result<()>
     where
         [(); 4 * (NR + 1)]:,
         [(); 4 * NK]:,
+        // [(); NK * NB]:,
     {
         // Circuit declaration
         let config = CircuitConfig::standard_recursion_config();
         let mut builder = CircuitBuilder::<F, D>::new(config);
 
-        let key_target: [ByteArrayTarget; 4 * NK] =
+        let key_target: [ByteArrayTarget; NK * NB] =
             array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
         let sbox_lut = sbox_lut(&mut builder);
         let mix_matrix = state_mix_matrix_bits(&mut builder);
         let expanded_key_target: [[ByteArrayTarget; 4]; 4 * (NR + 1)] =
-            builder.key_expansion::<NK, NR>(sbox_lut, key_target);
+            builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target);
 
         let input_state_target = builder.add_virtual_state();
 
@@ -663,16 +664,6 @@ }
 
         let expanded_key: [[u8; 4]; 4 * (NR + 1)] = key_expansion::<NK, NB, NR>(&key);
         let native_ciphertext = encrypt_block::<NR>(&input_state, &expanded_key);
-
-        // AES-128:
-        // let expected_ciphertext: [[u8; 4]; 4] = [
-        //     [0x39, 0x02, 0xdc, 0x19],
-        //     [0x25, 0xdc, 0x11, 0x6a],
-        //     [0x84, 0x09, 0x85, 0x0b],
-        //     [0x1d, 0xfb, 0x97, 0x32],
-        // ];
-        // // sanity check
-        // assert_eq!(native_ciphertext, expected_ciphertext);
 
         let mut pw = PartialWitness::<F>::new();
 

M src/native.rs → src/native.rs

@@ -16,8 +16,7 @@ fn flatten_state(s: State) -> [u8; 16] {
     let mut r = [0u8; 16];
     for i in 0..4 {
         for j in 0..4 {
-            // TODO
-            r[i * 4 + j] = s[i][j];
+            r[i * 4 + j] = s[j][i];
         }
     }
     r
@@ -209,17 +208,91 @@ 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
             0x4f, 0x3c,
         ];
         let expected: [u8; 16] = [
-            0x39, 0x02, 0xdc, 0x19, 0x25, 0xdc, 0x11, 0x6a, 0x84, 0x09, 0x85, 0x0b, 0x1d, 0xfb,
-            0x97, 0x32,
+            0x39, 0x25, 0x84, 0x1d, 0x02, 0xdc, 0x09, 0xfb, 0xdc, 0x11, 0x85, 0x97, 0x19, 0x6a,
+            0x0b, 0x32,
         ];
         let expanded_key = key_expansion::<4, 4, 10>(&key);
         let c = encrypt_block::<10>(&input, &expanded_key);
         assert_eq!(flatten_state(c), expected);
     }
 
+    use aes::{
+        cipher::{BlockDecrypt, BlockEncrypt, BlockSizeUser, KeyInit, generic_array::GenericArray},
+        {Aes128, Aes192, Aes256},
+    };
+    use rand::RngExt;
     /// test with loop of random-values against a rust well-tested library
     #[test]
     fn test_encrypt_block_external_lib() {
-        // TODO random-values loop against a rust well-tested library
+        let mut rng = rand::rng();
+
+        // AES-128
+        {
+            const NK: usize = 4;
+            const NB: usize = 4;
+            const NR: usize = 10;
+
+            for _ in 0..10 {
+                let key: [u8; NK * NB] = array::from_fn(|_| rng.random());
+                let block: [u8; 16] = array::from_fn(|_| rng.random());
+                let block = GenericArray::from(block);
+                test_encrypt_block_external_lib_op::<NK, NB, NR, Aes128>(key, block);
+            }
+        }
+        // AES-192
+        {
+            const NK: usize = 6;
+            const NB: usize = 4;
+            const NR: usize = 12;
+
+            for _ in 0..10 {
+                let key: [u8; NK * NB] = array::from_fn(|_| rng.random());
+                let block: [u8; 16] = array::from_fn(|_| rng.random());
+                let block = GenericArray::from(block);
+                test_encrypt_block_external_lib_op::<NK, NB, NR, Aes192>(key, block);
+            }
+        }
+        // AES-256
+        {
+            const NK: usize = 8;
+            const NB: usize = 4;
+            const NR: usize = 14;
+
+            for _ in 0..10 {
+                let key: [u8; NK * NB] = array::from_fn(|_| rng.random());
+                let block: [u8; 16] = array::from_fn(|_| rng.random());
+                let block = GenericArray::from(block);
+                test_encrypt_block_external_lib_op::<NK, NB, NR, Aes256>(key, block);
+            }
+        }
+    }
+    fn test_encrypt_block_external_lib_op<const NK: usize, const NB: usize, const NR: usize, C>(
+        key_arr: [u8; NK * NB],
+        block: GenericArray<u8, C::BlockSize>,
+    ) where
+        C: BlockEncrypt + BlockDecrypt + BlockSizeUser + KeyInit,
+        [(); 4 * (NR + 1)]:,
+    {
+        let cipher = C::new_from_slice(&key_arr).unwrap();
+
+        let mut block_arr: [u8; 16] = [0; 16];
+        block_arr.clone_from_slice(block.as_slice());
+        let mut block = block.clone();
+        let block_copy = block.clone();
+
+        // encrypt block in-place
+        cipher.encrypt_block(&mut block);
+        let binding = block.clone();
+        let expected_cipher = binding.as_slice();
+
+        // and decrypt it back
+        cipher.decrypt_block(&mut block);
+        assert_eq!(block, block_copy);
+
+        // now use our native implementation
+        let expanded_key = key_expansion::<NK, NB, NR>(&key_arr);
+        let c = encrypt_block::<NR>(&block_arr, &expanded_key);
+        // check native implementation vs external lib
+        assert_eq!(flatten_state(c), expected_cipher);
     }
 }