From 393eaa61cae7ed7ac91945f3cd9a6b47993d2d0f Mon Sep 17 00:00:00 2001 From: John Coffey Date: Fri, 16 Oct 2026 08:08:54 -0700 Subject: [PATCH] Re-queue DKIM rotation and certificate renewal a domain is missing DKIM rotation and ACME renewal are queued only when a domain is saved into automatic mode, and each run queues the next. A domain whose first run never happened, or whose chain broke, kept its keys forever with nothing scheduled: on one server nine automatic domains had no DkimManagement task and keys with no nextTransitionAt. - The task scheduler checks daily, and five minutes after start, that every domain with automatic DKIM has a DkimManagement task and every domain with automatic certificates an AcmeRenewal task, and queues any that are missing. Only the node with the task scheduler role runs it. - A task counts as present in any state, failed included, so a domain whose task keeps failing isn't queued again every day. - DKIM: only when DNS publishes DKIM, or when the domain has no keys at all. Without publishing, a run neither rotates nor queues a follow-up. - ACME: a renewal that finds a still-valid certificate reschedules itself for when it falls due (#93), so queuing one orders nothing early. - Each re-queued task is logged as task-manager.task-queued, with the domain and "Re-queued a missing domain task". Tests: automation_tests' new orphaned-key part drops a domain's task and its keys' schedule, then checks the domain gets a task, the same keys get a schedule again, and a second check queues nothing. --- .../services/src/task_manager/domain_check.rs | 264 ++++++++++++++++++ crates/services/src/task_manager/mod.rs | 1 + crates/services/src/task_manager/scheduler.rs | 29 ++ tests/src/automation/dkim.rs | 118 ++++++++ 4 files changed, 412 insertions(+) create mode 100644 crates/services/src/task_manager/domain_check.rs diff --git a/crates/services/src/task_manager/domain_check.rs b/crates/services/src/task_manager/domain_check.rs new file mode 100644 index 0000000..167dab5 --- /dev/null +++ b/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, + pub acme: HashSet, +} + +/// 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 { + // Registry::list() returns no Domain objects, so query their ids + let domain_ids = server + .registry() + .query::>(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 { + 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_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> { + 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 { + server + .registry() + .query::>( + 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 { + // 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::(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) -> 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))); + } +} diff --git a/crates/services/src/task_manager/mod.rs b/crates/services/src/task_manager/mod.rs index d8c1976..b4bdb7b 100644 --- a/crates/services/src/task_manager/mod.rs +++ b/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; diff --git a/crates/services/src/task_manager/scheduler.rs b/crates/services/src/task_manager/scheduler.rs index a32a2d6..1c07d98 100644 --- a/crates/services/src/task_manager/scheduler.rs +++ b/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) { 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) { 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", } } } diff --git a/tests/src/automation/dkim.rs b/tests/src/automation/dkim.rs index e645382..475490f 100644 --- a/tests/src/automation/dkim.rs +++ b/tests/src/automation/dkim.rs @@ -2,6 +2,8 @@ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * 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("admin@example.org"); @@ -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("admin@example.org"); + 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::().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("admin@example.org"); @@ -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| { -- 2.54.0