From 4f8ffd4fb328bdbacaa8ea516efc79adb44c1ca8 Mon Sep 17 00:00:00 2001
From: Bruno Valente <140794260+bruno-valante@users.noreply.github.com>
Date: Thu, 5 Sep 2024 18:33:09 +0800
Subject: [PATCH 1/2] Fix Critical Security Vulnerability Caused by Circom
 Operator Misuse

---
 circuits/ctr.circom           |  73 +++++++++++++----
 circuits/finite_field.circom  | 149 ++++++++++++++++++++++++++++++++++
 circuits/key_expansion.circom |  10 +--
 circuits/mix_columns.circom   |  18 ++--
 circuits/sbox128.circom       |  27 ++----
 circuits/tbox.circom          |  38 ++-------
 circuits/utils.circom         |   5 +-
 7 files changed, 235 insertions(+), 85 deletions(-)
 create mode 100644 circuits/finite_field.circom

diff --git a/circuits/ctr.circom b/circuits/ctr.circom
index c388d99..96f05a0 100644
--- a/circuits/ctr.circom
+++ b/circuits/ctr.circom
@@ -2,6 +2,8 @@ pragma circom 2.1.9;
 
 include "cipher.circom";
 include "transformations.circom";
+include "circomlib/circuits/comparators.circom";
+include "circomlib/circuits/bitify.circom";
 
 template EncryptCTR(l,nk){
         signal input plainText[l];
@@ -9,6 +11,24 @@ template EncryptCTR(l,nk){
         signal input key[nk * 4];
         signal output cipher[l];
 
+        component checkPlainText[l];
+        for (var i = 0; i < l; i++) {
+                checkPlainText[i] = Num2Bits(8);
+                checkPlainText[i].in <== plainText[i];
+        }
+
+        component checkIv[16];
+        for (var i = 0; i < 16; i++) {
+                checkIv[i] = Num2Bits(8);
+                checkIv[i].in <== iv[i];
+        }
+
+        component checkKey[nk * 4];
+        for (var i = 0; i < nk * 4; i++) {
+                checkKey[i] = Num2Bits(8);
+                checkKey[i].in <== key[i];
+        }
+
         var n = l\16;
         if(l%16 > 0){
                 n = n + 1;
@@ -114,31 +134,52 @@ template AddCipher(){
     }
 }
 
+template ByteInc() {
+        signal input in;
+        signal input control;
+        signal output out;
+        signal output carry;
+
+        signal added;
+        added <== in + control;
+
+        signal addedDiff;
+        addedDiff <== added - 256;
+        carry <== IsZero()(addedDiff);
+
+        out <== added - carry * 256;
+}
+
 // converts iv to counter blocks
 // iv is 16 bytes
 template GenerateCounterBlocks(n){
         assert(n < 0xffffffff);
         signal input iv[16];
+        signal blockNonce[n][16];
         signal output counterBlocks[n][4][4];
-
-        var ivr[16] = iv;
-
         component toBlocks[n];
+        
+        component ivByteInc[n-1][16];
+
+
+        toBlocks[0] = ToBlocks(16);
+        toBlocks[0].stream <== iv;
+        counterBlocks[0] <== toBlocks[0].blocks[0];
+
+        for (var i = 1; i < n; i++) {
+                for (var j = 15; j >= 0; j--) {
+                        ivByteInc[i-1][j] = ByteInc();
+                        ivByteInc[i-1][j].in <== toBlocks[i-1].stream[j];
+                        if (j==15) {
+                                ivByteInc[i-1][j].control <== 1;
+                        } else {
+                                ivByteInc[i-1][j].control <== ivByteInc[i-1][j+1].carry;
+                        }
+                        blockNonce[i][j] <== ivByteInc[i-1][j].out;
+                }        
 
-        for (var i = 0; i < n; i++) {
                 toBlocks[i] = ToBlocks(16);
-                toBlocks[i].stream <-- ivr;
+                toBlocks[i].stream <== blockNonce[i];
                 counterBlocks[i] <== toBlocks[i].blocks[0];
-                ivr[15] = (ivr[15] + 1)%256;
-                if (ivr[15] == 0){
-                        ivr[14] = (ivr[14] + 1)%256;
-                        if (ivr[14] == 0){
-                                ivr[13] = (ivr[13] + 1)%256;
-                                if (ivr[13] == 0){
-                                        ivr[12] = (ivr[12] + 1)%256;
-                                }
-                        }
-                }
-
         }
 }
\ No newline at end of file
diff --git a/circuits/finite_field.circom b/circuits/finite_field.circom
new file mode 100644
index 0000000..9916350
--- /dev/null
+++ b/circuits/finite_field.circom
@@ -0,0 +1,149 @@
+pragma circom 2.1.9;
+
+include "circomlib/circuits/bitify.circom";
+
+template FieldAddConst(c) {
+    signal input in[8];
+    signal input control;
+    signal output out[8];
+
+    for (var i=0; i<8; i++) {
+        if(c & (1<<i) != 0) {
+            out[i] <== in[i] + control - 2 * in[i] * control;
+        } else {
+            out[i] <== in[i];
+        }
+    }
+}
+
+template FieldMul2() {
+    signal input in;
+    signal output out;
+
+    signal inBits[8];
+    inBits <== Num2Bits(8)(in);
+
+    component reduce = FieldAddConst(0x1b);
+    reduce.in[0] <== 0;
+    for (var i = 1; i < 8; i++) {
+        reduce.in[i] <== inBits[i-1];
+    }
+    reduce.control <== inBits[7];
+    out <== Bits2Num(8)(reduce.out);
+}
+
+template FieldMul3() {
+    signal input in;
+    signal output out;
+
+    signal inBits[8] <== Num2Bits(8)(in);
+
+    component reduce = FieldAddConst(0x1b);
+    reduce.in[0] <== inBits[0];
+    for (var i = 1; i < 8; i++) {
+        reduce.in[i] <== inBits[i-1] + inBits[i] - 2 * inBits[i-1] * inBits[i];
+    }
+    reduce.control <== inBits[7];
+    out <== Bits2Num(8)(reduce.out);
+}
+
+template IsOdd(nBits) {
+    signal input in;
+    signal output out;
+    if (nBits == 1) {
+        out <== in;
+    } else {
+        signal bits[nBits] <== Num2Bits(nBits)(in);
+        out <== bits[0];
+    }
+}
+
+template FieldMul() {
+    signal input a;
+    signal input b;
+    signal inBits[2][8];
+    signal output out;
+
+    inBits[0] <== Num2Bits(8)(a);
+    inBits[1] <== Num2Bits(8)(b);
+
+    var power[15] = [0x1, 0x2, 0x4, 0x8, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d, 0x9a];
+
+    signal mulMatrix[8][8];
+    var outLinesLc[8];
+    for (var i = 0; i < 8; i++) {
+        outLinesLc[i] = 0;
+    }
+    for (var i = 0; i < 8; i++) {
+        for (var j = 0; j < 8; j++) {
+            mulMatrix[i][j] <== inBits[0][i] * inBits[1][j];
+            for (var t = 0; t < 8; t++) {
+                if (power[i+j] & (1 << t) != 0) {
+                    outLinesLc[t] += mulMatrix[i][j];
+                }
+            }
+        }    
+    }
+    signal outBitsUnreduced[8];
+    signal outBits[8];
+    for (var i = 0; i < 8; i++) {
+        outBitsUnreduced[i] <== outLinesLc[i];
+        outBits[i] <== IsOdd(6)(outBitsUnreduced[i]);
+    }
+
+    out <== Bits2Num(8)(outBits);
+}
+
+template FieldInv() {
+    signal input in;
+    signal output out;
+
+    var inv[256] = [0x00, 0x01, 0x8d, 0xf6, 0xcb, 0x52, 0x7b, 0xd1, 0xe8, 0x4f, 0x29, 0xc0, 0xb0, 0xe1, 0xe5, 0xc7,
+                    0x74, 0xb4, 0xaa, 0x4b, 0x99, 0x2b, 0x60, 0x5f, 0x58, 0x3f, 0xfd, 0xcc, 0xff, 0x40, 0xee, 0xb2,
+                    0x3a, 0x6e, 0x5a, 0xf1, 0x55, 0x4d, 0xa8, 0xc9, 0xc1, 0x0a, 0x98, 0x15, 0x30, 0x44, 0xa2, 0xc2,
+                    0x2c, 0x45, 0x92, 0x6c, 0xf3, 0x39, 0x66, 0x42, 0xf2, 0x35, 0x20, 0x6f, 0x77, 0xbb, 0x59, 0x19,
+                    0x1d, 0xfe, 0x37, 0x67, 0x2d, 0x31, 0xf5, 0x69, 0xa7, 0x64, 0xab, 0x13, 0x54, 0x25, 0xe9, 0x09,
+                    0xed, 0x5c, 0x05, 0xca, 0x4c, 0x24, 0x87, 0xbf, 0x18, 0x3e, 0x22, 0xf0, 0x51, 0xec, 0x61, 0x17,
+                    0x16, 0x5e, 0xaf, 0xd3, 0x49, 0xa6, 0x36, 0x43, 0xf4, 0x47, 0x91, 0xdf, 0x33, 0x93, 0x21, 0x3b,
+                    0x79, 0xb7, 0x97, 0x85, 0x10, 0xb5, 0xba, 0x3c, 0xb6, 0x70, 0xd0, 0x06, 0xa1, 0xfa, 0x81, 0x82,
+                    0x83, 0x7e, 0x7f, 0x80, 0x96, 0x73, 0xbe, 0x56, 0x9b, 0x9e, 0x95, 0xd9, 0xf7, 0x02, 0xb9, 0xa4,
+                    0xde, 0x6a, 0x32, 0x6d, 0xd8, 0x8a, 0x84, 0x72, 0x2a, 0x14, 0x9f, 0x88, 0xf9, 0xdc, 0x89, 0x9a,
+                    0xfb, 0x7c, 0x2e, 0xc3, 0x8f, 0xb8, 0x65, 0x48, 0x26, 0xc8, 0x12, 0x4a, 0xce, 0xe7, 0xd2, 0x62,
+                    0x0c, 0xe0, 0x1f, 0xef, 0x11, 0x75, 0x78, 0x71, 0xa5, 0x8e, 0x76, 0x3d, 0xbd, 0xbc, 0x86, 0x57,
+                    0x0b, 0x28, 0x2f, 0xa3, 0xda, 0xd4, 0xe4, 0x0f, 0xa9, 0x27, 0x53, 0x04, 0x1b, 0xfc, 0xac, 0xe6,
+                    0x7a, 0x07, 0xae, 0x63, 0xc5, 0xdb, 0xe2, 0xea, 0x94, 0x8b, 0xc4, 0xd5, 0x9d, 0xf8, 0x90, 0x6b,
+                    0xb1, 0x0d, 0xd6, 0xeb, 0xc6, 0x0e, 0xcf, 0xad, 0x08, 0x4e, 0xd7, 0xe3, 0x5d, 0x50, 0x1e, 0xb3,
+                    0x5b, 0x23, 0x38, 0x34, 0x68, 0x46, 0x03, 0x8c, 0xdd, 0x9c, 0x7d, 0xa0, 0xcd, 0x1a, 0x41, 0x1c];
+
+    out <-- inv[in];
+    signal checkRes <== FieldMul()(in, out);
+    signal isZeroIn <== IsZero()(in);
+    signal isZeroOut <== IsZero()(out);
+    signal checkZero <== isZeroIn * isZeroOut;
+    (1 - checkRes) * (1 - checkZero) === 0;
+}
+
+template AffineTransform() {
+    signal input inBits[8];
+    signal output outBits[8];
+
+    var matrix[8][8] = [[1, 0, 0, 0, 1, 1, 1, 1],
+                        [1, 1, 0, 0, 0, 1, 1, 1],
+                        [1, 1, 1, 0, 0, 0, 1, 1],
+                        [1, 1, 1, 1, 0, 0, 0, 1],
+                        [1, 1, 1, 1, 1, 0, 0, 0],
+                        [0, 1, 1, 1, 1, 1, 0, 0],
+                        [0, 0, 1, 1, 1, 1, 1, 0],
+                        [0, 0, 0, 1, 1, 1, 1, 1]];
+    var offset[8] = [1, 1, 0, 0, 0, 1, 1, 0];
+    for (var i = 0; i < 8; i++) {
+        var lc = 0;
+        for (var j = 0; j < 8; j++) {
+            if (matrix[i][j] == 1) {
+                lc += inBits[j];
+            }
+        }
+        lc += offset[i];
+        outBits[i] <== IsOdd(3)(lc);
+    }
+}
\ No newline at end of file
diff --git a/circuits/key_expansion.circom b/circuits/key_expansion.circom
index 81c8c14..2e28c32 100644
--- a/circuits/key_expansion.circom
+++ b/circuits/key_expansion.circom
@@ -59,7 +59,7 @@ template KeyExpansion(nk,nr) {
     
     for (var round = 1; round <= effectiveRounds; round++) {
         var outputWordLen = round == effectiveRounds ? 4 : nk; 
-        nextRound[round - 1] = NextRound(nk, outputWordLen);
+        nextRound[round - 1] = NextRound(nk, outputWordLen, round);
 
         for (var i = 0; i < nk; i++) {
             for (var j = 0; j < 4; j++) {
@@ -67,8 +67,6 @@ template KeyExpansion(nk,nr) {
             }
         }
 
-        nextRound[round - 1].round <== round;
-
         for (var i = 0; i < outputWordLen; i++) {
             for (var j = 0; j < 4; j++) {
                 keyExpanded[(round * nk) + i][j] <== nextRound[round - 1].nextKey[i][j];
@@ -79,9 +77,8 @@ template KeyExpansion(nk,nr) {
 
 // @param nk: number of keys which can be 4, 6, 8
 // @param o: number of output words which can be 4 or nk
-template NextRound(nk, o){
+template NextRound(nk, o, round){
     signal input key[nk][4]; 
-    signal input round;
     signal output nextKey[o][4];
 
     component rotateWord = Rotate(1, 4);
@@ -93,8 +90,7 @@ template NextRound(nk, o){
     substituteWord[0] = SubstituteWord();
     substituteWord[0].bytes <== rotateWord.rotated;
 
-    component rcon = RCon();
-    rcon.round <== round; 
+    component rcon = RCon(round);
 
     component xorWord[o + 1];
     xorWord[0] = XorWord();
diff --git a/circuits/mix_columns.circom b/circuits/mix_columns.circom
index ec8ff70..0a28804 100644
--- a/circuits/mix_columns.circom
+++ b/circuits/mix_columns.circom
@@ -55,10 +55,10 @@ template S0(){
     }
 
     num2bits[0] = Num2Bits(8);
-    num2bits[0].in <-- TBox(0, in[0]);
+    num2bits[0].in <== TBox(0)(in[0]);
 
     num2bits[1] = Num2Bits(8);
-    num2bits[1].in <-- TBox(1, in[1]);
+    num2bits[1].in <== TBox(1)(in[1]);
 
     xor[0] = XorBits();
     xor[0].a <== num2bits[0].out;
@@ -92,10 +92,10 @@ template S1(){
     num2bits[0].in <== in[0];
 
     num2bits[1] = Num2Bits(8);
-    num2bits[1].in <-- TBox(0, in[1]);
+    num2bits[1].in <== TBox(0)(in[1]);
 
     num2bits[2] = Num2Bits(8);
-    num2bits[2].in <-- TBox(1, in[2]);
+    num2bits[2].in <== TBox(1)(in[2]);
 
     num2bits[3] = Num2Bits(8);
     num2bits[3].in <== in[3];
@@ -134,10 +134,10 @@ template S2() {
     }
 
     num2bits[2] = Num2Bits(8);
-    num2bits[2].in <-- TBox(0, in[2]);
+    num2bits[2].in <== TBox(0)(in[2]);
 
     num2bits[3] = Num2Bits(8);
-    num2bits[3].in <-- TBox(1, in[3]);
+    num2bits[3].in <== TBox(1)(in[3]);
 
     xor[0] = XorBits();
     xor[0].a <== num2bits[0].out;
@@ -173,10 +173,10 @@ template S3() {
     }
 
     num2bits[0] = Num2Bits(8);
-    num2bits[0].in <-- TBox(1, in[0]);
+    num2bits[0].in <== TBox(1)(in[0]);
 
     num2bits[3] = Num2Bits(8);
-    num2bits[3].in <-- TBox(0, in[3]);
+    num2bits[3].in <== TBox(0)(in[3]);
 
     xor[0] = XorBits();
     xor[0].a <== num2bits[0].out;
@@ -187,7 +187,7 @@ template S3() {
     xor[1].b <== num2bits[2].out;
 
     xor[2] = XorBits();
-    xor[2].a <-- num2bits[3].out;
+    xor[2].a <== num2bits[3].out;
     xor[2].b <== xor[1].out;
 
     component b2n = Bits2Num(8);
diff --git a/circuits/sbox128.circom b/circuits/sbox128.circom
index c949f0e..40def67 100644
--- a/circuits/sbox128.circom
+++ b/circuits/sbox128.circom
@@ -1,29 +1,14 @@
 pragma circom 2.1.9;
 
+include "finite_field.circom";
 include "circomlib/circuits/comparators.circom";
 
 template SBox128() {
     signal input in;
     signal output out;
 
-    var sbox[16 * 16] = [
-        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
-    ];
-    
-    out <-- sbox[in];
-}
+    signal inv <== FieldInv()(in);
+    signal invBits[8] <== Num2Bits(8)(inv);
+    signal outBits[8] <== AffineTransform()(invBits);
+    out <== Bits2Num(8)(outBits);
+}
\ No newline at end of file
diff --git a/circuits/tbox.circom b/circuits/tbox.circom
index 976fbea..e4df15a 100644
--- a/circuits/tbox.circom
+++ b/circuits/tbox.circom
@@ -1,38 +1,18 @@
 pragma circom 2.1.9;
 
 include "transformations.circom";
+include "finite_field.circom";
 include "circomlib/circuits/bitify.circom";
 
 //tbox[0] =>> multiplication by 2
 //tbox[1] =>> multiplication by 3
-function TBox(index, subbyte) {
-    var tbox[2][256] = [
-        [
-            0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60,
-            62, 64, 66, 68, 70, 72, 74, 76, 78, 80, 82, 84, 86, 88, 90, 92, 94, 96, 98, 100, 102, 104, 106, 108, 110, 112, 114,
-            116, 118, 120, 122, 124, 126, 128, 130, 132, 134, 136, 138, 140, 142, 144, 146, 148, 150, 152, 154, 156, 158, 160,
-            162, 164, 166, 168, 170, 172, 174, 176, 178, 180, 182, 184, 186, 188, 190, 192, 194, 196, 198, 200, 202, 204, 206,
-            208, 210, 212, 214, 216, 218, 220, 222, 224, 226, 228, 230, 232, 234, 236, 238, 240, 242, 244, 246, 248, 250, 252,
-            254, 27, 25, 31, 29, 19, 17, 23, 21, 11, 9, 15, 13, 3, 1, 7, 5, 59, 57, 63, 61, 51, 49, 55, 53, 43, 41, 47, 45, 35,
-            33, 39, 37, 91, 89, 95, 93, 83, 81, 87, 85, 75, 73, 79, 77, 67, 65, 71, 69, 123, 121, 127, 125, 115, 113, 119, 117,
-            107, 105, 111, 109, 99, 97, 103, 101, 155, 153, 159, 157, 147, 145, 151, 149, 139, 137, 143, 141, 131, 129, 135, 133,
-            187, 185, 191, 189, 179, 177, 183, 181, 171, 169, 175, 173, 163, 161, 167, 165, 219, 217, 223, 221, 211, 209, 215,
-            213, 203, 201, 207, 205, 195, 193, 199, 197, 251, 249, 255, 253, 243, 241, 247, 245, 235, 233, 239, 237, 227, 225,
-            231, 229
-        ],
-        [
-            0, 3, 6, 5, 12, 15, 10, 9, 24, 27, 30, 29, 20, 23, 18, 17, 48, 51, 54, 53, 60, 63, 58, 57, 40, 43, 46, 45, 36, 39, 34,
-            33, 96, 99, 102, 101, 108, 111, 106, 105, 120, 123, 126, 125, 116, 119, 114, 113, 80, 83, 86, 85, 92, 95, 90, 89, 72,
-            75, 78, 77, 68, 71, 66, 65, 192, 195, 198, 197, 204, 207, 202, 201, 216, 219, 222, 221, 212, 215, 210, 209, 240, 243,
-            246, 245, 252, 255, 250, 249, 232, 235, 238, 237, 228, 231, 226, 225, 160, 163, 166, 165, 172, 175, 170, 169, 184,
-            187, 190, 189, 180, 183, 178, 177, 144, 147, 150, 149, 156, 159, 154, 153, 136, 139, 142, 141, 132, 135, 130, 129,
-            155, 152, 157, 158, 151, 148, 145, 146, 131, 128, 133, 134, 143, 140, 137, 138, 171, 168, 173, 174, 167, 164, 161,
-            162, 179, 176, 181, 182, 191, 188, 185, 186, 251, 248, 253, 254, 247, 244, 241, 242, 227, 224, 229, 230, 239, 236,
-            233, 234, 203, 200, 205, 206, 199, 196, 193, 194, 211, 208, 213, 214, 223, 220, 217, 218, 91, 88, 93, 94, 87, 84, 81,
-            82, 67, 64, 69, 70, 79, 76, 73, 74, 107, 104, 109, 110, 103, 100, 97, 98, 115, 112, 117, 118, 127, 124, 121, 122, 59,
-            56, 61, 62, 55, 52, 49, 50, 35, 32, 37, 38, 47, 44, 41, 42, 11, 8, 13, 14, 7, 4, 1, 2, 19, 16, 21, 22, 31, 28, 25, 26
-        ]
-    ];
+template TBox(index) {
+    signal input in;
+    signal output out;
 
-    return tbox[index][subbyte];
+    if (index == 0) {
+        out <== FieldMul2()(in);
+    } else if (index == 1) {
+        out <== FieldMul3()(in);
+    }
 }
\ No newline at end of file
diff --git a/circuits/utils.circom b/circuits/utils.circom
index d9b6bea..1a4ada5 100644
--- a/circuits/utils.circom
+++ b/circuits/utils.circom
@@ -48,8 +48,7 @@ template SubstituteWord() {
 }
 
 // Outputs a round constant for a given round number
-template RCon() {
-    signal input round;
+template RCon(round) {
     signal output out[4];
 
     assert(round > 0 && round <= 10);
@@ -67,7 +66,7 @@ template RCon() {
         [0x36, 0x00, 0x00, 0x00]
     ];
 
-    out <-- rcon[round-1];
+    out <== rcon[round-1];
 }
 
 

From f1668bd20f4fd36963a3376c4f456460f4447b73 Mon Sep 17 00:00:00 2001
From: Bruno Valente <140794260+bruno-valante@users.noreply.github.com>
Date: Thu, 5 Sep 2024 19:11:22 +0800
Subject: [PATCH 2/2] Add comments

---
 circuits/finite_field.circom | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/circuits/finite_field.circom b/circuits/finite_field.circom
index 9916350..f7cf239 100644
--- a/circuits/finite_field.circom
+++ b/circuits/finite_field.circom
@@ -2,13 +2,16 @@ pragma circom 2.1.9;
 
 include "circomlib/circuits/bitify.circom";
 
+// Finite field addition, the signal variable plus a compile-time constant
 template FieldAddConst(c) {
     signal input in[8];
+    // control bit, if 0, then do not perform addition
     signal input control;
     signal output out[8];
 
     for (var i=0; i<8; i++) {
         if(c & (1<<i) != 0) {
+            // XOR operation
             out[i] <== in[i] + control - 2 * in[i] * control;
         } else {
             out[i] <== in[i];
@@ -16,6 +19,9 @@ template FieldAddConst(c) {
     }
 }
 
+// Finite field multiplication by 2 operation for AES. This involves left-shifting 'input' by 1 (input << 1), 
+// and then XORing with 0x1B if the most significate bit is 1. This is because the irreducible polynomial 
+// for AES's finite field (GF(2^8)) is x^8 + x^4 + x^3 + x + 1.
 template FieldMul2() {
     signal input in;
     signal output out;
@@ -32,6 +38,8 @@ template FieldMul2() {
     out <== Bits2Num(8)(reduce.out);
 }
 
+// Finite field multiplication by 3 operation for AES. This involves (input << 1) ⊕ input and then XORing 
+// with 0x1B if the most significate bit is 1.
 template FieldMul3() {
     signal input in;
     signal output out;
@@ -47,6 +55,7 @@ template FieldMul3() {
     out <== Bits2Num(8)(reduce.out);
 }
 
+// Determine the parity (odd or even) of an integer that can be accommodated within 'nBits' bits.
 template IsOdd(nBits) {
     signal input in;
     signal output out;
@@ -58,6 +67,7 @@ template IsOdd(nBits) {
     }
 }
 
+// Finite field multiplication. 
 template FieldMul() {
     signal input a;
     signal input b;
@@ -67,6 +77,7 @@ template FieldMul() {
     inBits[0] <== Num2Bits(8)(a);
     inBits[1] <== Num2Bits(8)(b);
 
+    // List of finite field elements obtained by successively doubling, starting from 1.
     var power[15] = [0x1, 0x2, 0x4, 0x8, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d, 0x9a];
 
     signal mulMatrix[8][8];
@@ -74,6 +85,7 @@ template FieldMul() {
     for (var i = 0; i < 8; i++) {
         outLinesLc[i] = 0;
     }
+    // Apply elementary multiplication
     for (var i = 0; i < 8; i++) {
         for (var j = 0; j < 8; j++) {
             mulMatrix[i][j] <== inBits[0][i] * inBits[1][j];
@@ -88,12 +100,15 @@ template FieldMul() {
     signal outBits[8];
     for (var i = 0; i < 8; i++) {
         outBitsUnreduced[i] <== outLinesLc[i];
+        // Each element in 'outLinesLc' is incremented by a known constant number of 
+        // elements from 'mulMatrix', less than 31.
         outBits[i] <== IsOdd(6)(outBitsUnreduced[i]);
     }
 
     out <== Bits2Num(8)(outBits);
 }
 
+// Finite Field Inversion. Specially, if the input is 0, the output is also 0.
 template FieldInv() {
     signal input in;
     signal output out;
@@ -115,14 +130,19 @@ template FieldInv() {
                     0xb1, 0x0d, 0xd6, 0xeb, 0xc6, 0x0e, 0xcf, 0xad, 0x08, 0x4e, 0xd7, 0xe3, 0x5d, 0x50, 0x1e, 0xb3,
                     0x5b, 0x23, 0x38, 0x34, 0x68, 0x46, 0x03, 0x8c, 0xdd, 0x9c, 0x7d, 0xa0, 0xcd, 0x1a, 0x41, 0x1c];
 
+    // Obtain an unchecked result from a lookup table
     out <-- inv[in];
+    // Compute the product of the input and output, expected to be 1
     signal checkRes <== FieldMul()(in, out);
+    // For the special case when the input is 0, both input and output should be 0
     signal isZeroIn <== IsZero()(in);
     signal isZeroOut <== IsZero()(out);
     signal checkZero <== isZeroIn * isZeroOut;
+    // Ensure that either the product is 1 or both input and output are 0, satisfying at least one condition
     (1 - checkRes) * (1 - checkZero) === 0;
 }
 
+// AffineTransform required by the S-box computation.
 template AffineTransform() {
     signal input inBits[8];
     signal output outBits[8];