Feat: Added Header Masking

packages/circuits/README.md

@@ -29,6 +29,9 @@ include "@zk-email/circuits/email-verifier.circom";
   - `n`: Number of bits per chunk the RSA key is split into. Recommended to be 121.
   - `k`: Number of chunks the RSA key is split into. Recommended to be 17.
   - `ignoreBodyHashCheck`: Set 1 to skip body hash check in case data to prove/extract is only in the headers.
+  - `enableHeaderMasking`: Set 1 to turn on header masking.
+  - `enableBodyMasking`: Set 1 to turn on body masking.
+  - `removeSoftLineBreaks`: Set 1 to remove soft line breaks (`=\r\n`) from the email body.
 
   `Note`: We use these values for n and k because their product (n * k) needs to be more than 2048 (RSA constraint) and n has to be less than half of 255 to fit in a circom signal.
 
@@ -41,10 +44,14 @@ include "@zk-email/circuits/email-verifier.circom";
   - `emailBodyLength`: Length of the email body including the SHA-256 padding.
   - `bodyHashIndex`: Index of the body hash `bh` in the `emailHeader`.
   - `precomputedSHA[32]`: Precomputed SHA-256 hash of the email body till the bodyHashIndex.
+  - `headerMask[maxHeadersLength]`: Mask to be applied on the `emailHeader`.
+  - `bodyMask[maxBodyLength]`: Mask to be applied on the `emailBody`.
+  - `decodedEmailBody[maxBodyLength]`: Decoded email body after removing soft line breaks.
 
   **Output Signal** 
   - `pubkeyHash`: Poseidon hash of the pubkey - Poseidon(n/2)(n/2 chunks of pubkey with k*2 bits per chunk).
-
+  - `maskedHeader[maxHeadersLength]`: Masked email header.
+  - `maskedBody[maxBodyLength]`: Masked email body.
 <br/>
 
 ## **Libraries**
@@ -257,6 +264,33 @@ DigitBytesToInt: Converts a byte array representing digits to an integer.
   - `out`: The output integer after conversion.
 </details>
 
+<details>
+<summary>
+AssertBit: Asserts that a given input is binary.
+</summary>
+
+- **[Source](utils/bytes.circom#L1-L7)**
+- **Inputs**:
+  - `in`: An input signal, expected to be 0 or 1.
+- **Outputs**:
+  - None. This template will throw an assertion error if the input is not binary.
+
+</details>
+
+<details>
+<summary>
+ByteMask: Masks an input array using a binary mask array.
+</summary>
+
+- **[Source](utils/bytes.circom#L9-L25)**
+- **Parameters**:
+  - `maxLength`: The maximum length of the input and mask arrays.
+- **Inputs**:
+  - `in`: An array of signals representing the body to be masked.
+  - `mask`: An array of signals representing the binary mask.
+- **Outputs**:
+  - `out`: An array of signals representing the masked input.
+</details>
 
 ### `utils/constants.circom`
 
@@ -359,5 +393,20 @@ EmailNullifier: Calculates the email nullifier using Poseidon hash.
   - `out`: The email nullifier.
 </details>
 
+### `helpers/remove-soft-line-breaks.circom`
+
+<details>
+<summary>
+RemoveSoftLineBreaks: Verifies the removal of soft line breaks from an encoded input string.
+</summary>
 
+- **[Source](helpers/remove-soft-line-breaks.circom)**
+- **Parameters**:
+  - `maxLength`: The maximum length of the input strings.
+- **Inputs**:
+  - `encoded[maxLength]`: An array of ASCII values representing the input string with potential soft line breaks.
+  - `decoded[maxLength]`: An array of ASCII values representing the expected output after removing soft line breaks.
+- **Outputs**:
+  - `isValid`: A signal that is 1 if the decoded input correctly represents the encoded input with soft line breaks removed, 0 otherwise.
 
+</details>

packages/circuits/email-verifier.circom

@@ -22,8 +22,9 @@ include "./helpers/remove-soft-line-breaks.circom";
 /// @param n Number of bits per chunk the RSA key is split into. Recommended to be 121.
 /// @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 enableHeaderMasking Set 1 to turn on header masking.
 /// @param enableBodyMasking Set 1 to turn on body masking.
+/// @param removeSoftLineBreaks Set 1 to remove soft line breaks from the email body.
 /// @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.
@@ -36,8 +37,9 @@ include "./helpers/remove-soft-line-breaks.circom";
 /// @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.
+/// @output maskedHeader[maxHeadersLength] Masked email header.
 /// @output maskedBody[maxBodyLength] Masked email body.
-template EmailVerifier(maxHeadersLength, maxBodyLength, n, k, ignoreBodyHashCheck, removeSoftLineBreaks, enableBodyMasking) {
+template EmailVerifier(maxHeadersLength, maxBodyLength, n, k, ignoreBodyHashCheck, enableHeaderMasking, enableBodyMasking, removeSoftLineBreaks) {
     assert(maxHeadersLength % 64 == 0);
     assert(maxBodyLength % 64 == 0);
     assert(n * k > 2048); // to support 2048 bit RSA
@@ -89,6 +91,15 @@ template EmailVerifier(maxHeadersLength, maxBodyLength, n, k, ignoreBodyHashChec
     rsaVerifier.modulus <== pubkey;
     rsaVerifier.signature <== signature;
 
+    if (enableHeaderMasking == 1) {
+        signal input headerMask[maxHeadersLength];
+        signal output maskedHeader[maxHeadersLength];
+        component byteMask = ByteMask(maxHeadersLength);
+
+        byteMask.in <== emailHeader;
+        byteMask.mask <== headerMask;
+        maskedHeader <== byteMask.out;
+    }
 
     // Calculate the SHA256 hash of the body and verify it matches the hash in the header
     if (ignoreBodyHashCheck != 1) {
@@ -133,25 +144,22 @@ template EmailVerifier(maxHeadersLength, maxBodyLength, n, k, ignoreBodyHashChec
 
         if (removeSoftLineBreaks == 1) {
             signal input decodedEmailBodyIn[maxBodyLength];
-            signal output decodedEmailBodyOut[maxBodyLength];
             component qpEncodingChecker = RemoveSoftLineBreaks(maxBodyLength);
 
             qpEncodingChecker.encoded <== emailBody;
             qpEncodingChecker.decoded <== decodedEmailBodyIn;
 
             qpEncodingChecker.isValid === 1;
-
-            decodedEmailBodyOut <== qpEncodingChecker.decoded;
         }
 
         if (enableBodyMasking == 1) {
-            signal input mask[maxBodyLength];
+            signal input bodyMask[maxBodyLength];
             signal output maskedBody[maxBodyLength];
             component byteMask = ByteMask(maxBodyLength);
 
-            byteMask.body <== emailBody;
-            byteMask.mask <== mask;
-            maskedBody <== byteMask.maskedBody;
+            byteMask.in <== emailBody;
+            byteMask.mask <== bodyMask;
+            maskedBody <== byteMask.out;
         }
     }
 

packages/circuits/package.json

@@ -4,7 +4,7 @@
   "license": "MIT",
   "scripts": {
     "publish": "yarn npm publish --access=public",
-    "test": "NODE_OPTIONS=--max_old_space_size=8192 jest tests"
+    "test": "NODE_OPTIONS=--max_old_space_size=8192 jest --runInBand --detectOpenHandles --forceExit --verbose tests"
   },
   "dependencies": {
     "@zk-email/zk-regex-circom": "^2.1.0",

packages/circuits/tests/base64.test.ts

@@ -1,58 +1,57 @@
 import { wasm } from "circom_tester";
 import path from "path";
 
-
 describe("Base64 Lookup", () => {
-  jest.setTimeout(10 * 60 * 1000); // 10 minutes
-
-  let circuit: any;
-
-  beforeAll(async () => {
-    circuit = await wasm(
-      path.join(__dirname, "./test-circuits/base64-test.circom"),
-      {
-        recompile: true,
-        include: path.join(__dirname, "../../../node_modules"),
-        // output: path.join(__dirname, "./compiled-test-circuits"),
-      }
-    );
-  });
-
-  it("should decode valid base64 chars", async function () {
-    const inputs = [
-      [65, 0], // A
-      [90, 25], // Z
-      [97, 26], // a
-      [122, 51], // z
-      [48, 52], // 0
-      [57, 61], // 9
-      [43, 62], // +
-      [47, 63], // /
-      [61, 0], // =
-    ]
-
-    for (const [input, output] of inputs) {
-      const witness = await circuit.calculateWitness({
-        in: input
-      });
-      await circuit.checkConstraints(witness);
-      await circuit.assertOut(witness, { out: output })
-    }
-  });
-
-  it("should fail with invalid chars", async function () {
-    const inputs = [34, 64, 91, 44];
-
-    expect.assertions(inputs.length);
-    for (const input of inputs) {
-    try {
-      const witness = await circuit.calculateWitness({
-        in: input
-      });
-      await circuit.checkConstraints(witness);
-    } catch (error) {
-      expect((error as Error).message).toMatch("Assert Failed");
-    }
-    }
-  });
+    jest.setTimeout(30 * 60 * 1000); // 30 minutes
+
+    let circuit: any;
+
+    beforeAll(async () => {
+        circuit = await wasm(
+            path.join(__dirname, "./test-circuits/base64-test.circom"),
+            {
+                recompile: true,
+                include: path.join(__dirname, "../../../node_modules"),
+                // output: path.join(__dirname, "./compiled-test-circuits"),
+            }
+        );
+    });
+
+    it("should decode valid base64 chars", async function () {
+        const inputs = [
+            [65, 0], // A
+            [90, 25], // Z
+            [97, 26], // a
+            [122, 51], // z
+            [48, 52], // 0
+            [57, 61], // 9
+            [43, 62], // +
+            [47, 63], // /
+            [61, 0], // =
+        ];
+
+        for (const [input, output] of inputs) {
+            const witness = await circuit.calculateWitness({
+                in: input,
+            });
+            await circuit.checkConstraints(witness);
+            await circuit.assertOut(witness, { out: output });
+        }
+    });
+
+    it("should fail with invalid chars", async function () {
+        const inputs = [34, 64, 91, 44];
+
+        expect.assertions(inputs.length);
+        for (const input of inputs) {
+            try {
+                const witness = await circuit.calculateWitness({
+                    in: input,
+                });
+                await circuit.checkConstraints(witness);
+            } catch (error) {
+                expect((error as Error).message).toMatch("Assert Failed");
+            }
+        }
+    });
 });

packages/circuits/tests/body-masker.test.ts → packages/circuits/tests/byte-mask.test.ts

@@ -6,7 +6,7 @@ describe("ByteMask Circuit", () => {
 
     beforeAll(async () => {
         circuit = await wasm_tester(
-            path.join(__dirname, "./test-circuits/body-masker-test.circom"),
+            path.join(__dirname, "./test-circuits/byte-mask-test.circom"),
             {
                 recompile: true,
                 include: path.join(__dirname, "../../../node_modules"),
@@ -17,14 +17,14 @@ describe("ByteMask Circuit", () => {
 
     it("should mask the body correctly", async () => {
         const input = {
-            body: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+            in: [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],
+            out: [1, 0, 3, 0, 5, 0, 7, 0, 9, 0],
         });
     });
 

packages/circuits/tests/email-verifier-no-body.test.ts

@@ -0,0 +1,48 @@
+import fs from "fs";
+import { wasm as wasm_tester } from "circom_tester";
+import path from "path";
+import { DKIMVerificationResult } from "@zk-email/helpers/src/dkim";
+import { generateEmailVerifierInputsFromDKIMResult } from "@zk-email/helpers/src/input-generators";
+import { verifyDKIMSignature } from "@zk-email/helpers/src/dkim";
+
+describe("EmailVerifier : Without body check", () => {
+    jest.setTimeout(30 * 60 * 1000); // 30 minutes
+
+    let dkimResult: DKIMVerificationResult;
+    let circuit: any;
+
+    beforeAll(async () => {
+        const rawEmail = fs.readFileSync(
+            path.join(__dirname, "./test-emails/test.eml"),
+            "utf8"
+        );
+        dkimResult = await verifyDKIMSignature(rawEmail);
+
+        circuit = await wasm_tester(
+            path.join(
+                __dirname,
+                "./test-circuits/email-verifier-no-body-test.circom"
+            ),
+            {
+                recompile: true,
+                include: path.join(__dirname, "../../../node_modules"),
+                // output: path.join(__dirname, "./compiled-test-circuits"),
+            }
+        );
+    });
+
+    it("should verify email when ignore_body_hash_check is true", async function () {
+        // The result wont have shaPrecomputeSelector, maxHeadersLength, maxBodyLength, ignoreBodyHashCheck
+        const emailVerifierInputs = generateEmailVerifierInputsFromDKIMResult(
+            dkimResult,
+            {
+                maxHeadersLength: 640,
+                maxBodyLength: 768,
+                ignoreBodyHashCheck: true,
+            }
+        );
+
+        const witness = await circuit.calculateWitness(emailVerifierInputs);
+        await circuit.checkConstraints(witness);
+    });
+});