feat: generic chunked aes nivc (#44)

* feat: generic chunked AES NIVC * cleanup tests/builds/version * test: AES 2 chunk NIVC_FULL_2 * Update circuits/aes-gcm/nivc/aes-gctr-nivc.circom suggestion from Sambhav Co-authored-by: Sambhav Dusad <[email protected]> * fix: aes-gctr-nivc and test --------- Co-authored-by: Sambhav Dusad <[email protected]>
pluto · Nov 14, 2024 · 52d4bf6 · 52d4bf6
1 parent 73260ab
commit 52d4bf6
Show file tree

Hide file tree

Showing 6 changed files with 183 additions and 22 deletions.
diff --git a/builds/target_1024b/aes_gctr_nivc_1024b.circom b/builds/target_1024b/aes_gctr_nivc_1024b.circom
@@ -2,4 +2,4 @@ pragma circom 2.1.9;
 
 include "../../circuits/aes-gcm/nivc/aes-gctr-nivc.circom";
 
-component main { public [step_in] } = AESGCTRFOLD();
+component main { public [step_in] } = AESGCTRFOLD(1);
diff --git a/builds/target_512b/aes_gctr_nivc_512b.circom b/builds/target_512b/aes_gctr_nivc_512b.circom
@@ -2,4 +2,4 @@ pragma circom 2.1.9;
 
 include "../../circuits/aes-gcm/nivc/aes-gctr-nivc.circom";
 
-component main { public [step_in] } = AESGCTRFOLD();
+component main { public [step_in] } = AESGCTRFOLD(1);
diff --git a/circuits/aes-gcm/nivc/aes-gctr-nivc.circom b/circuits/aes-gcm/nivc/aes-gctr-nivc.circom
@@ -5,35 +5,57 @@ include "../../utils/array.circom";
 include "../../utils/hash.circom";
 
 // Compute AES-GCTR
-template AESGCTRFOLD() {
+template AESGCTRFOLD(NUM_CHUNKS) {
     signal input key[16];
     signal input iv[12];
     signal input aad[16];
+
     signal input ctr[4];
-    signal input plainText[16];
 
-    signal input cipherText[16];
+    signal input plainText[NUM_CHUNKS][16];
+    signal input cipherText[NUM_CHUNKS][16];
 
     signal input step_in[1];
     signal output step_out[1];
 
-    component aes     = AESGCTRFOLDABLE();
-    aes.key         <== key;
-    aes.iv          <== iv;
-    aes.aad         <== aad;
-    aes.plainText   <== plainText;
-    aes.lastCounter <== ctr;
-
-    signal ciphertext_equal_check[16];
-    for(var i = 0 ; i < 16 ; i++) {
-        ciphertext_equal_check[i] <== IsEqual()([aes.cipherText[i], cipherText[i]]);
-        ciphertext_equal_check[i] === 1;
+    component aes[NUM_CHUNKS];
+    for(var i = 0 ; i < NUM_CHUNKS ; i++) {
+        aes[i] = AESGCTRFOLDABLE();
+        if( i == 0) {
+            aes[i].plainText   <== plainText[i];
+            aes[i].lastCounter <== ctr;
+        } else {
+            aes[i].plainText   <== plainText[i];
+            aes[i].lastCounter <== aes[i - 1].counter;
+        }
+        aes[i].key         <== key;
+        aes[i].iv          <== iv;
+        aes[i].aad         <== aad;
+    }
+
+    signal ciphertext_equal_check[NUM_CHUNKS][16];
+    for(var i = 0 ; i < NUM_CHUNKS; i++) {
+        for(var j = 0 ; j < 16 ; j++) {
+            ciphertext_equal_check[i][j] <== IsEqual()([aes[i].cipherText[j], cipherText[i][j]]);
+            ciphertext_equal_check[i][j] === 1;
+        }
     }
 
 
-    var packedPlaintext = 0;
-    for(var i = 0 ; i < 16 ; i++) {
-        packedPlaintext += plainText[i] * 2**(8*i);
+    var packedPlaintext[NUM_CHUNKS];
+    for(var i = 0 ; i < NUM_CHUNKS ; i++) {
+        packedPlaintext[i] = 0;
+        for(var j = 0 ; j < 16 ; j++) {
+            packedPlaintext[i] += plainText[i][j] * 2**(8*j);
+        }
+    }
+    signal hash[NUM_CHUNKS];
+    for(var i = 0 ; i < NUM_CHUNKS ; i++) {
+        if(i == 0) {
+            hash[i] <== PoseidonChainer()([step_in[0],packedPlaintext[i]]);
+        } else {
+            hash[i] <== PoseidonChainer()([hash[i-1], packedPlaintext[i]]);
+        }
     }
-    step_out[0] <== PoseidonChainer()([step_in[0],packedPlaintext]);
+    step_out[0] <== hash[NUM_CHUNKS - 1];
 }
diff --git a/circuits/test/aes-gcm/nivc/aes-gctr-nivc.test.ts b/circuits/test/aes-gcm/nivc/aes-gctr-nivc.test.ts
@@ -21,6 +21,7 @@ describe("aes-gctr-nivc", () => {
         circuit_one_block = await circomkit.WitnessTester("aes-gcm-fold", {
             file: "aes-gcm/nivc/aes-gctr-nivc",
             template: "AESGCTRFOLD",
+            params: [1]
         });
 
         let key = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
@@ -40,6 +41,7 @@ describe("aes-gctr-nivc", () => {
         circuit_one_block = await circomkit.WitnessTester("aes-gcm-fold", {
             file: "aes-gcm/nivc/aes-gctr-nivc",
             template: "AESGCTRFOLD",
+            params: [1]
         });
 
         let key = [0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31, 0x31];
@@ -67,6 +69,7 @@ describe("aes-gctr-nivc", () => {
         circuit_one_block = await circomkit.WitnessTester("aes-gcm-fold", {
             file: "aes-gcm/nivc/aes-gctr-nivc",
             template: "AESGCTRFOLD",
+            params: [1]
         });
 
         const ctr = [0x00, 0x00, 0x00, 0x01];
@@ -80,6 +83,7 @@ describe("aes-gctr-nivc", () => {
         circuit_one_block = await circomkit.WitnessTester("aes-gcm-fold", {
             file: "aes-gcm/nivc/aes-gctr-nivc",
             template: "AESGCTRFOLD",
+            params: [1]
         });
 
         const ctr_0 = [0x00, 0x00, 0x00, 0x01];
@@ -90,4 +94,20 @@ describe("aes-gctr-nivc", () => {
         const witness_1 = await circuit_one_block.compute({ key: key, iv: iv, plainText: plainText2, aad: aad, ctr: ctr_1, cipherText: ct_part2, step_in: witness_0.step_out }, ["step_out"])
         assert.deepEqual(witness_1.step_out, PoseidonModular([BigInt(witness_0.step_out.toString()), bytesToBigInt(plainText2)]));
     });
+
+    let circuit_two_block: WitnessTester<["key", "iv", "plainText", "aad", "ctr", "cipherText", "step_in"], ["step_out"]>;
+    it("all correct for two folds at once", async () => {
+        circuit_two_block = await circomkit.WitnessTester("aes-gcm-fold", {
+            file: "aes-gcm/nivc/aes-gctr-nivc",
+            template: "AESGCTRFOLD",
+            params: [2]
+        });
+
+        const ctr_0 = [0x00, 0x00, 0x00, 0x01];
+        const step_in_0 = 0;
+
+        const witness = await circuit_two_block.compute({ key: key, iv: iv, aad: aad, ctr: ctr_0, plainText: [plainText1, plainText2], cipherText: [ct_part1, ct_part2], step_in: step_in_0 }, ["step_out"])
+        let hash_0 = PoseidonModular([step_in_0, bytesToBigInt(plainText1)]);
+        assert.deepEqual(witness.step_out, PoseidonModular([hash_0, bytesToBigInt(plainText2)]));
+    });
 });
diff --git a/circuits/test/full/full.test.ts b/circuits/test/full/full.test.ts
@@ -232,6 +232,7 @@ describe("NIVC_FULL", async () => {
         aesCircuit = await circomkit.WitnessTester("AESGCTRFOLD", {
             file: "aes-gcm/nivc/aes-gctr-nivc",
             template: "AESGCTRFOLD",
+            params: [1]
         });
         console.log("#constraints (AES-GCTR):", await aesCircuit.getConstraintCount());
 
@@ -323,4 +324,122 @@ describe("NIVC_FULL", async () => {
         console.log("finalValue", extractValue.step_out);
         assert.deepEqual(extractValue.step_out, final_value_hash);
     });
-});
+});
+
+
+describe("NIVC_FULL_2", async () => {
+    let aesCircuit: WitnessTester<["key", "iv", "aad", "ctr", "plainText", "cipherText", "step_in"], ["step_out"]>;
+    let httpCircuit: WitnessTester<["step_in", "data", "start_line_hash", "header_hashes", "body_hash"], ["step_out"]>;
+    let json_mask_object_circuit: WitnessTester<["step_in", "data", "key", "keyLen"], ["step_out"]>;
+    let json_mask_arr_circuit: WitnessTester<["step_in", "data", "index"], ["step_out"]>;
+    let extract_value_circuit: WitnessTester<["step_in", "data"], ["step_out"]>;
+
+    const MAX_NUMBER_OF_HEADERS = 2;
+
+    const DATA_BYTES = 320;
+    const MAX_STACK_HEIGHT = 5;
+
+    const MAX_KEY_LENGTH = 8;
+    const MAX_VALUE_LENGTH = 32;
+
+    before(async () => {
+        aesCircuit = await circomkit.WitnessTester("AESGCTRFOLD", {
+            file: "aes-gcm/nivc/aes-gctr-nivc",
+            template: "AESGCTRFOLD",
+            params: [2]
+        });
+        console.log("#constraints (AES-GCTR):", await aesCircuit.getConstraintCount());
+
+        httpCircuit = await circomkit.WitnessTester(`HttpNIVC`, {
+            file: "http/nivc/http_nivc",
+            template: "HttpNIVC",
+            params: [DATA_BYTES, MAX_NUMBER_OF_HEADERS],
+        });
+        console.log("#constraints (HttpNIVC):", await httpCircuit.getConstraintCount());
+
+        json_mask_object_circuit = await circomkit.WitnessTester(`JsonMaskObjectNIVC`, {
+            file: "json/nivc/masker",
+            template: "JsonMaskObjectNIVC",
+            params: [DATA_BYTES, MAX_STACK_HEIGHT, MAX_KEY_LENGTH],
+        });
+        console.log("#constraints (JSON-MASK-OBJECT):", await json_mask_object_circuit.getConstraintCount());
+
+        json_mask_arr_circuit = await circomkit.WitnessTester(`JsonMaskArrayIndexNIVC`, {
+            file: "json/nivc/masker",
+            template: "JsonMaskArrayIndexNIVC",
+            params: [DATA_BYTES, MAX_STACK_HEIGHT],
+        });
+        console.log("#constraints (JSON-MASK-ARRAY-INDEX):", await json_mask_arr_circuit.getConstraintCount());
+
+        extract_value_circuit = await circomkit.WitnessTester(`JsonMaskExtractFinal`, {
+            file: "json/nivc/extractor",
+            template: "MaskExtractFinal",
+            params: [DATA_BYTES, MAX_VALUE_LENGTH],
+        });
+        console.log("#constraints (JSON-MASK-EXTRACT-FINAL):", await extract_value_circuit.getConstraintCount());
+    });
+
+    it("NIVC_CHAIN_2", async () => {
+        // Run AES chain
+        let ctr = [0x00, 0x00, 0x00, 0x01];
+        const init_nivc_input = 0;
+
+        let pt = [http_response_plaintext.slice(0, 16), http_response_plaintext.slice(16, 32)];
+        let ct = [http_response_ciphertext.slice(0, 16), http_response_ciphertext.slice(16, 32)];
+        let aes_gcm = await aesCircuit.compute({ key: Array(16).fill(0), iv: Array(12).fill(0), ctr: ctr, plainText: pt, aad: Array(16).fill(0), cipherText: ct, step_in: init_nivc_input }, ["step_out"]);
+        let i = 0;
+        console.log("AES `step_out[", i, "]`: ", aes_gcm.step_out);
+        for (i = 1; i < (DATA_BYTES / (16 * 2)); i++) {
+            ctr[3] += 2; // This will work since we don't run a test that overlows a byte
+            let pt = [http_response_plaintext.slice(i * 32, i * 32 + 16), http_response_plaintext.slice(i * 32 + 16, i * 32 + 32)];
+            let ct = [http_response_ciphertext.slice(i * 32, i * 32 + 16), http_response_ciphertext.slice(i * 32 + 16, i * 32 + 32)];
+            aes_gcm = await aesCircuit.compute({ key: Array(16).fill(0), iv: Array(12).fill(0), ctr: ctr, plainText: pt, aad: Array(16).fill(0), cipherText: ct, step_in: aes_gcm.step_out }, ["step_out"]);
+            console.log("AES `step_out[", i, "]`: ", aes_gcm.step_out);
+        }
+        assert.deepEqual(http_response_hash, aes_gcm.step_out);
+
+        let http = await httpCircuit.compute({ step_in: aes_gcm.step_out, data: http_response_plaintext, start_line_hash: http_start_line_hash, header_hashes: [http_header_0_hash, http_header_1_hash], body_hash: http_body_mask_hash }, ["step_out"]);
+        console.log("HttpNIVC `step_out`:", http.step_out);
+
+        let key0 = [100, 97, 116, 97, 0, 0, 0, 0]; // "data"
+        let key0Len = 4;
+        let key1 = [105, 116, 101, 109, 115, 0, 0, 0]; // "items"
+        let key1Len = 5;
+        let key2 = [112, 114, 111, 102, 105, 108, 101, 0]; // "profile"
+        let key2Len = 7;
+        let key3 = [110, 97, 109, 101, 0, 0, 0, 0]; // "name"
+        let key3Len = 4;
+
+        let json_extract_key0 = await json_mask_object_circuit.compute({ step_in: http.step_out, data: http_body, key: key0, keyLen: key0Len }, ["step_out"]);
+        console.log("JSON Extract key0 `step_out`:", json_extract_key0.step_out);
+        assert.deepEqual(json_extract_key0.step_out, json_key0_mask_hash);
+
+        let json_extract_key1 = await json_mask_object_circuit.compute({ step_in: json_extract_key0.step_out, data: json_key0_mask, key: key1, keyLen: key1Len }, ["step_out"]);
+        assert.deepEqual(json_extract_key1.step_out, json_key1_mask_hash);
+        console.log("JSON Extract key1 `step_out`:", json_extract_key1.step_out);
+
+        let json_extract_arr = await json_mask_arr_circuit.compute({ step_in: json_extract_key1.step_out, data: json_key1_mask, index: 0 }, ["step_out"]);
+        assert.deepEqual(json_extract_arr.step_out, json_arr_mask_hash);
+        console.log("JSON Extract arr `step_out`:", json_extract_arr.step_out);
+
+        let json_extract_key2 = await json_mask_object_circuit.compute({ step_in: json_extract_arr.step_out, data: json_arr_mask, key: key2, keyLen: key2Len }, ["step_out"]);
+        assert.deepEqual(json_extract_key2.step_out, json_key2_mask_hash);
+        console.log("JSON Extract key2 `step_out`:", json_extract_key2.step_out);
+
+        let json_extract_key3 = await json_mask_object_circuit.compute({ step_in: json_extract_key2.step_out, data: json_key2_mask, key: key3, keyLen: key3Len }, ["step_out"]);
+        assert.deepEqual(json_extract_key3.step_out, json_key3_mask_hash);
+        console.log("JSON Extract key3 `step_out`:", json_extract_key3.step_out);
+
+        // TODO (autoparallel): we need to rethink extraction here.
+        let finalOutput = toByte("\"Taylor Swift\"");
+        let finalOutputPadded = finalOutput.concat(Array(Math.max(0, MAX_VALUE_LENGTH - finalOutput.length)).fill(0));
+        let final_value_hash = DataHasher(finalOutputPadded);
+        let extractValue = await extract_value_circuit.compute({ step_in: json_extract_key3.step_out, data: json_key3_mask }, ["step_out"]);
+        console.log("finalValue", extractValue.step_out);
+        assert.deepEqual(extractValue.step_out, final_value_hash);
+    });
+});
+
+
+
+
diff --git a/package.json b/package.json
@@ -1,7 +1,7 @@
 {
   "name": "web-prover-circuits",
   "description": "ZK Circuits for WebProofs",
-  "version": "0.5.2",
+  "version": "0.5.3",
   "license": "Apache-2.0",
   "repository": {
     "type": "git",
Original file line number	Diff line number	Diff line change
Expand Up		@@ -2,4 +2,4 @@ pragma circom 2.1.9;

		include "../../circuits/aes-gcm/nivc/aes-gctr-nivc.circom";

		component main { public [step_in] } = AESGCTRFOLD();
		component main { public [step_in] } = AESGCTRFOLD(1);