[Project] Lua_config: Add post init scripts

diff --git a/src/libserver/cfg_file.h b/src/libserver/cfg_file.h
index ed5e1325b9073244b24164357cd134d060fa2b3b..8af34a63eedf022788f1e369686d2ef918e61fa5 100644 (file)
--- a/src/libserver/cfg_file.h
+++ b/src/libserver/cfg_file.h
@@ -296,9 +296,14 @@ enum rspamd_action_flags {
 
 struct rspamd_action;
 
-struct rspamd_config_post_load_script {
+struct rspamd_config_cfg_lua_script {
        gint cbref;
-       struct rspamd_config_post_load_script *prev, *next;
+       struct rspamd_config_cfg_lua_script *prev, *next;
+};
+
+struct rspamd_config_post_init_script {
+       gint cbref;
+       struct rspamd_config_post_init_script *prev, *next;
 };
 
 struct rspamd_lang_detector;
@@ -447,9 +452,7 @@ struct rspamd_config {
 
        GList *classify_headers;                                                /**< list of headers using for statistics                               */
        struct module_s **compiled_modules;                             /**< list of compiled C modules                                                 */
-       struct worker_s **compiled_workers;                             /**< list of compiled C modules                                                 */
-       struct rspamd_config_post_load_script *finish_callbacks; /**< list of callbacks called on worker's termination  */
-       struct rspamd_log_format *log_format;                   /**< parsed log format                                                                  */
+       struct worker_s **compiled_workers;                             /**< list of compiled C modules                                                 */struct rspamd_log_format *log_format;                 /**< parsed log format                                                                  */
        gchar *log_format_str;                                                  /**< raw log format string                                                              */
 
        struct rspamd_external_libs_ctx *libs_ctx;              /**< context for external libraries                                             */
@@ -460,7 +463,9 @@ struct rspamd_config {
 
        GHashTable *trusted_keys;                                               /**< list of trusted public keys                                                */
 
-       struct rspamd_config_post_load_script *on_load; /**< list of scripts executed on config load                    */
+       struct rspamd_config_cfg_lua_script *on_load_scripts;   /**< list of scripts executed on workers load                   */
+       struct rspamd_config_cfg_lua_script *post_init_scripts; /**< list of scripts executed on workers load                   */
+       struct rspamd_config_cfg_lua_script *on_term_scripts; /**< list of callbacks called on worker's termination     */
 
        gchar *ssl_ca_path;                                                             /**< path to CA certs                                                                   */
        gchar *ssl_ciphers;                                                             /**< set of preferred ciphers                                                   */
@@ -529,13 +534,15 @@ enum rspamd_post_load_options {
        RSPAMD_CONFIG_INIT_VALIDATE = 1 << 3,
        RSPAMD_CONFIG_INIT_NO_TLD = 1 << 4,
        RSPAMD_CONFIG_INIT_PRELOAD_MAPS = 1 << 5,
+       RSPAMD_CONFIG_INIT_POST_LOAD_LUA = 1 << 6,
 };
 
 #define RSPAMD_CONFIG_LOAD_ALL (RSPAMD_CONFIG_INIT_URL| \
                RSPAMD_CONFIG_INIT_LIBS| \
                RSPAMD_CONFIG_INIT_SYMCACHE| \
                RSPAMD_CONFIG_INIT_VALIDATE| \
-               RSPAMD_CONFIG_INIT_PRELOAD_MAPS)
+               RSPAMD_CONFIG_INIT_PRELOAD_MAPS| \
+               RSPAMD_CONFIG_INIT_POST_LOAD_LUA)
 
 /**
  * Do post load actions for config

diff --git a/src/libserver/cfg_utils.c b/src/libserver/cfg_utils.c
index 0f38cc3892cd4b5447ead0372672d33ed7ba9049..3761d08619d0473bfe2df858fbf7bc63baf4c34d 100644 (file)
--- a/src/libserver/cfg_utils.c
+++ b/src/libserver/cfg_utils.c
@@ -245,18 +245,21 @@ rspamd_config_new (enum rspamd_config_init_flags flags)
 void
 rspamd_config_free (struct rspamd_config *cfg)
 {
-       struct rspamd_config_post_load_script *sc, *sctmp;
+       struct rspamd_config_cfg_lua_script *sc, *sctmp;
        struct rspamd_config_settings_elt *set, *stmp;
        struct rspamd_worker_log_pipe *lp, *ltmp;
 
-       DL_FOREACH_SAFE (cfg->finish_callbacks, sc, sctmp) {
+       /* Scripts part */
+       DL_FOREACH_SAFE (cfg->on_term_scripts, sc, sctmp) {
                luaL_unref (cfg->lua_state, LUA_REGISTRYINDEX, sc->cbref);
-               g_free (sc);
        }
 
-       DL_FOREACH_SAFE (cfg->on_load, sc, sctmp) {
+       DL_FOREACH_SAFE (cfg->on_load_scripts, sc, sctmp) {
+               luaL_unref (cfg->lua_state, LUA_REGISTRYINDEX, sc->cbref);
+       }
+
+       DL_FOREACH_SAFE (cfg->post_init_scripts, sc, sctmp) {
                luaL_unref (cfg->lua_state, LUA_REGISTRYINDEX, sc->cbref);
-               g_free (sc);
        }
 
        DL_FOREACH_SAFE (cfg->setting_ids, set, stmp) {
@@ -870,6 +873,10 @@ rspamd_config_post_load (struct rspamd_config *cfg,
                rspamd_map_preload (cfg);
        }
 
+       if (opts & RSPAMD_CONFIG_INIT_POST_LOAD_LUA) {
+               rspamd_lua_run_config_post_init (cfg->lua_state, cfg);
+       }
+
        return ret;
 }
 

diff --git a/src/lua/lua_common.c b/src/lua/lua_common.c
index f7793748ae0dc4ef2d550331482f6a40c59d8fde..6ee69c07201407b8efda50ba41c9ead98b0d156e 100644 (file)
--- a/src/lua/lua_common.c
+++ b/src/lua/lua_common.c
@@ -1842,13 +1842,13 @@ void
 rspamd_lua_run_postloads (lua_State *L, struct rspamd_config *cfg,
                struct event_base *ev_base, struct rspamd_worker *w)
 {
-       struct rspamd_config_post_load_script *sc;
+       struct rspamd_config_cfg_lua_script *sc;
        struct rspamd_config **pcfg;
        struct event_base **pev_base;
        struct rspamd_worker **pw;
 
        /* Execute post load scripts */
-       LL_FOREACH (cfg->on_load, sc) {
+       LL_FOREACH (cfg->on_load_scripts, sc) {
                struct thread_entry *thread = lua_thread_pool_get_for_config (cfg);
                thread->error_callback = rspamd_lua_run_postloads_error;
                L = thread->lua_state;
@@ -1871,6 +1871,30 @@ rspamd_lua_run_postloads (lua_State *L, struct rspamd_config *cfg,
 }
 
 
+void
+rspamd_lua_run_config_post_init (lua_State *L, struct rspamd_config *cfg)
+{
+       struct rspamd_config_cfg_lua_script *sc;
+       struct rspamd_config **pcfg;
+
+       LL_FOREACH (cfg->post_init_scripts, sc) {
+               lua_pushcfunction (L, &rspamd_lua_traceback);
+               gint err_idx = lua_gettop (L);
+
+               lua_rawgeti (L, LUA_REGISTRYINDEX, sc->cbref);
+               pcfg = lua_newuserdata (L, sizeof (*pcfg));
+               *pcfg = cfg;
+               rspamd_lua_setclass (L, "rspamd{config}", -1);
+
+               if (lua_pcall (L, 1, 0, err_idx) != 0) {
+                       msg_err_config ("cannot run config post init script: %s",
+                                       lua_tostring (L, -1));
+               }
+
+               lua_settop (L, err_idx - 1);
+       }
+}
+
 static void
 rspamd_lua_run_postloads_error (struct thread_entry *thread, int ret, const char *msg)
 {

diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h
index 776baf3b9a60d182f29dee1cccdec482544e4cc8..258c41491b7e45a48e31c6af6139bcd64c21ca16 100644 (file)
--- a/src/lua/lua_common.h
+++ b/src/lua/lua_common.h
@@ -408,7 +408,7 @@ void *rspamd_lua_check_udata_maybe (lua_State *L, gint pos, const gchar *classna
  * @param sc
  * @param task
  */
-void lua_call_finish_script (struct rspamd_config_post_load_script *sc,
+void lua_call_finish_script (struct rspamd_config_cfg_lua_script *sc,
                struct rspamd_task *task);
 
 /**
@@ -420,6 +420,8 @@ void lua_call_finish_script (struct rspamd_config_post_load_script *sc,
 void rspamd_lua_run_postloads (lua_State *L, struct rspamd_config *cfg,
                struct event_base *ev_base, struct rspamd_worker *w);
 
+void rspamd_lua_run_config_post_init (lua_State *L, struct rspamd_config *cfg);
+
 /**
  * Adds new destructor for a local function for specific pool
  * @param L

diff --git a/src/lua/lua_config.c b/src/lua/lua_config.c
index cff27ca1e0eb4a49b5e3fb72c43de1286108cc25..832e10614af7223cdb1f1f1a801369b87b05ff8f 100644 (file)
--- a/src/lua/lua_config.c
+++ b/src/lua/lua_config.c
@@ -553,6 +553,7 @@ LUA_FUNCTION_DEF (config, register_worker_script);
 /***
  * @method rspamd_config:add_on_load(function(cfg, ev_base, worker) ... end)
  * Registers the following script to be executed when configuration is completely loaded
+ * and the worker is already started (forked)
  * @param {function} script function to be executed
  * @example
 rspamd_config:add_on_load(function(cfg, ev_base, worker)
@@ -583,6 +584,14 @@ end)
  */
 LUA_FUNCTION_DEF (config, add_periodic);
 
+/***
+ * @method rspamd_config:add_post_init(function(cfg) ... end)
+ * Registers the following script to be executed when configuration is completely loaded
+ * @available 2.0+
+ * @param {function} script function to be executed
+ */
+LUA_FUNCTION_DEF (config, add_post_init);
+
 /***
  * @method rspamd_config:get_symbols_count()
  * Returns number of symbols registered in rspamd configuration
@@ -833,6 +842,7 @@ static const struct luaL_reg configlib_m[] = {
        LUA_INTERFACE_DEF (config, register_re_selector),
        LUA_INTERFACE_DEF (config, add_on_load),
        LUA_INTERFACE_DEF (config, add_periodic),
+       LUA_INTERFACE_DEF (config, add_post_init),
        LUA_INTERFACE_DEF (config, get_symbols_count),
        LUA_INTERFACE_DEF (config, get_symbols_cksum),
        LUA_INTERFACE_DEF (config, get_symbols_counters),
@@ -2894,16 +2904,35 @@ lua_config_add_on_load (lua_State *L)
 {
        LUA_TRACE_POINT;
        struct rspamd_config *cfg = lua_check_config (L, 1);
-       struct rspamd_config_post_load_script *sc;
+       struct rspamd_config_cfg_lua_script *sc;
+
+       if (cfg == NULL || lua_type (L, 2) != LUA_TFUNCTION) {
+               return luaL_error (L, "invalid arguments");
+       }
+
+       sc = rspamd_mempool_alloc0 (cfg->cfg_pool, sizeof (*sc));
+       lua_pushvalue (L, 2);
+       sc->cbref = luaL_ref (L, LUA_REGISTRYINDEX);
+       DL_APPEND (cfg->on_load_scripts, sc);
+
+       return 0;
+}
+
+static gint
+lua_config_add_post_init (lua_State *L)
+{
+       LUA_TRACE_POINT;
+       struct rspamd_config *cfg = lua_check_config (L, 1);
+       struct rspamd_config_cfg_lua_script *sc;
 
        if (cfg == NULL || lua_type (L, 2) != LUA_TFUNCTION) {
                return luaL_error (L, "invalid arguments");
        }
 
-       sc = g_malloc0 (sizeof (*sc));
+       sc = rspamd_mempool_alloc0 (cfg->cfg_pool, sizeof (*sc));
        lua_pushvalue (L, 2);
        sc->cbref = luaL_ref (L, LUA_REGISTRYINDEX);
-       DL_APPEND (cfg->on_load, sc);
+       DL_APPEND (cfg->post_init_scripts, sc);
 
        return 0;
 }
@@ -3315,13 +3344,13 @@ lua_config_register_finish_script (lua_State *L)
 {
        LUA_TRACE_POINT;
        struct rspamd_config *cfg = lua_check_config (L, 1);
-       struct rspamd_config_post_load_script *sc;
+       struct rspamd_config_cfg_lua_script *sc;
 
        if (cfg != NULL && lua_type (L, 2) == LUA_TFUNCTION) {
-               sc = g_malloc0 (sizeof (*sc));
+               sc = rspamd_mempool_alloc0 (cfg->cfg_pool, sizeof (*sc));
                lua_pushvalue (L, 2);
                sc->cbref = luaL_ref (L, LUA_REGISTRYINDEX);
-               DL_APPEND (cfg->finish_callbacks, sc);
+               DL_APPEND (cfg->on_term_scripts, sc);
        }
        else {
                return luaL_error (L, "invalid arguments");
@@ -4033,7 +4062,7 @@ luaopen_config (lua_State * L)
 }
 
 void
-lua_call_finish_script (struct rspamd_config_post_load_script *sc,
+lua_call_finish_script (struct rspamd_config_cfg_lua_script *sc,
                struct rspamd_task *task) {
 
        struct rspamd_task **ptask;

diff --git a/src/worker.c b/src/worker.c
index fc206cc3212a15c6edd590fb7d68483d6c1971f9..5a01c6bd3055dc7ad9ec35a6e89a53d0253ffb80 100644 (file)
--- a/src/worker.c
+++ b/src/worker.c
@@ -92,9 +92,9 @@ rspamd_worker_call_finish_handlers (struct rspamd_worker *worker)
        struct rspamd_task *task;
        struct rspamd_config *cfg = worker->srv->cfg;
        struct rspamd_abstract_worker_ctx *ctx;
-       struct rspamd_config_post_load_script *sc;
+       struct rspamd_config_cfg_lua_script *sc;
 
-       if (cfg->finish_callbacks) {
+       if (cfg->on_term_scripts) {
                ctx = worker->ctx;
                /* Create a fake task object for async events */
                task = rspamd_task_new (worker, cfg, NULL, NULL, ctx->ev_base);
@@ -106,7 +106,7 @@ rspamd_worker_call_finish_handlers (struct rspamd_worker *worker)
                                (event_finalizer_t) rspamd_task_free,
                                task);
 
-               DL_FOREACH (cfg->finish_callbacks, sc) {
+               DL_FOREACH (cfg->on_term_scripts, sc) {
                        lua_call_finish_script (sc, task);
                }