Clean up challenge material when a solve is cancelled

src/cleanup.rs

@@ -0,0 +1,122 @@
+//! Teardown of challenge material that survives cancellation.
+
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+
+type CleanupFn<T> = Arc<dyn Fn(T) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;
+
+/// Tracks what a solver has put in place and must take down.
+///
+/// The solver calls [`run`](Self::run) on every path it controls. If the
+/// solving future is dropped first (a caller-side `timeout`, a `select!`, a
+/// shutdown), whatever is still pending is cleaned up from `Drop` on a task
+/// spawned on the current Tokio runtime. That task is best-effort: it does not
+/// run if there is no runtime, or if the runtime shuts down before it does.
+pub(crate) struct CleanupGuard<T: Clone + Send + 'static> {
+    pending: Vec<T>,
+    cleanup: CleanupFn<T>,
+}
+
+impl<T: Clone + Send + 'static> CleanupGuard<T> {
+    pub(crate) fn new<F, Fut>(cleanup: F) -> Self
+    where
+        F: Fn(T) -> Fut + Send + Sync + 'static,
+        Fut: Future<Output = ()> + Send + 'static,
+    {
+        Self {
+            pending: Vec::new(),
+            cleanup: Arc::new(move |item| Box::pin(cleanup(item))),
+        }
+    }
+
+    pub(crate) fn push(&mut self, item: T) {
+        self.pending.push(item);
+    }
+
+    pub(crate) fn pending(&self) -> &[T] {
+        &self.pending
+    }
+
+    /// Cleans up every pending item, in order. An item leaves the list only
+    /// once its cleanup has completed, so a drop in the middle hands the
+    /// in-flight item and the rest over to `Drop`.
+    pub(crate) async fn run(&mut self) {
+        while let Some(item) = self.pending.first().cloned() {
+            (self.cleanup)(item).await;
+            self.pending.remove(0);
+        }
+    }
+}
+
+impl<T: Clone + Send + 'static> Drop for CleanupGuard<T> {
+    fn drop(&mut self) {
+        if self.pending.is_empty() {
+            return;
+        }
+        let Ok(runtime) = tokio::runtime::Handle::try_current() else {
+            return;
+        };
+        let pending = std::mem::take(&mut self.pending);
+        let cleanup = self.cleanup.clone();
+        runtime.spawn(async move {
+            for item in pending {
+                cleanup(item).await;
+            }
+        });
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Mutex;
+    use std::time::Duration;
+
+    use super::*;
+
+    fn recording_guard() -> (CleanupGuard<u32>, Arc<Mutex<Vec<u32>>>) {
+        let cleaned = Arc::new(Mutex::new(Vec::new()));
+        let sink = cleaned.clone();
+        let guard = CleanupGuard::new(move |item: u32| {
+            let sink = sink.clone();
+            async move { sink.lock().unwrap().push(item) }
+        });
+        (guard, cleaned)
+    }
+
+    #[tokio::test]
+    async fn run_cleans_every_item_in_order_and_empties_the_guard() {
+        let (mut guard, cleaned) = recording_guard();
+        guard.push(1);
+        guard.push(2);
+        guard.run().await;
+        assert!(guard.pending().is_empty());
+        drop(guard);
+        tokio::task::yield_now().await;
+        assert_eq!(*cleaned.lock().unwrap(), vec![1, 2]);
+    }
+
+    #[tokio::test]
+    async fn dropping_with_pending_items_spawns_their_cleanup() {
+        let (mut guard, cleaned) = recording_guard();
+        guard.push(7);
+        guard.push(8);
+        drop(guard);
+        tokio::time::timeout(Duration::from_secs(1), async {
+            while cleaned.lock().unwrap().len() < 2 {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .expect("spawned cleanup did not run");
+        assert_eq!(*cleaned.lock().unwrap(), vec![7, 8]);
+    }
+
+    #[test]
+    fn dropping_outside_a_runtime_does_not_panic() {
+        let (mut guard, cleaned) = recording_guard();
+        guard.push(1);
+        drop(guard);
+        assert!(cleaned.lock().unwrap().is_empty());
+    }
+}

src/lib.rs

@@ -5,6 +5,7 @@
 #![doc = include_str!("../README.md")]
 
 pub mod account_store;
+mod cleanup;
 pub mod dns;
 pub mod error;
 pub mod provider;

src/solver.rs

@@ -1,9 +1,11 @@
+use std::sync::Arc;
 use std::time::Duration;
 
 use instant_acme::{
     AuthorizationStatus, ChallengeType, Identifier, Order, OrderStatus, RetryPolicy,
 };
 
+use crate::cleanup::CleanupGuard;
 use crate::dns::wait_for_propagation;
 use crate::error::DnsError;
 use crate::provider::{DnsProvider, PropagationTiming};
@@ -42,6 +44,7 @@ pub(crate) trait ChallengeDriver {
     async fn finalize(&mut self) -> Result<(String, String), DnsError>;
 }
 
+#[derive(Clone)]
 struct PlacedRecord {
     fqdn: String,
     value: String,
@@ -50,17 +53,17 @@ struct PlacedRecord {
 /// End-to-end DNS-01 orchestrator. Wraps a `DnsProvider` and drives an
 /// `instant_acme::Order` from `Pending` to a usable certificate, calling
 /// `present` / `cleanup` on the provider as needed.
-pub struct Dns01Solver<P: DnsProvider> {
-    provider: P,
+pub struct Dns01Solver<P: DnsProvider + 'static> {
+    provider: Arc<P>,
     fallback_resolvers: Vec<String>,
     timing: Option<PropagationTiming>,
     skip_propagation: Option<Duration>,
 }
 
-impl<P: DnsProvider> Dns01Solver<P> {
+impl<P: DnsProvider + 'static> Dns01Solver<P> {
     pub fn new(provider: P) -> Self {
         Self {
-            provider,
+            provider: Arc::new(provider),
             fallback_resolvers: crate::dns::DEFAULT_RESOLVERS
                 .iter()
                 .map(|s| (*s).to_string())
@@ -98,8 +101,14 @@ impl<P: DnsProvider> Dns01Solver<P> {
     }
 
     /// Drive the order from Pending to a certificate. Returns
-    /// `(certificate_chain_pem, private_key_pem)`. Cleanup of placed TXT
-    /// records is best-effort and runs even on error paths.
+    /// `(certificate_chain_pem, private_key_pem)`.
+    ///
+    /// Placed TXT records are removed before returning, on success and on
+    /// error. If this future is dropped before completing (a caller-side
+    /// timeout, a `select!`, a shutdown), the records still in place are
+    /// removed from a task spawned on the current Tokio runtime. Cleanup is
+    /// best-effort either way: a failing `cleanup` is not reported, and the
+    /// spawned task does not run if the runtime shuts down first.
     pub async fn solve_and_finalize(&self, order: Order) -> Result<(String, String), DnsError> {
         let driver = InstantAcmeDriver::new(order);
         self.solve_with_driver(driver).await
@@ -109,27 +118,23 @@ impl<P: DnsProvider> Dns01Solver<P> {
         &self,
         mut driver: D,
     ) -> Result<(String, String), DnsError> {
-        let placed = match self.present_all(&mut driver).await {
-            Ok(p) => p,
-            Err(e) => return Err(e),
+        let mut placed = self.cleanup_guard();
+        let outcome = match self.present_then_ready(&mut driver, &mut placed).await {
+            Ok(()) => driver.finalize().await,
+            Err(e) => Err(e),
         };
-        let outcome = driver.finalize().await;
-        self.cleanup_all(&placed).await;
+        placed.run().await;
         outcome
     }
 
-    async fn present_all<D: ChallengeDriver>(
-        &self,
-        driver: &mut D,
-    ) -> Result<Vec<PlacedRecord>, DnsError> {
-        let mut placed: Vec<PlacedRecord> = Vec::new();
-        match self.present_then_ready(driver, &mut placed).await {
-            Ok(()) => Ok(placed),
-            Err(e) => {
-                self.cleanup_all(&placed).await;
-                Err(e)
+    fn cleanup_guard(&self) -> CleanupGuard<PlacedRecord> {
+        let provider = Arc::clone(&self.provider);
+        CleanupGuard::new(move |record: PlacedRecord| {
+            let provider = Arc::clone(&provider);
+            async move {
+                let _ = provider.cleanup(&record.fqdn, &record.value).await;
             }
-        }
+        })
     }
 
     /// Places every record first, then waits for propagation once, and only
@@ -138,7 +143,7 @@ impl<P: DnsProvider> Dns01Solver<P> {
     async fn present_then_ready<D: ChallengeDriver>(
         &self,
         driver: &mut D,
-        placed: &mut Vec<PlacedRecord>,
+        placed: &mut CleanupGuard<PlacedRecord>,
     ) -> Result<(), DnsError> {
         let mut tokens: Vec<usize> = Vec::new();
 
@@ -163,7 +168,7 @@ impl<P: DnsProvider> Dns01Solver<P> {
                 let timing = self.timing.unwrap_or_else(|| self.provider.timing());
                 let resolvers: Vec<&str> =
                     self.fallback_resolvers.iter().map(String::as_str).collect();
-                for record in placed.iter() {
+                for record in placed.pending() {
                     wait_for_propagation(&record.fqdn, &record.value, &resolvers, timing).await?;
                 }
             }
@@ -174,12 +179,6 @@ impl<P: DnsProvider> Dns01Solver<P> {
         }
         Ok(())
     }
-
-    async fn cleanup_all(&self, placed: &[PlacedRecord]) {
-        for record in placed {
-            let _ = self.provider.cleanup(&record.fqdn, &record.value).await;
-        }
-    }
 }
 
 /// Production `ChallengeDriver` that wraps an `instant_acme::Order`.
@@ -618,6 +617,52 @@ mod tests {
         );
     }
 
+    #[tokio::test(start_paused = true)]
+    async fn records_are_cleaned_up_when_the_solve_is_cancelled() {
+        use std::sync::Arc;
+
+        struct ArcProvider(Arc<FakeProvider>);
+        #[async_trait::async_trait]
+        impl DnsProvider for ArcProvider {
+            async fn present(&self, fqdn: &str, value: &str) -> Result<(), DnsError> {
+                self.0.present(fqdn, value).await
+            }
+            async fn cleanup(&self, fqdn: &str, value: &str) -> Result<(), DnsError> {
+                self.0.cleanup(fqdn, value).await
+            }
+        }
+
+        let provider = Arc::new(FakeProvider::new());
+        let driver = FakeDriver::new(vec![
+            Scripted::Pending {
+                domain: "a.example".into(),
+            },
+            Scripted::Pending {
+                domain: "b.example".into(),
+            },
+        ]);
+        // The propagation delay outlasts the caller's timeout, so the solve
+        // is dropped after both records are placed and before any cleanup.
+        let solver = Dns01Solver::new(ArcProvider(provider.clone()))
+            .skip_propagation_check(Duration::from_secs(600));
+        let cancelled =
+            tokio::time::timeout(Duration::from_secs(1), solver.solve_with_driver(driver)).await;
+        assert!(cancelled.is_err(), "the solve should have been cancelled");
+
+        // Let the cleanup task spawned on drop run; bounded so a missing
+        // cleanup fails the assertion below instead of hanging.
+        for _ in 0..100 {
+            if provider.cleaned_fqdns().len() >= 2 {
+                break;
+            }
+            tokio::task::yield_now().await;
+        }
+        assert_eq!(
+            provider.cleaned_fqdns(),
+            vec!["_acme-challenge.a.example", "_acme-challenge.b.example"]
+        );
+    }
+
     #[tokio::test(start_paused = true)]
     async fn skip_propagation_delay_elapses_before_set_ready() {
         use std::sync::Arc;