[Feature] Implement multi-flags fuzzy replies

4 files changed, 85 insertions, 11 deletions

diff --git a/src/fuzzy_storage.c b/src/fuzzy_storage.c
index cbea7d7f1..a10fd145d 100644
--- a/src/fuzzy_storage.c
+++ b/src/fuzzy_storage.c
@@ -685,6 +685,25 @@ rspamd_fuzzy_process_command (struct fuzzy_session *session)
 
 reply:
 	result.tag = cmd->tag;
+
+	if (session->epoch < RSPAMD_FUZZY_EPOCH11) {
+		/* We need to convert flags to legacy format */
+		guint32 flag = 0;
+
+		/* We select the least significant flag if multiple flags are set */
+		for (flag = 0; flag < 32; flag ++) {
+			if (result.flag & (1U << flag)) {
+				break;
+			}
+		}
+
+		if (flag == (1U << 31)) {
+			flag = 0;
+		}
+
+		result.flag = flag;
+	}
+
 	memcpy (&session->reply.rep, &result, sizeof (result));
 
 	rspamd_fuzzy_update_stats (session->ctx,
@@ -716,19 +735,32 @@ rspamd_fuzzy_command_valid (struct rspamd_fuzzy_cmd *cmd, gint r)
 {
 	enum rspamd_fuzzy_epoch ret = RSPAMD_FUZZY_EPOCH_MAX;
 
-	if (cmd->version == RSPAMD_FUZZY_VERSION) {
+	switch (cmd->version) {
+	case 4:
 		if (cmd->shingles_count > 0) {
 			if (r == sizeof (struct rspamd_fuzzy_shingle_cmd)) {
-				ret = RSPAMD_FUZZY_EPOCH9;
+				ret = RSPAMD_FUZZY_EPOCH11;
 			}
 		}
 		else {
 			if (r == sizeof (*cmd)) {
-				ret = RSPAMD_FUZZY_EPOCH9;
+				ret = RSPAMD_FUZZY_EPOCH11;
 			}
 		}
-	}
-	else if (cmd->version == 2) {
+		break;
+	case 3:
+		if (cmd->shingles_count > 0) {
+			if (r == sizeof (struct rspamd_fuzzy_shingle_cmd)) {
+				ret = RSPAMD_FUZZY_EPOCH10;
+			}
+		}
+		else {
+			if (r == sizeof (*cmd)) {
+				ret = RSPAMD_FUZZY_EPOCH10;
+			}
+		}
+		break;
+	case 2:
 		/*
 		 * rspamd 0.8 has slightly different tokenizer then it might be not
 		 * 100% compatible
@@ -741,6 +773,9 @@ rspamd_fuzzy_command_valid (struct rspamd_fuzzy_cmd *cmd, gint r)
 		else {
 			ret = RSPAMD_FUZZY_EPOCH8;
 		}
+		break;
+	default:
+		break;
 	}
 
 	return ret;
@@ -858,7 +893,7 @@ rspamd_fuzzy_cmd_from_wire (guchar *buf, guint buflen, struct fuzzy_session *s)
 			return FALSE;
 		}
 		/* Encrypted is epoch 10 at least */
-		s->epoch = RSPAMD_FUZZY_EPOCH10;
+		s->epoch = epoch;
 		break;
 	case sizeof (struct rspamd_fuzzy_encrypted_shingle_cmd):
 		s->cmd_type = CMD_ENCRYPTED_SHINGLE;
@@ -875,7 +910,7 @@ rspamd_fuzzy_cmd_from_wire (guchar *buf, guint buflen, struct fuzzy_session *s)
 			return FALSE;
 		}
 
-		s->epoch = RSPAMD_FUZZY_EPOCH10;
+		s->epoch = epoch;
 		break;
 	default:
 		msg_debug ("invalid fuzzy command of size %d received", buflen);
diff --git a/src/fuzzy_storage.h b/src/fuzzy_storage.h
index ff3438b38..aca6ab952 100644
--- a/src/fuzzy_storage.h
+++ b/src/fuzzy_storage.h
@@ -6,7 +6,7 @@
 #include "shingles.h"
 #include "cryptobox.h"
 
-#define RSPAMD_FUZZY_VERSION 3
+#define RSPAMD_FUZZY_VERSION 4
 #define RSPAMD_FUZZY_KEYLEN 8
 
 /* Commands for fuzzy storage */
@@ -22,7 +22,8 @@ enum rspamd_fuzzy_epoch {
 	RSPAMD_FUZZY_EPOCH6 = 0, /**< pre 0.6.x */
 	RSPAMD_FUZZY_EPOCH8, /**< 0.8 till 0.9 */
 	RSPAMD_FUZZY_EPOCH9, /**< 0.9 + */
-	RSPAMD_FUZZY_EPOCH10, /**< 1.0 + encryption */
+	RSPAMD_FUZZY_EPOCH10, /**< 1.0+ encryption */
+	RSPAMD_FUZZY_EPOCH11, /**< 1.3+ multiple flags */
 	RSPAMD_FUZZY_EPOCH_MAX
 };
 
diff --git a/src/libserver/fuzzy_backend.c b/src/libserver/fuzzy_backend.c
index 106e379f8..fa5518c2d 100644
--- a/src/libserver/fuzzy_backend.c
+++ b/src/libserver/fuzzy_backend.c
@@ -32,6 +32,7 @@ struct rspamd_fuzzy_backend {
 
 static const gdouble sql_sleep_time = 0.1;
 static const guint max_retries = 10;
+static const guint32 flags_mask = (1U << 31);
 
 #define msg_err_fuzzy_backend(...) rspamd_default_log_function (G_LOG_LEVEL_CRITICAL, \
         backend->pool->tag.tagname, backend->pool->tag.uid, \
@@ -526,6 +527,10 @@ rspamd_fuzzy_backend_check (struct rspamd_fuzzy_backend *backend,
 			rep.prob = 1.0;
 			rep.flag = sqlite3_column_int (
 					prepared_stmts[RSPAMD_FUZZY_BACKEND_CHECK].stmt, 2);
+
+			if (!(rep.flag & flags_mask)) {
+				rep.flag = (1U << rep.flag) | flags_mask;
+			}
 		}
 	}
 	else if (cmd->shingles_count > 0) {
@@ -609,6 +614,10 @@ rspamd_fuzzy_backend_check (struct rspamd_fuzzy_backend *backend,
 						rep.flag = sqlite3_column_int (
 								prepared_stmts[RSPAMD_FUZZY_BACKEND_GET_DIGEST_BY_ID].stmt,
 								3);
+
+						if (!(rep.flag & flags_mask)) {
+							rep.flag = (1U << rep.flag) | flags_mask;
+						}
 					}
 				}
 			}
@@ -664,6 +673,11 @@ rspamd_fuzzy_backend_add (struct rspamd_fuzzy_backend *backend,
 		return FALSE;
 	}
 
+	if (cmd->flag > 31) {
+		msg_err_fuzzy_backend ("flag more than 31 is no longer supported");
+		return FALSE;
+	}
+
 	rc = rspamd_fuzzy_backend_run_stmt (backend, FALSE,
 			RSPAMD_FUZZY_BACKEND_CHECK,
 			cmd->digest);
@@ -675,7 +689,7 @@ rspamd_fuzzy_backend_add (struct rspamd_fuzzy_backend *backend,
 				2);
 		rspamd_fuzzy_backend_cleanup_stmt (backend, RSPAMD_FUZZY_BACKEND_CHECK);
 
-		if (flag == cmd->flag) {
+		if ((flag & cmd->flag) == cmd->flag) {
 			/* We need to increase weight */
 			rc = rspamd_fuzzy_backend_run_stmt (backend, TRUE,
 					RSPAMD_FUZZY_BACKEND_UPDATE,
@@ -690,11 +704,28 @@ rspamd_fuzzy_backend_add (struct rspamd_fuzzy_backend *backend,
 		}
 		else {
 			/* We need to relearn actually */
+			if (flag & flags_mask) {
+				/* This is already new format */
+				flag |= cmd->flag;
+			}
+			else {
+				/* Convert to the new format */
+				if (flag > 31) {
+					msg_warn_fuzzy_backend ("storage had flag more than 31, remove "
+							"it");
+					flag = cmd->flag | flags_mask;
+				}
+				else {
+					flag = (1U << flag) | cmd->flag | flags_mask;
+				}
+			}
+
 			rc = rspamd_fuzzy_backend_run_stmt (backend, TRUE,
 					RSPAMD_FUZZY_BACKEND_UPDATE_FLAG,
 					(gint64) cmd->value,
 					(gint64) cmd->flag,
 					cmd->digest);
+
 			if (rc != SQLITE_OK) {
 				msg_warn_fuzzy_backend ("cannot update hash to %d -> "
 						"%*xs: %s", (gint) cmd->flag,
@@ -707,7 +738,7 @@ rspamd_fuzzy_backend_add (struct rspamd_fuzzy_backend *backend,
 		rspamd_fuzzy_backend_cleanup_stmt (backend, RSPAMD_FUZZY_BACKEND_CHECK);
 		rc = rspamd_fuzzy_backend_run_stmt (backend, FALSE,
 				RSPAMD_FUZZY_BACKEND_INSERT,
-				(gint) cmd->flag,
+				(gint) (1U << cmd->flag),
 				cmd->digest,
 				(gint64) cmd->value,
 				(gint64) timestamp);
diff --git a/src/plugins/fuzzy_check.c b/src/plugins/fuzzy_check.c
index bf4cd3a0a..c11ab40bf 100644
--- a/src/plugins/fuzzy_check.c
+++ b/src/plugins/fuzzy_check.c
@@ -188,6 +188,13 @@ parse_flags (struct fuzzy_rule *rule,
 
 			if (elt != NULL) {
 				map->fuzzy_flag = ucl_obj_toint (elt);
+
+				if (map->fuzzy_flag > 31) {
+					msg_err_config ("flags more than 31 are no longer "
+							"supported by rspamd");
+					return;
+				}
+
 				elt = ucl_object_lookup (val, "max_score");
 
 				if (elt != NULL) {