diff options
| -rw-r--r-- | src/libserver/worker_util.c | 2 | ||||
| -rw-r--r-- | src/lua/lua_common.c | 76 | ||||
| -rw-r--r-- | 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 */ |