[Project] Move stat tokens generation to Lua

diff --git a/lualib/lua_stat.lua b/lualib/lua_stat.lua
index 9b72a1f22bf1207bd3696d35b982814a52854127..0ced4b428c880079fbe28f943bcedf3ce087755b 100644 (file)
--- 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 d2772e9cae4c40f98b4c0debaa99138ab5fc3913..b5ddf143caf8d6ab8b4445e9b38c58a8ed100741 100644 (file)
--- 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 746199d45f829ac395fe9fde95cf9d40faff7e01..a547ca93a29d7b7dbd0ce1429988f7ec08080806 100644 (file)
--- 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;