feat: Better aes

circuits.json

@@ -3,5 +3,10 @@
     "file": "aes/polyval",
     "template": "POLYVAL",
     "params": [128]
+  },
+  "cipher_4": {
+    "file": "cipher",
+    "template": "Cipher",
+    "params": [4]
   }
 }

circuits/aes-ctr/cipher.circom

@@ -0,0 +1,106 @@
+pragma circom 2.1.9;
+
+include "key_expansion.circom";
+include "circomlib/circuits/comparators.circom";
+include "circomlib/circuits/bitify.circom";
+include "circomlib/circuits/gates.circom";
+include "transformations.circom";
+include "mix_columns.circom";
+
+// Cipher Process
+// nk: number of keys which can be 4, 6, 8
+// AES 128, 192, 256 have 10, 12, 14 rounds.
+// Input Block   Initial Round Key          Round Key             Final Round Key
+//     │                │                       │                       │
+//     ▼                ▼                       ▼                       ▼
+//  ┌─────────┐    ┌──────────┐ ┌────────┐ ┌──────────┐ ┌────────┐ ┌──────────┐
+//  │ Block   │──► │   Add    │ │  Sub   │ │   Mix    │ │  Sub   │ │   Add    │
+//  │         │    │  Round   │ │ Bytes  │ │ Columns  │ │ Bytes  │ │  Round   │
+//  │         │    │   Key    │ │        │ │          │ │        │ │   Key    │
+//  └─────────┘    └────┬─────┘ └───┬────┘ └────┬─────┘ └───┬────┘ └────┬─────┘
+//                      │           │           │           │           │
+//                      ▼           ▼           ▼           ▼           ▼
+//                 ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐
+//                 │ Round 0 │ │ Round 1 │ │ Round 2 │ │ Round   │ │  Final  │
+//                 │         │ │   to    │ │   to    │ │ Nr - 1  │ │  Round  │
+//                 │         │ │ Nr - 2  │ │ Nr - 1  │ │         │ │         │
+//                 └─────────┘ └─────────┘ └─────────┘ └─────────┘ └────┬────┘
+//                                                                      │
+//                                                                      ▼
+//                                                                 Ciphertext
+
+// @param nk: number of keys which can be 4, 6, 8
+// @inputs block: 4x4 matrix representing the input block
+// @inputs key: array of nk*4 bytes representing the key
+// @outputs cipher: 4x4 matrix representing the output block
+template Cipher(nk){
+        assert(nk == 4 || nk == 6 || nk == 8 );
+        signal input block[4][4];
+        signal input key[nk * 4];
+        signal output cipher[4][4];
+
+        var nr = Rounds(nk);
+
+        component keyExpansion = KeyExpansion(nk,nr);
+        keyExpansion.key <== key;
+
+        component addRoundKey[nr+1]; 
+        component subBytes[nr];
+        component shiftRows[nr];
+        component mixColumns[nr-1];
+
+        signal interBlock[nr][4][4];
+
+        addRoundKey[0] = AddRoundKey();
+        addRoundKey[0].state <== block;
+        for (var i = 0; i < 4; i++) {
+                addRoundKey[0].roundKey[i] <== keyExpansion.keyExpanded[i];
+        }
+
+        interBlock[0] <== addRoundKey[0].newState;
+        for (var i = 1; i < nr; i++) {
+                subBytes[i-1] = SubBlock();
+                subBytes[i-1].state <== interBlock[i-1];
+
+                shiftRows[i-1] = ShiftRows();
+                shiftRows[i-1].state <== subBytes[i-1].newState;
+
+                mixColumns[i-1] = MixColumns();
+                mixColumns[i-1].state <== shiftRows[i-1].newState;
+
+                addRoundKey[i] = AddRoundKey();
+                addRoundKey[i].state <== mixColumns[i-1].out;
+                 for (var j = 0; j < 4; j++) {
+                        addRoundKey[i].roundKey[j] <== keyExpansion.keyExpanded[j + (i * 4)];
+                }
+
+                interBlock[i] <== addRoundKey[i].newState;
+        }
+
+        subBytes[nr-1] = SubBlock();
+        subBytes[nr-1].state <== interBlock[nr-1];
+
+        shiftRows[nr-1] = ShiftRows();
+        shiftRows[nr-1].state <== subBytes[nr-1].newState;
+
+        addRoundKey[nr] = AddRoundKey();
+        addRoundKey[nr].state <== shiftRows[nr-1].newState;
+        for (var i = 0; i < 4; i++) {
+                addRoundKey[nr].roundKey[i] <== keyExpansion.keyExpanded[i + (nr * 4)];
+        }
+
+        cipher <== addRoundKey[nr].newState;
+}
+
+// @param nk: number of keys which can be 4, 6, 8
+// @returns number of rounds
+// AES 128, 192, 256 have 10, 12, 14 rounds.
+function Rounds (nk) {
+    if (nk == 4) {
+       return 10;
+    } else if (nk == 6) {
+        return 12;
+    } else {
+        return 14;
+    }
+}

circuits/aes-ctr/ctr.circom

@@ -0,0 +1,144 @@
+pragma circom 2.1.9;
+
+include "cipher.circom";
+include "transformations.circom";
+
+template EncryptCTR(l,nk){
+        signal input plainText[l];
+        signal input iv[16];
+        signal input key[nk * 4];
+        signal output cipher[l];
+
+        var n = l\16;
+        if(l%16 > 0){
+                n = n + 1;
+        }
+
+        component toBlocks = ToBlocks(l);
+        toBlocks.stream <== plainText;
+
+        component aes[n];
+
+        signal cipherBlocks[n][4][4];
+        component AddCipher[n];
+
+        component generateCtrBlocks = GenerateCounterBlocks(n);
+        generateCtrBlocks.iv <== iv;
+
+        for(var i = 0 ; i < n; i++){
+                aes[i] = Cipher(nk);
+                aes[i].key <== key;
+                aes[i].block <== generateCtrBlocks.counterBlocks[i];
+
+                AddCipher[i] = AddCipher();
+                AddCipher[i].state <== toBlocks.blocks[i];
+                AddCipher[i].cipher <== aes[i].cipher;
+
+                cipherBlocks[i] <== AddCipher[i].newState;
+        }
+
+        component toStream = ToStream(n,l);
+        toStream.blocks <== cipherBlocks;
+
+        cipher <== toStream.stream;
+}
+
+
+//convert stream of plain text to blocks of 16 bytes
+template ToBlocks(l){
+        signal input stream[l];
+
+        var n = l\16;
+        if(l%16 > 0){
+                n = n + 1;
+        }
+        signal output blocks[n][4][4];
+
+        var i, j, k;
+
+        for (var idx = 0; idx < l; idx++) {
+                blocks[i][k][j] <== stream[idx];
+                k = k + 1;
+                if (k == 4){
+                        k = 0;
+                        j = j + 1;
+                        if (j == 4){
+                                j = 0;
+                                i = i + 1;
+                        }
+                }
+        }
+
+        if (l%16 > 0){
+               blocks[i][k][j] <== 1;
+               k = k + 1;
+        }
+}
+
+// convert blocks of 16 bytes to stream of bytes
+template ToStream(n,l){
+        signal input blocks[n][4][4];
+
+        signal output stream[l];
+
+        var i, j, k;
+
+        while(i*16 + j*4 + k < l){
+                stream[i*16 + j*4 + k] <== blocks[i][k][j];
+                k = k + 1;
+                if (k == 4){
+                        k = 0;
+                        j = j + 1;
+                        if (j == 4){
+                                j = 0;
+                                i = i + 1;
+                        }
+                }
+        }
+}
+
+template AddCipher(){
+    signal input state[4][4];
+    signal input cipher[4][4];
+    signal output newState[4][4];
+
+    component xorbyte[4][4];
+
+    for (var i = 0; i < 4; i++) {
+        for (var j = 0; j < 4; j++) {
+            xorbyte[i][j] = XorByte();
+            xorbyte[i][j].a <== state[i][j];
+            xorbyte[i][j].b <== cipher[i][j];
+            newState[i][j] <== xorbyte[i][j].out;
+        }
+    }
+}
+
+// converts iv to counter blocks
+// iv is 16 bytes
+template GenerateCounterBlocks(n){
+        assert(n < 0xffffffff);
+        signal input iv[16];
+        signal output counterBlocks[n][4][4];
+
+        var ivr[16] = iv;
+
+        component toBlocks[n];
+
+        for (var i = 0; i < n; i++) {
+                toBlocks[i] = ToBlocks(16);
+                toBlocks[i].stream <-- ivr;
+                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;
+                                }
+                        }
+                }
+
+        }
+}

circuits/aes-ctr/key_expansion.circom

@@ -0,0 +1,123 @@
+pragma circom 2.1.9;
+
+include "sbox128.circom";
+include "utils.circom";
+
+// Key Expansion Process
+//
+// Original Key (Nk words)
+// ┌───┬───┬───┬───┐
+// │W0 │W1 │W2 │W3 │  (for AES-128, Nk=4)
+// └─┬─┴─┬─┴─┬─┴─┬─┘
+//   │   │   │   │
+//   ▼   ▼   ▼   ▼
+// ┌───────────────────────────────────────────────────────┐
+// │                 Key Expansion Process                 │
+// │                                                       │
+// │  ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐      │
+// │  │RotWord  │ │SubWord  │ │  XOR    │ │  XOR    │      │
+// │  │         │ │         │ │ Rcon(i) │ │ W[i-Nk] │      │
+// │  └────┬────┘ └────┬────┘ └────┬────┘ └────┬────┘      │
+// │       │           │           │           │           │
+// │       └───────────┴───────────┴───────────┘           │
+// │                       │                               │
+// │                       ▼                               │
+// │            ┌─────────────────────┐                    │
+// │            │  New Expanded Key   │                    │
+// │            │       Word          │                    │
+// │            └─────────────────────┘                    │
+// │                       │                               │
+// └───────────────────────┼───────────────────────────────┘
+//                         │
+//                         ▼
+//               Expanded Key Words
+//        ┌───┬───┬───┬───┬───┬───┬───┬───┐
+//        │W4 │W5 │W6 │W7 │W8 │W9 │...│W43│  (for AES-128, 44 words total)
+//        └───┴───┴───┴───┴───┴───┴───┴───┘
+
+
+// @param nk: number of keys which can be 4, 6, 8
+// @param nr: number of rounds which can be 10, 12, 14 for AES 128, 192, 256
+// @inputs key: array of nk*4 bytes representing the key
+// @outputs keyExpanded: array of (nr+1)*4 words i.e for AES 128, 192, 256 it will be 44, 52, 60 words
+template KeyExpansion(nk,nr) {
+    assert(nk == 4 || nk == 6 || nk == 8 );    
+    signal input key[nk * 4];
+
+    var totalWords = (4 * (nr + 1));
+    var effectiveRounds = nk == 4 ? 10 : totalWords\nk;
+
+    signal output keyExpanded[totalWords][4];
+
+    for (var i = 0; i < nk; i++) {
+        for (var j = 0; j < 4; j++) {
+            keyExpanded[i][j] <== key[(4 * i) + j];
+        }
+    }
+
+    component nextRound[effectiveRounds];
+
+    for (var round = 1; round <= effectiveRounds; round++) {
+        var outputWordLen = round == effectiveRounds ? 4 : nk; 
+        nextRound[round - 1] = NextRound(nk, outputWordLen);
+
+        for (var i = 0; i < nk; i++) {
+            for (var j = 0; j < 4; j++) {
+                nextRound[round - 1].key[i][j] <== keyExpanded[(round * nk) + i - nk][j];
+            }
+        }
+
+        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];
+            }
+        }
+    }
+}
+
+// @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){
+    signal input key[nk][4]; 
+    signal input round;
+    signal output nextKey[o][4];
+
+    component rotateWord = Rotate(1, 4);
+    for (var i = 0; i < 4; i++) {
+        rotateWord.bytes[i] <== key[nk - 1][i];
+    }
+
+    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[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];
+        }
+        xorWord[i+1].bytes2 <== key[i];
+
+        for (var j = 0; j < 4; j++) {
+            nextKey[i][j] <== xorWord[i+1].out[j];
+        }
+    }
+}
+
+