[Feature] Support protocol flags

diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index 8c5cd40806acbdbdcaef129459a5cdfb61e0683e..602878d9bfd7a4dfbf6725191fb834028028c133 100644 (file)
--- 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 f673fc74d299d7304851cc32fbb16f0a2cd6ca64..418af70d965d6edbf0ed249b8ecf5b5723dcdb14 100644 (file)
--- 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 d6a9fc7c92b3cac4e0ca487ac244d8e2184178b6..ec66febd4850978b12e49da237ae5ddcce4ecb85 100644 (file)
--- 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 0bc86533b15fbf69269a962a58964be30e0a9e00..cd33b43ae69aea9eb7a66478cc72a27b9f05d1d1 100644 (file)
--- 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) {