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

3 files changed, 29 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;
+		}
 	}