From 4950b1e115df10cd76c8684e084a4cb2550f20e1 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Fri, 6 May 2016 17:13:13 +0100 Subject: [PATCH] [Feature] Implement compare scripts setup --- src/rspamd_proxy.c | 122 ++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 115 insertions(+), 7 deletions(-) diff --git a/src/rspamd_proxy.c b/src/rspamd_proxy.c index 4bf8d2e18..5112e7328 100644 --- a/src/rspamd_proxy.c +++ b/src/rspamd_proxy.c @@ -101,6 +101,8 @@ struct rspamd_proxy_ctx { struct event rotate_ev; gdouble rotate_tm; lua_State *lua_state; + /* Array of callback functions called on end of scan to compare results */ + GArray *cmp_refs; }; enum rspamd_backend_flags { @@ -139,7 +141,7 @@ struct rspamd_proxy_session { static GQuark rspamd_proxy_quark (void) { - return g_quark_from_static_string ("http-proxy"); + return g_quark_from_static_string ("rspamd-proxy"); } static gboolean @@ -247,7 +249,7 @@ rspamd_proxy_parse_mirror (rspamd_mempool_t *pool, if (ucl_object_type (obj) != UCL_OBJECT) { g_set_error (err, rspamd_proxy_quark (), 100, - "upstream option must be an object"); + "mirror option must be an object"); return FALSE; } @@ -255,7 +257,7 @@ rspamd_proxy_parse_mirror (rspamd_mempool_t *pool, elt = ucl_object_lookup (obj, "name"); if (elt == NULL) { g_set_error (err, rspamd_proxy_quark (), 100, - "upstream option must have some name definition"); + "mirror option must have some name definition"); return FALSE; } @@ -271,7 +273,7 @@ rspamd_proxy_parse_mirror (rspamd_mempool_t *pool, if (up->key == NULL) { g_set_error (err, rspamd_proxy_quark (), 100, - "cannot read upstream key"); + "cannot read mirror key"); goto err; } @@ -281,7 +283,7 @@ rspamd_proxy_parse_mirror (rspamd_mempool_t *pool, if (elt == NULL) { g_set_error (err, rspamd_proxy_quark (), 100, - "upstream option must have some hosts definition"); + "mirror option must have some hosts definition"); goto err; } @@ -289,7 +291,7 @@ rspamd_proxy_parse_mirror (rspamd_mempool_t *pool, up->u = rspamd_upstreams_create (ctx->cfg->ups_ctx); if (!rspamd_upstreams_from_ucl (up->u, elt, 11333, NULL)) { g_set_error (err, rspamd_proxy_quark (), 100, - "upstream has bad hosts definition"); + "mirror has bad hosts definition"); goto err; } @@ -351,7 +353,7 @@ rspamd_proxy_parse_mirror (rspamd_mempool_t *pool, "must return function"); lua_settop (L, 0); - return FALSE; + goto err; } ref_idx = luaL_ref (L, LUA_REGISTRYINDEX); @@ -383,6 +385,99 @@ err: return FALSE; } +static gboolean +rspamd_proxy_parse_script (rspamd_mempool_t *pool, + const ucl_object_t *obj, + gpointer ud, + struct rspamd_rcl_section *section, + GError **err) +{ + const ucl_object_t *elt; + struct rspamd_proxy_ctx *ctx; + struct rspamd_rcl_struct_parser *pd = ud; + lua_State *L; + const gchar *lua_script; + gsize slen; + gint err_idx, ref_idx; + GString *tb = NULL; + struct stat st; + + ctx = pd->user_struct; + L = ctx->lua_state; + + if (ucl_object_type (obj) != UCL_STRING) { + g_set_error (err, rspamd_proxy_quark (), 100, + "script option must be a string with file or lua chunk"); + + return FALSE; + } + + lua_script = ucl_object_tolstring (elt, &slen); + lua_pushcfunction (L, &rspamd_lua_traceback); + err_idx = lua_gettop (L); + + if (stat (lua_script, &st) != -1) { + /* Load file */ + if (luaL_loadfile (L, lua_script) != 0) { + g_set_error (err, + rspamd_proxy_quark (), + EINVAL, + "cannot load lua parser script: %s", + lua_tostring (L, -1)); + lua_settop (L, 0); /* Error function */ + + goto err; + } + } + else { + /* Load data directly */ + if (luaL_loadbuffer (L, lua_script, slen, "proxy parser") != 0) { + g_set_error (err, + rspamd_proxy_quark (), + EINVAL, + "cannot load lua parser script: %s", + lua_tostring (L, -1)); + lua_settop (L, 0); /* Error function */ + + goto err; + } + } + + /* Now do it */ + if (lua_pcall (L, 0, 1, err_idx) != 0) { + tb = lua_touserdata (L, -1); + g_set_error (err, + rspamd_proxy_quark (), + EINVAL, + "cannot init lua parser script: %s", + tb->str); + g_string_free (tb, TRUE); + lua_settop (L, 0); + + goto err; + } + + if (!lua_isfunction (L, -1)) { + g_set_error (err, + rspamd_proxy_quark (), + EINVAL, + "cannot init lua parser script: " + "must return function"); + lua_settop (L, 0); + + goto err; + } + + ref_idx = luaL_ref (L, LUA_REGISTRYINDEX); + lua_settop (L, 0); + g_array_append_val (ctx->cmp_refs, ref_idx); + + return TRUE; + +err: + return FALSE; +} + gpointer init_rspamd_proxy (struct rspamd_config *cfg) { @@ -444,6 +539,15 @@ init_rspamd_proxy (struct rspamd_config *cfg) RSPAMD_CL_FLAG_MULTIPLE, "List of mirrors"); + rspamd_rcl_register_worker_option (cfg, + type, + "script", + rspamd_proxy_parse_script, + ctx, + 0, + RSPAMD_CL_FLAG_MULTIPLE, + "Compare script to be executed"); + return ctx; } @@ -652,6 +756,7 @@ proxy_backend_mirror_error_handler (struct rspamd_http_connection *conn, GError err->message); proxy_backend_close_connection (bk_conn); + REF_RELEASE (bk_conn->s); } static gint @@ -670,6 +775,7 @@ proxy_backend_mirror_finish_handler (struct rspamd_http_connection *conn, } proxy_backend_close_connection (bk_conn); + REF_RELEASE (bk_conn->s); return 0; } @@ -739,6 +845,7 @@ proxy_open_mirror_connections (struct rspamd_proxy_session *session) &session->ctx->io_tv, session->ctx->ev_base); g_ptr_array_add (session->mirror_conns, bk_conn); + REF_RETAIN (session); msg_info_session ("send request to %s", m->name); } } @@ -982,6 +1089,7 @@ start_rspamd_proxy (struct rspamd_worker *worker) ctx->ev_base, worker->srv->cfg); ctx->lua_state = worker->srv->cfg->lua_state; + ctx->cmp_refs = g_array_new (FALSE, FALSE, sizeof (gint)); double_to_tv (ctx->timeout, &ctx->io_tv); rspamd_map_watch (worker->srv->cfg, ctx->ev_base, ctx->resolver); -- 2.39.5