[Feature] Allow to check rspamd_text using maps

diff --git a/src/lua/lua_map.c b/src/lua/lua_map.c
index f6fde2124e31f3e8921d1231fc80ed7c488685f3..bd5dbbc36f98a875136ba4f4c3777bb6ff9e1c13 100644 (file)
--- a/src/lua/lua_map.c
+++ b/src/lua/lua_map.c
@@ -481,6 +481,26 @@ lua_config_add_map (lua_State *L)
        return 1;
 }
 
+static const gchar *
+lua_map_process_string_key (lua_State *L, gint pos, gsize *len)
+{
+       struct rspamd_lua_text *t;
+
+       if (lua_type (L, pos) == LUA_TSTRING) {
+               return lua_tolstring (L, pos, len);
+       }
+       else if (lua_type (L, pos) == LUA_TUSERDATA) {
+               t = lua_check_text (L, pos);
+
+               if (t) {
+                       *len = t->len;
+                       return t->start;
+               }
+       }
+
+       return NULL;
+}
+
 /* Radix and hash table functions */
 static gint
 lua_map_get_key (lua_State * L)
@@ -530,14 +550,14 @@ lua_map_get_key (lua_State * L)
                        }
                }
                else if (map->type == RSPAMD_LUA_MAP_SET) {
-                       key = lua_tostring (L, 2);
+                       key = lua_map_process_string_key (L, 2, &len);
 
                        if (key) {
                                ret = g_hash_table_lookup (map->data.hash, key) != NULL;
                        }
                }
                else if (map->type == RSPAMD_LUA_MAP_REGEXP) {
-                       key = lua_tolstring (L, 2, &len);
+                       key = lua_map_process_string_key (L, 2, &len);
 
                        if (key) {
                                value = rspamd_match_regexp_map (map->data.re_map, key, len);
@@ -550,7 +570,7 @@ lua_map_get_key (lua_State * L)
                }
                else if (map->type == RSPAMD_LUA_MAP_HASH) {
                        /* key-value map */
-                       key = lua_tostring (L, 2);
+                       key = lua_map_process_string_key (L, 2, &len);
 
                        if (key) {
                                value = g_hash_table_lookup (map->data.hash, key);