Lazy DKIM key parsing (closes #1211)

crates/common/src/config/smtp/auth.rs

@@ -13,13 +13,13 @@ use mail_auth::{
 };
 use mail_parser::decoders::base64::base64_decode;
 use utils::config::{
-    utils::{AsKey, ParseValue},
     Config,
+    utils::{AsKey, ParseValue},
 };
 
 use crate::{
     config::CONNECTION_VARS,
-    expr::{self, if_block::IfBlock, tokenizer::TokenMap, Constant, ConstantValue},
+    expr::{self, Constant, ConstantValue, if_block::IfBlock, tokenizer::TokenMap},
 };
 
 use super::*;
@@ -31,9 +31,20 @@ pub struct MailAuthConfig {
     pub spf: SpfAuthConfig,
     pub dmarc: DmarcAuthConfig,
     pub iprev: IpRevAuthConfig,
+    pub signatures: AHashMap<String, Arc<ArcSwap<LazySignature>>>,
+}
 
-    pub signers: AHashMap<String, Arc<DkimSigner>>,
-    pub sealers: AHashMap<String, Arc<ArcSealer>>,
+#[allow(clippy::large_enum_variant)]
+pub enum LazySignature {
+    Resolved(ResolvedSignature),
+    Pending(Config),
+    Failed,
+}
+
+#[derive(Clone)]
+pub struct ResolvedSignature {
+    pub signer: Arc<DkimSigner>,
+    pub sealer: Arc<ArcSealer>,
 }
 
 #[derive(Clone)]
@@ -150,8 +161,7 @@ impl Default for MailAuthConfig {
                     "relaxed",
                 ),
             },
-            signers: Default::default(),
-            sealers: Default::default(),
+            signatures: Default::default(),
         }
     }
 }
@@ -193,23 +203,40 @@ impl MailAuthConfig {
             .unwrap_or(true);
 
         // Parse signatures
-        for id in config
-            .sub_keys("signature", ".algorithm")
-            .map(|k| k.to_string())
-            .collect::<Vec<_>>()
-        {
-            let id = id.to_string();
-            if let Some((signer, sealer)) = build_signature(config, &id) {
-                mail_auth.signers.insert(id.clone(), Arc::new(signer));
-                mail_auth.sealers.insert(id, Arc::new(sealer));
+        let mut signatures: AHashMap<&str, Config> = AHashMap::new();
+        let mut current_id = None;
+        for (k, v) in config.keys.iter() {
+            if let Some(prefix) = k.strip_prefix("signature.") {
+                if let Some(id) = prefix.strip_suffix(".algorithm") {
+                    current_id = Some(id);
+                }
+                #[allow(clippy::unwrap_or_default)]
+                if let Some(current_id) = current_id {
+                    signatures
+                        .entry(current_id)
+                        .or_insert_with(Config::default)
+                        .keys
+                        .insert(k.to_string(), v.to_string());
+                }
+            } else if !signatures.is_empty() {
+                break;
             }
         }
+        mail_auth.signatures = signatures
+            .into_iter()
+            .map(|(id, config)| {
+                (
+                    id.to_string(),
+                    Arc::new(ArcSwap::from_pointee(LazySignature::Pending(config))),
+                )
+            })
+            .collect();
 
         mail_auth
     }
 }
 
-fn build_signature(config: &mut Config, id: &str) -> Option<(DkimSigner, ArcSealer)> {
+pub fn build_signature(config: &mut Config, id: &str) -> Option<(DkimSigner, ArcSealer)> {
     match config.property_require::<Algorithm>(("signature", id, "algorithm"))? {
         Algorithm::RsaSha256 => {
             let pk = config

crates/common/src/core.rs

@@ -6,32 +6,32 @@
 
 use std::{sync::Arc, time::Duration};
 
-use directory::{backend::internal::manage::ManageDirectory, Directory, QueryBy, Type};
+use directory::{Directory, QueryBy, Type, backend::internal::manage::ManageDirectory};
 use jmap_proto::types::{
     blob::BlobId, collection::Collection, property::Property, state::StateChange,
 };
 use sieve::Sieve;
 use store::{
+    BitmapKey, BlobClass, BlobStore, Deserialize, FtsStore, InMemoryStore, IndexKey, IterateParams,
+    LogKey, Serialize, Store, U32_LEN, ValueKey,
     dispatch::DocumentSet,
     roaring::RoaringBitmap,
     write::{
-        key::DeserializeBigEndian, log::ChangeLogBuilder, now, BatchBuilder, BitmapClass, BlobOp,
-        DirectoryClass, QueueClass, TagValue, ValueClass,
+        BatchBuilder, BitmapClass, BlobOp, DirectoryClass, QueueClass, TagValue, ValueClass,
+        key::DeserializeBigEndian, log::ChangeLogBuilder, now,
     },
-    BitmapKey, BlobClass, BlobStore, Deserialize, FtsStore, InMemoryStore, IndexKey, IterateParams,
-    LogKey, Serialize, Store, ValueKey, U32_LEN,
 };
 use trc::AddContext;
 use utils::BlobHash;
 
 use crate::{
+    ImapId, Inner, MailboxState, Server,
     auth::{AccessToken, ResourceToken, TenantInfo},
     config::smtp::{
-        auth::{ArcSealer, DkimSigner},
+        auth::{ArcSealer, DkimSigner, LazySignature, ResolvedSignature, build_signature},
         queue::RelayHost,
     },
     ipc::StateEvent,
-    ImapId, Inner, MailboxState, Server,
 };
 
 impl Server {
@@ -110,40 +110,51 @@ impl Server {
         })
     }
 
-    pub fn get_arc_sealer(&self, name: &str, session_id: u64) -> Option<&ArcSealer> {
-        self.core
-            .smtp
-            .mail_auth
-            .sealers
-            .get(name)
-            .map(|s| s.as_ref())
-            .or_else(|| {
-                trc::event!(
-                    Arc(trc::ArcEvent::SealerNotFound),
-                    Id = name.to_string(),
-                    SpanId = session_id,
-                );
+    pub fn get_arc_sealer(&self, name: &str, session_id: u64) -> Option<Arc<ArcSealer>> {
+        self.resolve_signature(name).map(|s| s.sealer).or_else(|| {
+            trc::event!(
+                Arc(trc::ArcEvent::SealerNotFound),
+                Id = name.to_string(),
+                SpanId = session_id,
+            );
 
-                None
-            })
+            None
+        })
     }
 
-    pub fn get_dkim_signer(&self, name: &str, session_id: u64) -> Option<&DkimSigner> {
-        self.core
-            .smtp
-            .mail_auth
-            .signers
-            .get(name)
-            .map(|s| s.as_ref())
-            .or_else(|| {
-                trc::event!(
-                    Dkim(trc::DkimEvent::SignerNotFound),
-                    Id = name.to_string(),
-                    SpanId = session_id,
-                );
+    pub fn get_dkim_signer(&self, name: &str, session_id: u64) -> Option<Arc<DkimSigner>> {
+        self.resolve_signature(name).map(|s| s.signer).or_else(|| {
+            trc::event!(
+                Dkim(trc::DkimEvent::SignerNotFound),
+                Id = name.to_string(),
+                SpanId = session_id,
+            );
 
-                None
-            })
+            None
+        })
+    }
+
+    fn resolve_signature(&self, name: &str) -> Option<ResolvedSignature> {
+        let lazy_resolver_ = self.core.smtp.mail_auth.signatures.get(name)?;
+        match lazy_resolver_.load().as_ref() {
+            LazySignature::Resolved(resolved_signature) => Some(resolved_signature.clone()),
+            LazySignature::Pending(config) => {
+                let mut config = config.clone();
+                if let Some((signer, sealer)) = build_signature(&mut config, name) {
+                    let resolved = ResolvedSignature {
+                        signer: Arc::new(signer),
+                        sealer: Arc::new(sealer),
+                    };
+                    lazy_resolver_.store(Arc::new(LazySignature::Resolved(resolved.clone())));
+                    Some(resolved)
+                } else {
+                    config.log_errors();
+                    lazy_resolver_.store(Arc::new(LazySignature::Failed));
+                    None
+                }
+            }
+            LazySignature::Failed => None,
+        }
     }
 
     pub fn get_trusted_sieve_script(&self, name: &str, session_id: u64) -> Option<&Arc<Sieve>> {

tests/src/smtp/inbound/sign.rs

@@ -16,9 +16,9 @@ use store::Stores;
 use utils::config::Config;
 
 use crate::smtp::{
+    DnsCache, TempDir, TestSMTP,
     inbound::TestMessage,
     session::{TestSession, VerifyResponse},
-    DnsCache, TempDir, TestSMTP,
 };
 use smtp::core::Session;