From 3c4fbc745b1f34ad505749f5110722d7a13f0cee Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Fri, 16 Nov 2018 14:27:02 +0000
Subject: [PATCH] [Project] Move stat tokens generation to Lua

---
 lualib/lua_stat.lua         | 37 ++++++++++++++++++++
 src/libstat/stat_config.c   | 67 +++++++++++++++++++++++++++++++++++++
 src/libstat/stat_internal.h |  3 ++
 3 files changed, 107 insertions(+)

diff --git a/lualib/lua_stat.lua b/lualib/lua_stat.lua
index 9b72a1f22..0ced4b428 100644
--- a/lualib/lua_stat.lua
+++ b/lualib/lua_stat.lua
@@ -23,6 +23,7 @@ local logger = require "rspamd_logger"
 local sqlite3 = require "rspamd_sqlite3"
 local util = require "rspamd_util"
 local lua_redis = require "lua_redis"
+local lua_util = require "lua_util"
 local exports = {}
 
 local N = "stat_tools" -- luacheck: ignore (maybe unused)
@@ -516,4 +517,40 @@ end
 
 exports.redis_classifier_from_sqlite = redis_classifier_from_sqlite
 
+-- Reads statistics config and return preprocessed table
+local function process_stat_config(cfg)
+  local opts_section = cfg:get_all_opt('options')
+
+  -- Check if we have a dedicated section for statistics
+  if opts_section.statistics then
+    opts_section = opts_section.statistics
+  end
+
+  -- Default
+  local res_config = {
+    classify_headers = {
+      "User-Agent",
+      "X-Mailer",
+      "Content-Type",
+      "X-MimeOLE",
+    },
+    classify_images = true,
+    classify_mime_info = true,
+    classify_headers = true,
+  }
+
+  res_config = lua_util.override_defaults(res_config, opts_section)
+
+  return res_config
+end
+
+
+exports.gen_stat_tokens = function(cfg)
+  local stat_config = process_stat_config(cfg)
+
+  return function(task)
+    return get_stat_tokens(task, stat_config)
+  end
+end
+
 return exports
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;
-- 
2.39.5