feat: public key resolution in parsed email

Cargo.toml

@@ -41,3 +41,4 @@ wasm-bindgen-futures = "0.4.45"
 js-sys = "0.3.72"
 serde-wasm-bindgen = "0.6.5"
 rand = "0.8.5"
+base64 = "0.22.1"

src/cryptos.rs

@@ -1,10 +1,19 @@
 //! Cryptographic functions.
 
+#[cfg(target_arch = "wasm32")]
+use crate::EmailHeaders;
 use crate::{field_to_hex, hex_to_field};
+use anyhow::Result;
 use ethers::types::Bytes;
 use halo2curves::ff::Field;
 use poseidon_rs::{poseidon_bytes, poseidon_fields, Fr, PoseidonError};
 use rand_core::RngCore;
+#[cfg(target_arch = "wasm32")]
+use regex::Regex;
+#[cfg(target_arch = "wasm32")]
+use rsa::pkcs8::DecodePublicKey;
+#[cfg(target_arch = "wasm32")]
+use rsa::traits::PublicKeyParts;
 use serde::{
     de::{self, Visitor},
     Deserialize, Deserializer, Serialize, Serializer,
@@ -612,3 +621,82 @@ pub fn calculate_account_salt(email_addr: &str, account_code: &str) -> String {
     // Convert account salt to hexadecimal representation
     field_to_hex(&account_salt.0)
 }
+
+#[cfg(target_arch = "wasm32")]
+/// Fetches the public key from DNS records using the DKIM signature in the email headers.
+///
+/// # Arguments
+///
+/// * `email_headers` - An `EmailHeaders` object containing the headers of the email.
+///
+/// # Returns
+///
+/// A `Result` containing a vector of bytes representing the public key, or an error if the key is not found.
+pub async fn fetch_public_key(email_headers: EmailHeaders) -> Result<Vec<u8>> {
+    let mut selector = String::new();
+    let mut domain = String::new();
+
+    // Extract the selector and domain from the DKIM-Signature header
+    if let Some(headers) = email_headers.get_header("DKIM-Signature") {
+        if let Some(header) = headers.first() {
+            let s_re = Regex::new(r"s=([^;]+);").unwrap();
+            let d_re = Regex::new(r"d=([^;]+);").unwrap();
+
+            selector = s_re
+                .captures(header)
+                .and_then(|cap| cap.get(1))
+                .map_or("", |m| m.as_str())
+                .to_string();
+            domain = d_re
+                .captures(header)
+                .and_then(|cap| cap.get(1))
+                .map_or("", |m| m.as_str())
+                .to_string();
+        }
+    }
+
+    println!("Selector: {}, Domain: {}", selector, domain);
+
+    // Fetch the DNS TXT record for the domain key
+    let response = reqwest::get(format!(
+        "https://dns.google/resolve?name={}._domainkey.{}&type=TXT",
+        selector, domain
+    ))
+    .await?;
+    let data: serde_json::Value = response.json().await?;
+
+    // Extract the 'p' value from the Answer section
+    let mut p_value = None;
+    if let Some(answers) = data.get("Answer").and_then(|a| a.as_array()) {
+        for answer in answers {
+            if let Some(data) = answer.get("data").and_then(|d| d.as_str()) {
+                let parts: Vec<&str> = data.split(';').collect();
+                for part in parts {
+                    let key_value: Vec<&str> = part.trim().split('=').collect();
+                    if key_value.len() == 2 && key_value[0].trim() == "p" {
+                        p_value = Some(key_value[1].trim().to_string());
+                        break;
+                    }
+                }
+            }
+        }
+    }
+
+    if let Some(public_key_b64) = p_value {
+        // Decode the base64 string to get the public key bytes
+        let public_key_bytes = base64::decode(public_key_b64)?;
+
+        // Load the public key from DER format
+        let public_key = rsa::RsaPublicKey::from_public_key_der(&public_key_bytes)?;
+
+        // Extract the modulus from the public key
+        let modulus = public_key.n();
+
+        // Convert the modulus to a byte array in big-endian order
+        let modulus_bytes = modulus.to_bytes_be();
+
+        Ok(modulus_bytes)
+    } else {
+        Err(anyhow::anyhow!("Public key not found"))
+    }
+}

src/parse_email.rs

@@ -2,13 +2,17 @@
 
 use std::collections::HashMap;
 
+#[cfg(target_arch = "wasm32")]
+use crate::cryptos::fetch_public_key;
 use anyhow::Result;
 use cfdkim::canonicalize_signed_email;
 #[cfg(not(target_arch = "wasm32"))]
 use cfdkim::resolve_public_key;
 use hex;
 use itertools::Itertools;
 use mailparse::{parse_mail, ParsedMail};
+#[cfg(target_arch = "wasm32")]
+use regex::Regex;
 use rsa::traits::PublicKeyParts;
 use serde::{Deserialize, Serialize};
 use zk_regex_apis::extract_substrs::{
@@ -51,6 +55,10 @@ impl ParsedEmail {
         // Initialize a logger for the function scope.
         let logger = slog::Logger::root(slog::Discard, slog::o!());
 
+        // Extract all headers
+        let parsed_mail = parse_mail(raw_email.as_bytes())?;
+        let headers: EmailHeaders = EmailHeaders::new_from_mail(&parsed_mail);
+
         // Resolve the public key from the raw email bytes.
         #[cfg(not(target_arch = "wasm32"))]
         let public_key = match resolve_public_key(&logger, raw_email.as_bytes()).await? {
@@ -59,16 +67,12 @@ impl ParsedEmail {
         };
 
         #[cfg(target_arch = "wasm32")]
-        let public_key = Vec::new();
+        let public_key = fetch_public_key(headers.clone()).await?;
 
         // Canonicalize the signed email to separate the header, body, and signature.
         let (canonicalized_header, canonicalized_body, signature_bytes) =
             canonicalize_signed_email(raw_email.as_bytes())?;
 
-        // Extract all headers
-        let parsed_mail = parse_mail(raw_email.as_bytes())?;
-        let headers: EmailHeaders = EmailHeaders::new_from_mail(&parsed_mail);
-
         // Construct the `ParsedEmail` instance.
         let parsed_email = ParsedEmail {
             canonicalized_header: String::from_utf8(canonicalized_header)?, // Convert bytes to string, may return an error if not valid UTF-8.

src/wasm.rs

@@ -1,40 +1,108 @@
+#[cfg(target_arch = "wasm32")]
 use js_sys::Promise;
+#[cfg(target_arch = "wasm32")]
 use rand::rngs::OsRng;
+#[cfg(target_arch = "wasm32")]
 use serde_wasm_bindgen::to_value;
+#[cfg(target_arch = "wasm32")]
 use wasm_bindgen::prelude::*;
 
+#[cfg(target_arch = "wasm32")]
 use crate::{hex_to_field, AccountCode, AccountSalt, PaddedEmailAddr, ParsedEmail};
 
 #[wasm_bindgen]
 #[allow(non_snake_case)]
+#[cfg(target_arch = "wasm32")]
+/// Parses a raw email string into a structured `ParsedEmail` object.
+///
+/// This function utilizes the `ParsedEmail::new_from_raw_email` method to parse the email,
+/// and then serializes the result for JavaScript interoperability.
+///
+/// # Arguments
+///
+/// * `raw_email` - A `String` representing the raw email to be parsed.
+///
+/// # Returns
+///
+/// A `Promise` that resolves with the serialized `ParsedEmail` or rejects with an error message.
 pub async fn parseEmail(raw_email: String) -> Promise {
-    let parsed_email = ParsedEmail::new_from_raw_email(&raw_email)
-        .await
-        .expect("Failed to parse email");
-    let parsed_email = to_value(&parsed_email).expect("Failed to serialize ParsedEmail");
-    Promise::resolve(&parsed_email)
+    match ParsedEmail::new_from_raw_email(&raw_email).await {
+        Ok(parsed_email) => match to_value(&parsed_email) {
+            Ok(serialized_email) => Promise::resolve(&serialized_email),
+            Err(_) => Promise::reject(&JsValue::from_str("Failed to serialize ParsedEmail")),
+        },
+        Err(_) => Promise::reject(&JsValue::from_str("Failed to parse email")),
+    }
 }
 
 #[wasm_bindgen]
 #[allow(non_snake_case)]
-pub async fn generateAccountCode() -> JsValue {
-    to_value(&AccountCode::new(OsRng)).expect("Failed to serialize AccountCode")
+#[cfg(target_arch = "wasm32")]
+/// Generates a new `AccountCode` using a secure random number generator.
+///
+/// This function creates a new `AccountCode` and serializes it for JavaScript interoperability.
+///
+/// # Returns
+///
+/// A `Promise` that resolves with the serialized `AccountCode` or rejects with an error message.
+pub async fn generateAccountCode() -> Promise {
+    match to_value(&AccountCode::new(OsRng)) {
+        Ok(serialized_code) => Promise::resolve(&serialized_code),
+        Err(_) => Promise::reject(&JsValue::from_str("Failed to serialize AccountCode")),
+    }
 }
 
 #[wasm_bindgen]
 #[allow(non_snake_case)]
-pub async fn generateAccountSalt(email_addr: String, account_code: String) -> JsValue {
+#[cfg(target_arch = "wasm32")]
+/// Generates an `AccountSalt` using a padded email address and an account code.
+///
+/// This function converts the email address to a padded format, parses the account code,
+/// and generates an `AccountSalt`, which is then serialized for JavaScript interoperability.
+///
+/// # Arguments
+///
+/// * `email_addr` - A `String` representing the email address.
+/// * `account_code` - A `String` representing the account code in hexadecimal format.
+///
+/// # Returns
+///
+/// A `Promise` that resolves with the serialized `AccountSalt` or rejects with an error message.
+pub async fn generateAccountSalt(email_addr: String, account_code: String) -> Promise {
     let email_addr = PaddedEmailAddr::from_email_addr(&email_addr);
-    let account_code =
-        AccountCode::from(hex_to_field(&account_code).expect("Failed to parse AccountCode"));
-    let account_salt =
-        AccountSalt::new(&email_addr, account_code).expect("Failed to generate AccountSalt");
-    to_value(&account_salt).expect("Failed to serialize AccountSalt")
+    let account_code = match hex_to_field(&account_code) {
+        Ok(field) => AccountCode::from(field),
+        Err(_) => return Promise::reject(&JsValue::from_str("Failed to parse AccountCode")),
+    };
+    let account_salt = match AccountSalt::new(&email_addr, account_code) {
+        Ok(salt) => salt,
+        Err(_) => return Promise::reject(&JsValue::from_str("Failed to generate AccountSalt")),
+    };
+    match to_value(&account_salt) {
+        Ok(serialized_salt) => Promise::resolve(&serialized_salt),
+        Err(_) => Promise::reject(&JsValue::from_str("Failed to serialize AccountSalt")),
+    }
 }
 
 #[wasm_bindgen]
 #[allow(non_snake_case)]
-pub async fn padEmailAddr(email_addr: String) -> JsValue {
+#[cfg(target_arch = "wasm32")]
+/// Pads an email address to a fixed length format.
+///
+/// This function converts the email address to a padded format and serializes it
+/// for JavaScript interoperability.
+///
+/// # Arguments
+///
+/// * `email_addr` - A `String` representing the email address to be padded.
+///
+/// # Returns
+///
+/// A `Promise` that resolves with the serialized padded email address or rejects with an error message.
+pub async fn padEmailAddr(email_addr: String) -> Promise {
     let padded_email_addr = PaddedEmailAddr::from_email_addr(&email_addr);
-    to_value(&padded_email_addr).expect("Failed to serialize padded_email_addr")
+    match to_value(&padded_email_addr) {
+        Ok(serialized_addr) => Promise::resolve(&serialized_addr),
+        Err(_) => Promise::reject(&JsValue::from_str("Failed to serialize padded_email_addr")),
+    }
 }