Merge pull request #9 from zkemail/feat/partial-hashing

Feat/partial hashing

examples/partial_hash/src/main.nr

@@ -1,10 +1,10 @@
 use dep::zkemail::{
-    KEY_LIMBS_2048, dkim::verify_dkim_2048,
-    standard_outputs
+    KEY_LIMBS_2048, dkim::verify_dkim_2048, get_body_hash_by_index, base64::body_hash_base64_decode,
+    partial_hash::partial_sha256_var_end, standard_outputs
 };
 
-global MAX_EMAIL_HEADER_LENGTH: u32 = 1024;
-global MAX_EMAIL_BODY_LENGTH: u32 = 1536;
+global MAX_EMAIL_HEADER_LENGTH: u32 = 512;
+global MAX_PARTIAL_EMAIL_BODY_LENGTH: u32 = 192;
 
 /**
  * Verify an arbitrary email signed by a 1024-bit RSA DKIM signature
@@ -22,20 +22,35 @@ global MAX_EMAIL_BODY_LENGTH: u32 = 1536;
  *         1: Pedersen hash of DKIM signature (email nullifier)
  */
 fn main(
+    body_hash_index: u32,
     header: [u8; MAX_EMAIL_HEADER_LENGTH],
     header_length: u32,
+    body: [u8; MAX_PARTIAL_EMAIL_BODY_LENGTH], // use partial body length instead of full body length
+    body_length: u32,
+    partial_body_hash: [u32; 8],
+    partial_body_length: u32,
     pubkey: [Field; KEY_LIMBS_2048],
     pubkey_redc: [Field; KEY_LIMBS_2048],
     signature: [Field; KEY_LIMBS_2048]
 ) -> pub [Field; 2] {
     // check the body and header lengths are within bounds
-    assert(header_length <= MAX_EMAIL_HEADER_LENGTH);
-
+    assert(header_length <= MAX_EMAIL_HEADER_LENGTH, "Email header length exceeds maximum length");
+    assert(partial_body_length <= MAX_PARTIAL_EMAIL_BODY_LENGTH, "Partial email body length exceeds maximum length");
+
     // verify the dkim signature over the header
     verify_dkim_2048(header, header_length, pubkey, pubkey_redc, signature);
 
-    // check the recipient email address
-    // todo
+    // manually extract the body hash from the header
+    let body_hash_encoded = get_body_hash_by_index(header, body_hash_index);
+    let signed_body_hash: [u8; 32] = body_hash_base64_decode(body_hash_encoded);
+
+    // finish the partial hash
+    let computed_body_hash = partial_sha256_var_end(partial_body_hash, body, partial_body_length as u64, body_length as u64);    
+
+    // check the body hashes match
+    assert(
+        signed_body_hash == computed_body_hash, "Sha256 hash computed over body does not match DKIM-signed header"
+    );
 
     // hash the pubkey and signature for the standard outputs
     standard_outputs(pubkey, signature)

examples/verify_email_1024_bit_dkim/src/main.nr

@@ -1,10 +1,11 @@
 use dep::zkemail::{
-    KEY_LIMBS_1024, dkim::verify_dkim_1024, body_hash::{get_body_hash_by_index, compare_body_hash},
+    KEY_LIMBS_1024, dkim::verify_dkim_1024, get_body_hash_by_index, base64::body_hash_base64_decode,
     standard_outputs
 };
+use dep::std::hash::sha256_var;
 
-global MAX_EMAIL_HEADER_LENGTH: u32 = 1024;
-global MAX_EMAIL_BODY_LENGTH: u32 = 1536;
+global MAX_EMAIL_HEADER_LENGTH: u32 = 512;
+global MAX_EMAIL_BODY_LENGTH: u32 = 1024;
 
 /**
  * Verify an arbitrary email signed by a 1024-bit RSA DKIM signature
@@ -38,11 +39,19 @@ fn main(
     // verify the dkim signature over the header
     verify_dkim_1024(header, header_length, pubkey, pubkey_redc, signature);
 
-    // manually extract the body hash from the header
+    // extract the body hash from the header
     let body_hash_encoded = get_body_hash_by_index(header, body_hash_index);
 
-    // compare the retrieved body hash to the computed body hash
-    compare_body_hash(body_hash_encoded, body, body_length);
+    // base64 decode the body hash
+    let signed_body_hash: [u8; 32] = body_hash_base64_decode(body_hash_encoded);
+
+    // hash the asserted body
+    let computed_body_hash: [u8; 32] = sha256_var(body, body_length as u64);
+
+    // compare the body hashes
+    assert(
+        signed_body_hash == computed_body_hash, "SHA256 hash computed over body does not match body hash found in DKIM-signed header"
+    );
 
     // hash the pubkey and signature for the standard outputs
     standard_outputs(pubkey, signature)

examples/verify_email_2048_bit_dkim/src/main.nr

@@ -1,10 +1,11 @@
 use dep::zkemail::{
-    KEY_LIMBS_2048, dkim::verify_dkim_2048, body_hash::{get_body_hash_by_index, compare_body_hash},
+    KEY_LIMBS_2048, dkim::verify_dkim_2048, get_body_hash_by_index, base64::body_hash_base64_decode,
     standard_outputs
 };
+use dep::std::hash::sha256_var;
 
-global MAX_EMAIL_HEADER_LENGTH: u32 = 1024;
-global MAX_EMAIL_BODY_LENGTH: u32 = 1536;
+global MAX_EMAIL_HEADER_LENGTH: u32 = 512;
+global MAX_EMAIL_BODY_LENGTH: u32 = 1024;
 
 /**
  * Verify an arbitrary email signed by a 1024-bit RSA DKIM signature
@@ -38,11 +39,19 @@ fn main(
     // verify the dkim signature over the header
     verify_dkim_2048(header, header_length, pubkey, pubkey_redc, signature);
 
-    // manually extract the body hash from the header
+    // extract the body hash from the header
     let body_hash_encoded = get_body_hash_by_index(header, body_hash_index);
 
-    // compare the retrieved body hash to the computed body hash
-    compare_body_hash(body_hash_encoded, body, body_length);
+    // base64 decode the body hash
+    let signed_body_hash: [u8; 32] = body_hash_base64_decode(body_hash_encoded);
+
+    // hash the asserted body
+    let computed_body_hash: [u8; 32] = sha256_var(body, body_length as u64);
+
+    // compare the body hashes
+    assert(
+        signed_body_hash == computed_body_hash, "SHA256 hash computed over body does not match body hash found in DKIM-signed header"
+    );
 
     // hash the pubkey and signature for the standard outputs
     standard_outputs(pubkey, signature)

js/src/index.ts

@@ -4,28 +4,30 @@ import {
   MAX_HEADER_PADDED_BYTES,
   generatePartialSHA,
   sha256Pad,
+  findIndexInUint8Array,
 } from "@zk-email/helpers";
 import {
   DKIMVerificationResult,
   verifyDKIMSignature,
 } from "@zk-email/helpers/dist/dkim";
 import * as NoirBignum from "@mach-34/noir-bignum-paramgen";
+import { u8ToU32 } from "./utils";
 
 // This file is essentially https://github.com/zkemail/zk-email-verify/blob/main/packages/helpers/src/input-generators.ts
 // with a few modifications for noir input generation
 // also removes some of the unused functionality like masking
 
 type CircuitInput = {
-  emailHeader: string[];
-  emailHeaderLength: string;
+  header: string[];
+  header_length: string;
   pubkey: string[];
-  redcParams?: string[];
+  pubkey_redc: string[];
   signature: string[];
-  emailBody?: string[];
-  emailBodyLength?: string;
-  precomputedSHA?: string[];
-  bodyHashIndex?: string;
-  decodedEmailBodyIn?: string[];
+  body?: string[];
+  body_length?: string;
+  partial_body_length?: string;
+  partial_body_hash?: string[];
+  body_hash_index?: string;
 };
 
 type InputGenerationArgs = {
@@ -97,13 +99,13 @@ export function generateEmailVerifierInputsFromDKIMResult(
   );
 
   const circuitInputs: CircuitInput = {
-    emailHeader: Uint8ArrayToCharArray(messagePadded), // Packed into 1 byte signals
+    header: Uint8ArrayToCharArray(messagePadded), // Packed into 1 byte signals
     // modified from original: can use exact email header length
-    emailHeaderLength: headers.length.toString(),
+    header_length: headers.length.toString(),
     // modified from original: use noir bignum to format
     pubkey: NoirBignum.bnToLimbStrArray(publicKey),
     // not in original: add barrett reduction param for efficient rsa sig verification
-    redcParams: NoirBignum.bnToRedcLimbStrArray(publicKey),
+    pubkey_redc: NoirBignum.bnToRedcLimbStrArray(publicKey),
     // modified from original: use noir bignum to format
     signature: NoirBignum.bnToLimbStrArray(signature),
   };
@@ -128,77 +130,37 @@ export function generateEmailVerifierInputsFromDKIMResult(
       Math.max(maxBodyLength, bodySHALength)
     );
 
-    const { precomputedSha, bodyRemaining } = generatePartialSHA({
+    let { precomputedSha, bodyRemaining, bodyRemainingLength } = generatePartialSHA({
       body: bodyPadded,
       bodyLength: bodyPaddedLen,
       selectorString: params.shaPrecomputeSelector,
       maxRemainingBodyLength: maxBodyLength,
     });
+
+    // idk why this gets out of sync, todo: fix
+    if (params.shaPrecomputeSelector && bodyRemaining.length != bodyRemainingLength) {
+      bodyRemaining = bodyRemaining.slice(0, bodyRemainingLength);
+    }
 
-    // modified from original: can use exact email body length
-    // since circom needs 64 byte chunks and noir doesnt
-    circuitInputs.emailBodyLength = body.length.toString();
-    circuitInputs.precomputedSHA = Uint8ArrayToCharArray(precomputedSha);
-    circuitInputs.bodyHashIndex = bodyHashIndex.toString();
-    circuitInputs.emailBody = Uint8ArrayToCharArray(bodyRemaining);
-
-    if (params.removeSoftLineBreaks) {
-      circuitInputs.decodedEmailBodyIn = removeSoftLineBreaks(
-        circuitInputs.emailBody
-      );
+    // can use exact body lengths
+    circuitInputs.body_length = body.length.toString();
+    circuitInputs.body_hash_index = bodyHashIndex.toString();
+    circuitInputs.body = Uint8ArrayToCharArray(bodyRemaining);
+
+    if (params.shaPrecomputeSelector) {
+      // can use exact body lengths
+      const selector = new TextEncoder().encode(params.shaPrecomputeSelector);
+      const selectorIndex = findIndexInUint8Array(body, selector);
+      const shaCutoffIndex = Math.floor(selectorIndex / 64) * 64;
+      const remainingBodyLength = body.length - shaCutoffIndex;
+      circuitInputs.partial_body_length = remainingBodyLength.toString();
+
+      // format back into u32 so noir doesn't have to do it
+      circuitInputs.partial_body_hash = Array.from(
+        u8ToU32(precomputedSha)
+      ).map((x) => x.toString());
     }
   }
 
   return circuitInputs;
 }
-
-/**
- * Rename inputs for Noir format
- * @todo handle optional values
- *
- * @param inputs - the inputs to convert to Noir format
- * @param exactLength - whether to have exact length for header or (default) keep 0-padding
- * @returns - the inputs as the NoirJS witness simulator expects them
- */
-export function toNoirInputs(inputs: CircuitInput, exactLength = false) {
-  return {
-    body_hash_index: inputs.bodyHashIndex!,
-    header: exactLength
-      ? inputs.emailHeader.slice(0, Number(inputs.emailHeaderLength))!
-      : inputs.emailHeader!,
-    body: exactLength
-      ? inputs.emailBody!.slice(0, Number(inputs.emailBodyLength))!
-      : inputs.emailBody!,
-    body_length: inputs.emailBodyLength!,
-    header_length: inputs.emailHeaderLength!,
-    pubkey: inputs.pubkey!,
-    pubkey_redc: inputs.redcParams!,
-    signature: inputs.signature!,
-  };
-}
-
-/**
- * Format circuit inputs for a Prover.toml file
- *
- * @param inputs - the inputs to convert to Prover.toml format
- * @param exactLength - whether toNoirInputs should have exact length for header or keep 0-padding
- * @returns - the inputs as bb cli expects them to appear in a Prover.toml file
- */
-export function toProverToml(
-  inputs: CircuitInput,
-  exactLength = false
-): string {
-  const formatted = toNoirInputs(inputs, exactLength);
-  const lines: string[] = [];
-  for (const [key, value] of Object.entries(formatted)) {
-    let valueStr = "";
-    if (Array.isArray(value)) {
-      const valueStrArr = value.map((val) => `'${val}'`);
-      valueStr = `[${valueStrArr.join(", ")}]`;
-    } else {
-      valueStr = `'${value}'`;
-    }
-    lines.push(`${key} = ${valueStr}`);
-  }
-  return lines.join("\n");
-}

js/src/utils.ts

@@ -0,0 +1,46 @@
+/**
+ * Transforms a u32 array to a u8 array
+ * @dev sha-utils in zk-email-verify encodes partial hash as u8 array but noir expects u32
+ *      transform back to keep upstream code but not have noir worry about transformation
+ * 
+ * @param input - the input to convert to 32 bit array
+ * @returns - the input as a 32 bit array
+ */
+export function u8ToU32(input: Uint8Array): Uint32Array {
+  const out = new Uint32Array(input.length / 4);
+  for (let i = 0; i < out.length; i++) {
+    out[i] =
+      (input[i * 4 + 0] << 24) |
+      (input[i * 4 + 1] << 16) |
+      (input[i * 4 + 2] << 8) |
+      (input[i * 4 + 3] << 0);
+  }
+  return out;
+}
+
+/**
+ * Format circuit inputs for a Prover.toml file
+ *
+ * @param inputs - the inputs to convert to Prover.toml format
+ * @param exactLength - whether toNoirInputs should have exact length for header or keep 0-padding
+ * @returns - the inputs as bb cli expects them to appear in a Prover.toml file
+ */
+export function toProverToml(inputs: any): string {
+  const lines: string[] = [];
+  const structs: string[] = [];
+  for (const [key, value] of Object.entries(inputs)) {
+    if (Array.isArray(value)) {
+      const valueStrArr = value.map((val) => `'${val}'`);
+      lines.push(`${key} = [${valueStrArr.join(", ")}]`);
+    } else if (typeof value === "string") {
+      lines.push(`${key} = '${value}'`);
+    } else {
+      let values = "";
+      for (const [k, v] of Object.entries(value!)) {
+        values = values.concat(`${k} = '${v}'\n`);
+      }
+      structs.push(`[${key}]\n${values}`);
+    }
+  }
+  return lines.concat(structs).join("\n");
+}