config: fix cancellation safety issue with async lua calls

This commit addresses a tricky interaction with lua's garbage collection
and cancellation safety of async lua call handling.

In Rust, when a future is cancelled it is simply not called any more,
and any resources it owns will be released when it is dropped.

In mlua's Rust bindings, async calls are futures that are wrapped
up into state whose ownership is managed in the lua context
stack, which in turn has managed via its own garbage collection
process.

In KumoMTA, we implement a pool of Lua contexts which we keep
around for reuse.  Our logic will return our LuaConfig object
back to this pool automatically when it is dropped.

In a situation where tokio timeouts or select! are used with
a future that transitively contains a lua async call and that
call is cancelled, that can leave async locks/semaphores owned
in the un-collected objects in the lua stack.

For async locks/semaphores that were in the locked state, this
can leave them locked until the owning lua context is expired
from the pool, either due to age, capacity or use count being hit.

That's bad because it can block certain processing flows for 5 minutes
(the default `kumo.set_max_lua_context_age`).

What this commit does is switch away from automatically returning
the LuaConfig to the pool on Drop.  In that case the lua context
is destroy on Drop, immediately allowing destructors to release
any/all owned resources.

We do still want to benefit from the pool, so the bulk of
this commit is changing the consumers of load_config to
explicitly call `config.put()` to put the object back into
the pool, but only in the success case.

This does penalize flows that grab a config and then
`kumo.reject` as their final status. At the present time, the
majority of kumomta deployments are very much focused on (relatively)
trusted senders and an outbound use case, so this isn't a terrible
situation for the time being.

crates/config/src/lib.rs

@@ -87,17 +87,10 @@ pub struct LuaConfig {
     inner: Option<LuaConfigInner>,
 }
 
-impl Drop for LuaConfig {
-    fn drop(&mut self) {
-        if let Some(inner) = self.inner.take() {
-            pool_put(inner);
-        }
-    }
-}
-
 pub async fn set_policy_path(path: PathBuf) -> anyhow::Result<()> {
     POLICY_FILE.lock().replace(path);
-    load_config().await?;
+    let config = load_config().await?;
+    config.put();
     Ok(())
 }
 
@@ -223,6 +216,29 @@ impl LuaConfig {
         }
     }
 
+    /// Explicitly put the config object back into its containing pool.
+    /// Ideally we'd do this automatically when the object is dropped,
+    /// but lua's garbage collection makes this problematic:
+    /// if a future whose graph contains an async lua call within
+    /// this config object is cancelled (eg: simply stopped without
+    /// calling it again), and the config object is not explicitly garbage
+    /// collected, any futures and data owned by any dependencies of
+    /// the cancelled future remain alive until the next gc run,
+    /// which can cause things like async locks and semaphores to
+    /// have a lifetime extended by the maximum age of the lua context.
+    ///
+    /// The combat this, consumers of LuaConfig should explicitly
+    /// call `config.put()` after successfully using the config
+    /// object.
+    ///
+    /// Or framing it another way: consumers must not call `config.put()`
+    /// if a transitive dep might have been cancelled.
+    pub fn put(mut self) {
+        if let Some(inner) = self.inner.take() {
+            pool_put(inner);
+        }
+    }
+
     pub async fn async_call_callback<A: IntoLuaMulti + Clone, R: FromLuaMulti + Default>(
         &mut self,
         sig: &CallbackSignature<A, R>,

crates/kumo-server-common/src/http_server/auth.rs

@@ -51,22 +51,24 @@ impl AuthKind {
 
     async fn validate_impl(&self) -> anyhow::Result<bool> {
         let mut config = load_config().await?;
-        match self {
-            Self::TrustedIp(_) => Ok(true),
+        let result = match self {
+            Self::TrustedIp(_) => true,
             Self::Basic { user, password } => {
                 let sig = CallbackSignature::<(String, Option<String>), bool>::new(
                     "http_server_validate_auth_basic",
                 );
-                Ok(config
+                config
                     .async_call_callback(&sig, (user.to_string(), password.clone()))
-                    .await?)
+                    .await?
             }
             Self::Bearer { token } => {
                 let sig =
                     CallbackSignature::<String, bool>::new("http_server_validate_auth_bearer");
-                Ok(config.async_call_callback(&sig, token.to_string()).await?)
+                config.async_call_callback(&sig, token.to_string()).await?
             }
-        }
+        };
+        config.put();
+        Ok(result)
     }
 
     async fn lookup_cache(&self) -> Option<Result<bool, String>> {

crates/kumo-server-common/src/lib.rs

@@ -325,7 +325,11 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> {
 
             config
                 .convert_args_and_call_callback(&sig, &self.args)
-                .await
+                .await?;
+
+            config.put();
+
+            Ok(())
         }
     }
 

crates/kumod/src/http_server/inject_v1.rs

@@ -695,6 +695,8 @@ async fn inject_v1_impl(
             }
         }
     }
+    config.put();
+
     Ok(Json(InjectV1Response {
         success_count,
         fail_count,

crates/kumod/src/logging/disposition.rs

@@ -78,7 +78,10 @@ pub async fn log_disposition(args: LogDisposition<'_>) {
 
                     let enqueue: bool =
                         match lua_config.async_call_callback(&log_sig, msg.clone()).await {
-                            Ok(b) => b,
+                            Ok(b) => {
+                                lua_config.put();
+                                b
+                            }
                             Err(err) => {
                                 tracing::error!(
                                     "error while calling {name} event for log filter: {err:#}"

crates/kumod/src/logging/hooks.rs

@@ -170,6 +170,7 @@ impl LogHookState {
             let enqueue: bool = lua_config
                 .async_call_callback(&SHOULD_ENQ_LOG_RECORD_SIG, (msg.clone(), name))
                 .await?;
+            lua_config.put();
 
             if enqueue {
                 let queue_name = msg.get_queue_name()?;

crates/kumod/src/main.rs

@@ -293,6 +293,8 @@ async fn perform_init(opts: Opt) -> anyhow::Result<()> {
             .context("start_spool")?;
     }
 
+    config.put();
+
     Ok(())
 }
 

crates/kumod/src/mod_kumo.rs

@@ -159,6 +159,7 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> {
                 crate::queue::Queue::call_get_queue_config(&queue_name, &mut config)
                     .await
                     .map_err(any_err)?;
+            config.put();
             lua.to_value(&queue_config)
         })?,
     )?;

crates/kumod/src/queue.rs

@@ -1090,6 +1090,8 @@ impl Queue {
         let queue_config = Self::call_get_queue_config(&name, &mut config).await?;
 
         let pool = EgressPool::resolve(queue_config.egress_pool.as_deref(), &mut config).await?;
+        config.put();
+
         let rr = ArcSwap::new(EgressPoolRoundRobin::new(&pool).into());
 
         let activity = Activity::get(format!("Queue {name}"))?;
@@ -1916,6 +1918,8 @@ impl Queue {
         config
             .async_call_callback(&THROTTLE_INSERT_READY_SIG, msg.clone())
             .await?;
+        config.put();
+
         if let Some(due) = msg.get_due() {
             let now = Utc::now();
             if due > now {
@@ -2849,6 +2853,7 @@ impl QueueManager {
 
                 match result {
                     Ok(()) => {
+                        config.put();
                         let queue_name_after = msg.get_queue_name()?;
                         if queue_name != queue_name_after {
                             // We want to avoid the normal due-time adjustment

crates/kumod/src/ready_queue.rs

@@ -322,6 +322,7 @@ impl ReadyQueueManager {
                 tracing::error!("Error while calling get_egress_path_config: {err:#}");
                 err
             })?;
+        config.put();
 
         Ok(ReadyQueueConfig {
             name,

crates/kumod/src/smtp_dispatcher.rs

@@ -943,6 +943,9 @@ impl QueueDispatcher for SmtpDispatcher {
                         ),
                     )
                     .await;
+                if rewritten_code.is_ok() {
+                    config.put();
+                }
 
                 match rewritten_code {
                     Ok(Some(code)) if code != response.code => {

crates/kumod/src/smtp_server.rs

@@ -635,6 +635,7 @@ impl SmtpServer {
 
         let value = match value {
             Ok(v) => {
+                config.put();
                 SmtpServerTraceManager::submit(|| SmtpServerTraceEvent {
                     conn_meta: self.meta.clone_inner(),
                     payload: SmtpServerTraceEventPayload::Callback {
@@ -1011,6 +1012,7 @@ impl SmtpServer {
         let name = sig.name();
         match config.async_call_callback(sig, args).await {
             Ok(r) => {
+                config.put();
                 SmtpServerTraceManager::submit(|| SmtpServerTraceEvent {
                     conn_meta: self.meta.clone_inner(),
                     payload: SmtpServerTraceEventPayload::Callback {
@@ -2065,12 +2067,15 @@ impl QueueDispatcher for DeferredSmtpInjectionDispatcher {
             .async_call_callback(&DEFERRED_SMTP_SERVER_MSG_INJECT, (msg.clone(), meta))
             .await
         {
-            Ok(_) => Response {
-                code: 250,
-                enhanced_code: None,
-                content: "ok".to_string(),
-                command: None,
-            },
+            Ok(_) => {
+                config.put();
+                Response {
+                    code: 250,
+                    enhanced_code: None,
+                    content: "ok".to_string(),
+                    command: None,
+                }
+            }
             Err(err) => {
                 if let Some(rej) = RejectError::from_anyhow(&err) {
                     Response {

crates/tsa-daemon/src/main.rs

@@ -74,6 +74,7 @@ async fn perform_init() -> anyhow::Result<()> {
         .async_call_callback(&tsa_init_sig, ())
         .await
         .context("in tsa_init event")?;
+    config.put();
 
     spawn_shaping_updater()?;
 

crates/tsa-daemon/src/shaping_config.rs

@@ -15,6 +15,7 @@ pub async fn load_shaping() -> anyhow::Result<Arc<Shaping>> {
         .async_call_callback_non_default(&sig, ())
         .await
         .context("in tsa_load_shaping_data event")?;
+    config.put();
     Ok(Arc::new(shaping))
 }