From 00cdf1d7e70add06541099594e3d4eafe5af7f2f Mon Sep 17 00:00:00 2001
From: lonerapier <lonerapier@proton.me>
Date: Thu, 5 Dec 2024 14:04:35 +0530
Subject: [PATCH] better check

---
 circuits/chacha20/nivc/chacha20_nivc.circom  | 26 ++++++++++----------
 circuits/test/chacha20/chacha20-nivc.test.ts |  8 ++----
 circuits/test/full/full.test.ts              |  1 -
 3 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/circuits/chacha20/nivc/chacha20_nivc.circom b/circuits/chacha20/nivc/chacha20_nivc.circom
index b28e287..2daba75 100644
--- a/circuits/chacha20/nivc/chacha20_nivc.circom
+++ b/circuits/chacha20/nivc/chacha20_nivc.circom
@@ -35,8 +35,7 @@ template ChaCha20_NIVC(N) {
   // in => N 32-bit words => N 4 byte words
   signal input plainText[N][32];
   // out => N 32-bit words => N 4 byte words
-  signal input cipherText[N][32];
-  signal input length;
+  signal input cipherText[N*4];
 
   signal input step_in[1];
   signal output step_out[1];
@@ -116,20 +115,21 @@ template ChaCha20_NIVC(N) {
     }
   }
 
-  signal ciphertext_equal_check[N][32];
-  signal index_less_than_length[32*N];
-  signal ciphertext_not_equal[32 * N];
-  for(var i = 0 ; i < N; i++) {
-    for(var j = 0 ; j < 32 ; j++) {
-      var byteIndex = i*4 + j\8;
-      index_less_than_length[i*32 + j] <== LessThan(15)([byteIndex, length]);
-      ciphertext_not_equal[i*32 + j] <== IsEqual()([computedCipherText[i][j], cipherText[i][j]]);
-      ciphertext_equal_check[i][j] <== (1 - ciphertext_not_equal[i*32 + j]) * index_less_than_length[i*32 + j];
-      // 0 means ciphertext is equal and index < length
-      ciphertext_equal_check[i][j] === 0;
+  component toCiphertextBytes[N];
+  signal bigEndianCiphertext[N*4];
+  for (var i = 0 ; i < N ; 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++) {
+      bigEndianCiphertext[i*4 + j] <== toCiphertextBytes[i].words[j];
     }
   }
 
+  signal paddedCiphertextCheck <== IsEqualArrayPaddedLHS(N*4)([cipherText, bigEndianCiphertext]);
+  paddedCiphertextCheck === 1;
+
   component toBytes[N];
   signal bigEndianPlaintext[N*4];
   for(var i = 0 ; i < N; i++) {
diff --git a/circuits/test/chacha20/chacha20-nivc.test.ts b/circuits/test/chacha20/chacha20-nivc.test.ts
index 079cd3c..d6f6358 100644
--- a/circuits/test/chacha20/chacha20-nivc.test.ts
+++ b/circuits/test/chacha20/chacha20-nivc.test.ts
@@ -5,8 +5,8 @@ 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", "length", "step_in"], ["step_out"]>;
         it("should perform encryption", async () => {
             circuit = await circomkit.WitnessTester(`ChaCha20`, {
                 file: "chacha20/nivc/chacha20_nivc",
@@ -58,7 +58,6 @@ describe("chacha20-nivc", () => {
                 counter: counterBits,
                 cipherText: ciphertextBits,
                 plainText: plaintextBits,
-                length: plaintextBytes.length,
                 step_in: 0
             }, (["step_out"]));
             assert.deepEqual(w.step_out, DataHasher(plaintextBytes));
@@ -66,7 +65,6 @@ describe("chacha20-nivc", () => {
     });
 
     describe("padded plaintext", () => {
-        let circuit: WitnessTester<["key", "nonce", "counter", "plainText", "cipherText", "length", "step_in"], ["step_out"]>;
         it("should perform encryption", async () => {
             circuit = await circomkit.WitnessTester(`ChaCha20`, {
                 file: "chacha20/nivc/chacha20_nivc",
@@ -109,16 +107,14 @@ describe("chacha20-nivc", () => {
             let totalLength = 128;
             let paddedPlaintextBytes = plaintextBytes.concat(Array(totalLength - plaintextBytes.length).fill(0));
             let paddedCiphertextBytes = ciphertextBytes.concat(Array(totalLength - ciphertextBytes.length).fill(0));
-            const ciphertextBits = toInput(Buffer.from(paddedCiphertextBytes))
             const plaintextBits = toInput(Buffer.from(paddedPlaintextBytes))
             const counterBits = uintArray32ToBits([1])[0]
             let w = await circuit.compute({
                 key: toInput(Buffer.from(keyBytes)),
                 nonce: toInput(Buffer.from(nonceBytes)),
                 counter: counterBits,
-                cipherText: ciphertextBits,
+                cipherText: paddedCiphertextBytes,
                 plainText: plaintextBits,
-                length: plaintextBytes.length,
                 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
index 4791812..a7f34dd 100644
--- 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