From 7890a62a8a76a7711e5dbdb994f8dbbe5c0c1872 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Fri, 21 Oct 2016 17:36:12 +0100
Subject: [PATCH] [Feature] Add a very basic interface to access workers data
 from on_load

---
 src/libserver/worker_util.c |  2 +-
 src/lua/lua_common.c        | 76 ++++++++++++++++++++++++++++++++++++-
 src/lua/lua_common.h        |  2 +-
 3 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/src/libserver/worker_util.c b/src/libserver/worker_util.c
index 16e08eb31..6261e1991 100644
--- a/src/libserver/worker_util.c
+++ b/src/libserver/worker_util.c
@@ -306,7 +306,7 @@ rspamd_prepare_worker (struct rspamd_worker *worker, const char *name,
 	if (load_lua) {
 		struct rspamd_config *cfg = worker->srv->cfg;
 
-		rspamd_lua_run_postloads (cfg->lua_state, cfg, ev_base);
+		rspamd_lua_run_postloads (cfg->lua_state, cfg, ev_base, worker);
 	}
 
 	return ev_base;
diff --git a/src/lua/lua_common.c b/src/lua/lua_common.c
index 3f36306e2..f899f68ed 100644
--- a/src/lua/lua_common.c
+++ b/src/lua/lua_common.c
@@ -26,6 +26,19 @@ const luaL_reg null_reg[] = {
 	{NULL, NULL}
 };
 
+
+LUA_FUNCTION_DEF (worker, get_name);
+LUA_FUNCTION_DEF (worker, get_index);
+LUA_FUNCTION_DEF (worker, get_pid);
+
+const luaL_reg worker_reg[] = {
+	LUA_INTERFACE_DEF (worker, get_name),
+	LUA_INTERFACE_DEF (worker, get_index),
+	LUA_INTERFACE_DEF (worker, get_pid),
+	{"__tostring", rspamd_lua_class_tostring},
+	{NULL, NULL}
+};
+
 static GQuark
 lua_error_quark (void)
 {
@@ -277,6 +290,7 @@ rspamd_lua_init ()
 	lua_rawset (L, -3);
 	lua_pop (L, 1);
 
+	rspamd_lua_new_class (L, "rspamd{worker}", worker_reg);
 	rspamd_lua_add_preload (L, "ucl", luaopen_ucl);
 
 	if (luaL_dostring (L, rspamadm_script_global_functions) != 0) {
@@ -1146,11 +1160,12 @@ lua_check_ev_base (lua_State * L, gint pos)
 
 gboolean
 rspamd_lua_run_postloads (lua_State *L, struct rspamd_config *cfg,
-		struct event_base *ev_base)
+		struct event_base *ev_base, struct rspamd_worker *w)
 {
 	struct rspamd_config_post_load_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) {
@@ -1163,7 +1178,11 @@ rspamd_lua_run_postloads (lua_State *L, struct rspamd_config *cfg,
 		*pev_base = ev_base;
 		rspamd_lua_setclass (cfg->lua_state, "rspamd{ev_base}", -1);
 
-		if (lua_pcall (cfg->lua_state, 2, 0, 0) != 0) {
+		pw = lua_newuserdata (cfg->lua_state, sizeof (*pw));
+		*pw = w;
+		rspamd_lua_setclass (cfg->lua_state, "rspamd{worker}", -1);
+
+		if (lua_pcall (cfg->lua_state, 3, 0, 0) != 0) {
 			msg_err_config ("error executing post load code: %s",
 					lua_tostring (cfg->lua_state, -1));
 			lua_pop (cfg->lua_state, 1);
@@ -1174,3 +1193,56 @@ rspamd_lua_run_postloads (lua_State *L, struct rspamd_config *cfg,
 
 	return TRUE;
 }
+
+static struct rspamd_worker *
+lua_check_worker (lua_State *L, gint pos)
+{
+	void *ud = rspamd_lua_check_udata (L, pos, "rspamd{worker}");
+	luaL_argcheck (L, ud != NULL, pos, "'worker' expected");
+	return ud ? *((struct rspamd_worker **)ud) : NULL;
+}
+
+static gint
+lua_worker_get_name (lua_State *L)
+{
+	struct rspamd_worker *w = lua_check_worker (L, 1);
+
+	if (w) {
+		lua_pushstring (L, g_quark_to_string (w->type));
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 1;
+}
+
+static gint
+lua_worker_get_index (lua_State *L)
+{
+	struct rspamd_worker *w = lua_check_worker (L, 1);
+
+	if (w) {
+		lua_pushnumber (L, w->index);
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 1;
+}
+
+static gint
+lua_worker_get_pid (lua_State *L)
+{
+	struct rspamd_worker *w = lua_check_worker (L, 1);
+
+	if (w) {
+		lua_pushnumber (L, w->pid);
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 1;
+}
diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h
index ae0376328..a2b886098 100644
--- a/src/lua/lua_common.h
+++ b/src/lua/lua_common.h
@@ -371,7 +371,7 @@ void lua_call_finish_script (lua_State *L, struct
  * @param ev_base
  */
 gboolean rspamd_lua_run_postloads (lua_State *L, struct rspamd_config *cfg,
-		struct event_base *ev_base);
+		struct event_base *ev_base, struct rspamd_worker *w);
 
 #endif /* WITH_LUA */
 #endif /* RSPAMD_LUA_H */
-- 
2.39.5