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

@@ -11,9 +11,9 @@ plonk::circuit_builder::CircuitBuilder,
 };
 
 use crate::{
-    D,
     constants::{RCON, SBOX},
-    native::{State, rot_word, shift_rows},
+    native::{rot_word, shift_rows, State},
+    D,
 };
 
 #[derive(Debug, Copy, Clone)]
@@ -27,7 +27,7 @@ /// Adds state target.
     fn add_virtual_state(&mut self) -> StateTarget;
 
     /// AES cipher as in spec.
-    fn aes_cipher<const NR: usize>(
+    fn encrypt_block<const NR: usize>(
         &mut self,
         sbox_lut_idx: usize,
         mix_matrix: [[ByteArrayTarget; 4]; 4],
@@ -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
@@ -105,12 +105,12 @@ array::from_fn(|_| array::from_fn(|_| self.add_virtual_bool_target_safe()))
         }))
     }
 
-    fn aes_cipher<const NR: usize>(
+    fn encrypt_block<const NR: usize>(
         &mut self,
         sbox_lut_idx: usize,
-        mix_matrix: [[ByteArrayTarget; 4]; 4],
+        mix_matrix: [[ByteArrayTarget; 4]; 4], // constant mix_matrix
         s: StateTarget,
-        w: [[ByteArrayTarget; 4]; 4 * (NR + 1)],
+        w: [[ByteArrayTarget; 4]; 4 * (NR + 1)], // expanded key
     ) -> StateTarget {
         let mut s = s;
         s = self.state_add_round_key(&w[0..4], s);
@@ -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]);
@@ -399,10 +399,10 @@ use rand::RngExt;
 
     use crate::{
         circuit::{
-            ByteArrayTarget, CircuitBuilderAESState, PartialWitnessAESState,
-            PartialWitnessByteArray, sbox_lut, state_mix_matrix_bits,
+            sbox_lut, state_mix_matrix_bits, ByteArrayTarget, CircuitBuilderAESState,
+            PartialWitnessAESState, PartialWitnessByteArray,
         },
-        native::{State, encrypt_block, key_expansion, mix_columns, sub_bytes},
+        native::{encrypt_block, key_expansion, mix_columns, sub_bytes, State},
     };
 
     use super::D;
@@ -418,8 +418,9 @@ let sbox_lut = sbox_lut(&mut builder);
         let out_state_target = builder.state_sub_bytes(sbox_lut, state_target);
 
         let data = builder.build::<PoseidonGoldilocksConfig>();
-        let mut rng = rand::rng();
 
+        // set values to circuit
+        let mut rng = rand::rng();
         let test_states: [State; 10] = array::from_fn(|_| array::from_fn(|_| rng.random()));
 
         test_states.into_iter().try_for_each(|s| {
@@ -445,8 +446,9 @@ let mix_matrix = state_mix_matrix_bits(&mut builder);
         let out_state_target = builder.state_mix_columns(mix_matrix, state_target);
 
         let data = builder.build::<PoseidonGoldilocksConfig>();
+
+        // set values to circuit
         let mut rng = rand::rng();
-
         let test_states: [State; 10] = array::from_fn(|_| array::from_fn(|_| rng.random()));
 
         test_states.into_iter().try_for_each(|s| {
@@ -474,6 +476,7 @@ let a_times_b = builder.gf_2_8_mul(a, b);
 
         let data = builder.build::<PoseidonGoldilocksConfig>();
 
+        // set values to circuit
         let test_values = [
             (0x57, 0x01, 0x57),
             (0x57, 0x02, 0xae),
@@ -485,7 +488,6 @@ (0x57, 0x40, 0x1c),
             (0x57, 0x80, 0x38),
             (0x57, 0x13, 0xfe),
         ];
-
         test_values
             .into_iter()
             .map(|(a, b, c)| {
@@ -506,58 +508,61 @@ data.verify(proof)
             })
     }
 
-    /// Analogous tests to `native::tests::test_key_expansion`.
     #[test]
-    fn test_key_expansion_128() -> Result<()> {
-        // Circuit declaration
-        let config = CircuitConfig::standard_recursion_config();
-        let mut builder = CircuitBuilder::<F, D>::new(config);
-
-        let key_target: [ByteArrayTarget; _] =
-            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::<4, 10>(sbox_lut, key_target);
-
-        let data = builder.build::<PoseidonGoldilocksConfig>();
-
+    fn test_key_expansion() -> Result<()> {
         // AES-128
         let key = [
             0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
             0x4f, 0x3c,
         ];
-        let w = key_expansion::</*Nk,Nb,Nr*/ 4, 4, 10>(&key);
+        test_key_expansion_op::</*Nk,Nb,Nr*/ 4, 4, 10>(key)?;
 
-        let mut pw = PartialWitness::<F>::new();
+        // AES-196
+        let key = [
+            0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
+            0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
+        ];
+        test_key_expansion_op::</*Nk,Nb,Nr*/ 6, 4, 12>(key)?;
 
-        std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
+        // AES-256
+        let key = [
+            0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
+            0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
+            0x09, 0x14, 0xdf, 0xf4,
+        ];
+        test_key_expansion_op::</*Nk,Nb,Nr*/ 8, 4, 14>(key)?;
 
-        std::iter::zip(expanded_key_target, w).try_for_each(|(t, v)| {
-            std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))
-        })?;
-
-        let proof = data.prove(pw)?;
-        data.verify(proof)
+        Ok(())
     }
 
-    #[test]
-    fn test_key_expansion_196() -> Result<()> {
+    /// Analogous tests to `native::tests::test_key_expansion`.
+    fn test_key_expansion_op<const NK: usize, const NB: usize, const NR: usize>(
+        key: [u8; NK * NB],
+    ) -> Result<()>
+    where
+        [(); 4 * (NR + 1)]:,
+        [(); 4 * NK]:,
+    {
         // 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::<6, 12>(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: {}",
+            NK,
+            NB,
+            NR,
+            builder.num_gates()
+        );
         let data = builder.build::<PoseidonGoldilocksConfig>();
 
-        // AES-196
-        let key = [
-            0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
-            0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
-        ];
-        let w = key_expansion::</*Nk,Nb,Nr*/ 6, 4, 12>(&key);
+        // set values to circuit
+        let w = key_expansion::<NK, NB, NR>(&key);
 
         let mut pw = PartialWitness::<F>::new();
 
@@ -571,101 +576,106 @@ let proof = data.prove(pw)?;
         data.verify(proof)
     }
 
+    /// test against native version
     #[test]
-    fn test_key_expansion_256() -> Result<()> {
-        // Circuit declaration
-        let config = CircuitConfig::standard_recursion_config();
-        let mut builder = CircuitBuilder::<F, D>::new(config);
-
-        let key_target: [ByteArrayTarget; _] =
-            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::<8, 14>(sbox_lut, key_target);
+    fn test_encrypt_block_test_vector() -> Result<()> {
+        // AES-128
+        let input_state: [u8; 16] = [
+            0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
+            0x07, 0x34,
+        ];
+        let key = [
+            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
+            0x4f, 0x3c,
+        ];
+        test_encrypt_block_test_vector_op::<4, 4, 10>(input_state, key)?;
 
-        let data = builder.build::<PoseidonGoldilocksConfig>();
+        // AES-196
+        let input_state: [u8; 16] = [
+            0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
+            0x07, 0x34,
+        ];
+        let key = [
+            0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
+            0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
+        ];
+        test_encrypt_block_test_vector_op::<6, 4, 12>(input_state, key)?;
 
         // AES-256
+        let input_state: [u8; 16] = [
+            0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
+            0x07, 0x34,
+        ];
         let key = [
             0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
             0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
             0x09, 0x14, 0xdf, 0xf4,
         ];
-        let w = key_expansion::</*Nk,Nb,Nr*/ 8, 4, 14>(&key);
-
-        let mut pw = PartialWitness::<F>::new();
-
-        std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
+        test_encrypt_block_test_vector_op::<8, 4, 14>(input_state, key)?;
 
-        std::iter::zip(expanded_key_target, w).try_for_each(|(t, v)| {
-            std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))
-        })?;
-
-        let proof = data.prove(pw)?;
-        data.verify(proof)
+        Ok(())
     }
 
-    /// test against paper test vector (appendix B)
-    #[test]
-    fn test_encrypt_block_test_vector() -> Result<()> {
+    fn test_encrypt_block_test_vector_op<const NK: usize, const NB: usize, const NR: usize>(
+        input_state: [u8; 16],
+        key: [u8; NK * NB],
+    ) -> 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; _] =
+        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 = builder.key_expansion::<4, 10>(sbox_lut, key_target);
+        let expanded_key_target: [[ByteArrayTarget; 4]; 4 * (NR + 1)] =
+            builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target);
 
         let input_state_target = builder.add_virtual_state();
 
-        let output = builder.aes_cipher(
+        let output = builder.encrypt_block(
             sbox_lut,
             mix_matrix,
             input_state_target,
             expanded_key_target,
         );
 
+        println!(
+            "encrypt_block (NK:{}, NB:{}, NR:{}) num_gates: {}",
+            NK,
+            NB,
+            NR,
+            builder.num_gates()
+        );
         let data = builder.build::<PoseidonGoldilocksConfig>();
 
-        let input_state: [[u8; _]; _] = [
-            [0x32, 0x43, 0xf6, 0xa8],
-            [0x88, 0x5a, 0x30, 0x8d],
-            [0x31, 0x31, 0x98, 0xa2],
-            [0xe0, 0x37, 0x07, 0x34],
-        ];
+        // set values to circuit
+        let mut input_state_matrix: State = [[0; 4]; 4];
+        for i in 0..4 {
+            for j in 0..4 {
+                input_state_matrix[i][j] = input_state[i + 4 * j];
+            }
+        }
 
-        let key = [
-            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
-            0x4f, 0x3c,
-        ];
-
-        let expanded_key = key_expansion::<4, 4, 10>(&key);
-
-        let expected_result: [[u8; _]; _] = [
-            [0x39, 0x02, 0xdc, 0x19],
-            [0x25, 0xdc, 0x11, 0x6a],
-            [0x84, 0x09, 0x85, 0x0b],
-            [0x1d, 0xfb, 0x97, 0x32],
-        ];
-
-        // TODO
-        let expected_result = encrypt_block::<10>(
-            &array::from_fn(|i| input_state[i % 4][i / 4]),
-            &expanded_key,
-        );
+        let expanded_key: [[u8; 4]; 4 * (NR + 1)] = key_expansion::<NK, NB, NR>(&key);
+        let native_ciphertext = encrypt_block::<NR>(&input_state, &expanded_key);
 
         let mut pw = PartialWitness::<F>::new();
 
         std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
 
-        pw.set_target_state(input_state_target, input_state)?;
+        pw.set_target_state(input_state_target, input_state_matrix)?;
 
         std::iter::zip(expanded_key_target, expanded_key).try_for_each(|(t, v)| {
             std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))
         })?;
 
-        std::iter::zip(output.0, expected_result).try_for_each(|(o, e)| {
+        std::iter::zip(output.0, native_ciphertext).try_for_each(|(o, e)| {
             std::iter::zip(o, e).try_for_each(|(o, e)| pw.set_byte_array_target(o, e))
         })?;
 

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
@@ -112,11 +111,7 @@ }
 
 pub(crate) fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
     key: &[u8; NK * NB],
-) -> [[u8; 4]; 4 * (NR + 1)]
-where
-    [(); 4 * (NR + 1)]:,
-    [(); 4 * 4 * (NR + 1)]:,
-{
+) -> [[u8; 4]; 4 * (NR + 1)] {
     let mut w = [[0u8; 4]; 4 * (NR + 1)]; // expanded key
 
     for i in 0..NK {
@@ -208,25 +203,96 @@ let input: [u8; 16] = [
             0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
             0x07, 0x34,
         ];
-        let key = [
+        let key: [u8; 16] = [
             0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
             0x4f, 0x3c,
         ];
+        let expected: [u8; 16] = [
+            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);
-
-        let expected: [u8; 16] = [
-            0x39, 0x02, 0xdc, 0x19, 0x25, 0xdc, 0x11, 0x6a, 0x84, 0x09, 0x85, 0x0b, 0x1d, 0xfb,
-            0x97, 0x32,
-        ];
-
         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);
     }
 }