From 136a8c897d400bad8e2068a910e8246ebecc2cfc Mon Sep 17 00:00:00 2001
From: Saleel <saleel@saleel.xyz>
Date: Fri, 11 Oct 2024 00:59:50 +0530
Subject: [PATCH] feat: enable dns archiver flag in tryVerifyDKIM

---
 packages/helpers/src/dkim/index.ts | 88 ++++++++++++++++++++++--------
 1 file changed, 66 insertions(+), 22 deletions(-)

diff --git a/packages/helpers/src/dkim/index.ts b/packages/helpers/src/dkim/index.ts
index 4bd714a17..8ee07b36f 100644
--- a/packages/helpers/src/dkim/index.ts
+++ b/packages/helpers/src/dkim/index.ts
@@ -1,7 +1,7 @@
-import { pki } from 'node-forge';
-import { DkimVerifier } from '../lib/mailauth/dkim-verifier';
-import { writeToStream } from '../lib/mailauth/tools';
-import sanitizers from './sanitizers';
+import { pki } from "node-forge";
+import { DkimVerifier } from "../lib/mailauth/dkim-verifier";
+import { writeToStream } from "../lib/mailauth/tools";
+import sanitizers from "./sanitizers";
 
 // `./mailauth` is modified version of https://github.com/postalsys/mailauth
 // Main modification are including emailHeaders in the DKIM result, making it work in the browser, add types
@@ -26,32 +26,36 @@ export interface DKIMVerificationResult {
  * @param email Entire email data as a string or buffer
  * @param domain Domain to verify DKIM signature for. If not provided, the domain is extracted from the `From` header
  * @param enableSanitization If true, email will be applied with various sanitization to try and pass DKIM verification
+ * @param enableZKEmailDNSArchiver If provided, this public (modulus as bigint) key will be used instead of the one in the email
  * @returns
  */
 export async function verifyDKIMSignature(
   email: Buffer | string,
-  domain: string = '',
+  domain: string = "",
   enableSanitization: boolean = true,
+  enableZKEmailDNSArchiver: bigint | null = null
 ): Promise<DKIMVerificationResult> {
   const emailStr = email.toString();
 
-  let dkimResult = await tryVerifyDKIM(email, domain);
+  let dkimResult = await tryVerifyDKIM(email, domain, enableZKEmailDNSArchiver);
 
   // If DKIM verification fails, try again after sanitizing email
   let appliedSanitization;
-  if (dkimResult.status.comment === 'bad signature' && enableSanitization) {
+  if (dkimResult.status.comment === "bad signature" && enableSanitization) {
     const results = await Promise.all(
-      sanitizers.map((sanitize) => tryVerifyDKIM(sanitize(emailStr), domain).then((result) => ({
-        result,
-        sanitizer: sanitize.name,
-      }))),
+      sanitizers.map((sanitize) =>
+        tryVerifyDKIM(sanitize(emailStr), domain).then((result) => ({
+          result,
+          sanitizer: sanitize.name,
+        }))
+      )
     );
 
-    const passed = results.find((r) => r.result.status.result === 'pass');
+    const passed = results.find((r) => r.result.status.result === "pass");
 
     if (passed) {
       console.log(
-        `DKIM: Verification passed after applying sanitization "${passed.sanitizer}"`,
+        `DKIM: Verification passed after applying sanitization "${passed.sanitizer}"`
       );
       dkimResult = passed.result;
       appliedSanitization = passed.sanitizer;
@@ -68,16 +72,16 @@ export async function verifyDKIMSignature(
     bodyHash,
   } = dkimResult;
 
-  if (result !== 'pass') {
+  if (result !== "pass") {
     throw new Error(
-      `DKIM signature verification failed for domain ${signingDomain}. Reason: ${comment}`,
+      `DKIM signature verification failed for domain ${signingDomain}. Reason: ${comment}`
     );
   }
 
   const pubKeyData = pki.publicKeyFromPem(publicKey.toString());
 
   return {
-    signature: BigInt(`0x${Buffer.from(signature, 'base64').toString('hex')}`),
+    signature: BigInt(`0x${Buffer.from(signature, "base64").toString("hex")}`),
     headers: status.signedHeaders,
     body,
     bodyHash,
@@ -91,28 +95,68 @@ export async function verifyDKIMSignature(
   };
 }
 
-async function tryVerifyDKIM(email: Buffer | string, domain: string = '') {
-  const dkimVerifier = new DkimVerifier({});
+async function tryVerifyDKIM(
+  email: Buffer | string,
+  domain: string = "",
+  enableZKEmailDNSArchiver: bigint | null = null
+) {
+  const dkimVerifier = new DkimVerifier({
+    ...(enableZKEmailDNSArchiver && {
+      resolver: async (name: string, type: string) => {
+        if (type !== "TXT") {
+          throw new Error(
+            `ZK Email Archive only supports TXT records - got ${type}`
+          );
+        }
+        const ZKEMAIL_DNS_ARCHIVER_API = "https://archive.prove.email/api/key";
+
+        // Get domain from full dns record name - $selector._domainkey.$domain.com
+        const domain = name.split(".").slice(-2).join(".");
+        const selector = name.split(".")[0];
+
+        const queryUrl = new URL(ZKEMAIL_DNS_ARCHIVER_API);
+        queryUrl.searchParams.set("domain", domain);
+
+        const resp = await fetch(queryUrl);
+        const data = await resp.json();
+
+        const dkimRecord = data.find(
+          (record: any) => record.selector === selector
+        );
+
+        if (!dkimRecord) {
+          throw new Error(
+            `DKIM record not found for domain ${domain} and selector ${selector} in ZK Email Archive.`
+          );
+        }
+
+        console.log("dkimRecord", dkimRecord.value);
+
+        return [dkimRecord.value];
+      },
+    }),
+  });
+
   await writeToStream(dkimVerifier, email as any);
 
   let domainToVerifyDKIM = domain;
   if (!domainToVerifyDKIM) {
     if (dkimVerifier.headerFrom.length > 1) {
       throw new Error(
-        'Multiple From header in email and domain for verification not specified',
+        "Multiple From header in email and domain for verification not specified"
       );
     }
 
-    domainToVerifyDKIM = dkimVerifier.headerFrom[0].split('@')[1];
+    domainToVerifyDKIM = dkimVerifier.headerFrom[0].split("@")[1];
   }
 
   const dkimResult = dkimVerifier.results.find(
-    (d: any) => d.signingDomain === domainToVerifyDKIM,
+    (d: any) => d.signingDomain === domainToVerifyDKIM
   );
 
   if (!dkimResult) {
     throw new Error(
-      `DKIM signature not found for domain ${domainToVerifyDKIM}`,
+      `DKIM signature not found for domain ${domainToVerifyDKIM}`
     );
   }