Merge working tests (#73)

* stuff * Working tests for aes-gcm and gctr * Remove logs --------- Co-authored-by: devloper <[email protected]>
pluto · Sep 17, 2024 · 3c5d812 · 3c5d812
1 parent b5ee77c
commit 3c5d812
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 41 deletions.
diff --git a/circuits/aes-gcm/aes-gcm.circom b/circuits/aes-gcm/aes-gcm.circom
@@ -47,7 +47,6 @@ template AESGCM(l) {
 
     // Step 2: Define a block, J0 with 96 bits of iv and 32 bits of 0s
     // you can of the 96bits as a nonce and the 32 bits of 0s as an integer counter
-    // TODO(WJ 2024-09-16): make this a block of bytes not bits
     component J0builder = ToBlocks(16);
     for (var i = 0; i < 12; i++) {
         J0builder.stream[i] <== iv[i];
@@ -57,13 +56,15 @@ template AESGCM(l) {
     }
     component J0WordIncrementer = IncrementWord();
     J0WordIncrementer.in <== J0builder.blocks[0][3];
+
+    component J0WordIncrementer2 = IncrementWord();
+    J0WordIncrementer2.in <== J0WordIncrementer.out;
+
     signal J0[4][4];
     for (var i = 0; i < 3; i++) {
         J0[i] <== J0builder.blocks[0][i];
     }
-    // TODO(WJ 2024-09-16): maybe need to increment this again before passing to gctr. Check section 7.3 of nist spec
-    J0[3] <== J0WordIncrementer.out;
-
+    J0[3] <== J0WordIncrementer2.out;
 
     // Step 3: Let C = GCTRK(inc32(J0), P)
     component gctr = GCTR(l, 4);

diff --git a/circuits/aes-gcm/gctr.circom b/circuits/aes-gcm/gctr.circom
@@ -47,21 +47,20 @@ template GCTR(INPUT_LEN, nk) {
     // total number of bits in the plaintext blocks
     var bytesExcludingLastBlock = 16 * (nBlocks);
 
+    log("INPUT_LEN");
+    log(INPUT_LEN);
+    log("nBlocks");
     log(nBlocks);
+    log("lastBlockSize");
     log(lastBlockSize);
+    log("bytesExcludingLastBlock");
     log(bytesExcludingLastBlock);
-    // TODO(WJ 2024-09-16): Get this assert to pass
-    assert(INPUT_LEN == nBlocks * 16 + lastBlockSize);
 
-    // last block of plaintext
-    signal tempLastBlock[lastBlockSize];
-    for (var i = 0; i < lastBlockSize; i++) {
-        tempLastBlock[i] <== plainText[bytesExcludingLastBlock + i];
-    }
+    assert(INPUT_LEN == nBlocks * 16 + lastBlockSize);
 
     // generate plaintext blocks
     // note to not use the last block of plaintext
-    // because it will be padded byt the toBlocks components
+    // because it will be padded by the toBlocks components
     component plainTextBlocks = ToBlocks(INPUT_LEN);
     plainTextBlocks.stream <== plainText;
 
@@ -75,6 +74,8 @@ template GCTR(INPUT_LEN, nk) {
     // component to increment the last word of the counter block
     component inc32[nBlocks];
     // For i = 2 to nBlocks, let CBi = inc32(CBi-1).
+
+    // TODO: Actually test me on a block larger than 16 bytes. 
     for (var i = 1; i < nBlocks; i++) {
         inc32[i] = IncrementWord();
         inc32[i].in <== CounterBlocks[i - 1][3]; // idea: use the counterblock here directly so that we don't need to use this toCounterblock thing
@@ -87,13 +88,14 @@ template GCTR(INPUT_LEN, nk) {
         CounterBlocks[i][3] <== inc32[i].out;
     }
 
-
     // Convert blocks of 16 bytes to stream
     component toStream = ToStream(nBlocks, bytesExcludingLastBlock);
     // Step 2: Encrypt each counter block with the key
     component aes[nBlocks+1]; // +1 for the last block
     component AddCipher[nBlocks];
 
+    // NOTE: All this code does for one block is encrypt and xor, 
+    // which is identical to CTR. 
     for (var i = 0; i < nBlocks; i++) {
         // encrypt counter block
         aes[i] = Cipher(nk);
@@ -112,32 +114,26 @@ template GCTR(INPUT_LEN, nk) {
     // Step 3: Handle the last block separately
     // Y* = X* ⊕ MSBlen(X*) (CIPH_K (CB_n*))
 
-    // encrypt the last counter block
+    // TODO: When we only have one block, this double Cipher's. We shouldnnt do this when l % 16 == 0
+    // encrypt the last counter block 
     aes[nBlocks] = Cipher(nk);
     aes[nBlocks].key <== key;
     aes[nBlocks].block <== CounterBlocks[nBlocks-1];
+    component aesCipherToStream = ToStream(1, 16);
+    aesCipherToStream.blocks[0] <== aes[nBlocks].cipher;
 
     // XOR the cipher with the last chunk of unpadded plaintext
-    component aesCipherToStream = ToStream(1, 16);
     component addLastCipher = XorMultiple(2, lastBlockSize);
-    aesCipherToStream.blocks[0] <== aes[nBlocks].cipher;
     for (var i = 0; i < lastBlockSize; i++) {
         // convert cipher to stream
         addLastCipher.inputs[0][i] <== aesCipherToStream.stream[i];
-        addLastCipher.inputs[1][i] <== tempLastBlock[i];
+        addLastCipher.inputs[1][i] <== plainText[bytesExcludingLastBlock + i];
     }
 
-    /// set the last block in the toStream component
-
-
     for (var i = 0; i < bytesExcludingLastBlock; i++) {
-        // cipher text is of size inpute  lenght 
-        // bitblock = 128 * nblocks
-        // input length = (nblocks-1) * 128 + lastBlockSize
-        // setting all but the last block
-        cipherText[i] <== toStream.stream[i]; // out of bounds error
+        cipherText[i] <== toStream.stream[i];
     }
-     
+
     for (var i = 0; i < lastBlockSize; i++) {
         cipherText[bytesExcludingLastBlock + i] <== addLastCipher.out[i];
     }

diff --git a/circuits/test/aes-gcm.test.ts b/circuits/test/aes-gcm.test.ts
@@ -1,6 +1,6 @@
 import { assert } from "chai";
 import { WitnessTester } from "circomkit";
-import { circomkit } from "./common";
+import { circomkit, hexBytesToBigInt } from "./common";
 
 describe("aes-gcm", () => {
   let circuit: WitnessTester<["key", "iv", "plainText", "aad"], ["cipherText", "tag"]>;
@@ -23,7 +23,7 @@ describe("aes-gcm", () => {
   it("should have correct output", async () => {
     const witness = await circuit.compute({ key: key, iv: iv, plainText: plainText, aad: aad }, ["cipherText", "authTag"])
 
-    assert.deepEqual(witness.cipherText, expected_output)
+    assert.deepEqual(witness.cipherText, hexBytesToBigInt(expected_output))
   });
 
 });

diff --git a/circuits/test/common/index.ts b/circuits/test/common/index.ts
@@ -8,6 +8,13 @@ export const circomkit = new Circomkit({
 
 export { WitnessTester };
 
+export function hexBytesToBigInt(hexBytes: number[]): any[] {
+  return hexBytes.map(byte => {
+      let n = BigInt(byte);
+      return n;
+  });
+}
+
 export function hexToBitArray(hex: string): number[] {
   // Remove '0x' prefix if present and ensure lowercase
   hex = hex.replace(/^0x/i, "").toLowerCase();

diff --git a/circuits/test/gctr.test.ts b/circuits/test/gctr.test.ts
@@ -1,27 +1,28 @@
 import { WitnessTester } from "circomkit";
-import { circomkit } from "./common";
+import { circomkit, hexBytesToBigInt } from "./common";
+import { assert } from "chai";
 
 // signal input key[nk * 4];
 // signal input initialCounterBlock[128];
 // signal input plainText[INPUT_LEN];
 // signal output cipherText[INPUT_LEN];
 describe("GCTR", () => {
-    let circuit: WitnessTester<["key", "initialCounterBlock", "plainText"], ["cipherText"]>;
+    let circuit: WitnessTester<["plainText", "initialCounterBlock", "key"], ["cipherText"]>;
     it("should encrypt the plaintext", async () => {
         circuit = await circomkit.WitnessTester(`GCTR`, {
             file: "aes-gcm/gctr",
             template: "GCTR",
-            params: [129, 4],
+            params: [16, 4],
         });
-        await circuit.expectPass(
-            {
-                key: [0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c],
-                initialCounterBlock: [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],
-                plainText: [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],
-            },
-            {
-                cipherText: [0x69, 0xc4, 0xe0, 0xd8, 0x6a, 0x7b, 0x04, 0x30, 0xd8, 0xcd, 0xb7, 0x80, 0x70, 0xb4, 0xc5, 0x5a],
-            }
-        );
+
+        // GOOD TEST CASE.
+        const key = [0xca, 0xaa, 0x3f, 0x6f, 0xd3, 0x18, 0x22, 0xed, 0x2d, 0x21, 0x25, 0xf2, 0x25, 0xb0, 0x16, 0x9f];
+        const column_wise_icb = [0x7f,0x48,0x12,0x00,0x6d,0x3e,0xfa,0x00,0x90,0x8c,0x55,0x00,0x41,0x14,0x2a,0x02];
+        const pt = [0x84, 0xc9, 0x07, 0xb1, 0x1a, 0xe3, 0xb7, 0x9f,0xc4, 0x45, 0x1d, 0x1b, 0xf1, 0x7f, 0x4a, 0x99];
+        const ct = [0xfd, 0xb4, 0xaa, 0xfa, 0x35, 0x19, 0xd3, 0xc0,0x55, 0xbe, 0x8b, 0x34, 0x77, 0x64, 0xea, 0x33];
+
+        const witness = await circuit.compute({ key: key, initialCounterBlock: column_wise_icb, plainText: pt }, ["cipherText"])
+
+        assert.deepEqual(witness.cipherText, hexBytesToBigInt(ct))
     });
 });