commit e2600fe: [Feature] Allow to output group results

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2019-08-12 13:37:24 +0100
URL: https://github.com/rspamd/rspamd/commit/e2600fef6aac4e838871783e41ab934ab43e2e5b

[Feature] Allow to output group results
Issue: #2985

---
 src/libserver/protocol.c | 36 ++++++++++++++++++++++++++++++++++--
 src/libserver/task.h     |  4 +++-
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index e62333262..a51cbf17f 100644
--- a/src/libserver/protocol.c
+++ b/src/libserver/protocol.c
@@ -380,6 +380,7 @@ rspamd_protocol_handle_flag (struct rspamd_task *task, const gchar *str,
 	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);
+	CHECK_PROTOCOL_FLAG("groups", RSPAMD_TASK_PROTOCOL_FLAG_GROUPS);
 
 	if (!known) {
 		msg_warn_protocol ("unknown flag: %*s", (gint)len, str);
@@ -1108,8 +1109,8 @@ rspamd_metric_symbol_ucl (struct rspamd_task *task, struct rspamd_symbol_result
 	}
 
 	if (description) {
-		ucl_object_insert_key (obj, ucl_object_fromstring (
-				description), "description", 0, false);
+		ucl_object_insert_key (obj, ucl_object_fromstring (description),
+				"description", 0, false);
 	}
 
 	if (sym->options != NULL) {
@@ -1125,6 +1126,19 @@ rspamd_metric_symbol_ucl (struct rspamd_task *task, struct rspamd_symbol_result
 	return obj;
 }
 
+static ucl_object_t *
+rspamd_metric_group_ucl (struct rspamd_task *task,
+		struct rspamd_symbols_group *gr, gdouble score)
+{
+	ucl_object_t *obj = NULL;
+
+	obj = ucl_object_typed_new (UCL_OBJECT);
+	ucl_object_insert_key (obj, ucl_object_fromdouble (score),
+			"score", 0, false);
+
+	return obj;
+}
+
 static ucl_object_t *
 rspamd_scan_result_ucl (struct rspamd_task *task,
 						struct rspamd_scan_result *mres, ucl_object_t *top)
@@ -1189,6 +1203,7 @@ rspamd_scan_result_ucl (struct rspamd_task *task,
 
 	/* Now handle symbols */
 	if (task->cmd != CMD_CHECK) {
+		/* For checkv2 we insert symbols as a separate object */
 		obj = ucl_object_typed_new (UCL_OBJECT);
 	}
 
@@ -1200,12 +1215,29 @@ rspamd_scan_result_ucl (struct rspamd_task *task,
 	})
 
 	if (task->cmd != CMD_CHECK) {
+		/* For checkv2 we insert symbols as a separate object */
 		ucl_object_insert_key (top, obj, "symbols", 0, false);
 	}
 	else {
+		/* For legacy check we just insert it as "default" all together */
 		ucl_object_insert_key (top, obj, DEFAULT_METRIC, 0, false);
 	}
 
+	/* Handle groups if needed */
+	if (task->protocol_flags & RSPAMD_TASK_PROTOCOL_FLAG_GROUPS) {
+		struct rspamd_symbols_group *gr;
+		gdouble gr_score;
+
+		obj = ucl_object_typed_new (UCL_OBJECT);
+
+		kh_foreach (mres->sym_groups, gr, gr_score,{
+			sobj = rspamd_metric_group_ucl (task, gr, gr_score);
+			ucl_object_insert_key (obj, sobj, gr->name, 0, false);
+		});
+
+		ucl_object_insert_key (top, obj, "groups", 0, false);
+	}
+
 	return obj;
 }
 
diff --git a/src/libserver/task.h b/src/libserver/task.h
index 9086578d7..28e0dc070 100644
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -128,7 +128,9 @@ enum rspamd_task_stage {
 #define RSPAMD_TASK_PROTOCOL_FLAG_EXT_URLS (1u << 4u)
 /* Client allows body block (including headers in no FLAG_MILTER) */
 #define RSPAMD_TASK_PROTOCOL_FLAG_BODY_BLOCK (1u << 5u)
-#define RSPAMD_TASK_PROTOCOL_FLAG_MAX_SHIFT (5u)
+/* Emit groups information */
+#define RSPAMD_TASK_PROTOCOL_FLAG_GROUPS (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)->cmd == CMD_CHECK_SPAMC))