From: Vsevolod Stakhov Date: Sat, 24 Jun 2017 11:32:16 +0000 (+0100) Subject: [Feature] Allow to customize spam header, remove existing spam headers X-Git-Tag: 1.6.2~80 X-Git-Url: https://source.dussan.org/?a=commitdiff_plain;h=32330869f5d61f085e97c056418626f4f6a934fa;p=rspamd.git [Feature] Allow to customize spam header, remove existing spam headers --- diff --git a/src/libserver/milter.c b/src/libserver/milter.c index 82996d01d..65b0e6a54 100644 --- a/src/libserver/milter.c +++ b/src/libserver/milter.c @@ -1344,15 +1344,59 @@ rspamd_milter_update_userdata (struct rspamd_milter_session *session, } static void +rspamd_milter_remove_header_safe (struct rspamd_milter_session *session, + const gchar *key, gint nhdr) +{ + gint saved_nhdr, i; + gpointer found; + GString *hname, *hvalue; + struct rspamd_milter_private *priv = session->priv; + + found = g_hash_table_lookup (priv->headers, key); + + if (found) { + saved_nhdr = GPOINTER_TO_INT (found); + + hname = g_string_new (key); + hvalue = g_string_new (""); + + if (nhdr >= 1) { + rspamd_milter_send_action (session, + RSPAMD_MILTER_CHGHEADER, + nhdr, hname, hvalue); + } + else if (nhdr == 0) { + /* We need to clear all headers */ + for (i = 1; i <= saved_nhdr; i ++) { + rspamd_milter_send_action (session, + RSPAMD_MILTER_CHGHEADER, + i, hname, hvalue); + } + } + else { + /* Remove from the end */ + if (nhdr >= -saved_nhdr) { + rspamd_milter_send_action (session, + RSPAMD_MILTER_CHGHEADER, + saved_nhdr + nhdr + 1, hname, hvalue); + } + } + + g_string_free (hname, TRUE); + g_string_free (hvalue, TRUE); + } +} + +/* + * Returns `TRUE` if action has been processed internally by this function + */ +static gboolean rspamd_milter_process_milter_block (struct rspamd_milter_session *session, - const ucl_object_t *obj) + const ucl_object_t *obj, gint action) { const ucl_object_t *elt, *cur, *cur_elt; ucl_object_iter_t it; - struct rspamd_milter_private *priv = session->priv; GString *hname, *hvalue; - gint nhdr, saved_nhdr, i; - gpointer found; if (obj && ucl_object_type (obj) == UCL_OBJECT) { elt = ucl_object_lookup (obj, "remove_headers"); @@ -1365,42 +1409,9 @@ rspamd_milter_process_milter_block (struct rspamd_milter_session *session, while ((cur = ucl_object_iterate (elt, &it, true)) != NULL) { if (ucl_object_type (cur) == UCL_INT) { - nhdr = ucl_object_toint (cur); - - found = g_hash_table_lookup (priv->headers, - ucl_object_key (cur)); - - if (found) { - saved_nhdr = GPOINTER_TO_INT (found); - - hname = g_string_new (ucl_object_key (cur)); - hvalue = g_string_new (""); - - if (nhdr >= 1) { - rspamd_milter_send_action (session, - RSPAMD_MILTER_CHGHEADER, - nhdr, hname, hvalue); - } - else if (nhdr == 0) { - /* We need to clear all headers */ - for (i = 1; i <= saved_nhdr; i ++) { - rspamd_milter_send_action (session, - RSPAMD_MILTER_CHGHEADER, - i, hname, hvalue); - } - } - else { - /* Remove from the end */ - if (nhdr >= -saved_nhdr) { - rspamd_milter_send_action (session, - RSPAMD_MILTER_CHGHEADER, - saved_nhdr + nhdr + 1, hname, hvalue); - } - } - - g_string_free (hname, TRUE); - g_string_free (hvalue, TRUE); - } + rspamd_milter_remove_header_safe (session, + ucl_object_key (cur), + ucl_object_toint (cur)); } } } @@ -1429,6 +1440,50 @@ rspamd_milter_process_milter_block (struct rspamd_milter_session *session, } } } + + if (action == METRIC_ACTION_ADD_HEADER) { + elt = ucl_object_lookup (obj, "spam_header"); + + if (elt) { + if (ucl_object_type (elt) == UCL_STRING) { + rspamd_milter_remove_header_safe (session, + milter_ctx->spam_header, + 0); + + hname = g_string_new (milter_ctx->spam_header); + hvalue = g_string_new (ucl_object_tostring (elt)); + rspamd_milter_send_action (session, RSPAMD_MILTER_CHGHEADER, + (guint32)1, hname, hvalue); + g_string_free (hname, TRUE); + g_string_free (hvalue, TRUE); + rspamd_milter_send_action (session, RSPAMD_MILTER_ACCEPT); + + return TRUE; + } + else if (ucl_object_type (elt) == UCL_OBJECT) { + it = NULL; + + while ((cur = ucl_object_iterate (elt, &it, true)) != NULL) { + rspamd_milter_remove_header_safe (session, + ucl_object_key (cur), + 0); + + hname = g_string_new (ucl_object_key (cur)); + hvalue = g_string_new (ucl_object_tostring (cur)); + rspamd_milter_send_action (session, RSPAMD_MILTER_CHGHEADER, + (guint32) 1, hname, hvalue); + g_string_free (hname, TRUE); + g_string_free (hvalue, TRUE); + } + + rspamd_milter_send_action (session, RSPAMD_MILTER_ACCEPT); + + return TRUE; + } + } + } + + return FALSE; } void @@ -1440,6 +1495,7 @@ rspamd_milter_send_task_results (struct rspamd_milter_session *session, gint action = METRIC_ACTION_REJECT; rspamd_fstring_t *xcode = NULL, *rcode = NULL, *reply = NULL; GString *hname, *hvalue; + gboolean processed = FALSE; if (results == NULL) { msg_err_milter ("cannot find scan results, tempfail"); @@ -1477,7 +1533,7 @@ rspamd_milter_send_task_results (struct rspamd_milter_session *session, elt = ucl_object_lookup (results, "milter"); if (elt) { - rspamd_milter_process_milter_block (session, elt); + processed = rspamd_milter_process_milter_block (session, elt, action); } /* DKIM-Signature */ @@ -1493,6 +1549,10 @@ rspamd_milter_send_task_results (struct rspamd_milter_session *session, g_string_free (hvalue, TRUE); } + if (processed) { + goto cleanup; + } + switch (action) { case METRIC_ACTION_REJECT: rcode = rspamd_fstring_new_init (RSPAMD_MILTER_RCODE_REJECT, @@ -1540,11 +1600,13 @@ rspamd_milter_send_task_results (struct rspamd_milter_session *session, break; case METRIC_ACTION_ADD_HEADER: + /* Remove existing headers */ + rspamd_milter_remove_header_safe (session, + milter_ctx->spam_header, + 0); + hname = g_string_new (milter_ctx->spam_header); - /* TODO: Perhaps, we can customize it as well */ hvalue = g_string_new ("Yes"); - - /* TODO: We to track headers to make it work fine */ rspamd_milter_send_action (session, RSPAMD_MILTER_CHGHEADER, (guint32)1, hname, hvalue); g_string_free (hname, TRUE); @@ -1559,6 +1621,7 @@ rspamd_milter_send_task_results (struct rspamd_milter_session *session, break; } +cleanup: rspamd_fstring_free (rcode); rspamd_fstring_free (xcode); rspamd_fstring_free (reply);