commit f7200dd: [Minor] Milter: Add support of quarantine and discard actions

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2019-01-26 14:22:37 +0000
URL: https://github.com/rspamd/rspamd/commit/f7200dd41947e09e42d6528b16b659532c3c7b5a

[Minor] Milter: Add support of quarantine and discard actions

---
 src/libserver/milter.c   | 29 ++++++++++++++++++++++++-----
 src/libserver/milter.h   |  2 ++
 src/libserver/protocol.c | 11 +++++++++++
 src/rspamd_proxy.c       |  1 +
 4 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/src/libserver/milter.c b/src/libserver/milter.c
index 16d6252e9..b871c2347 100644
--- a/src/libserver/milter.c
+++ b/src/libserver/milter.c
@@ -25,6 +25,7 @@
 #include "libutil/http.h"
 #include "libutil/http_private.h"
 #include "libserver/protocol_internal.h"
+#include "libserver/cfg_file_private.h"
 #include "libmime/filter.h"
 #include "libserver/worker_util.h"
 #include "utlist.h"
@@ -1557,7 +1558,7 @@ rspamd_milter_remove_header_safe (struct rspamd_milter_session *session,
  */
 static gboolean
 rspamd_milter_process_milter_block (struct rspamd_milter_session *session,
-		const ucl_object_t *obj, gint action)
+		const ucl_object_t *obj, struct rspamd_action *action)
 {
 	const ucl_object_t *elt, *cur, *cur_elt;
 	ucl_object_iter_t it;
@@ -1731,7 +1732,7 @@ rspamd_milter_process_milter_block (struct rspamd_milter_session *session,
 		}
 	}
 
-	if (action == METRIC_ACTION_ADD_HEADER) {
+	if (action->action_type == METRIC_ACTION_ADD_HEADER) {
 		elt = ucl_object_lookup (obj, "spam_header");
 
 		if (elt) {
@@ -1783,7 +1784,7 @@ rspamd_milter_send_task_results (struct rspamd_milter_session *session,
 	const ucl_object_t *elt;
 	struct rspamd_milter_private *priv = session->priv;
 	const gchar *str_action;
-	gint action = METRIC_ACTION_REJECT;
+	struct rspamd_action *action;
 	rspamd_fstring_t *xcode = NULL, *rcode = NULL, *reply = NULL;
 	GString *hname, *hvalue;
 	gboolean processed = FALSE;
@@ -1805,7 +1806,14 @@ rspamd_milter_send_task_results (struct rspamd_milter_session *session,
 	}
 
 	str_action = ucl_object_tostring (elt);
-	rspamd_action_from_str (str_action, &action);
+	action = rspamd_config_get_action (milter_ctx->cfg, str_action);
+
+	if (action == NULL) {
+		msg_err_milter ("action %s has not been registered", str_action);
+		rspamd_milter_send_action (session, RSPAMD_MILTER_TEMPFAIL);
+
+		goto cleanup;
+	}
 
 	elt = ucl_object_lookup (results, "messages");
 	if (elt) {
@@ -1860,7 +1868,7 @@ rspamd_milter_send_task_results (struct rspamd_milter_session *session,
 		goto cleanup;
 	}
 
-	switch (action) {
+	switch (action->action_type) {
 	case METRIC_ACTION_REJECT:
 		if (priv->discard_on_reject) {
 			rspamd_milter_send_action (session, RSPAMD_MILTER_DISCARD);
@@ -1939,6 +1947,17 @@ rspamd_milter_send_task_results (struct rspamd_milter_session *session,
 		rspamd_milter_send_action (session, RSPAMD_MILTER_ACCEPT);
 		break;
 
+	case METRIC_ACTION_MILTER_QUARANTINE:
+		/* TODO: be more flexible about SMTP messages */
+		rspamd_milter_send_action (session, RSPAMD_MILTER_QUARANTINE,
+				RSPAMD_MILTER_QUARANTINE_MESSAGE);
+
+		/* Quarantine also requires accept action, all hail Sendmail */
+		rspamd_milter_send_action (session, RSPAMD_MILTER_ACCEPT);
+		break;
+	case METRIC_ACTION_MILTER_DISCARD:
+		rspamd_milter_send_action (session, RSPAMD_MILTER_DISCARD);
+		break;
 	case METRIC_ACTION_GREYLIST:
 	case METRIC_ACTION_NOACTION:
 	default:
diff --git a/src/libserver/milter.h b/src/libserver/milter.h
index 35dc8f908..7f375b018 100644
--- a/src/libserver/milter.h
+++ b/src/libserver/milter.h
@@ -43,12 +43,14 @@ enum rspamd_milter_reply {
 struct rspamd_email_address;
 struct event_base;
 struct rspamd_http_message;
+struct rspamd_config;
 
 struct rspamd_milter_context {
 	const gchar *spam_header;
 	const gchar *client_ca_name;
 	const gchar *reject_message;
 	void *sessions_cache;
+	struct rspamd_config *cfg;
 	gboolean discard_on_reject;
 	gboolean quarantine_on_reject;
 };
diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index 713c54499..d8c790e8f 100644
--- a/src/libserver/protocol.c
+++ b/src/libserver/protocol.c
@@ -968,6 +968,17 @@ rspamd_metric_result_ucl (struct rspamd_task *task,
 				"subject", 0, false);
 		}
 	}
+	if (action->flags & RSPAMD_ACTION_MILTER) {
+		/* Treat milter action specially */
+		if (action->action_type == METRIC_ACTION_MILTER_DISCARD) {
+			ucl_object_insert_key (obj, ucl_object_fromstring ("discard"),
+					"reject", 0, false);
+		}
+		else if (action->action_type == METRIC_ACTION_MILTER_QUARANTINE) {
+			ucl_object_insert_key (obj, ucl_object_fromstring ("quarantine"),
+					"reject", 0, false);
+		}
+	}
 
 	/* Now handle symbols */
 	if (task->cmd == CMD_CHECK_V2) {
diff --git a/src/rspamd_proxy.c b/src/rspamd_proxy.c
index aff866bb4..9ea2b0d74 100644
--- a/src/rspamd_proxy.c
+++ b/src/rspamd_proxy.c
@@ -2245,6 +2245,7 @@ start_rspamd_proxy (struct rspamd_worker *worker) {
 	ctx->milter_ctx.sessions_cache = ctx->sessions_cache;
 	ctx->milter_ctx.client_ca_name = ctx->client_ca_name;
 	ctx->milter_ctx.reject_message = ctx->reject_message;
+	ctx->milter_ctx.cfg = ctx->cfg;
 	rspamd_milter_init_library (&ctx->milter_ctx);
 
 	rspamd_lua_run_postloads (ctx->cfg->lua_state, ctx->cfg, ctx->ev_base,