Lua dataset/v6

src/datasets.c

@@ -861,13 +861,15 @@ int DatasetsInit(void)
     uint32_t default_hashsize = 0;
     GetDefaultMemcap(&default_memcap, &default_hashsize);
     if (datasets != NULL) {
+        SCLogDebug("datasets %p", datasets);
         int list_pos = 0;
         ConfNode *iter = NULL;
         TAILQ_FOREACH(iter, &datasets->head, next) {
             if (iter->name == NULL) {
                 list_pos++;
                 continue;
             }
+            SCLogDebug("datasets %p name %s", datasets, iter->name);
 
             char save[PATH_MAX] = "";
             char load[PATH_MAX] = "";

src/detect-lua-extensions.c

@@ -28,6 +28,7 @@
 
 #include "threads.h"
 #include "decode.h"
+#include "datasets.h"
 
 #include "detect.h"
 #include "detect-parse.h"
@@ -414,6 +415,90 @@ static int LuaSetFlowint(lua_State *luastate)
     return 0;
 }
 
+struct LuaDataset {
+    Dataset *set;
+};
+
+static int LuaDatasetGC(lua_State *luastate)
+{
+    SCLogDebug("GC running");
+    struct LuaDataset *s = (struct LuaDataset *)lua_touserdata(luastate, 1);
+    SCLogDebug("deref %s", s->set->name);
+    s->set = NULL;
+    return 0;
+}
+
+static int LuaDatasetGetRef(lua_State *luastate)
+{
+    if (!lua_isstring(luastate, 1)) {
+        LUA_ERROR("1st arg is not a string");
+    }
+
+    const char *name = lua_tostring(luastate, 1);
+    if (name == NULL) {
+        LUA_ERROR("null string");
+    }
+
+    Dataset *dataset = DatasetFind(name, DATASET_TYPE_STRING);
+    if (dataset == NULL) {
+        LUA_ERROR("dataset not found");
+    }
+
+    struct LuaDataset *s = (struct LuaDataset *)lua_newuserdata(luastate, sizeof(*s));
+    if (s == NULL) {
+        LUA_ERROR("failed to get userdata");
+    }
+    s->set = dataset;
+
+    // TODO not sure why we need to do this here. Would expect it to be
+    // a registration time thing.
+    lua_newtable(luastate);
+    lua_pushcfunction(luastate, LuaDatasetGC);
+    lua_setfield(luastate, -2, "__gc");
+    lua_setmetatable(luastate, -2);
+    return 1;
+}
+
+static int LuaDatasetAdd(lua_State *luastate)
+{
+    struct LuaDataset *s = (struct LuaDataset *)lua_touserdata(luastate, 1);
+    if (s == NULL) {
+        LUA_ERROR("dataset is not initialized");
+    }
+    if (!lua_isstring(luastate, 2)) {
+        LUA_ERROR("1st arg is not a string");
+    }
+    if (!lua_isnumber(luastate, 3)) {
+        LUA_ERROR("2nd arg is not a number");
+    }
+
+    const uint8_t *str = (const uint8_t *)lua_tostring(luastate, 2);
+    if (str == NULL) {
+        LUA_ERROR("1st arg is not null string");
+    }
+
+    uint32_t str_len = lua_tonumber(luastate, 3);
+
+    int r = DatasetAdd(s->set, (const uint8_t *)str, str_len);
+    /* return value through luastate, as a luanumber */
+    lua_pushnumber(luastate, (lua_Number)r);
+    return 1;
+}
+
+// clang-format off
+static const struct luaL_Reg datasetlib [] = {
+    { "get_ref", LuaDatasetGetRef },
+    { "add", LuaDatasetAdd },
+    { NULL, NULL }
+};
+// clang-format on
+
+static void LuaDatasetRegister(lua_State *luastate)
+{
+    luaL_newlib(luastate, datasetlib);
+    lua_setglobal(luastate, "dataset");
+}
+
 static int LuaIncrFlowint(lua_State *luastate)
 {
     uint32_t idx;
@@ -578,6 +663,8 @@ int LuaRegisterExtensions(lua_State *lua_state)
     lua_pushcfunction(lua_state, LuaGetByteVar);
     lua_setglobal(lua_state, "SCByteVarGet");
 
+    LuaDatasetRegister(lua_state);
+
     LuaRegisterFunctions(lua_state);
     LuaRegisterHttpFunctions(lua_state);
     LuaRegisterDnsFunctions(lua_state);

src/detect-lua.c

@@ -131,18 +131,18 @@ void DetectLuaRegister(void)
 
 #if 0
 /** \brief dump stack from lua state to screen */
-void LuaDumpStack(lua_State *state)
+static void LuaDumpStack(lua_State *state, const char *prefix)
 {
     int size = lua_gettop(state);
-    int i;
+    printf("%s: size %d\n", prefix, size);
 
-    for (i = 1; i <= size; i++) {
+    for (int i = 1; i <= size; i++) {
         int type = lua_type(state, i);
-        printf("Stack size=%d, level=%d, type=%d, ", size, i, type);
+        printf("- %s: Stack size=%d, level=%d, type=%d, ", prefix, size, i, type);
 
         switch (type) {
             case LUA_TFUNCTION:
-                printf("function %s", lua_tostring(state, i) ? "true" : "false");
+                printf("function %s", lua_tostring(state, i));
                 break;
             case LUA_TBOOLEAN:
                 printf("bool %s", lua_toboolean(state, i) ? "true" : "false");
@@ -525,6 +525,18 @@ static void *DetectLuaThreadInit(void *data)
         goto error;
     }
 
+    /* thread_init call */
+    lua_getglobal(t->luastate, "thread_init");
+    if (lua_isfunction(t->luastate, -1)) {
+        if (lua_pcall(t->luastate, 0, 0, 0) != 0) {
+            SCLogError("couldn't run script 'thread_init' function: %s",
+                    lua_tostring(t->luastate, -1));
+            goto error;
+        }
+    } else {
+        lua_pop(t->luastate, 1);
+    }
+
     return (void *)t;
 
 error: