Add AssertBytes template to validate byte inputs in CustomHasher circuit (#51)

circuits/circuits/tests/utils/packBytesAndPoseidon_tester.circom

@@ -0,0 +1,5 @@
+pragma circom  2.1.6;
+
+include "../../utils/passport/customHashers.circom";
+
+component main = PackBytesAndPoseidon(16);

circuits/circuits/utils/passport/customHashers.circom

@@ -2,6 +2,7 @@ pragma circom 2.1.9;
 include "../crypto/bigInt/bigIntFunc.circom";
 include "circomlib/circuits/poseidon.circom";
 include "@openpassport/zk-email-circuits/utils/bytes.circom";
+include "circomlib/circuits/comparators.circom";
 
 /// @notice CutomHasher circuit - used to Poseidon up to 256 signals
 /// @param k Number of signals to hash
@@ -52,9 +53,36 @@ template CustomHasher(k) {
 /// @param in Input array
 /// @param out Output hash
 template PackBytesAndPoseidon(k) {
-    var packed_length = computeIntChunkLength(k);
-
     signal input in[k];
+
+    AssertBytes(k)(in);
+
+    var packed_length = computeIntChunkLength(k);
     signal packed[packed_length] <== PackBytes(k)(in);
     signal output out <== CustomHasher(packed_length)(packed);
-}
+}
+/// @title AssertBytes
+/// @notice Asserts that every element in the input array is a valid byte (i.e., less than 256).
+/// @param k The number of elements in the input array.
+/// @param in The input array containing byte values to be validated.
+/// @dev This template uses a chain of LessThan components to check that each byte is below 256. The results are cumulatively multiplied in checkArray, and the final constraint (checkArray[k-1] === 1) enforces that all bytes meet the condition.
+template AssertBytes(k) {
+    signal input in[k];
+    signal checkArray[k];
+    component lessThan256[k];
+
+    for (var i = 0; i < k; i++) {
+        if (i == 0) {
+            lessThan256[i] = LessThan(8);
+            lessThan256[i].in[0] <== in[i];
+            lessThan256[i].in[1] <== 256;
+            checkArray[i] <== lessThan256[i].out;
+        } else {
+            lessThan256[i] = LessThan(8);
+            lessThan256[i].in[0] <== in[i];
+            lessThan256[i].in[1] <== 256;
+            checkArray[i] <== lessThan256[i].out * checkArray[i-1];
+        }
+    }
+    checkArray[k-1] === 1;
+}

circuits/tests/other_circuits/custom_hasher.test.ts

@@ -2,18 +2,30 @@ import { expect } from 'chai';
 import path from 'path';
 import { wasm as wasm_tester } from 'circom_tester';
 import { formatInput } from '../../../common/src/utils/circuits/generateInputs';
-import { customHasher } from '../../../common/src/utils/hash';
+import { customHasher, packBytesAndPoseidon } from '../../../common/src/utils/hash';
 
 describe('CustomHasher', function () {
   this.timeout(0);
-  let circuit;
+  let circuitCustomHasher;
+  let circuitPackBytesAndPoseidon;
 
   this.beforeAll(async () => {
-    const circuitPath = path.resolve(
+    const circuitPathCustomHasher = path.resolve(
       __dirname,
       '../../circuits/tests/utils/customHasher_tester.circom'
     );
-    circuit = await wasm_tester(circuitPath, {
+    const circuitPathPackBytesAndPoseidon = path.resolve(
+      __dirname,
+      '../../circuits/tests/utils/packBytesAndPoseidon_tester.circom'
+    );
+    circuitCustomHasher = await wasm_tester(circuitPathCustomHasher, {
+      include: [
+        'node_modules',
+        './node_modules/@zk-kit/binary-merkle-root.circom/src',
+        './node_modules/circomlib/circuits',
+      ],
+    });
+    circuitPackBytesAndPoseidon = await wasm_tester(circuitPathPackBytesAndPoseidon, {
       include: [
         'node_modules',
         './node_modules/@zk-kit/binary-merkle-root.circom/src',
@@ -35,12 +47,55 @@ describe('CustomHasher', function () {
     const inputs = { in: formatInput(randomNumbers) };
 
     it('customHasher output should be the same between circom and js implementation', async () => {
-      const witness = await circuit.calculateWitness(inputs, true);
-      const hashValueCircom = (await circuit.getOutput(witness, ['out'])).out;
+      const witness = await circuitCustomHasher.calculateWitness(inputs, true);
+      const hashValueCircom = (await circuitCustomHasher.getOutput(witness, ['out'])).out;
       console.log('\x1b[34m', 'hashValueCircom: ', hashValueCircom, '\x1b[0m');
       const hashValueJs = customHasher(randomNumbers);
       console.log('\x1b[34m', 'hashValueJs: ', hashValueJs, '\x1b[0m');
       expect(BigInt(hashValueCircom).toString()).to.equal(BigInt(hashValueJs).toString());
     });
   });
+
+  describe('packBytesAndPoseidon', async () => {
+    let randomNumbers = Array(16)
+      .fill(0)
+      .map(() => {
+        return Math.floor(Math.random() * 256); // Random number between 0 and 255
+      });
+    const inputs = { in: formatInput(randomNumbers) };
+    console.log('\x1b[34m', 'inputs: ', inputs, '\x1b[0m');
+
+    it('packBytesAndPoseidon output should be the same between circom and js implementation', async () => {
+      const witness = await circuitPackBytesAndPoseidon.calculateWitness(inputs, true);
+      const hashValueCircom = (await circuitPackBytesAndPoseidon.getOutput(witness, ['out'])).out;
+      console.log('\x1b[34m', 'hashValueCircom: ', hashValueCircom, '\x1b[0m');
+      const hashValueJs = packBytesAndPoseidon(randomNumbers);
+      console.log('\x1b[34m', 'hashValueJs: ', hashValueJs, '\x1b[0m');
+      expect(BigInt(hashValueCircom).toString()).to.equal(BigInt(hashValueJs).toString());
+    });
+
+    it('packBytesAndPoseidon should fail, inputs[0] = 256', async () => {
+      try {
+        const wrongRandomNumbers = [...randomNumbers];
+        wrongRandomNumbers[0] = 256;
+        const wrongInputs = { in: formatInput(wrongRandomNumbers) };
+        const witness = await circuitPackBytesAndPoseidon.calculateWitness(wrongInputs, true);
+        expect.fail('Expected an error but none was thrown.');
+      } catch (error) {
+        expect(error.message).to.include('Assert Failed');
+      }
+    });
+
+    it('packBytesAndPoseidon should fail, inputs[15] = 256', async () => {
+      try {
+        const wrongRandomNumbers = [...randomNumbers];
+        wrongRandomNumbers[15] = 256;
+        const wrongInputs = { in: formatInput(wrongRandomNumbers) };
+        const witness = await circuitPackBytesAndPoseidon.calculateWitness(wrongInputs, true);
+        expect.fail('Expected an error but none was thrown.');
+      } catch (error) {
+        expect(error.message).to.include('Assert Failed');
+      }
+    });
+  });
 });