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) {
@@ -145,13 +156,13 @@ template EmailVerifier(maxHeadersLength, maxBodyLength, n, k, ignoreBodyHashChec
         }
 
         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/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.test.ts

@@ -265,7 +265,7 @@ describe("EmailVerifier : With body masking", () => {
         circuit = await wasm_tester(
             path.join(
                 __dirname,
-                "./test-circuits/email-verifier-with-mask-test.circom"
+                "./test-circuits/email-verifier-with-body-mask-test.circom"
             ),
             {
                 recompile: true,
@@ -287,7 +287,7 @@ describe("EmailVerifier : With body masking", () => {
                 maxBodyLength: 768,
                 ignoreBodyHashCheck: false,
                 enableBodyMasking: true,
-                mask: mask.map((value) => (value ? 1 : 0)),
+                bodyMask: mask.map((value) => (value ? 1 : 0)),
             }
         );
 
@@ -344,3 +344,56 @@ describe("EmailVerifier : With soft line breaks", () => {
         await circuit.checkConstraints(witness);
     });
 });
+
+describe("EmailVerifier : With header masking", () => {
+    jest.setTimeout(10 * 60 * 1000); // 10 minutes
+
+    let dkimResult: DKIMVerificationResult;
+    let circuit: any;
+
+    beforeAll(async () => {
+        const rawEmail = fs.readFileSync(
+            path.join(__dirname, "./test-emails/test.eml")
+        );
+        dkimResult = await verifyDKIMSignature(rawEmail);
+
+        circuit = await wasm_tester(
+            path.join(
+                __dirname,
+                "./test-circuits/email-verifier-with-header-mask-test.circom"
+            ),
+            {
+                recompile: true,
+                include: path.join(__dirname, "../../../node_modules"),
+                output: path.join(__dirname, "./compiled-test-circuits"),
+            }
+        );
+    });
+
+    it("should verify email with header masking", async function () {
+        const mask = Array.from({ length: 640 }, (_, i) =>
+            i > 25 && i < 50 ? 1 : 0
+        );
+
+        const emailVerifierInputs = generateEmailVerifierInputsFromDKIMResult(
+            dkimResult,
+            {
+                maxHeadersLength: 640,
+                maxBodyLength: 768,
+                ignoreBodyHashCheck: false,
+                enableHeaderMasking: true,
+                headerMask: mask.map((value) => (value ? 1 : 0)),
+            }
+        );
+
+        const expectedMaskedHeader = emailVerifierInputs.emailHeader!.map(
+            (byte, i) => (mask[i] === 1 ? byte : 0)
+        );
+
+        const witness = await circuit.calculateWitness(emailVerifierInputs);
+        await circuit.checkConstraints(witness);
+        await circuit.assertOut(witness, {
+            maskedHeader: expectedMaskedHeader,
+        });
+    });
+});

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

@@ -2,4 +2,4 @@ pragma circom 2.1.6;
 
 include "../../email-verifier.circom";
 
-component main { public [ pubkey ] } = EmailVerifier(640, 768, 121, 17, 1, 0, 0);
+component main { public [ pubkey ] } = EmailVerifier(640, 768, 121, 17, 1, 0, 0, 0);

packages/circuits/tests/test-circuits/email-verifier-test.circom

@@ -2,4 +2,4 @@ pragma circom 2.1.6;
 
 include "../../email-verifier.circom";
 
-component main { public [ pubkey ] } = EmailVerifier(640, 768, 121, 17, 0, 0, 0);
+component main { public [ pubkey ] } = EmailVerifier(640, 768, 121, 17, 0, 0, 0, 0);

packages/circuits/tests/test-circuits/email-verifier-with-mask-test.circom → packages/circuits/tests/test-circuits/email-verifier-with-body-mask-test.circom

@@ -2,4 +2,4 @@ pragma circom 2.1.6;
 
 include "../../email-verifier.circom";
 
-component main { public [ pubkey ] } = EmailVerifier(640, 768, 121, 17, 0, 0, 1);
+component main { public [ pubkey ] } = EmailVerifier(640, 768, 121, 17, 0, 0, 1, 0);

packages/circuits/tests/test-circuits/email-verifier-with-header-mask-test.circom

@@ -0,0 +1,5 @@
+pragma circom 2.1.6;
+
+include "../../email-verifier.circom";
+
+component main { public [ pubkey ] } = EmailVerifier(640, 768, 121, 17, 0, 1, 0, 0);

packages/circuits/tests/test-circuits/email-verifier-with-soft-line-breaks-test.circom

@@ -2,4 +2,4 @@ pragma circom 2.1.6;
 
 include "../../email-verifier.circom";
 
-component main { public [ pubkey ] } = EmailVerifier(640, 1408, 121, 17, 0, 1, 0);
+component main { public [ pubkey ] } = EmailVerifier(640, 1408, 121, 17, 0, 0, 0, 1);

packages/circuits/utils/bytes.circom

@@ -157,29 +157,29 @@ template AssertBit() {
     in * (in - 1) === 0;
 }
 
-// The ByteMask 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 ByteMask template masks an input array using a binary mask array.
+// Each element in the input 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.
+// - maxLength: The maximum length of the input 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:
-// - maskedBody: An array of signals representing the masked body.
-template ByteMask(maxBodyLength) {
-    signal input body[maxBodyLength];
-    signal input mask[maxBodyLength];
-    signal output maskedBody[maxBodyLength];
+// - out: An array of signals representing the masked input.
+template ByteMask(maxLength) {
+    signal input in[maxLength];
+    signal input mask[maxLength];
+    signal output out[maxLength];
 
-    component bit_check[maxBodyLength];
+    component bit_check[maxLength];
 
-    for (var i = 0; i < maxBodyLength; i++) {
+    for (var i = 0; i < maxLength; i++) {
         bit_check[i] = AssertBit();
         bit_check[i].in <== mask[i];
-        maskedBody[i] <== body[i] * mask[i];
+        out[i] <== in[i] * mask[i];
     }
 }

packages/helpers/src/input-generators.ts

@@ -13,17 +13,20 @@ type CircuitInput = {
   precomputedSHA?: string[];
   bodyHashIndex?: string;
   decodedEmailBodyIn?: string[];
-  mask?: number[];
+  headerMask?: number[];
+  bodyMask?: number[];
 };
 
 type InputGenerationArgs = {
   ignoreBodyHashCheck?: boolean;
+  enableHeaderMasking?: boolean;
   enableBodyMasking?: boolean;
   shaPrecomputeSelector?: string;
   maxHeadersLength?: number; // Max length of the email header including padding
   maxBodyLength?: number; // Max length of the email body after shaPrecomputeSelector including padding
   removeSoftLineBreaks?: boolean;
-  mask?: number[];
+  headerMask?: number[];
+  bodyMask?: number[];
 };
 
 function removeSoftLineBreaks(body: string[]): string[] {
@@ -95,6 +98,10 @@ export function generateEmailVerifierInputsFromDKIMResult(
     signature: toCircomBigIntBytes(signature),
   };
 
+  if (params.enableHeaderMasking) {
+    circuitInputs.headerMask = params.headerMask;
+  }
+
   if (!params.ignoreBodyHashCheck) {
     if (!body || !bodyHash) {
       throw new Error(
@@ -132,7 +139,7 @@ export function generateEmailVerifierInputsFromDKIMResult(
     }
 
     if (params.enableBodyMasking) {
-      circuitInputs.mask = params.mask;
+      circuitInputs.bodyMask = params.bodyMask;
     }
   }