[Project] Allow to set custom Lua handlers for main commands

src/lua/lua_worker.c

@@ -25,6 +25,7 @@
 #endif
 
 #include <sys/wait.h>
+#include <src/libserver/rspamd_control.h>
 
 /***
  * @module rspamd_worker
@@ -42,6 +43,7 @@ LUA_FUNCTION_DEF (worker, is_scanner);
 LUA_FUNCTION_DEF (worker, is_primary_controller);
 LUA_FUNCTION_DEF (worker, spawn_process);
 LUA_FUNCTION_DEF (worker, get_mem_stats);
+LUA_FUNCTION_DEF (worker, add_control_handler);
 
 const luaL_reg worker_reg[] = {
 		LUA_INTERFACE_DEF (worker, get_name),
@@ -53,6 +55,7 @@ const luaL_reg worker_reg[] = {
 		LUA_INTERFACE_DEF (worker, is_scanner),
 		LUA_INTERFACE_DEF (worker, is_primary_controller),
 		LUA_INTERFACE_DEF (worker, get_mem_stats),
+		LUA_INTERFACE_DEF (worker, add_control_handler),
 		{"__tostring", rspamd_lua_class_tostring},
 		{NULL, NULL}
 };
@@ -240,6 +243,145 @@ lua_worker_is_primary_controller (lua_State *L)
 	return 1;
 }
 
+struct rspamd_control_cbdata {
+	lua_State *L;
+	rspamd_mempool_t *pool;
+	struct rspamd_worker *w;
+	struct rspamd_config *cfg;
+	struct ev_loop *event_loop;
+	struct rspamd_async_session *session;
+	enum rspamd_control_type cmd;
+	gint cbref;
+	gint fd;
+};
+
+static gboolean
+lua_worker_control_fin_session (void *ud)
+{
+	struct rspamd_control_reply rep;
+	struct rspamd_control_cbdata *cbd = (struct rspamd_control_cbdata *)ud;
+	rspamd_mempool_t *pool;
+	lua_State *L;
+
+	L = cbd->L;
+	pool = cbd->pool;
+
+	memset (&rep, 0, sizeof (rep));
+	rep.type = cbd->cmd;
+
+	if (write (cbd->fd, &rep, sizeof (rep)) != sizeof (rep)) {
+		msg_err_pool ("cannot write reply to the control socket: %s",
+				strerror (errno));
+	}
+
+	return TRUE;
+}
+
+static void
+lua_worker_control_session_dtor (void *ud)
+{
+	struct rspamd_control_cbdata *cbd = (struct rspamd_control_cbdata *)ud;
+
+	rspamd_mempool_delete (cbd->pool);
+}
+
+static gboolean
+lua_worker_control_handler (struct rspamd_main *rspamd_main,
+							struct rspamd_worker *worker,
+							gint fd,
+							gint attached_fd,
+							struct rspamd_control_command *cmd,
+							gpointer ud)
+{
+	struct rspamd_async_session *session, **psession;
+	struct rspamd_control_cbdata *cbd = (struct rspamd_control_cbdata *)ud;
+	rspamd_mempool_t *pool;
+	lua_State *L;
+	gint err_idx;
+
+	L = cbd->L;
+	pool = cbd->pool;
+	session = rspamd_session_create (cbd->pool,
+			lua_worker_control_fin_session,
+			NULL,
+			lua_worker_control_session_dtor,
+			cbd);
+	cbd->session = session;
+	cbd->fd = fd;
+
+	lua_pushcfunction (L, &rspamd_lua_traceback);
+	err_idx = lua_gettop (L);
+	lua_rawgeti (L, LUA_REGISTRYINDEX, cbd->cbref);
+	psession = lua_newuserdata (L, sizeof (*psession));
+	rspamd_lua_setclass (L, "rspamd{session}", -1);
+	*psession = session;
+
+	if (lua_pcall (L, 1, 0, err_idx) != 0) {
+		msg_err_pool ("cannot init lua parser script: %s", lua_tostring (L, -1));
+		lua_settop (L, err_idx - 1);
+
+		struct rspamd_control_reply rep;
+
+		memset (&rep, 0, sizeof (rep));
+		rep.type = cbd->cmd;
+		rep.reply.monitored_change.status = -1;
+
+		if (write (fd, &rep, sizeof (rep)) != sizeof (rep)) {
+			msg_err_pool ("cannot write reply to the control socket: %s",
+					strerror (errno));
+		}
+
+		rspamd_session_destroy (session);
+	}
+	else {
+		lua_settop (L, err_idx - 1);
+		rspamd_session_pending (session);
+	}
+
+	return TRUE;
+}
+
+static gint
+lua_worker_add_control_handler (lua_State *L)
+{
+	struct rspamd_worker *w = lua_check_worker (L, 1);
+	struct rspamd_config *cfg = lua_check_config (L, 2);
+	struct ev_loop *event_loop = lua_check_ev_base (L, 3);
+	const gchar *cmd_name = luaL_checkstring (L, 4);
+	enum rspamd_control_type cmd;
+	struct rspamd_control_cbdata *cbd;
+
+	if (w && cfg && event_loop && cmd_name && lua_isfunction (L, 5)) {
+		cmd = rspamd_control_command_from_string (cmd_name);
+
+		if (cmd == RSPAMD_CONTROL_MAX) {
+			return luaL_error (L, "invalid command type: %s", cmd_name);
+		}
+
+		rspamd_mempool_t *pool = rspamd_mempool_new (
+				rspamd_mempool_suggest_size (), "lua_control");
+		cbd = rspamd_mempool_alloc0 (pool, sizeof (*cbd));
+		cbd->pool = pool;
+		cbd->event_loop = event_loop;
+		cbd->w = w;
+		cbd->cfg = cfg;
+		cbd->cmd = cmd;
+		cbd->L = L;
+		/* Refcount callback */
+		lua_pushvalue (L, 5);
+		cbd->cbref = luaL_ref (L, LUA_REGISTRYINDEX);
+
+		rspamd_control_worker_add_cmd_handler (w, cmd, lua_worker_control_handler,
+				cbd);
+	}
+	else {
+		return luaL_error (L, "invalid arguments, need worker, cfg, "
+						"ev_loop, cmd_name and callback function");
+	}
+
+	return 0;
+}
+
 #ifdef WITH_JEMALLOC
 static void
 lua_worker_jemalloc_stats_cb (void *ud, const char *msg)