fix: chacha circuit (#74)

lonerapier · Autoparallel · web-flow · commit bfc120be2457 · 2024-12-05T08:46:40.000-07:00
* add correct padding check for AES

* add ciphertext padding check for chacha

* update package version

* better check

* pass plaintext as bytes

* fix test

* remove length

* revert aes

* remove aes

* refactor: use bytes for chacha input

---------

Co-authored-by: Colin Roberts &lt;colin@autoparallel.xyz&gt;
diff --git a/builds/target_1024b/chacha20_nivc_1024.circom b/builds/target_1024b/chacha20_nivc_1024.circom
@@ -2,4 +2,4 @@ pragma circom 2.1.9;
 
 include "../../circuits/chacha20/nivc/chacha20_nivc.circom";
 
-component main { public [step_in] } = ChaCha20_NIVC(256);
+component main { public [step_in] } = ChaCha20_NIVC(1024);
diff --git a/builds/target_512b/chacha20_nivc_512b.circom b/builds/target_512b/chacha20_nivc_512b.circom
@@ -2,4 +2,4 @@ pragma circom 2.1.9;
 
 include "../../circuits/chacha20/nivc/chacha20_nivc.circom";
 
-component main { public [step_in] } = ChaCha20_NIVC(128);
+component main { public [step_in] } = ChaCha20_NIVC(512);
diff --git a/circuits/chacha20/nivc/chacha20_nivc.circom b/circuits/chacha20/nivc/chacha20_nivc.circom
@@ -22,8 +22,8 @@ include "../../utils/array.circom";
 // +---+---+---+---+
 // | # | N | N | N |
 // +---+---+---+---+
-// paramaterized by n which is the number of 32-bit words to encrypt
-template ChaCha20_NIVC(N) {
+// paramaterized by `DATA_BYTES` which is the plaintext length in bytes
+template ChaCha20_NIVC(DATA_BYTES) {
   // key => 8 32-bit words = 32 bytes
   signal input key[8][32];
   // nonce => 3 32-bit words = 12 bytes
@@ -33,13 +33,23 @@ template ChaCha20_NIVC(N) {
 
   // the below can be both ciphertext or plaintext depending on the direction
   // in => N 32-bit words => N 4 byte words
-  signal input plainText[N][32];
+  signal input plainText[DATA_BYTES];
   // out => N 32-bit words => N 4 byte words
-  signal input cipherText[N][32];
+  signal input cipherText[DATA_BYTES];
 
   signal input step_in[1];
   signal output step_out[1];
 
+  signal plaintextBits[DATA_BYTES / 4][32];
+  component toBits[DATA_BYTES / 4];
+  for (var i = 0 ; i < DATA_BYTES / 4 ; i++) {
+    toBits[i] = fromWords32ToLittleEndian();
+    for (var j = 0 ; j < 4 ; j++) {
+      toBits[i].words[j] <== plainText[i*4 + j];
+    }
+    plaintextBits[i] <== toBits[i].data;
+  }
+
   var tmp[16][32] = [
     [
       // constant 0x61707865
@@ -88,24 +98,24 @@ template ChaCha20_NIVC(N) {
 
   // do the ChaCha20 rounds
   // rounds opperates on 4 words at a time
-  component rounds[N/16];
-  component xors[N];
-  component counter_adder[N/16 - 1];
+  component rounds[DATA_BYTES / 64];
+  component xors[DATA_BYTES];
+  component counter_adder[DATA_BYTES / 64 - 1];
 
-  signal computedCipherText[N][32];
+  signal computedCipherText[DATA_BYTES / 4][32];
 
-  for(i = 0; i < N/16; i++) {
+  for(i = 0; i < DATA_BYTES / 64; i++) {
     rounds[i] = Round();
     rounds[i].in <== tmp;
     // XOR block with input
     for(j = 0; j < 16; j++) {
       xors[i*16 + j] = XorBits(32);
-      xors[i*16 + j].a <== plainText[i*16 + j];
+      xors[i*16 + j].a <== plaintextBits[i*16 + j];
       xors[i*16 + j].b <== rounds[i].out[j];
       computedCipherText[i*16 + j] <== xors[i*16 + j].out;
     }
 
-    if(i < N/16 - 1) {
+    if(i < DATA_BYTES / 64 - 1) {
       counter_adder[i] = AddBits(32);
       counter_adder[i].a <== tmp[12];
       counter_adder[i].b <== one;
@@ -115,25 +125,21 @@ template ChaCha20_NIVC(N) {
     }
   }
 
-  signal ciphertext_equal_check[N][32];
-  for(var i = 0 ; i < N; i++) {
-    for(var j = 0 ; j < 32 ; j++) {
-      ciphertext_equal_check[i][j] <== IsEqual()([computedCipherText[i][j], cipherText[i][j]]);
-      ciphertext_equal_check[i][j] === 1;
-    }
-  }
-
-  component toBytes[N];
-  signal bigEndianPlaintext[N*4];
-  for(var i = 0 ; i < N; i++) {
-    toBytes[i] = fromLittleEndianToWords32();
-    for(var j = 0 ; j < 32 ; j++) {
-      toBytes[i].data[j] <== plainText[i][j];
+  component toCiphertextBytes[DATA_BYTES / 4];
+  signal bigEndianCiphertext[DATA_BYTES];
+  for (var i = 0 ; i < DATA_BYTES / 4 ; i++) {
+    toCiphertextBytes[i] = fromLittleEndianToWords32();
+    for (var j = 0 ; j < 32 ; j++) {
+      toCiphertextBytes[i].data[j] <== computedCipherText[i][j];
     }
-    for(var j = 0; j < 4; j++) {
-      bigEndianPlaintext[i*4 + j] <== toBytes[i].words[j];
+    for (var j = 0 ; j < 4 ; j++) {
+      bigEndianCiphertext[i*4 + j] <== toCiphertextBytes[i].words[j];
     }
   }
-  signal data_hash <== DataHasher(N*4)(bigEndianPlaintext);
+
+  signal paddedCiphertextCheck <== IsEqualArrayPaddedLHS(DATA_BYTES)([cipherText, bigEndianCiphertext]);
+  paddedCiphertextCheck === 1;
+
+  signal data_hash <== DataHasher(DATA_BYTES)(plainText);
   step_out[0] <== data_hash;
 }
diff --git a/circuits/test/chacha20/chacha20-nivc.test.ts b/circuits/test/chacha20/chacha20-nivc.test.ts
@@ -1,41 +1,41 @@
 import { WitnessTester } from "circomkit";
-import { circomkit, toUint32Array, uintArray32ToBits } from "../common";
+import { circomkit, toByte, toUint32Array, uintArray32ToBits } from "../common";
 import { DataHasher } from "../common/poseidon";
 import { assert } from "chai";
 
 
 describe("chacha20-nivc", () => {
+    let circuit: WitnessTester<["key", "nonce", "counter", "plainText", "cipherText", "step_in"], ["step_out"]>;
     describe("16 block test", () => {
-        let circuit: WitnessTester<["key", "nonce", "counter", "plainText", "cipherText", "step_in"], ["step_out"]>;
         it("should perform encryption", async () => {
             circuit = await circomkit.WitnessTester(`ChaCha20`, {
                 file: "chacha20/nivc/chacha20_nivc",
                 template: "ChaCha20_NIVC",
-                params: [16] // number of 32-bit words in the key, 32 * 16 = 512 bits
+                params: [64] // number of bytes for plaintext
             });
             // Test case from RCF https://www.rfc-editor.org/rfc/rfc7539.html#section-2.4.2
-            // the input encoding here is not the most intuitive. inputs are serialized as little endian. 
+            // the input encoding here is not the most intuitive. inputs are serialized as little endian.
             // i.e. "e4e7f110" is serialized as "10 f1 e7 e4". So the way i am reading in inputs is
-            // to ensure that every 32 bit word is byte reversed before being turned into bits. 
+            // to ensure that every 32 bit word is byte reversed before being turned into bits.
             // i think this should be easy when we compute witness in rust.
             let keyBytes = [
-                    0x00, 0x01, 0x02, 0x03,
-                    0x04, 0x05, 0x06, 0x07,
-                    0x08, 0x09, 0x0a, 0x0b,
-                    0x0c, 0x0d, 0x0e, 0x0f,
-                    0x10, 0x11, 0x12, 0x13,
-                    0x14, 0x15, 0x16, 0x17,
-                    0x18, 0x19, 0x1a, 0x1b,
-                    0x1c, 0x1d, 0x1e, 0x1f
-                ];
+                0x00, 0x01, 0x02, 0x03,
+                0x04, 0x05, 0x06, 0x07,
+                0x08, 0x09, 0x0a, 0x0b,
+                0x0c, 0x0d, 0x0e, 0x0f,
+                0x10, 0x11, 0x12, 0x13,
+                0x14, 0x15, 0x16, 0x17,
+                0x18, 0x19, 0x1a, 0x1b,
+                0x1c, 0x1d, 0x1e, 0x1f
+            ];
 
             let nonceBytes =
                 [
                     0x00, 0x00, 0x00, 0x00,
                     0x00, 0x00, 0x00, 0x4a,
                     0x00, 0x00, 0x00, 0x00
                 ];
-            let plaintextBytes = 
+            let plaintextBytes =
                 [
                     0x4c, 0x61, 0x64, 0x69, 0x65, 0x73, 0x20, 0x61, 0x6e, 0x64, 0x20, 0x47, 0x65, 0x6e, 0x74, 0x6c,
                     0x65, 0x6d, 0x65, 0x6e, 0x20, 0x6f, 0x66, 0x20, 0x74, 0x68, 0x65, 0x20, 0x63, 0x6c, 0x61, 0x73,
@@ -48,21 +48,75 @@ describe("chacha20-nivc", () => {
                     0xe9, 0x7e, 0x7a, 0xec, 0x1d, 0x43, 0x60, 0xc2, 0x0a, 0x27, 0xaf, 0xcc, 0xfd, 0x9f, 0xae, 0x0b,
                     0xf9, 0x1b, 0x65, 0xc5, 0x52, 0x47, 0x33, 0xab, 0x8f, 0x59, 0x3d, 0xab, 0xcd, 0x62, 0xb3, 0x57,
                     0x16, 0x39, 0xd6, 0x24, 0xe6, 0x51, 0x52, 0xab, 0x8f, 0x53, 0x0c, 0x35, 0x9f, 0x08, 0x61, 0xd8
-                ]; 
-            const ciphertextBits = toInput(Buffer.from(ciphertextBytes))
-            const plaintextBits = toInput(Buffer.from(plaintextBytes))
-			const counterBits = uintArray32ToBits([1])[0]
-			let w = await circuit.compute({
-				key: toInput(Buffer.from(keyBytes)),
-				nonce: toInput(Buffer.from(nonceBytes)),
-				counter: counterBits,
-				cipherText: ciphertextBits,
-                plainText: plaintextBits,
+                ];
+            const counterBits = uintArray32ToBits([1])[0]
+            let w = await circuit.compute({
+                key: toInput(Buffer.from(keyBytes)),
+                nonce: toInput(Buffer.from(nonceBytes)),
+                counter: counterBits,
+                cipherText: ciphertextBytes,
+                plainText: plaintextBytes,
                 step_in: 0
-			}, (["step_out"]));
+            }, (["step_out"]));
             assert.deepEqual(w.step_out, DataHasher(plaintextBytes));
         });
     });
+
+    describe("padded plaintext", () => {
+        it("should perform encryption", async () => {
+            circuit = await circomkit.WitnessTester(`ChaCha20`, {
+                file: "chacha20/nivc/chacha20_nivc",
+                template: "ChaCha20_NIVC",
+                params: [128] // number of bytes in plaintext
+            });
+            // Test case from RCF https://www.rfc-editor.org/rfc/rfc7539.html#section-2.4.2
+            // the input encoding here is not the most intuitive. inputs are serialized as little endian.
+            // i.e. "e4e7f110" is serialized as "10 f1 e7 e4". So the way i am reading in inputs is
+            // to ensure that every 32 bit word is byte reversed before being turned into bits.
+            // i think this should be easy when we compute witness in rust.
+            let keyBytes = [
+                0x00, 0x01, 0x02, 0x03,
+                0x04, 0x05, 0x06, 0x07,
+                0x08, 0x09, 0x0a, 0x0b,
+                0x0c, 0x0d, 0x0e, 0x0f,
+                0x10, 0x11, 0x12, 0x13,
+                0x14, 0x15, 0x16, 0x17,
+                0x18, 0x19, 0x1a, 0x1b,
+                0x1c, 0x1d, 0x1e, 0x1f
+            ];
+
+            let nonceBytes =
+                [
+                    0x00, 0x00, 0x00, 0x00,
+                    0x00, 0x00, 0x00, 0x4a,
+                    0x00, 0x00, 0x00, 0x00
+                ];
+            let plaintextBytes =
+                toByte("Ladies and Gentlemen of the class of '99: If I could offer you only one tip ");
+
+            let ciphertextBytes =
+                [
+                    0x6e, 0x2e, 0x35, 0x9a, 0x25, 0x68, 0xf9, 0x80, 0x41, 0xba, 0x07, 0x28, 0xdd, 0x0d, 0x69, 0x81,
+                    0xe9, 0x7e, 0x7a, 0xec, 0x1d, 0x43, 0x60, 0xc2, 0x0a, 0x27, 0xaf, 0xcc, 0xfd, 0x9f, 0xae, 0x0b,
+                    0xf9, 0x1b, 0x65, 0xc5, 0x52, 0x47, 0x33, 0xab, 0x8f, 0x59, 0x3d, 0xab, 0xcd, 0x62, 0xb3, 0x57,
+                    0x16, 0x39, 0xd6, 0x24, 0xe6, 0x51, 0x52, 0xab, 0x8f, 0x53, 0x0c, 0x35, 0x9f, 0x08, 0x61, 0xd8,
+                    0x07, 0xca, 0x0d, 0xbf, 0x50, 0x0d, 0x6a, 0x61, 0x56, 0xa3, 0x8e, 0x08
+                ];
+            let totalLength = 128;
+            let paddedPlaintextBytes = plaintextBytes.concat(Array(totalLength - plaintextBytes.length).fill(0));
+            let paddedCiphertextBytes = ciphertextBytes.concat(Array(totalLength - ciphertextBytes.length).fill(0));
+            const counterBits = uintArray32ToBits([1])[0]
+            let w = await circuit.compute({
+                key: toInput(Buffer.from(keyBytes)),
+                nonce: toInput(Buffer.from(nonceBytes)),
+                counter: counterBits,
+                cipherText: paddedCiphertextBytes,
+                plainText: paddedPlaintextBytes,
+                step_in: 0
+            }, (["step_out"]));
+            assert.deepEqual(w.step_out, DataHasher(paddedPlaintextBytes));
+        });
+    });
 });
 
 
diff --git a/circuits/test/full/full.test.ts b/circuits/test/full/full.test.ts
@@ -2,7 +2,6 @@ import { assert } from "chai";
 import { circomkit, WitnessTester, toByte, uintArray32ToBits, toUint32Array } from "../common";
 import { DataHasher } from "../common/poseidon";
 import { toInput } from "../chacha20/chacha20-nivc.test";
-import { buffer } from "stream/consumers";
 
 // HTTP/1.1 200 OK
 // content-type: application/json; charset=utf-8
@@ -355,7 +354,7 @@ describe("NIVC_FULL_CHACHA", async () => {
         chacha20Circuit = await circomkit.WitnessTester("CHACHA20", {
             file: "chacha20/nivc/chacha20_nivc",
             template: "ChaCha20_NIVC",
-            params: [80] // 80 * 32 = 2560 bits / 8 = 320 bytes
+            params: [320]
         });
         console.log("#constraints (CHACHA20):", await chacha20Circuit.getConstraintCount());
 
@@ -392,11 +391,9 @@ describe("NIVC_FULL_CHACHA", async () => {
         const init_nivc_input = 0;
         // Run ChaCha20
         const counterBits = uintArray32ToBits([1])[0]
-        const ptIn = toInput(Buffer.from(http_response_plaintext));
-        const ctIn = toInput(Buffer.from(chacha20_http_response_ciphertext));
         const keyIn = toInput(Buffer.from(Array(32).fill(0)));
         const nonceIn = toInput(Buffer.from([0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00]));
-        let chacha20 = await chacha20Circuit.compute({ key: keyIn, nonce: nonceIn, counter: counterBits, plainText: ptIn, cipherText: ctIn, step_in: init_nivc_input }, ["step_out"]);
+        let chacha20 = await chacha20Circuit.compute({ key: keyIn, nonce: nonceIn, counter: counterBits, plainText: http_response_plaintext, cipherText: chacha20_http_response_ciphertext, step_in: init_nivc_input }, ["step_out"]);
         console.log("ChaCha20 `step_out`:", chacha20.step_out);
         assert.deepEqual(http_response_hash, chacha20.step_out);
 
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,7 +1,7 @@
 {
   "name": "web-prover-circuits",
   "description": "ZK Circuits for WebProofs",
-  "version": "0.5.9",
+  "version": "0.5.10",
   "license": "Apache-2.0",
   "repository": {
     "type": "git",

Original file line number	Diff line number	Diff line change
`@@ -2,4 +2,4 @@ pragma circom 2.1.9;`
`2`	`2`
`3`	`3`	`include "../../circuits/chacha20/nivc/chacha20_nivc.circom";`
`4`	`4`
`5`		`-component main { public [step_in] } = ChaCha20_NIVC(256);`
	`5`	`+component main { public [step_in] } = ChaCha20_NIVC(1024);`
Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,7 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "web-prover-circuits",`
`3`	`3`	`"description": "ZK Circuits for WebProofs",`
`4`		`- "version": "0.5.9",`
	`4`	`+ "version": "0.5.10",`
`5`	`5`	`"license": "Apache-2.0",`
`6`	`6`	`"repository": {`
`7`	`7`	`"type": "git",`