smartmta/rust/crates/mailer-security/src/dkim.rs

use mail_auth::common::crypto::{Ed25519Key, RsaKey, Sha256};
use mail_auth::common::headers::HeaderWriter;
use mail_auth::dkim::{Canonicalization, DkimSigner};
use mail_auth::{AuthenticatedMessage, DkimOutput, DkimResult, MessageAuthenticator};
use rustls_pki_types::{PrivateKeyDer, PrivatePkcs1KeyDer, pem::PemObject};
use serde::{Deserialize, Serialize};

use crate::error::{Result, SecurityError};

/// Result of DKIM verification.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DkimVerificationResult {
    /// Whether the DKIM signature is valid.
    pub is_valid: bool,
    /// The signing domain (d= tag).
    pub domain: Option<String>,
    /// The selector (s= tag).
    pub selector: Option<String>,
    /// Result status: "pass", "fail", "permerror", "temperror", "none".
    pub status: String,
    /// Human-readable details.
    pub details: Option<String>,
}

/// Convert raw `mail-auth` DKIM outputs to our serializable results.
///
/// This is used internally by `verify_dkim` and by the compound `verify_email_security`.
pub fn dkim_outputs_to_results(dkim_outputs: &[DkimOutput<'_>]) -> Vec<DkimVerificationResult> {
    if dkim_outputs.is_empty() {
        return vec![DkimVerificationResult {
            is_valid: false,
            domain: None,
            selector: None,
            status: "none".to_string(),
            details: Some("No DKIM signatures found".to_string()),
        }];
    }

    dkim_outputs
        .iter()
        .map(|output| {
            let (is_valid, status, details) = match output.result() {
                DkimResult::Pass => (true, "pass", None),
                DkimResult::Neutral(err) => (false, "neutral", Some(err.to_string())),
                DkimResult::Fail(err) => (false, "fail", Some(err.to_string())),
                DkimResult::PermError(err) => (false, "permerror", Some(err.to_string())),
                DkimResult::TempError(err) => (false, "temperror", Some(err.to_string())),
                DkimResult::None => (false, "none", None),
            };

            let (domain, selector) = output
                .signature()
                .map(|sig| (Some(sig.d.clone()), Some(sig.s.clone())))
                .unwrap_or((None, None));

            DkimVerificationResult {
                is_valid,
                domain,
                selector,
                status: status.to_string(),
                details,
            }
        })
        .collect()
}

/// Verify DKIM signatures on a raw email message.
///
/// Uses the `mail-auth` crate which performs full RFC 6376 verification
/// including DNS lookups for the public key.
pub async fn verify_dkim(
    raw_message: &[u8],
    authenticator: &MessageAuthenticator,
) -> Result<Vec<DkimVerificationResult>> {
    let message = AuthenticatedMessage::parse(raw_message)
        .ok_or_else(|| SecurityError::Parse("Failed to parse email for DKIM verification".into()))?;

    let dkim_outputs = authenticator.verify_dkim(&message).await;
    Ok(dkim_outputs_to_results(&dkim_outputs))
}

/// Sign a raw email message with DKIM (RSA-SHA256).
///
/// * `raw_message` - The raw RFC 5322 message bytes
/// * `domain` - The signing domain (d= tag)
/// * `selector` - The DKIM selector (s= tag)
/// * `private_key_pem` - RSA private key in PEM format (PKCS#1 or PKCS#8)
///
/// Returns the DKIM-Signature header string to prepend to the message.
pub fn sign_dkim(
    raw_message: &[u8],
    domain: &str,
    selector: &str,
    private_key_pem: &str,
) -> Result<String> {
    // Try PKCS#1 PEM first, then PKCS#8
    let key_der = PrivatePkcs1KeyDer::from_pem_slice(private_key_pem.as_bytes())
        .map(PrivateKeyDer::Pkcs1)
        .or_else(|_| {
            // Try PKCS#8
            rustls_pki_types::PrivatePkcs8KeyDer::from_pem_slice(private_key_pem.as_bytes())
                .map(PrivateKeyDer::Pkcs8)
        })
        .map_err(|e| SecurityError::Key(format!("Failed to parse private key PEM: {}", e)))?;

    let rsa_key = RsaKey::<Sha256>::from_key_der(key_der)
        .map_err(|e| SecurityError::Key(format!("Failed to load RSA key: {}", e)))?;

    let signature = DkimSigner::from_key(rsa_key)
        .domain(domain)
        .selector(selector)
        .headers(["From", "To", "Subject", "Date", "Message-ID", "MIME-Version", "Content-Type"])
        .header_canonicalization(Canonicalization::Relaxed)
        .body_canonicalization(Canonicalization::Relaxed)
        .sign(raw_message)
        .map_err(|e| SecurityError::Dkim(format!("DKIM signing failed: {}", e)))?;

    Ok(signature.to_header())
}

/// Sign a raw email message with DKIM using Ed25519-SHA256 (RFC 8463).
///
/// * `raw_message` - The raw RFC 5322 message bytes
/// * `domain` - The signing domain (d= tag)
/// * `selector` - The DKIM selector (s= tag)
/// * `private_key_pkcs8_der` - Ed25519 private key in PKCS#8 DER format
///
/// Returns the DKIM-Signature header string to prepend to the message.
pub fn sign_dkim_ed25519(
    raw_message: &[u8],
    domain: &str,
    selector: &str,
    private_key_pkcs8_der: &[u8],
) -> Result<String> {
    let ed_key = Ed25519Key::from_pkcs8_maybe_unchecked_der(private_key_pkcs8_der)
        .map_err(|e| SecurityError::Key(format!("Failed to load Ed25519 key: {}", e)))?;

    let signature = DkimSigner::from_key(ed_key)
        .domain(domain)
        .selector(selector)
        .headers(["From", "To", "Subject", "Date", "Message-ID", "MIME-Version", "Content-Type"])
        .header_canonicalization(Canonicalization::Relaxed)
        .body_canonicalization(Canonicalization::Relaxed)
        .sign(raw_message)
        .map_err(|e| SecurityError::Dkim(format!("Ed25519 DKIM signing failed: {}", e)))?;

    Ok(signature.to_header())
}

/// Sign a raw email message with DKIM, auto-selecting RSA or Ed25519 based on `key_type`.
///
/// * `key_type` - `"rsa"` (default) or `"ed25519"`
/// * For RSA: `private_key_pem` is a PEM-encoded RSA key
/// * For Ed25519: `private_key_pem` is a PEM-encoded PKCS#8 Ed25519 key
pub fn sign_dkim_auto(
    raw_message: &[u8],
    domain: &str,
    selector: &str,
    private_key_pem: &str,
    key_type: &str,
) -> Result<String> {
    match key_type.to_lowercase().as_str() {
        "ed25519" => {
            // Parse PEM to DER for Ed25519
            let der = rustls_pki_types::PrivatePkcs8KeyDer::from_pem_slice(private_key_pem.as_bytes())
                .map_err(|e| SecurityError::Key(format!("Failed to parse Ed25519 PEM: {}", e)))?;
            sign_dkim_ed25519(raw_message, domain, selector, der.secret_pkcs8_der())
        }
        _ => sign_dkim(raw_message, domain, selector, private_key_pem),
    }
}

/// Generate a DKIM DNS TXT record value for a given public key.
///
/// Returns the value for a TXT record at `{selector}._domainkey.{domain}`.
/// `key_type` should be `"rsa"` or `"ed25519"`.
pub fn dkim_dns_record_value(public_key_pem: &str) -> String {
    // Extract the base64 content from PEM
    let key_b64: String = public_key_pem
        .lines()
        .filter(|line| !line.starts_with("-----"))
        .collect::<Vec<_>>()
        .join("");

    format!("v=DKIM1; h=sha256; k=rsa; p={}", key_b64)
}

/// Generate a DKIM DNS TXT record value with explicit key type.
///
/// * `key_type` - `"rsa"` or `"ed25519"`
pub fn dkim_dns_record_value_typed(public_key_pem: &str, key_type: &str) -> String {
    let key_b64: String = public_key_pem
        .lines()
        .filter(|line| !line.starts_with("-----"))
        .collect::<Vec<_>>()
        .join("");

    let k = match key_type.to_lowercase().as_str() {
        "ed25519" => "ed25519",
        _ => "rsa",
    };

    format!("v=DKIM1; h=sha256; k={}; p={}", k, key_b64)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_dkim_dns_record_value() {
        let pem = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBg==\n-----END PUBLIC KEY-----";
        let record = dkim_dns_record_value(pem);
        assert!(record.starts_with("v=DKIM1; h=sha256; k=rsa; p="));
        assert!(record.contains("MIIBIjANBg=="));
    }

    #[test]
    fn test_sign_dkim_invalid_key() {
        let result = sign_dkim(b"From: test@example.com\r\n\r\nBody", "example.com", "mta", "not a key");
        assert!(result.is_err());
    }

    #[test]
    fn test_sign_dkim_ed25519() {
        // Generate an Ed25519 key pair using mail-auth
        let pkcs8_der = Ed25519Key::generate_pkcs8().expect("generate ed25519 key");
        let ed_key = Ed25519Key::from_pkcs8_der(&pkcs8_der).expect("parse ed25519 key");
        let _pub_key = ed_key.public_key();

        let msg = b"From: test@example.com\r\nTo: rcpt@example.com\r\nSubject: Test\r\n\r\nBody";
        let result = sign_dkim_ed25519(msg, "example.com", "ed25519sel", &pkcs8_der);
        assert!(result.is_ok());
        let header = result.unwrap();
        assert!(header.contains("a=ed25519-sha256"));
        assert!(header.contains("d=example.com"));
        assert!(header.contains("s=ed25519sel"));
    }

    #[test]
    fn test_sign_dkim_auto_dispatches() {
        // RSA with invalid key should error
        let msg = b"From: test@example.com\r\n\r\nBody";
        let result = sign_dkim_auto(msg, "example.com", "mta", "not a key", "rsa");
        assert!(result.is_err());

        // Ed25519 with invalid PEM should error
        let result = sign_dkim_auto(msg, "example.com", "mta", "not a key", "ed25519");
        assert!(result.is_err());
    }

    #[test]
    fn test_dkim_dns_record_value_typed() {
        let pem = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBg==\n-----END PUBLIC KEY-----";
        let rsa_record = dkim_dns_record_value_typed(pem, "rsa");
        assert!(rsa_record.contains("k=rsa"));

        let ed_record = dkim_dns_record_value_typed(pem, "ed25519");
        assert!(ed_record.contains("k=ed25519"));
    }
}