[Project] Move stat tokens generation to Lua

2 files changed, 70 insertions, 0 deletions

diff --git a/src/libstat/stat_config.c b/src/libstat/stat_config.c
index d2772e9ca..b5ddf143c 100644
--- a/src/libstat/stat_config.c
+++ b/src/libstat/stat_config.c
@@ -164,6 +164,67 @@ rspamd_stat_init (struct rspamd_config *cfg, struct event_base *ev_base)
 	stat_ctx->classifiers = g_ptr_array_new ();
 	stat_ctx->async_elts = g_queue_new ();
 	stat_ctx->ev_base = ev_base;
+	stat_ctx->lua_stat_tokens_ref = -1;
+
+	/* Interact with lua_stat */
+	if (luaL_dostring (L, "return require \"lua_stat\"") != 0) {
+		msg_err_config ("cannot require lua_stat: %s",
+				lua_tostring (L, -1));
+	}
+	else {
+		if (lua_type (L, -1) != LUA_TTABLE) {
+			msg_err_config ("lua stat must return "
+							"table and not %s",
+					lua_typename (L, lua_type (L, -1)));
+		}
+		else {
+			lua_pushstring (L, "gen_stat_tokens");
+			lua_gettable (L, -2);
+
+			if (lua_type (L, -1) != LUA_TFUNCTION) {
+				msg_err_config ("gen_stat_tokens must return "
+								"function and not %s",
+						lua_typename (L, lua_type (L, -1)));
+			}
+			else {
+				/* Call this function to obtain closure */
+				gint err_idx, ret;
+				GString *tb;
+				struct rspamd_config **pcfg;
+
+				lua_pushcfunction (L, &rspamd_lua_traceback);
+				err_idx = lua_gettop (L);
+				lua_pushvalue (L, err_idx - 1);
+
+				pcfg = lua_newuserdata (L, sizeof (*pcfg));
+				*pcfg = cfg;
+				rspamd_lua_setclass (L, "rspamd{config}", -1);
+
+				if ((ret = lua_pcall (L, 1, 1, err_idx)) != 0) {
+					tb = lua_touserdata (L, -1);
+					msg_err_config ("call to cleanup_rules lua "
+									"script failed (%d): %v", ret, tb);
+
+					if (tb) {
+						g_string_free (tb, TRUE);
+					}
+				}
+				else {
+					if (lua_type (L, -1) != LUA_TFUNCTION) {
+						msg_err_config ("gen_stat_tokens invocation must return "
+										"function and not %s",
+								lua_typename (L, lua_type (L, -1)));
+					}
+					else {
+						stat_ctx->lua_stat_tokens_ref = luaL_ref (L, LUA_REGISTRYINDEX);
+					}
+				}
+			}
+		}
+	}
+
+	/* Cleanup mess */
+	lua_settop (L, 0);
 
 	/* Create statfiles from the classifiers */
 	cur = cfg->classifiers;
@@ -349,6 +410,12 @@ rspamd_stat_close (void)
 	g_queue_free (stat_ctx->async_elts);
 	g_ptr_array_free (st_ctx->statfiles, TRUE);
 	g_ptr_array_free (st_ctx->classifiers, TRUE);
+
+	if (st_ctx->lua_stat_tokens_ref != -1) {
+		luaL_unref (st_ctx->cfg->lua_state, LUA_REGISTRYINDEX,
+				st_ctx->lua_stat_tokens_ref);
+	}
+
 	g_free (st_ctx);
 
 	/* Set global var to NULL */
diff --git a/src/libstat/stat_internal.h b/src/libstat/stat_internal.h
index 746199d45..a547ca93a 100644
--- a/src/libstat/stat_internal.h
+++ b/src/libstat/stat_internal.h
@@ -86,6 +86,9 @@ struct rspamd_stat_ctx {
 	GPtrArray *classifiers; /* struct rspamd_classifier */
 	GQueue *async_elts; /* struct rspamd_stat_async_elt */
 	struct rspamd_config *cfg;
+
+	gint lua_stat_tokens_ref;
+
 	/* Global tokenizer */
 	struct rspamd_stat_tokenizer *tokenizer;
 	gpointer tkcf;