[Project] Lua_config: Add post init scripts

6 files changed, 94 insertions, 25 deletions

diff --git a/src/libserver/cfg_file.h b/src/libserver/cfg_file.h
index ed5e1325b..8af34a63e 100644
--- 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 0f38cc389..3761d0861 100644
--- 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 f7793748a..6ee69c072 100644
--- 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 776baf3b9..258c41491 100644
--- 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 cff27ca1e..832e10614 100644
--- 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)
@@ -584,6 +585,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
  * @return {number} number of symbols registered in the 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 fc206cc32..5a01c6bd3 100644
--- 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);
 		}