From b59fa94fa9f3953e6ba568fa787c0dfc53032f71 Mon Sep 17 00:00:00 2001
From: ayman <aymanshaik1015@gmail.com>
Date: Thu, 25 Jul 2024 21:39:38 +0530
Subject: [PATCH] feat: add aes256 support

---
 circuits/aes_128.circom       |  0
 circuits/key_expansion.circom | 13 ++++--
 tests/key_expansion.test.ts   | 78 +++++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+), 4 deletions(-)
 delete mode 100644 circuits/aes_128.circom

diff --git a/circuits/aes_128.circom b/circuits/aes_128.circom
deleted file mode 100644
index e69de29..0000000
diff --git a/circuits/key_expansion.circom b/circuits/key_expansion.circom
index b17b65c..d395bb9 100644
--- a/circuits/key_expansion.circom
+++ b/circuits/key_expansion.circom
@@ -13,7 +13,7 @@ template KeyExpansion(nk, nr) {
     var totalWords = (4 * (nr + 1));
     var effectiveRounds = (totalWords % nk == 0 ? totalWords - nk : totalWords) \ nk;
     var leftoverWords = totalWords - (effectiveRounds * nk);
-    
+
     signal output keyExpanded[totalWords][4];
 
     for (var i = 0; i < nk; i++) {
@@ -55,21 +55,26 @@ template NextRound(nk, o){
         rotateWord.bytes[i] <== key[nk - 1][i];
     }
     
-    component substituteWord = SubstituteWord();
-    substituteWord.bytes <== rotateWord.rotated;
+    component substituteWord[2];
+    substituteWord[0] = SubstituteWord();
+    substituteWord[0].bytes <== rotateWord.rotated;
 
     component rcon = RCon();
     rcon.round <== round; 
 
     component xorWord[o + 1];
     xorWord[0] = XorWord();
-    xorWord[0].bytes1 <== substituteWord.substituted;
+    xorWord[0].bytes1 <== substituteWord[0].substituted;
     xorWord[0].bytes2 <== rcon.out;
 
     for (var i = 0; i < o; i++) {
         xorWord[i+1] = XorWord();
         if (i == 0) {
             xorWord[i+1].bytes1 <== xorWord[0].out;
+        } else if(nk == 8 && i == 4) {
+            substituteWord[1] = SubstituteWord();
+            substituteWord[1].bytes <== nextKey[i - 1];
+            xorWord[i+1].bytes1 <== substituteWord[1].substituted;
         } else {
             xorWord[i+1].bytes1 <== nextKey[i-1];
         }
diff --git a/tests/key_expansion.test.ts b/tests/key_expansion.test.ts
index 6b142b2..2c47838 100644
--- a/tests/key_expansion.test.ts
+++ b/tests/key_expansion.test.ts
@@ -127,4 +127,82 @@ describe("KeyExpansion", () => {
 
     await circuit.expectPass({ key }, { keyExpanded });
   });
+
+  it("should compute correctly for aes256", async () => {
+    const circuit: WitnessTester<["key"], ["keyExpanded"]> = await circomkit.WitnessTester(`SubBytes`, {
+      file: "key_expansion",
+      template: "KeyExpansion",
+      params: [8, 14],
+    });
+    console.log("#constraints:", await circuit.getConstraintCount());
+    const 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,
+    ];
+
+    const keyExpanded = [
+      [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],
+      [0x9b, 0xa3, 0x54, 0x11],
+      [0x8e, 0x69, 0x25, 0xaf],
+      [0xa5, 0x1a, 0x8b, 0x5f],
+      [0x20, 0x67, 0xfc, 0xde],
+      [0xa8, 0xb0, 0x9c, 0x1a],
+      [0x93, 0xd1, 0x94, 0xcd],
+      [0xbe, 0x49, 0x84, 0x6e],
+      [0xb7, 0x5d, 0x5b, 0x9a],
+      [0xd5, 0x9a, 0xec, 0xb8],
+      [0x5b, 0xf3, 0xc9, 0x17],
+      [0xfe, 0xe9, 0x42, 0x48],
+      [0xde, 0x8e, 0xbe, 0x96],
+      [0xb5, 0xa9, 0x32, 0x8a],
+      [0x26, 0x78, 0xa6, 0x47],
+      [0x98, 0x31, 0x22, 0x29],
+      [0x2f, 0x6c, 0x79, 0xb3],
+      [0x81, 0x2c, 0x81, 0xad],
+      [0xda, 0xdf, 0x48, 0xba],
+      [0x24, 0x36, 0x0a, 0xf2],
+      [0xfa, 0xb8, 0xb4, 0x64],
+      [0x98, 0xc5, 0xbf, 0xc9],
+      [0xbe, 0xbd, 0x19, 0x8e],
+      [0x26, 0x8c, 0x3b, 0xa7],
+      [0x09, 0xe0, 0x42, 0x14],
+      [0x68, 0x00, 0x7b, 0xac],
+      [0xb2, 0xdf, 0x33, 0x16],
+      [0x96, 0xe9, 0x39, 0xe4],
+      [0x6c, 0x51, 0x8d, 0x80],
+      [0xc8, 0x14, 0xe2, 0x04],
+      [0x76, 0xa9, 0xfb, 0x8a],
+      [0x50, 0x25, 0xc0, 0x2d],
+      [0x59, 0xc5, 0x82, 0x39],
+      [0xde, 0x13, 0x69, 0x67],
+      [0x6c, 0xcc, 0x5a, 0x71],
+      [0xfa, 0x25, 0x63, 0x95],
+      [0x96, 0x74, 0xee, 0x15],
+      [0x58, 0x86, 0xca, 0x5d],
+      [0x2e, 0x2f, 0x31, 0xd7],
+      [0x7e, 0x0a, 0xf1, 0xfa],
+      [0x27, 0xcf, 0x73, 0xc3],
+      [0x74, 0x9c, 0x47, 0xab],
+      [0x18, 0x50, 0x1d, 0xda],
+      [0xe2, 0x75, 0x7e, 0x4f],
+      [0x74, 0x01, 0x90, 0x5a],
+      [0xca, 0xfa, 0xaa, 0xe3],
+      [0xe4, 0xd5, 0x9b, 0x34],
+      [0x9a, 0xdf, 0x6a, 0xce],
+      [0xbd, 0x10, 0x19, 0x0d],
+      [0xfe, 0x48, 0x90, 0xd1],
+      [0xe6, 0x18, 0x8d, 0x0b],
+      [0x04, 0x6d, 0xf3, 0x44],
+      [0x70, 0x6c, 0x63, 0x1e],
+    ];
+
+    await circuit.expectPass({ key }, { keyExpanded });
+  });
 });