commit f46d3ab: [Feature] Support protocol flags

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2019-07-16 15:32:17 +0100
URL: https://github.com/rspamd/rspamd/commit/f46d3ab320e9a0638f4e96826a315608c65ea6c7

[Feature] Support protocol flags

---
 src/libserver/protocol.c          | 91 +++++++++++++++++++++++++++++++++++++++
 src/libserver/protocol_internal.h |  1 +
 src/libserver/task.h              | 10 ++++-
 src/lua/lua_task.c                |  2 +-
 4 files changed, 102 insertions(+), 2 deletions(-)

diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index 8c5cd4080..602878d9b 100644
--- a/src/libserver/protocol.c
+++ b/src/libserver/protocol.c
@@ -338,6 +338,93 @@ rspamd_protocol_process_recipients (struct rspamd_task *task,
 	}
 }
 
+#define COMPARE_FLAG_LIT(lit) (len == sizeof(lit) - 1 && memcmp ((lit), str, len) == 0)
+#define CHECK_PROTOCOL_FLAG(lit, fl) do { \
+	if (!known && COMPARE_FLAG_LIT(lit)) { \
+		task->protocol_flags |= (fl); \
+		known = TRUE; \
+		msg_debug_protocol ("add protocol flag %s", lit); \
+	} \
+} while (0)
+#define CHECK_TASK_FLAG(lit, fl) do { \
+	if (!known && COMPARE_FLAG_LIT(lit)) { \
+		task->flags |= (fl); \
+		known = TRUE; \
+		msg_debug_protocol ("add task flag %s", lit); \
+	} \
+} while (0)
+
+static void
+rspamd_protocol_handle_flag (struct rspamd_task *task, const gchar *str,
+		gsize len)
+{
+	gboolean known = FALSE;
+
+	CHECK_TASK_FLAG("pass_all", RSPAMD_TASK_FLAG_PASS_ALL);
+	CHECK_TASK_FLAG("no_log", RSPAMD_TASK_FLAG_NO_LOG);
+	CHECK_TASK_FLAG("skip", RSPAMD_TASK_FLAG_NO_LOG);
+	CHECK_TASK_FLAG("no_stat", RSPAMD_TASK_FLAG_NO_STAT);
+	CHECK_TASK_FLAG("ssl", RSPAMD_TASK_FLAG_SSL);
+
+	CHECK_PROTOCOL_FLAG("milter", RSPAMD_TASK_PROTOCOL_FLAG_MILTER);
+	CHECK_PROTOCOL_FLAG("zstd", RSPAMD_TASK_PROTOCOL_FLAG_COMPRESSED);
+	CHECK_PROTOCOL_FLAG("ext_urls", RSPAMD_TASK_PROTOCOL_FLAG_EXT_URLS);
+	CHECK_PROTOCOL_FLAG("body_block", RSPAMD_TASK_PROTOCOL_FLAG_BODY_BLOCK);
+
+	if (!known) {
+		msg_warn_protocol ("unknown flag: %*s", (gint)len, str);
+	}
+}
+
+#undef COMPARE_FLAG
+#undef CHECK_PROTOCOL_FLAG
+
+static void
+rspamd_protocol_process_flags (struct rspamd_task *task, const rspamd_ftok_t *hdr)
+{
+	enum {
+		skip_spaces,
+		read_flag,
+	} state = skip_spaces;
+	const gchar *p, *end, *start;
+
+	p = hdr->begin;
+	end = hdr->begin + hdr->len;
+	start = NULL;
+
+	while (p < end) {
+		switch (state) {
+		case skip_spaces:
+			if (g_ascii_isspace (*p)) {
+				p ++;
+			}
+			else {
+				state = read_flag;
+				start = p;
+			}
+			break;
+		case read_flag:
+			if (*p == ',') {
+				if (p > start) {
+					rspamd_protocol_handle_flag (task, start, p - start);
+				}
+				start = NULL;
+				state = skip_spaces;
+				p ++;
+			}
+			else {
+				p ++;
+			}
+			break;
+		}
+	}
+
+	/* Check remainder */
+	if (start && end > start && state == read_flag) {
+		rspamd_protocol_handle_flag (task, start, end - start);
+	}
+}
+
 #define IF_HEADER(name) \
 	srch.begin = (name); \
 	srch.len = sizeof (name) - 1; \
@@ -407,6 +494,10 @@ rspamd_protocol_handle_headers (struct rspamd_task *task,
 							hv_tok);
 					msg_debug_protocol ("read filename header, value: %s", task->msg.fpath);
 				}
+				IF_HEADER (FLAGS_HEADER) {
+					msg_debug_protocol ("read flags header, value: %T", hv_tok);
+					rspamd_protocol_process_flags (task, hv_tok);
+				}
 				break;
 			case 'q':
 			case 'Q':
diff --git a/src/libserver/protocol_internal.h b/src/libserver/protocol_internal.h
index f673fc74d..418af70d9 100644
--- a/src/libserver/protocol_internal.h
+++ b/src/libserver/protocol_internal.h
@@ -84,6 +84,7 @@ extern "C" {
 #define MTA_NAME_HEADER "MTA-Name"
 #define MILTER_HEADER "Milter"
 #define FILENAME_HEADER "Filename"
+#define FLAGS_HEADER "Flags"
 #define CERT_ISSUER_HEADER "TLS-Cert-Issuer"
 #define MAILER_HEADER "Mailer"
 #define RAW_DATA_HEADER "Raw"
diff --git a/src/libserver/task.h b/src/libserver/task.h
index d6a9fc7c9..ec66febd4 100644
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -118,13 +118,21 @@ enum rspamd_task_stage {
 #define RSPAMD_TASK_FLAG_MESSAGE_REWRITE (1u << 24u)
 #define RSPAMD_TASK_FLAG_MAX_SHIFT (24u)
 
+/* Spamc message */
 #define RSPAMD_TASK_PROTOCOL_FLAG_SPAMC (1u << 0u)
+/* Request has a JSON control block */
 #define RSPAMD_TASK_PROTOCOL_FLAG_HAS_CONTROL (1u << 1u)
+/* Request has been done by a local client */
 #define RSPAMD_TASK_PROTOCOL_FLAG_LOCAL_CLIENT (1u << 2u)
+/* Request has been sent via milter */
 #define RSPAMD_TASK_PROTOCOL_FLAG_MILTER (1u << 3u)
+/* Compress protocol reply */
 #define RSPAMD_TASK_PROTOCOL_FLAG_COMPRESSED (1u << 4u)
+/* Include all URLs */
 #define RSPAMD_TASK_PROTOCOL_FLAG_EXT_URLS (1u << 5u)
-#define RSPAMD_TASK_PROTOCOL_FLAG_MAX_SHIFT (5u)
+/* Client allows body block (including headers in no FLAG_MILTER) */
+#define RSPAMD_TASK_PROTOCOL_FLAG_BODY_BLOCK (1u << 6u)
+#define RSPAMD_TASK_PROTOCOL_FLAG_MAX_SHIFT (6u)
 
 #define RSPAMD_TASK_IS_SKIPPED(task) (((task)->flags & RSPAMD_TASK_FLAG_SKIP))
 #define RSPAMD_TASK_IS_SPAMC(task) (((task)->protocol_flags & RSPAMD_TASK_PROTOCOL_FLAG_SPAMC))
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 0bc86533b..cd33b43ae 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -4819,7 +4819,7 @@ lua_task_get_flags (lua_State *L)
 
 		flags = task->flags;
 
-		for (i = 0; i < RSPAMD_TASK_FLAG_MAX_SHIFT; i ++) {
+		for (i = 0; i <= RSPAMD_TASK_FLAG_MAX_SHIFT; i ++) {
 			bit = (1U << i);
 
 			if (flags & bit) {