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

use mail_auth::common::crypto::{RsaKey, Sha256};
use mail_auth::common::headers::HeaderWriter;
use mail_auth::dkim::{Canonicalization, DkimSigner};
use mail_auth::{AuthenticatedMessage, 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>,
}

/// 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;

    let mut results = Vec::new();

    if dkim_outputs.is_empty() {
        results.push(DkimVerificationResult {
            is_valid: false,
            domain: None,
            selector: None,
            status: "none".to_string(),
            details: Some("No DKIM signatures found".to_string()),
        });
        return Ok(results);
    }

    for output in &dkim_outputs {
        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));

        results.push(DkimVerificationResult {
            is_valid,
            domain,
            selector,
            status: status.to_string(),
            details,
        });
    }

    Ok(results)
}

/// 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())
}

/// Generate a DKIM DNS TXT record value for a given public key.
///
/// Returns the value for a TXT record at `{selector}._domainkey.{domain}`.
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)
}

#[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());
    }
}