[Feature] Do not classify message if some class is missing

4 files changed, 47 insertions, 0 deletions

diff --git a/src/libstat/backends/mmaped_file.c b/src/libstat/backends/mmaped_file.c
index eba730cab..2065edd5c 100644
--- a/src/libstat/backends/mmaped_file.c
+++ b/src/libstat/backends/mmaped_file.c
@@ -871,6 +871,13 @@ rspamd_mmaped_file_process_tokens (struct rspamd_task *task, GPtrArray *tokens,
 		tok->values[id] = rspamd_mmaped_file_get_block (mf, h1, h2);
 	}
 
+	if (mf->cf->is_spam) {
+		task->flags |= RSPAMD_TASK_FLAG_HAS_SPAM_TOKENS;
+	}
+	else {
+		task->flags |= RSPAMD_TASK_FLAG_HAS_HAM_TOKENS;
+	}
+
 	return TRUE;
 }
 
diff --git a/src/libstat/backends/redis_backend.c b/src/libstat/backends/redis_backend.c
index 3b1667e1c..d66b6d68e 100644
--- a/src/libstat/backends/redis_backend.c
+++ b/src/libstat/backends/redis_backend.c
@@ -828,9 +828,24 @@ rspamd_redis_processed (redisAsyncContext *c, gpointer r, gpointer priv)
 
 						processed ++;
 					}
+
+					if (rt->stcf->is_spam) {
+						task->flags |= RSPAMD_TASK_FLAG_HAS_SPAM_TOKENS;
+					}
+					else {
+						task->flags |= RSPAMD_TASK_FLAG_HAS_HAM_TOKENS;
+					}
+				}
+				else {
+					msg_err_task ("got invalid length of reply vector from redis: "
+							"%d, expected: %d",
+							(gint)reply->elements,
+							(gint)task->tokens->len);
 				}
 			}
 			else {
+				msg_err_task ("got invalid reply from redis: %d",
+						reply->type);
 			}
 
 			msg_debug_task ("received tokens for %s: %d processed, %d found",
diff --git a/src/libstat/backends/sqlite3_backend.c b/src/libstat/backends/sqlite3_backend.c
index 36bbaa20e..f90e6b93a 100644
--- a/src/libstat/backends/sqlite3_backend.c
+++ b/src/libstat/backends/sqlite3_backend.c
@@ -723,6 +723,13 @@ rspamd_sqlite3_process_tokens (struct rspamd_task *task,
 		else {
 			tok->values[id] = 0.0;
 		}
+
+		if (rt->cf->is_spam) {
+			task->flags |= RSPAMD_TASK_FLAG_HAS_SPAM_TOKENS;
+		}
+		else {
+			task->flags |= RSPAMD_TASK_FLAG_HAS_HAM_TOKENS;
+		}
 	}
 
 
diff --git a/src/libstat/stat_process.c b/src/libstat/stat_process.c
index 6b1b16aed..1761dc7a3 100644
--- a/src/libstat/stat_process.c
+++ b/src/libstat/stat_process.c
@@ -321,6 +321,24 @@ rspamd_stat_classifiers_process (struct rspamd_stat_ctx *st_ctx,
 	guint i;
 	struct rspamd_classifier *cl;
 
+	if (st_ctx->classifiers->len == 0) {
+		return;
+	}
+
+	/*
+	 * Do not classify a message if some class is missing
+	 */
+	if (!(task->flags & RSPAMD_TASK_FLAG_HAS_SPAM_TOKENS)) {
+		msg_warn_task ("skip statistics as SPAM class is missing");
+
+		return;
+	}
+	if (!(task->flags & RSPAMD_TASK_FLAG_HAS_HAM_TOKENS)) {
+		msg_warn_task ("skip statistics as HAM class is missing");
+
+		return;
+	}
+
 	for (i = 0; i < st_ctx->classifiers->len; i++) {
 		cl = g_ptr_array_index (st_ctx->classifiers, i);
 		g_assert (cl != NULL);