[Feature] Allow custom milter quarantine and tempfail messages

Issue: #4892
Closes: #4892

3 files changed, 73 insertions, 20 deletions

diff --git a/src/libserver/milter.c b/src/libserver/milter.c
index dd64822c4..f35278a0e 100644
--- a/src/libserver/milter.c
+++ b/src/libserver/milter.c
@@ -1288,9 +1288,10 @@ rspamd_milter_send_action(struct rspamd_milter_session *session,
 	rspamd_fstring_t *reply = NULL;
 	gsize len;
 	GString *name, *value;
-	const char *reason, *body_str;
+	const char *body_str;
 	struct rspamd_milter_outbuf *obuf;
 	struct rspamd_milter_private *priv = session->priv;
+	const rspamd_fstring_t *reason;
 
 	va_start(ap, act);
 	cmd = act;
@@ -1307,16 +1308,21 @@ rspamd_milter_send_action(struct rspamd_milter_session *session,
 		SET_COMMAND(cmd, 0, reply, pos);
 		break;
 	case RSPAMD_MILTER_QUARANTINE:
-		reason = va_arg(ap, const char *);
+		reason = va_arg(ap, const rspamd_fstring_t *);
 
-		if (reason == NULL) {
-			reason = "";
+		if (reason != NULL) {
+			len = reason->len;
+			msg_debug_milter("send quarantine action %*s", (int) len, reason->str);
+			SET_COMMAND(cmd, len + 1, reply, pos);
+			memcpy(pos, reason->str, len);
+			pos[len] = '\0';
+		}
+		else {
+			msg_debug_milter("send quarantine action with no specific string");
+			SET_COMMAND(cmd, 1, reply, pos);
+			*pos = '\0';
 		}
 
-		len = strlen(reason);
-		msg_debug_milter("send quarantine action %s", reason);
-		SET_COMMAND(cmd, len + 1, reply, pos);
-		memcpy(pos, reason, len + 1);
 		break;
 	case RSPAMD_MILTER_ADDHEADER:
 		name = va_arg(ap, GString *);
@@ -2119,9 +2125,18 @@ void rspamd_milter_send_task_results(struct rspamd_milter_session *session,
 			rspamd_milter_send_action(session, RSPAMD_MILTER_DISCARD);
 		}
 		else if (priv->quarantine_on_reject) {
-			/* TODO: be more flexible about SMTP messages */
-			rspamd_milter_send_action(session, RSPAMD_MILTER_QUARANTINE,
-									  RSPAMD_MILTER_QUARANTINE_MESSAGE);
+			if (!reply) {
+				if (milter_ctx->quarantine_message == NULL) {
+					reply = rspamd_fstring_new_init(
+						RSPAMD_MILTER_QUARANTINE_MESSAGE,
+						sizeof(RSPAMD_MILTER_QUARANTINE_MESSAGE) - 1);
+				}
+				else {
+					reply = rspamd_fstring_new_init(milter_ctx->quarantine_message,
+													strlen(milter_ctx->quarantine_message));
+				}
+				rspamd_milter_send_action(session, RSPAMD_MILTER_QUARANTINE, reply);
+			}
 
 			/* Quarantine also requires accept action, all hail Sendmail */
 			rspamd_milter_send_action(session, RSPAMD_MILTER_ACCEPT);
@@ -2154,8 +2169,15 @@ void rspamd_milter_send_task_results(struct rspamd_milter_session *session,
 										sizeof(RSPAMD_MILTER_XCODE_TEMPFAIL) - 1);
 
 		if (!reply) {
-			reply = rspamd_fstring_new_init(RSPAMD_MILTER_TEMPFAIL_MESSAGE,
-											sizeof(RSPAMD_MILTER_TEMPFAIL_MESSAGE) - 1);
+			if (milter_ctx->tempfail_message == NULL) {
+				reply = rspamd_fstring_new_init(
+					RSPAMD_MILTER_TEMPFAIL_MESSAGE,
+					sizeof(RSPAMD_MILTER_TEMPFAIL_MESSAGE) - 1);
+			}
+			else {
+				reply = rspamd_fstring_new_init(milter_ctx->tempfail_message,
+												strlen(milter_ctx->tempfail_message));
+			}
 		}
 
 		rspamd_milter_set_reply(session, rcode, xcode, reply);
@@ -2193,9 +2215,18 @@ void rspamd_milter_send_task_results(struct rspamd_milter_session *session,
 		break;
 
 	case METRIC_ACTION_QUARANTINE:
-		/* TODO: be more flexible about SMTP messages */
-		rspamd_milter_send_action(session, RSPAMD_MILTER_QUARANTINE,
-								  RSPAMD_MILTER_QUARANTINE_MESSAGE);
+		if (!reply) {
+			if (milter_ctx->quarantine_message == NULL) {
+				reply = rspamd_fstring_new_init(
+					RSPAMD_MILTER_QUARANTINE_MESSAGE,
+					sizeof(RSPAMD_MILTER_QUARANTINE_MESSAGE) - 1);
+			}
+			else {
+				reply = rspamd_fstring_new_init(milter_ctx->quarantine_message,
+												strlen(milter_ctx->quarantine_message));
+			}
+			rspamd_milter_send_action(session, RSPAMD_MILTER_QUARANTINE, reply);
+		}
 
 		/* Quarantine also requires accept action, all hail Sendmail */
 		rspamd_milter_send_action(session, RSPAMD_MILTER_ACCEPT);
diff --git a/src/libserver/milter.h b/src/libserver/milter.h
index 840f69126..5b4b750e2 100644
--- a/src/libserver/milter.h
+++ b/src/libserver/milter.h
@@ -1,11 +1,11 @@
-/*-
- * Copyright 2017 Vsevolod Stakhov
+/*
+ * Copyright 2024 Vsevolod Stakhov
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
- *   http://www.apache.org/licenses/LICENSE-2.0
+ *    http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
@@ -55,6 +55,8 @@ struct rspamd_milter_context {
 	const char *spam_header;
 	const char *client_ca_name;
 	const char *reject_message;
+	const char *quarantine_message;
+	const char *tempfail_message;
 	void *sessions_cache;
 	struct rspamd_config *cfg;
 	gboolean discard_on_reject;
diff --git a/src/rspamd_proxy.c b/src/rspamd_proxy.c
index 39df7c98d..415d22f07 100644
--- a/src/rspamd_proxy.c
+++ b/src/rspamd_proxy.c
@@ -150,8 +150,10 @@ struct rspamd_proxy_ctx {
 	char *spam_header;
 	/* CA name that can be used for client certificates */
 	char *client_ca_name;
-	/* Milter rejection message */
+	/* Milter messages */
 	char *reject_message;
+	char *quarantine_message;
+	char *tempfail_message;
 	/* Sessions cache */
 	void *sessions_cache;
 	struct rspamd_milter_context milter_ctx;
@@ -874,6 +876,22 @@ init_rspamd_proxy(struct rspamd_config *cfg)
 									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, reject_message),
 									  0,
 									  "Use custom rejection message");
+	rspamd_rcl_register_worker_option(cfg,
+									  type,
+									  "quarantine_message",
+									  rspamd_rcl_parse_struct_string,
+									  ctx,
+									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, quarantine_message),
+									  0,
+									  "Use custom quarantine message");
+	rspamd_rcl_register_worker_option(cfg,
+									  type,
+									  "tempfail_message",
+									  rspamd_rcl_parse_struct_string,
+									  ctx,
+									  G_STRUCT_OFFSET(struct rspamd_proxy_ctx, tempfail_message),
+									  0,
+									  "Use custom tempfail message");
 
 	return ctx;
 }
@@ -2435,6 +2453,8 @@ 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.quarantine_message = ctx->quarantine_message;
+	ctx->milter_ctx.tempfail_message = ctx->tempfail_message;
 	ctx->milter_ctx.cfg = ctx->cfg;
 	rspamd_milter_init_library(&ctx->milter_ctx);