Merge pull request #203 from zkemail/feat/add-body-masking

Feat: Added a Body masking Template

packages/circuits/email-verifier.circom

@@ -9,6 +9,7 @@ include "./lib/sha.circom";
 include "./utils/array.circom";
 include "./utils/regex.circom";
 include "./utils/hash.circom";
+include "./utils/bytes.circom";
 include "./helpers/remove-soft-line-breaks.circom";
 
 
@@ -22,6 +23,7 @@ include "./helpers/remove-soft-line-breaks.circom";
 /// @param k Number of chunks the RSA key is split into. Recommended to be 17.
 /// @param ignoreBodyHashCheck Set 1 to skip body hash check in case data to prove/extract is only in the headers.
 /// @param removeSoftLineBreaks Set 1 to remove soft line breaks from the email body.
+/// @param enableBodyMasking Set 1 to turn on body masking.
 /// @input emailHeader[maxHeadersLength] Email headers that are signed (ones in `DKIM-Signature` header) as ASCII int[], padded as per SHA-256 block size.
 /// @input emailHeaderLength Length of the email header including the SHA-256 padding.
 /// @input pubkey[k] RSA public key split into k chunks of n bits each.
@@ -31,9 +33,11 @@ include "./helpers/remove-soft-line-breaks.circom";
 /// @input bodyHashIndex Index of the body hash `bh` in the emailHeader.
 /// @input precomputedSHA[32] Precomputed SHA-256 hash of the email body till the bodyHashIndex.
 /// @input decodedEmailBodyIn[maxBodyLength] Decoded email body without soft line breaks.
+/// @input mask[maxBodyLength] Mask for the email body.
 /// @output pubkeyHash Poseidon hash of the pubkey - Poseidon(n/2)(n/2 chunks of pubkey with k*2 bits per chunk).
 /// @output decodedEmailBodyOut[maxBodyLength] Decoded email body with soft line breaks removed.
-template EmailVerifier(maxHeadersLength, maxBodyLength, n, k, ignoreBodyHashCheck, removeSoftLineBreaks) {
+/// @output maskedBody[maxBodyLength] Masked email body.
+template EmailVerifier(maxHeadersLength, maxBodyLength, n, k, ignoreBodyHashCheck, removeSoftLineBreaks, enableBodyMasking) {
     assert(maxHeadersLength % 64 == 0);
     assert(maxBodyLength % 64 == 0);
     assert(n * k > 2048); // to support 2048 bit RSA
@@ -139,8 +143,17 @@ template EmailVerifier(maxHeadersLength, maxBodyLength, n, k, ignoreBodyHashChec
 
             decodedEmailBodyOut <== qpEncodingChecker.decoded;
         }
-    }
 
+        if (enableBodyMasking == 1) {
+            signal input mask[maxBodyLength];
+            signal output maskedBody[maxBodyLength];
+            component byteMask = ByteMask(maxBodyLength);
+
+            byteMask.body <== emailBody;
+            byteMask.mask <== mask;
+            maskedBody <== byteMask.maskedBody;
+        }
+    }
 
     // Calculate the Poseidon hash of DKIM public key as output
     // This can be used to verify (by verifier/contract) the pubkey used in the proof without needing the full key

packages/circuits/tests/body-masker.test.ts

@@ -0,0 +1,47 @@
+import { wasm as wasm_tester } from "circom_tester";
+import path from "path";
+
+describe("ByteMask Circuit", () => {
+    let circuit: any;
+
+    beforeAll(async () => {
+        circuit = await wasm_tester(
+            path.join(__dirname, "./test-circuits/body-masker-test.circom"),
+            {
+                recompile: true,
+                include: path.join(__dirname, "../../../node_modules"),
+                output: path.join(__dirname, "./compiled-test-circuits"),
+            }
+        );
+    });
+
+    it("should mask the body correctly", async () => {
+        const input = {
+            body: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+            mask: [1, 0, 1, 0, 1, 0, 1, 0, 1, 0],
+        };
+
+        const witness = await circuit.calculateWitness(input);
+        await circuit.checkConstraints(witness);
+        await circuit.assertOut(witness, {
+            maskedBody: [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, {
+                maskedBody: [1, 0, 3, 0, 5, 0, 7, 0, 9, 0],
+            });
+        } catch (error) {
+            expect(error).toBeTruthy();
+        }
+    });
+});