From 786e053ba54874e0768ffdbcce11a4292b887ff4 Mon Sep 17 00:00:00 2001
From: Wez Furlong <wez@wezfurlong.org>
Date: Wed, 16 Aug 2023 16:21:51 -0700
Subject: [PATCH] catch kumo.on called inside kumo.on

Event handlers must be registered at the file scope and not from
within another event handler in order for events to be consistently
triggered and handled.

In particular, the `init` event is only ever triggered once on
server startup.  If other handlers are registered from within the
init event, those will only ever fire when we re-use that original
lua context, which is good for a limited number of uses before it
is aged out of the resource pool.  That might work under very
limited or lightweight testing scenarios, but will otherwise
result in very inconsistent behavior.

We now catch and prevent, with a very visible error, attempting
to call `kumo.on` from within an event handler.
---
 crates/config/src/lib.rs             | 14 ++++++++++++++
 crates/kumo-server-common/src/lib.rs | 10 ++++++++++
 2 files changed, 24 insertions(+)

diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs
index 9826f7f7..cc46f84f 100644
--- a/crates/config/src/lib.rs
+++ b/crates/config/src/lib.rs
@@ -209,6 +209,15 @@ pub fn register(func: RegisterFunc) {
 }
 
 impl LuaConfig {
+    fn set_current_event(&mut self, name: &str) -> mlua::Result<()> {
+        self.inner
+            .as_mut()
+            .unwrap()
+            .lua
+            .globals()
+            .set("_KUMO_CURRENT_EVENT", name.to_string())
+    }
+
     /// Call a callback registered via `on`.
     pub async fn async_call_callback<
         'lua,
@@ -222,6 +231,7 @@ impl LuaConfig {
     ) -> anyhow::Result<R> {
         let name = name.as_ref();
         let decorated_name = format!("kumomta-on-{}", name);
+        self.set_current_event(name)?;
         match self
             .inner
             .as_mut()
@@ -246,6 +256,7 @@ impl LuaConfig {
     ) -> anyhow::Result<R> {
         let name = name.as_ref();
         let decorated_name = format!("kumomta-on-{}", name);
+        self.set_current_event(name)?;
         match self
             .inner
             .as_mut()
@@ -270,6 +281,7 @@ impl LuaConfig {
     ) -> anyhow::Result<Option<R>> {
         let name = name.as_ref();
         let decorated_name = format!("kumomta-on-{}", name);
+        self.set_current_event(name)?;
         let lua = self.inner.as_mut().unwrap();
         let opt_func: mlua::Value = lua.lua.named_registry_value(&decorated_name)?;
 
@@ -308,6 +320,7 @@ impl LuaConfig {
     ) -> anyhow::Result<RegistryKey> {
         let name = name.as_ref();
         let decorated_name = format!("kumomta-on-{}", name);
+        self.set_current_event(name)?;
 
         let inner = self.inner.as_mut().unwrap();
 
@@ -353,6 +366,7 @@ impl LuaConfig {
     ) -> anyhow::Result<R> {
         let name = name.as_ref();
         let decorated_name = format!("kumomta-on-{}", name);
+        self.set_current_event(name);
         match self
             .inner
             .as_mut()
diff --git a/crates/kumo-server-common/src/lib.rs b/crates/kumo-server-common/src/lib.rs
index a3d9b922..c7d613ad 100644
--- a/crates/kumo-server-common/src/lib.rs
+++ b/crates/kumo-server-common/src/lib.rs
@@ -36,6 +36,16 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> {
         lua.create_function(move |lua, (name, func): (String, Function)| {
             let decorated_name = format!("kumomta-on-{}", name);
 
+            if let Ok(current_event) = lua.globals().get::<_, String>("_KUMO_CURRENT_EVENT") {
+                return Err(mlua::Error::external(format!(
+                    "Attempting to register an event handler via \
+                    `kumo.on('{name}', ...)` from within the event handler \
+                    '{current_event}'. You must move your event handler registration \
+                    so that it is setup directly when the policy is loaded \
+                    in order for it to consistently trigger and handle events."
+                )));
+            }
+
             let existing: Value = lua.named_registry_value(&decorated_name)?;
             match existing {
                 Value::Nil => {}