Skip to content

Commit

Permalink
Merge pull request #216 from zkemail/feat/header-masking
Browse files Browse the repository at this point in the history
Feat: Added Header Masking
  • Loading branch information
Divide-By-0 authored Sep 13, 2024
2 parents d718290 + 58f7d73 commit fc99497
Show file tree
Hide file tree
Showing 21 changed files with 573 additions and 340 deletions.
51 changes: 50 additions & 1 deletion packages/circuits/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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**
Expand Down Expand Up @@ -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`

Expand Down Expand Up @@ -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>
26 changes: 17 additions & 9 deletions packages/circuits/email-verifier.circom
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
}
}

Expand Down
2 changes: 1 addition & 1 deletion packages/circuits/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
105 changes: 52 additions & 53 deletions packages/circuits/tests/base64.test.ts
Original file line number Diff line number Diff line change
@@ -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");
}
}
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
Expand All @@ -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],
});
});

Expand Down
48 changes: 48 additions & 0 deletions packages/circuits/tests/email-verifier-no-body.test.ts
Original file line number Diff line number Diff line change
@@ -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);
});
});
Loading

0 comments on commit fc99497

Please sign in to comment.