From 70174a623cb3716f8c6217b4ccc6049502cdfeb0 Mon Sep 17 00:00:00 2001
From: shreyas-londhe <shreyas_londhe@outlook.com>
Date: Sun, 14 Jul 2024 01:22:31 +0200
Subject: [PATCH] fix: added a assert-bit check in body-masker

---
 packages/circuits/helpers/body-masker.circom | 26 ++++++++++++++++++++
 packages/circuits/tests/body-masker.test.ts  | 17 +++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/packages/circuits/helpers/body-masker.circom b/packages/circuits/helpers/body-masker.circom
index f3178ddc..474aa237 100644
--- a/packages/circuits/helpers/body-masker.circom
+++ b/packages/circuits/helpers/body-masker.circom
@@ -1,11 +1,37 @@
 pragma circom 2.1.6;
 
+// Asserts that a given input is binary.
+//
+// Inputs:
+// - in: an input signal, expected to be 0 or 1.
+template AssertBit() {
+  signal input in;
+  in * (in - 1) === 0;
+}
+
+// The BodyMasker template masks an input body array using a binary mask array.
+// Each element in the body array is multiplied by the corresponding element in the mask array.
+// The mask array is validated to ensure all elements are binary (0 or 1).
+//
+// Parameters:
+// - maxBodyLength: The maximum length of the body and mask arrays.
+//
+// Inputs:
+// - body: An array of signals representing the body to be masked.
+// - mask: An array of signals representing the binary mask.
+//
+// Outputs:
+// - masked_body: An array of signals representing the masked body.
 template BodyMasker(maxBodyLength) {
     signal input body[maxBodyLength];
     signal input mask[maxBodyLength];
     signal output masked_body[maxBodyLength];
 
+    component bit_check[maxBodyLength];
+
     for (var i = 0; i < maxBodyLength; i++) {
+        bit_check[i] = AssertBit();
+        bit_check[i].in <== mask[i];
         masked_body[i] <== body[i] * mask[i];
     }
 }
\ No newline at end of file
diff --git a/packages/circuits/tests/body-masker.test.ts b/packages/circuits/tests/body-masker.test.ts
index 7e966a0c..aea8f3a7 100644
--- a/packages/circuits/tests/body-masker.test.ts
+++ b/packages/circuits/tests/body-masker.test.ts
@@ -27,4 +27,21 @@ describe("BodyMasker Circuit", () => {
             masked_body: [1, 0, 3, 0, 5, 0, 7, 0, 9, 0],
         });
     });
+
+    it("should fail if mask has non-bit numbers", async () => {
+        const input = {
+            body: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+            mask: [1, 2, 1, 0, 1, 0, 1, 0, 1, 0], // Mask with non-bit number (2)
+        };
+
+        try {
+            const witness = await circuit.calculateWitness(input);
+            await circuit.checkConstraints(witness);
+            await circuit.assertOut(witness, {
+                masked_body: [1, 0, 3, 0, 5, 0, 7, 0, 9, 0],
+            });
+        } catch (error) {
+            expect(error).toBeTruthy();
+        }
+    });
 });