From b2e34f1884ed8c6c0346aaa007a1da70bad0a561 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Wed, 1 Jun 2016 21:41:20 +0200 Subject: [PATCH] [Feature] Add configuration knobs for in and out parser scripts --- src/rspamd_proxy.c | 165 +++++++++++++++++++++++++++++++-------------- 1 file changed, 113 insertions(+), 52 deletions(-) diff --git a/src/rspamd_proxy.c b/src/rspamd_proxy.c index 2fb8d0b50..24c0fd4bf 100644 --- a/src/rspamd_proxy.c +++ b/src/rspamd_proxy.c @@ -65,6 +65,8 @@ struct rspamd_http_upstream { gchar *name; struct upstream_list *u; struct rspamd_cryptobox_pubkey *key; + gint parser_from_ref; + gint parser_to_ref; }; struct rspamd_http_mirror { @@ -73,7 +75,8 @@ struct rspamd_http_mirror { struct upstream_list *u; struct rspamd_cryptobox_pubkey *key; gdouble prob; - gint parser_ref; + gint parser_from_ref; + gint parser_to_ref; }; static const guint64 rspamd_rspamd_proxy_magic = 0xcdeb4fd1fc351980ULL; @@ -125,6 +128,8 @@ struct rspamd_proxy_backend_connection { struct rspamd_proxy_session *s; gint backend_sock; enum rspamd_backend_flags flags; + gint parser_from_ref; + gint parser_to_ref; }; struct rspamd_proxy_session { @@ -233,6 +238,101 @@ err: return FALSE; } +static gboolean +rspamd_proxy_parse_lua_parser (lua_State *L, const ucl_object_t *obj, + gint *ref_from, gint *ref_to, GError **err) +{ + const gchar *lua_script; + gsize slen; + gint err_idx, ref_idx; + GString *tb = NULL; + gboolean has_ref = FALSE; + + g_assert (obj != NULL); + g_assert (ref_from != NULL); + g_assert (ref_to != NULL); + + *ref_from = -1; + *ref_to = -1; + + lua_script = ucl_object_tolstring (obj, &slen); + lua_pushcfunction (L, &rspamd_lua_traceback); + err_idx = lua_gettop (L); + + /* Load data */ + 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 */ + + return FALSE; + } + + /* 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); + + return FALSE; + } + + if (lua_istable (L, -1)) { + /* + * We have a table, so we check for two keys: + * 'from' -> function + * 'to' -> function + * + * From converts parent request to a client one + * To converts client request to a parent one + */ + lua_pushstring (L, "from"); + lua_gettable (L, -2); + + if (lua_isfunction (L, -1)) { + ref_idx = luaL_ref (L, LUA_REGISTRYINDEX); + *ref_from = ref_idx; + has_ref = TRUE; + } + + lua_pushstring (L, "to"); + lua_gettable (L, -2); + + if (lua_isfunction (L, -1)) { + ref_idx = luaL_ref (L, LUA_REGISTRYINDEX); + *ref_to = ref_idx; + has_ref = TRUE; + } + } + else 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); + + return FALSE; + } + else { + /* Just parser from protocol */ + ref_idx = luaL_ref (L, LUA_REGISTRYINDEX); + *ref_from = ref_idx; + lua_settop (L, 0); + has_ref = TRUE; + } + + return has_ref; +} + static gboolean rspamd_proxy_parse_mirror (rspamd_mempool_t *pool, const ucl_object_t *obj, @@ -266,7 +366,8 @@ rspamd_proxy_parse_mirror (rspamd_mempool_t *pool, up = g_slice_alloc0 (sizeof (*up)); up->name = g_strdup (ucl_object_tostring (elt)); - up->parser_ref = -1; + up->parser_to_ref = -1; + up->parser_from_ref = -1; elt = ucl_object_lookup (obj, "key"); if (elt != NULL) { @@ -312,55 +413,10 @@ rspamd_proxy_parse_mirror (rspamd_mempool_t *pool, */ elt = ucl_object_lookup (obj, "parser"); if (elt) { - const gchar *lua_script; - gsize slen; - gint err_idx, ref_idx; - GString *tb = NULL; - - lua_script = ucl_object_tolstring (elt, &slen); - lua_pushcfunction (L, &rspamd_lua_traceback); - err_idx = lua_gettop (L); - - /* Load data */ - 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 */ - + if (!rspamd_proxy_parse_lua_parser (L, elt, &up->parser_from_ref, + &up->parser_to_ref, err)) { 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); - up->parser_ref = ref_idx; - lua_settop (L, 0); } elt = ucl_object_lookup_any (obj, "settings", "settings_id", NULL); @@ -382,8 +438,11 @@ err: rspamd_pubkey_unref (up->key); } - if (up->parser_ref != -1) { - luaL_unref (L, LUA_REGISTRYINDEX, up->parser_ref); + if (up->parser_from_ref != -1) { + luaL_unref (L, LUA_REGISTRYINDEX, up->parser_from_ref); + } + if (up->parser_to_ref != -1) { + luaL_unref (L, LUA_REGISTRYINDEX, up->parser_to_ref); } g_slice_free1 (sizeof (*up), up); @@ -846,7 +905,7 @@ proxy_backend_mirror_finish_handler (struct rspamd_http_connection *conn, session = bk_conn->s; if (!proxy_backend_parse_results (session, bk_conn, session->ctx->lua_state, - -1, msg->body_buf.begin, msg->body_buf.len)) { + bk_conn->parser_from_ref, msg->body_buf.begin, msg->body_buf.len)) { msg_warn_session ("cannot parse results from the mirror backend %s:%s", bk_conn->name, rspamd_inet_address_to_string (rspamd_upstream_addr (bk_conn->up))); @@ -888,6 +947,8 @@ proxy_open_mirror_connections (struct rspamd_proxy_session *session) bk_conn->up = rspamd_upstream_get (m->u, RSPAMD_UPSTREAM_ROUND_ROBIN, NULL, 0); + bk_conn->parser_from_ref = m->parser_from_ref; + bk_conn->parser_to_ref = m->parser_to_ref; if (bk_conn->up == NULL) { msg_err_session ("cannot select upstream for %s", m->name); -- 2.39.5