crates/services/src/task_manager/domain_check.rs

@@ -0,0 +1,264 @@
+/*
+ * SPDX-FileCopyrightText: 2026 Coffey Labs LLC
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+//! Keeps automatic DKIM rotation and certificate renewal running. Both
+//! tasks are only queued when a domain is saved into automatic mode, and
+//! each run queues the next one. A domain whose first run was missed (or
+//! whose chain was broken) never got another, so its keys stopped rotating.
+//! This check runs on the task scheduler and queues the missing task.
+
+use common::Server;
+use registry::{
+    schema::{
+        enums::{DnsRecordType, TaskType},
+        prelude::{ObjectType, Property},
+        structs::{
+            CertificateManagement, DkimManagement, DnsManagement, Domain, Task,
+            TaskDomainManagement, TaskStatus,
+        },
+    },
+    types::EnumImpl,
+};
+use std::collections::HashSet;
+use store::{
+    IterateParams, U64_LEN, ValueKey,
+    registry::RegistryQuery,
+    write::{BatchBuilder, TaskQueueClass, ValueClass, key::DeserializeBigEndian},
+};
+use trc::{AddContext, TaskManagerEvent};
+use types::id::Id;
+
+/// Domains that already have a task of each kind, whatever its state
+/// (pending, retrying, running or failed).
+#[derive(Default)]
+pub struct QueuedDomainTasks {
+    pub dkim: HashSet<Id>,
+    pub acme: HashSet<Id>,
+}
+
+/// Queues the DKIM and ACME tasks that automatic domains are missing.
+/// Returns how many were queued.
+pub async fn check_domain_tasks(server: &Server) -> trc::Result<usize> {
+    // Registry::list() returns no Domain objects, so query their ids
+    let domain_ids = server
+        .registry()
+        .query::<Vec<Id>>(RegistryQuery::new(ObjectType::Domain))
+        .await
+        .caused_by(trc::location!())?;
+
+    check_tasks_for_domains(server, &domain_ids).await
+}
+
+/// The same check for the given domains only.
+pub async fn check_tasks_for_domains(server: &Server, domain_ids: &[Id]) -> trc::Result<usize> {
+    let queued = queued_domain_tasks(server).await?;
+    let mut batch = BatchBuilder::new();
+    let mut total = 0;
+
+    for &domain_id in domain_ids {
+        let Some(domain) = server
+            .registry()
+            .object::<Domain>(domain_id)
+            .await
+            .caused_by(trc::location!())?
+        else {
+            continue;
+        };
+
+        for task in missing_tasks(server, domain_id, &domain, &queued).await? {
+            trc::event!(
+                TaskManager(TaskManagerEvent::TaskQueued),
+                Type = task.object_type().as_str(),
+                Domain = domain.name.clone(),
+                Details = "Re-queued a missing domain task",
+            );
+            batch.schedule_task(task);
+            total += 1;
+        }
+    }
+
+    if !batch.is_empty() {
+        server
+            .store()
+            .write(batch.build_all())
+            .await
+            .caused_by(trc::location!())?;
+        server.notify_task_queue();
+    }
+
+    Ok(total)
+}
+
+async fn missing_tasks(
+    server: &Server,
+    domain_id: Id,
+    domain: &Domain,
+    queued: &QueuedDomainTasks,
+) -> trc::Result<Vec<Task>> {
+    let mut tasks = Vec::new();
+
+    // Without published DNS a DKIM run can't rotate anything, and it doesn't
+    // queue a follow-up; it's only needed when the domain has no keys yet.
+    if needs_dkim_task(domain, queued.dkim.contains(&domain_id))
+        && (publishes_dkim(domain) || !has_dkim_keys(server, domain_id).await?)
+    {
+        tasks.push(Task::DkimManagement(TaskDomainManagement {
+            domain_id,
+            status: TaskStatus::now(),
+        }));
+    }
+
+    // A renewal that finds a certificate still valid reschedules itself for
+    // when it falls due (renew.rs), so queuing one now orders nothing early.
+    if needs_acme_task(domain, queued.acme.contains(&domain_id)) {
+        tasks.push(Task::AcmeRenewal(TaskDomainManagement {
+            domain_id,
+            status: TaskStatus::now(),
+        }));
+    }
+
+    Ok(tasks)
+}
+
+pub fn needs_dkim_task(domain: &Domain, has_task: bool) -> bool {
+    !has_task && matches!(domain.dkim_management, DkimManagement::Automatic(_))
+}
+
+pub fn needs_acme_task(domain: &Domain, has_task: bool) -> bool {
+    !has_task
+        && matches!(
+            domain.certificate_management,
+            CertificateManagement::Automatic(_)
+        )
+}
+
+pub fn publishes_dkim(domain: &Domain) -> bool {
+    matches!(
+        &domain.dns_management,
+        DnsManagement::Automatic(details) if details.publish_records.contains(&DnsRecordType::Dkim)
+    )
+}
+
+async fn has_dkim_keys(server: &Server, domain_id: Id) -> trc::Result<bool> {
+    server
+        .registry()
+        .query::<Vec<Id>>(
+            RegistryQuery::new(ObjectType::DkimSignature)
+                .equal(Property::DomainId, domain_id.document_id()),
+        )
+        .await
+        .map(|ids| !ids.is_empty())
+        .caused_by(trc::location!())
+}
+
+async fn queued_domain_tasks(server: &Server) -> trc::Result<QueuedDomainTasks> {
+    // The due index carries each task's type, so only the few domain tasks
+    // are read in full. Failed tasks sit at the end of it (due u64::MAX).
+    let mut ids = Vec::new();
+    server
+        .store()
+        .iterate(
+            IterateParams::new(
+                ValueKey::from(ValueClass::TaskQueue(TaskQueueClass::Due { id: 0, due: 1 })),
+                ValueKey::from(ValueClass::TaskQueue(TaskQueueClass::Due {
+                    id: u64::MAX,
+                    due: u64::MAX,
+                })),
+            )
+            .ascending(),
+            |key, value| {
+                let task_type = value.deserialize_be_u16(0).ok().and_then(TaskType::from_id);
+                if matches!(
+                    task_type,
+                    Some(TaskType::DkimManagement | TaskType::AcmeRenewal)
+                ) {
+                    ids.push(key.deserialize_be_u64(U64_LEN)?);
+                }
+                Ok(true)
+            },
+        )
+        .await
+        .caused_by(trc::location!())?;
+
+    let mut queued = QueuedDomainTasks::default();
+    for id in ids {
+        match server
+            .store()
+            .get_value::<Task>(ValueKey::from(ValueClass::TaskQueue(
+                TaskQueueClass::Task { id },
+            )))
+            .await
+            .caused_by(trc::location!())?
+        {
+            Some(Task::DkimManagement(task)) => {
+                queued.dkim.insert(task.domain_id);
+            }
+            Some(Task::AcmeRenewal(task)) => {
+                queued.acme.insert(task.domain_id);
+            }
+            _ => {}
+        }
+    }
+
+    Ok(queued)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use registry::schema::structs::{
+        CertificateManagementProperties, DkimManagementProperties, DnsManagementProperties,
+    };
+    use registry::types::map::Map;
+
+    fn domain(cert: bool, dkim: bool, dns_dkim: Option<bool>) -> Domain {
+        Domain {
+            certificate_management: if cert {
+                CertificateManagement::Automatic(CertificateManagementProperties::default())
+            } else {
+                CertificateManagement::Manual
+            },
+            dkim_management: if dkim {
+                DkimManagement::Automatic(DkimManagementProperties::default())
+            } else {
+                DkimManagement::Manual
+            },
+            dns_management: match dns_dkim {
+                None => DnsManagement::Manual,
+                Some(with_dkim) => DnsManagement::Automatic(DnsManagementProperties {
+                    publish_records: Map::new(if with_dkim {
+                        vec![DnsRecordType::Dkim]
+                    } else {
+                        vec![DnsRecordType::Mx]
+                    }),
+                    ..Default::default()
+                }),
+            },
+            ..Default::default()
+        }
+    }
+
+    #[test]
+    fn dkim_task_only_for_automatic_domains_without_one() {
+        assert!(needs_dkim_task(&domain(false, true, None), false));
+        assert!(!needs_dkim_task(&domain(false, true, None), true));
+        assert!(!needs_dkim_task(&domain(false, false, None), false));
+    }
+
+    #[test]
+    fn acme_task_only_for_automatic_domains_without_one() {
+        assert!(needs_acme_task(&domain(true, false, None), false));
+        assert!(!needs_acme_task(&domain(true, false, None), true));
+        assert!(!needs_acme_task(&domain(false, false, None), false));
+    }
+
+    #[test]
+    fn publishes_dkim_needs_automatic_dns_with_dkim_records() {
+        assert!(publishes_dkim(&domain(false, true, Some(true))));
+        assert!(!publishes_dkim(&domain(false, true, Some(false))));
+        assert!(!publishes_dkim(&domain(false, true, None)));
+    }
+}

crates/services/src/task_manager/mod.rs

@@ -22,6 +22,7 @@ pub mod alarm;
 pub mod destroy_account;
 pub mod dkim;
 pub mod dns;
+pub mod domain_check; // inbuxa: re-queue missing DKIM/ACME tasks
 pub mod imip;
 pub mod inbuxa_restore; // inbuxa: undelete
 pub mod index;

crates/services/src/task_manager/scheduler.rs

@@ -48,8 +48,14 @@ enum Event {
     EvaluateAlerts,
     // inbuxa: settings-reorg: probe the other nodes' ports
     ProbePeerPorts,
+    // inbuxa: re-queue DKIM rotation and ACME renewal tasks a domain is missing
+    CheckDomainTasks,
 }
 
+// inbuxa: first check a few minutes after start, then daily
+const DOMAIN_TASK_CHECK_DELAY: Duration = Duration::from_secs(5 * 60);
+const DOMAIN_TASK_CHECK_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
+
 /// When the next metric-history tick is due (MON-4), read from the registry
 /// so a change needs no reload (MON-3).
 async fn metrics_collection_delay(server: &common::Server) -> Duration {
@@ -146,6 +152,11 @@ pub fn spawn_task_scheduler(inner: Arc<Inner>) {
                 Event::StoreMetrics,
             );
 
+            // inbuxa: automatic domains keep their DKIM and ACME tasks
+            queue.schedule(
+                Instant::now() + DOMAIN_TASK_CHECK_DELAY,
+                Event::CheckDomainTasks,
+            );
         }
 
 
@@ -400,6 +411,23 @@ pub fn spawn_task_scheduler(inner: Arc<Inner>) {
                             server.store_metrics().await;
                         });
                     }
+                    // inbuxa: on the task scheduler only, like the tasks it queues
+                    Event::CheckDomainTasks => {
+                        queue.schedule(
+                            Instant::now() + DOMAIN_TASK_CHECK_INTERVAL,
+                            Event::CheckDomainTasks,
+                        );
+                        if roles.task_scheduler {
+                            let server = server.clone();
+                            tokio::spawn(async move {
+                                if let Err(err) =
+                                    super::domain_check::check_domain_tasks(&server).await
+                                {
+                                    trc::error!(err.details("Failed to check domain tasks"));
+                                }
+                            });
+                        }
+                    }
                     Event::TrainSpamClassifier => {
                         if let Some(train_frequency) = server
                             .core
@@ -497,6 +525,7 @@ impl Event {
             Event::StoreMetrics => "storeMetrics",
             Event::EvaluateAlerts => "evaluateAlerts",
             Event::ProbePeerPorts => "probePeerPorts",
+            Event::CheckDomainTasks => "checkDomainTasks",
         }
     }
 }

tests/src/automation/dkim.rs

@@ -2,6 +2,8 @@
  * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
  *
  * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ *
+ * Modified by Coffey Labs in 2026 for INBUXA.
  */
 
 use crate::utils::{account::Account, server::TestServer};
@@ -40,6 +42,7 @@ pub async fn test(test: &TestServer) {
     })
     .await;
     automatic_to_manual_dns_test(test).await;
+    orphaned_keys_test(test).await;
 
     println!("Running DKIM Management tests...");
     let account = test.account("[email protected]");
@@ -525,6 +528,100 @@ async fn unscheduled_keys_test(
     account.registry_destroy_all(ObjectType::DnsServer).await;
 }
 
+// inbuxa: keys left with no rotation schedule and no DKIM task (a domain
+// switched to automatic before scheduling existed) are picked up by the
+// scheduler's domain task check, and a second check queues nothing more
+async fn orphaned_keys_test(test: &TestServer) {
+    println!("Running DKIM orphaned key tests...");
+    let account = test.account("[email protected]");
+    DNS_RECORDS.lock().unwrap().clear();
+    account.dkim_signatures().await.assert_total(0, 0);
+    let dns_server_id = account.create_memory_dns_server().await;
+    let domain = "dkim-orphaned.org";
+    let domain_id = account
+        .registry_create_object(Domain {
+            name: domain.to_string(),
+            certificate_management: CertificateManagement::Manual,
+            dkim_management: dkim_management(LONG_ROTATION_MS),
+            dns_management: dns_management(dns_server_id),
+            ..Default::default()
+        })
+        .await;
+    let initial = account
+        .wait_for_dkim("active keys with a rotation schedule", |signatures| {
+            signatures.total() == 2
+                && signatures.stage_count(DkimRotationStage::Active) == 2
+                && signatures
+                    .keys()
+                    .all(|key| key.next_transition_at.is_some())
+        })
+        .await;
+    account.wait_for_dkim_task_count(domain_id, 1).await;
+
+    // Orphan the keys: no task and no schedule
+    account.registry_destroy_all(ObjectType::Task).await;
+    for (id, _) in account.registry_get_all::<DkimSignature>().await {
+        account
+            .registry_update_object(
+                ObjectType::DkimSignature,
+                id,
+                json!({ Property::NextTransitionAt: null }),
+            )
+            .await;
+    }
+    account
+        .wait_for_dkim("keys without a rotation schedule", |signatures| {
+            signatures
+                .keys()
+                .all(|key| key.next_transition_at.is_none())
+        })
+        .await;
+    account.wait_for_dkim_task_count(domain_id, 0).await;
+
+    // The check queues a DKIM task, which schedules the same keys
+    assert_eq!(
+        services::task_manager::domain_check::check_tasks_for_domains(&test.server, &[domain_id])
+            .await
+            .unwrap(),
+        1
+    );
+    let rescheduled = account
+        .wait_for_dkim("the same keys rescheduled", |signatures| {
+            signatures.total() == 2
+                && signatures.stage_count(DkimRotationStage::Active) == 2
+                && signatures
+                    .keys()
+                    .all(|key| key.next_transition_at.is_some())
+        })
+        .await;
+    assert_eq!(rescheduled.v1_rsa[0].selector, initial.v1_rsa[0].selector);
+    assert_eq!(
+        rescheduled.v1_ed25519[0].selector,
+        initial.v1_ed25519[0].selector
+    );
+    account.wait_for_dkim_task_count(domain_id, 1).await;
+
+    // The follow-up task counts as queued
+    assert_eq!(
+        services::task_manager::domain_check::check_tasks_for_domains(&test.server, &[domain_id])
+            .await
+            .unwrap(),
+        0
+    );
+    account.wait_for_dkim_task_count(domain_id, 1).await;
+
+    // Cleanup
+    account.registry_destroy_all(ObjectType::Task).await;
+    account
+        .registry_destroy_all(ObjectType::DkimSignature)
+        .await;
+    account
+        .registry_destroy(ObjectType::Domain, [domain_id])
+        .await
+        .assert_destroyed(&[domain_id]);
+    account.registry_destroy_all(ObjectType::DnsServer).await;
+}
+
 async fn automatic_to_manual_dns_test(test: &TestServer) {
     println!("Running DKIM automatic to manual DNS tests...");
     let account = test.account("[email protected]");
@@ -776,6 +873,27 @@ impl Account {
         panic!("DKIM signatures did not reach {expected}: {signatures:#?}");
     }
 
+    // inbuxa: settles on exactly `count` DKIM tasks for the domain
+    async fn wait_for_dkim_task_count(&self, domain_id: Id, count: usize) {
+        let mut last = None;
+        for _ in 0..60 {
+            if let Some(tasks) = self.tasks().await {
+                let found = tasks
+                    .iter()
+                    .filter(|task| {
+                        matches!(&task.task, Task::DkimManagement(task) if task.domain_id == domain_id)
+                    })
+                    .count();
+                if found == count {
+                    return;
+                }
+                last = Some(found);
+            }
+            tokio::time::sleep(std::time::Duration::from_millis(250)).await;
+        }
+        panic!("Expected {count} DKIM tasks for the domain, found {last:?}");
+    }
+
     async fn wait_for_no_dkim_tasks(&self, domain_id: Id) {
         for _ in 0..60 {
             if self.tasks().await.is_some_and(|tasks| {