From 81eb3c648eda46f1f3fa2e6110f16f7022232b28 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Mon, 13 Mar 2017 13:46:40 +0000 Subject: [PATCH] [Feature] Allow to cache and use flexible protocol reply --- src/libserver/protocol.c | 157 ++++++++++++++++++++++++++------------- src/libserver/protocol.h | 19 ++++- src/lua/lua_util.c | 5 +- 3 files changed, 127 insertions(+), 54 deletions(-) diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c index 3a03c3b50..b64f06d97 100644 --- a/src/libserver/protocol.c +++ b/src/libserver/protocol.c @@ -1064,80 +1064,133 @@ rspamd_protocol_output_profiling (struct rspamd_task *task, ucl_object_insert_key (top, prof, "profile", 0, false); } +struct rspamd_saved_protocol_reply { + ucl_object_t *obj; + enum rspamd_protocol_flags flags; +}; + +static void +rspamd_protocol_cached_dtor (gpointer p) +{ + struct rspamd_saved_protocol_reply *cached = p; + + ucl_object_unref (cached->obj); +} + ucl_object_t * -rspamd_protocol_write_ucl (struct rspamd_task *task) +rspamd_protocol_write_ucl (struct rspamd_task *task, + enum rspamd_protocol_flags flags) { struct rspamd_metric_result *metric_res; ucl_object_t *top = NULL, *obj; GHashTableIter hiter; GString *dkim_sig; const ucl_object_t *rmilter_reply; + struct rspamd_saved_protocol_reply *cached; + static const gchar *varname = "cached_reply"; gpointer h, v; - g_hash_table_iter_init (&hiter, task->results); - top = ucl_object_typed_new (UCL_OBJECT); - /* Convert results to an ucl object */ - while (g_hash_table_iter_next (&hiter, &h, &v)) { - metric_res = (struct rspamd_metric_result *)v; - obj = rspamd_metric_result_ucl (task, metric_res); - ucl_object_insert_key (top, obj, h, 0, false); + /* Check for cached reply */ + cached = rspamd_mempool_get_variable (task->task_pool, varname); + + if (cached) { + top = cached->obj; + /* + * Update flags: we don't need to set more flags than we need, + * so just xor previous flags and current flags and then or them + * to the cached one + */ + flags ^= cached->flags; + cached->flags |= flags; + + } + else { + top = ucl_object_typed_new (UCL_OBJECT); + cached = rspamd_mempool_alloc (task->task_pool, sizeof (*cached)); + cached->obj = top; + cached->flags = flags; + rspamd_mempool_set_variable (task->task_pool, varname, + cached, rspamd_protocol_cached_dtor); + } + + if (flags & RSPAMD_PROTOCOL_METRICS) { + g_hash_table_iter_init (&hiter, task->results); + /* Convert results to an ucl object */ + while (g_hash_table_iter_next (&hiter, &h, &v)) { + metric_res = (struct rspamd_metric_result *)v; + obj = rspamd_metric_result_ucl (task, metric_res); + ucl_object_insert_key (top, obj, h, 0, false); + } } - if (G_UNLIKELY (task->cfg->compat_messages)) { - const ucl_object_t *cur; - ucl_object_t *msg_object; - ucl_object_iter_t iter = NULL; + if (flags & RSPAMD_PROTOCOL_MESSAGES) { + if (G_UNLIKELY (task->cfg->compat_messages)) { + const ucl_object_t *cur; + ucl_object_t *msg_object; + ucl_object_iter_t iter = NULL; - msg_object = ucl_object_typed_new (UCL_ARRAY); + msg_object = ucl_object_typed_new (UCL_ARRAY); - while ((cur = ucl_object_iterate (task->messages, &iter, true)) != NULL) { - if (cur->type == UCL_STRING) { - ucl_array_append (msg_object, ucl_object_ref (cur)); + while ((cur = ucl_object_iterate (task->messages, &iter, true)) != NULL) { + if (cur->type == UCL_STRING) { + ucl_array_append (msg_object, ucl_object_ref (cur)); + } } - } - ucl_object_insert_key (top, msg_object, "messages", 0, false); - } - else { - ucl_object_insert_key (top, ucl_object_ref (task->messages), - "messages", 0, false); + ucl_object_insert_key (top, msg_object, "messages", 0, false); + } + else { + ucl_object_insert_key (top, ucl_object_ref (task->messages), + "messages", 0, false); + } } - if (task->cfg->log_urls || (task->flags & RSPAMD_TASK_FLAG_EXT_URLS)) { - if (g_hash_table_size (task->urls) > 0) { - ucl_object_insert_key (top, rspamd_urls_tree_ucl (task->urls, - task), "urls", 0, false); - } - if (g_hash_table_size (task->emails) > 0) { - ucl_object_insert_key (top, rspamd_emails_tree_ucl (task->emails, task), - "emails", 0, false); + if (flags & RSPAMD_PROTOCOL_URLS) { + if (task->cfg->log_urls || (task->flags & RSPAMD_TASK_FLAG_EXT_URLS)) { + if (g_hash_table_size (task->urls) > 0) { + ucl_object_insert_key (top, rspamd_urls_tree_ucl (task->urls, + task), "urls", 0, false); + } + if (g_hash_table_size (task->emails) > 0) { + ucl_object_insert_key (top, rspamd_emails_tree_ucl (task->emails, task), + "emails", 0, false); + } } } - if (G_UNLIKELY (RSPAMD_TASK_IS_PROFILING (task))) { - rspamd_protocol_output_profiling (task, top); + if (flags & RSPAMD_PROTOCOL_EXTRA) { + if (G_UNLIKELY (RSPAMD_TASK_IS_PROFILING (task))) { + rspamd_protocol_output_profiling (task, top); + } } - ucl_object_insert_key (top, ucl_object_fromstring (task->message_id), - "message-id", 0, false); + if (flags & RSPAMD_PROTOCOL_BASIC) { + ucl_object_insert_key (top, ucl_object_fromstring (task->message_id), + "message-id", 0, false); + } - dkim_sig = rspamd_mempool_get_variable (task->task_pool, "dkim-signature"); + if (flags & RSPAMD_PROTOCOL_DKIM) { + dkim_sig = rspamd_mempool_get_variable (task->task_pool, "dkim-signature"); - if (dkim_sig) { - GString *folded_header = rspamd_header_value_fold ("DKIM-Signature", - dkim_sig->str, 80, task->nlines_type); - ucl_object_insert_key (top, - ucl_object_fromstring_common (folded_header->str, - folded_header->len, UCL_STRING_RAW), - "dkim-signature", 0, false); - g_string_free (folded_header, TRUE); + if (dkim_sig) { + GString *folded_header = rspamd_header_value_fold ("DKIM-Signature", + dkim_sig->str, 80, task->nlines_type); + ucl_object_insert_key (top, + ucl_object_fromstring_common (folded_header->str, + folded_header->len, UCL_STRING_RAW), + "dkim-signature", 0, false); + g_string_free (folded_header, TRUE); + } } - rmilter_reply = rspamd_mempool_get_variable (task->task_pool, "rmilter-reply"); + if (flags & RSPAMD_PROTOCOL_RMILTER) { + rmilter_reply = rspamd_mempool_get_variable (task->task_pool, + "rmilter-reply"); - if (rmilter_reply) { - ucl_object_insert_key (top, ucl_object_ref (rmilter_reply), - "rmilter", 0, false); + if (rmilter_reply) { + ucl_object_insert_key (top, ucl_object_ref (rmilter_reply), + "rmilter", 0, false); + } } return top; @@ -1153,7 +1206,7 @@ rspamd_protocol_http_reply (struct rspamd_http_message *msg, gpointer h, v; ucl_object_t *top = NULL; rspamd_fstring_t *reply; - gint action; + gint action, flags = RSPAMD_PROTOCOL_DEFAULT; /* Write custom headers */ g_hash_table_iter_init (&hiter, task->reply_headers); @@ -1163,7 +1216,11 @@ rspamd_protocol_http_reply (struct rspamd_http_message *msg, rspamd_http_message_add_header (msg, hn->begin, hv->begin); } - top = rspamd_protocol_write_ucl (task); + if (task->cfg->log_urls || (task->flags & RSPAMD_TASK_FLAG_EXT_URLS)) { + flags |= RSPAMD_PROTOCOL_URLS; + } + + top = rspamd_protocol_write_ucl (task, flags); if (!(task->flags & RSPAMD_TASK_FLAG_NO_LOG)) { rspamd_roll_history_update (task->worker->srv->history, task); @@ -1200,8 +1257,6 @@ rspamd_protocol_http_reply (struct rspamd_http_message *msg, } } - ucl_object_unref (top); - if ((task->flags & RSPAMD_TASK_FLAG_COMPRESSED) && rspamd_libs_reset_compression (task->cfg->libs_ctx)) { /* We can compress output */ diff --git a/src/libserver/protocol.h b/src/libserver/protocol.h index 164a53d2f..c266c2f99 100644 --- a/src/libserver/protocol.h +++ b/src/libserver/protocol.h @@ -68,13 +68,30 @@ gboolean rspamd_protocol_handle_request (struct rspamd_task *task, void rspamd_protocol_http_reply (struct rspamd_http_message *msg, struct rspamd_task *task); +enum rspamd_protocol_flags { + RSPAMD_PROTOCOL_BASIC = 1 << 0, + RSPAMD_PROTOCOL_METRICS = 1 << 1, + RSPAMD_PROTOCOL_MESSAGES = 1 << 2, + RSPAMD_PROTOCOL_RMILTER = 1 << 3, + RSPAMD_PROTOCOL_DKIM = 1 << 4, + RSPAMD_PROTOCOL_URLS = 1 << 5, + RSPAMD_PROTOCOL_EXTRA = 1 << 6, +}; + +#define RSPAMD_PROTOCOL_DEFAULT (RSPAMD_PROTOCOL_BASIC| \ + RSPAMD_PROTOCOL_METRICS| \ + RSPAMD_PROTOCOL_MESSAGES| \ + RSPAMD_PROTOCOL_RMILTER| \ + RSPAMD_PROTOCOL_DKIM| \ + RSPAMD_PROTOCOL_EXTRA) /** * Write reply to ucl object filling log buffer * @param task * @param logbuf * @return */ -ucl_object_t * rspamd_protocol_write_ucl (struct rspamd_task *task); +ucl_object_t * rspamd_protocol_write_ucl (struct rspamd_task *task, + enum rspamd_protocol_flags flags); /** * Write reply for specified task command diff --git a/src/lua/lua_util.c b/src/lua/lua_util.c index 0b4d8dd9c..064aa53f8 100644 --- a/src/lua/lua_util.c +++ b/src/lua/lua_util.c @@ -624,7 +624,7 @@ lua_util_task_fin (struct rspamd_task *task, void *ud) { ucl_object_t **target = ud; - *target = rspamd_protocol_write_ucl (task); + *target = rspamd_protocol_write_ucl (task, RSPAMD_PROTOCOL_DEFAULT); rdns_resolver_release (task->resolver->r); return TRUE; @@ -669,7 +669,8 @@ lua_util_process_message (lua_State *L) ucl_object_unref (res); } else { - ucl_object_push_lua (L, rspamd_protocol_write_ucl (task), + ucl_object_push_lua (L, + rspamd_protocol_write_ucl (task, RSPAMD_PROTOCOL_DEFAULT), true); rdns_resolver_release (task->resolver->r); rspamd_session_destroy (task->s); -- 2.39.5