From 6b0164eb234fc3e61cd425e33d6ba5c9dd89a157 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Mon, 15 Jul 2019 13:57:15 +0100
Subject: [PATCH] [Minor] Neural: Use tableshape to verify profiles

---
 src/plugins/lua/neural.lua | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/plugins/lua/neural.lua b/src/plugins/lua/neural.lua
index 4e1b7057c..09bf86e08 100644
--- a/src/plugins/lua/neural.lua
+++ b/src/plugins/lua/neural.lua
@@ -27,6 +27,7 @@ local lua_util = require "lua_util"
 local fun = require "fun"
 local lua_settings = require "lua_settings"
 local meta_functions = require "lua_meta"
+local ts = require("tableshape").types
 local N = "neural"
 
 -- Module vars
@@ -50,6 +51,14 @@ local default_options = {
   symbol_ham = 'NEURAL_HAM',
 }
 
+local redis_profile_schema = ts.shape{
+  digest = ts.string,
+  symbols = ts.array_of(ts.string),
+  version = ts.number,
+  redis_key = ts.string,
+  distance = ts.number:is_optional(),
+}
+
 -- Rule structure:
 -- * static config fields (see `default_options`)
 -- * prefix - name or defined prefix
@@ -797,7 +806,7 @@ local function load_new_ann(rule, ev_base, set, profile, min_diff)
 
             local ucl = require "ucl"
             local profile_serialized = ucl.to_format(profile, 'json-compact', true)
-            profile.ann = ann -- To avoid serialization
+            set.ann.ann = ann -- To avoid serialization
 
             local function rank_cb(_, _)
               -- TODO: maybe add some logging
@@ -993,7 +1002,14 @@ local function load_ann_profile(element)
         ucl_err)
     return nil
   else
-    return parser:get_object()
+    local profile = parser:get_object()
+    local checked,schema_err = redis_profile_schema:transform(profile)
+    if not checked then
+      rspamd_logger.errx(rspamd_config, "cannot parse profile schema: %s", schema_err)
+
+      return nil
+    end
+    return checked
   end
 end
 
-- 
2.39.5