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

diff --git a/src/lua/lua_worker.c b/src/lua/lua_worker.c
index c3c9e729e4fa2180331140b865c57806c882090f..1768b1c2a3f4b97dd1a3490d4e42cd07a7e5244c 100644 (file)
--- a/src/lua/lua_worker.c
+++ b/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)