[Minor] Add function to get a lua_reference from a string returning a function

4 files changed, 71 insertions, 47 deletions

diff --git a/src/libserver/cfg_rcl.c b/src/libserver/cfg_rcl.c
index d91ebf3ae..717b16bea 100644
--- a/src/libserver/cfg_rcl.c
+++ b/src/libserver/cfg_rcl.c
@@ -1220,7 +1220,7 @@ rspamd_rcl_classifier_handler (rspamd_mempool_t *pool,
 	gboolean res = TRUE;
 	struct rspamd_rcl_section *stat_section;
 	struct rspamd_tokenizer_config *tkcf = NULL;
-	lua_State *L;
+	lua_State *L = cfg->lua_state;
 
 	g_assert (key != NULL);
 	ccf = rspamd_config_new_classifier (cfg, NULL);
@@ -1303,59 +1303,24 @@ rspamd_rcl_classifier_handler (rspamd_mempool_t *pool,
 
 	if (val) {
 		LL_FOREACH (val, cur) {
-			if (ucl_object_type (cur) == UCL_STRING) {
+			if (ucl_object_type(cur) == UCL_STRING) {
 				const gchar *lua_script;
 				gsize slen;
-				gint err_idx, ref_idx;
-
-				lua_script = ucl_object_tolstring (cur, &slen);
-				L = cfg->lua_state;
-				lua_pushcfunction (L, &rspamd_lua_traceback);
-				err_idx = lua_gettop (L);
-
+				gint ref_idx;
 
-				/* Load file */
-				if (luaL_loadbuffer (L, lua_script, slen, "learn_condition") != 0) {
-					g_set_error (err,
-							CFG_RCL_ERROR,
-							EINVAL,
-							"cannot load lua condition script: %s",
-							lua_tostring (L, -1));
-					lua_settop (L, 0); /* Error function */
+				lua_script = ucl_object_tolstring(cur, &slen);
+				ref_idx = rspamd_lua_function_ref_from_str(L,
+						lua_script, slen, err);
 
+				if (ref_idx == LUA_NOREF) {
 					return FALSE;
 				}
 
-				/* Now do it */
-				if (lua_pcall (L, 0, 1, err_idx) != 0) {
-					g_set_error (err,
-							CFG_RCL_ERROR,
-							EINVAL,
-							"cannot init lua condition script: %s",
-							lua_tostring (L, -1));
-					lua_settop (L, 0);
-
-					return FALSE;
-				}
-
-				if (!lua_isfunction (L, -1)) {
-					g_set_error (err,
-							CFG_RCL_ERROR,
-							EINVAL,
-							"cannot init lua condition script: "
-							"must return function");
-					lua_settop (L, 0);
-
-					return FALSE;
-				}
-
-				ref_idx = luaL_ref (L, LUA_REGISTRYINDEX);
-				rspamd_lua_add_ref_dtor (L, cfg->cfg_pool, ref_idx);
-				ccf->learn_conditions = rspamd_mempool_glist_append (
+				rspamd_lua_add_ref_dtor(L, cfg->cfg_pool, ref_idx);
+				ccf->learn_conditions = rspamd_mempool_glist_append(
 						cfg->cfg_pool,
 						ccf->learn_conditions,
 						GINT_TO_POINTER (ref_idx));
-				lua_settop (L, 0);
 			}
 		}
 	}
diff --git a/src/libstat/stat_process.c b/src/libstat/stat_process.c
index d8b0a3349..8ac4e499e 100644
--- a/src/libstat/stat_process.c
+++ b/src/libstat/stat_process.c
@@ -551,7 +551,7 @@ rspamd_stat_classifiers_learn (struct rspamd_stat_ctx *st_ctx,
 		while (cur) {
 			cb_ref = GPOINTER_TO_INT (cur->data);
 
-			lua_settop (L, 0);
+			gint old_top = lua_gettop (L);
 			lua_rawgeti (L, LUA_REGISTRYINDEX, cb_ref);
 			/* Push task and two booleans: is_spam and is_unlearn */
 			ptask = lua_newuserdata (L, sizeof (*ptask));
@@ -576,13 +576,13 @@ rspamd_stat_classifiers_learn (struct rspamd_stat_ctx *st_ctx,
 									lua_tostring (L, 2));
 						}
 
-						lua_settop (L, 0);
+						lua_settop (L, old_top);
 						break;
 					}
 				}
 			}
 
-			lua_settop (L, 0);
+			lua_settop (L, old_top);
 			cur = g_list_next (cur);
 		}
 
diff --git a/src/lua/lua_common.c b/src/lua/lua_common.c
index 06720c9f2..5d874d507 100644
--- a/src/lua/lua_common.c
+++ b/src/lua/lua_common.c
@@ -2292,6 +2292,56 @@ rspamd_lua_require_function (lua_State *L, const gchar *modname,
 	return FALSE;
 }
 
+gint
+rspamd_lua_function_ref_from_str (lua_State *L, const gchar *str, gsize slen,
+								  GError **err)
+{
+	gint err_idx, ref_idx;
+
+	lua_pushcfunction (L, &rspamd_lua_traceback);
+	err_idx = lua_gettop (L);
+
+	/* Load file */
+	if (luaL_loadbuffer (L, str, slen, "lua_embedded_str") != 0) {
+		g_set_error (err,
+				lua_error_quark(),
+				EINVAL,
+				"cannot load lua script: %s",
+				lua_tostring (L, -1));
+		lua_settop (L, err_idx - 1); /* Error function */
+
+		return LUA_NOREF;
+	}
+
+	/* Now call it */
+	if (lua_pcall (L, 0, 1, err_idx) != 0) {
+		g_set_error (err,
+				lua_error_quark(),
+				EINVAL,
+				"cannot init lua script: %s",
+				lua_tostring (L, -1));
+		lua_settop (L, err_idx - 1);
+
+		return LUA_NOREF;
+	}
+
+	if (!lua_isfunction (L, -1)) {
+		g_set_error (err,
+				lua_error_quark(),
+				EINVAL,
+				"cannot init lua script: "
+				"must return function");
+		lua_settop (L, err_idx - 1);
+
+		return LUA_NOREF;
+	}
+
+	ref_idx = luaL_ref (L, LUA_REGISTRYINDEX);
+	lua_settop (L, err_idx - 1);
+
+	return ref_idx;
+}
+
 
 gboolean
 rspamd_lua_try_load_redis (lua_State *L, const ucl_object_t *obj,
diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h
index d9e4fcaec..b929ab864 100644
--- a/src/lua/lua_common.h
+++ b/src/lua/lua_common.h
@@ -566,6 +566,15 @@ void rspamd_lua_add_ref_dtor (lua_State *L, rspamd_mempool_t *pool,
 							  gint ref);
 
 /**
+ * Returns a lua reference from a function like string, e.g. `return function(...) end`
+ * @param L
+ * @param str
+ * @return
+ */
+gint rspamd_lua_function_ref_from_str (lua_State *L, const gchar *str, gsize slen,
+									   GError **err);
+
+/**
 * Tries to load some module using `require` and get some method from it
 * @param L
 * @param modname