From b0dc1504eb3f788c698001a36b320bdd41a5d287 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Sat, 29 Jun 2019 17:24:54 +0100 Subject: [PATCH] [Project] Add kann load/save methods --- src/lua/lua_common.h | 1 + src/lua/lua_kann.c | 181 ++++++++++++++++++++++++++++++++++++++++++- src/lua/lua_text.c | 7 +- 3 files changed, 187 insertions(+), 2 deletions(-) diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h index 4c0605782..32b17a2fc 100644 --- a/src/lua/lua_common.h +++ b/src/lua/lua_common.h @@ -68,6 +68,7 @@ struct rspamd_lua_ip { #define RSPAMD_TEXT_FLAG_OWN (1u << 0u) #define RSPAMD_TEXT_FLAG_MMAPED (1u << 1u) #define RSPAMD_TEXT_FLAG_WIPE (1u << 2u) +#define RSPAMD_TEXT_FLAG_SYSMALLOC (1u << 3u) struct rspamd_lua_text { const gchar *start; guint len; diff --git a/src/lua/lua_kann.c b/src/lua/lua_kann.c index 836b2968b..3d50cc587 100644 --- a/src/lua/lua_kann.c +++ b/src/lua/lua_kann.c @@ -140,6 +140,20 @@ static luaL_reg rspamd_kann_new_f[] = { {NULL, NULL}, }; +LUA_FUNCTION_DEF (kann, load); +LUA_FUNCTION_DEF (kann, destroy); +LUA_FUNCTION_DEF (kann, save); +LUA_FUNCTION_DEF (kann, train); +LUA_FUNCTION_DEF (kann, forward); + +static luaL_reg rspamd_kann_m[] = { + LUA_INTERFACE_DEF (kann, save), + LUA_INTERFACE_DEF (kann, train), + LUA_INTERFACE_DEF (kann, forward), + {"__gc", lua_kann_destroy}, + {NULL, NULL}, +}; + static int rspamd_kann_table_to_flags (lua_State *L, int table_pos) { @@ -227,6 +241,11 @@ lua_load_kann (lua_State * L) luaL_register (L, NULL, rspamd_kann_new_f); lua_settable (L, -3); + /* Load ann from memory or file */ + lua_pushstring (L, "load"); + lua_pushcfunction (L, lua_kann_load); + lua_settable (L, -3); + return 1; } @@ -238,12 +257,20 @@ lua_check_kann_node (lua_State *L, int pos) return ud ? *((kad_node_t **)ud) : NULL; } +static kann_t * +lua_check_kann (lua_State *L, int pos) +{ + void *ud = rspamd_lua_check_udata (L, pos, KANN_NETWORK_CLASS); + luaL_argcheck (L, ud != NULL, pos, "'kann' expected"); + return ud ? *((kann_t **)ud) : NULL; +} + void luaopen_kann (lua_State *L) { /* Metatables */ rspamd_lua_new_class (L, KANN_NODE_CLASS, NULL); /* TODO: add methods */ lua_pop (L, 1); /* No need in metatable... */ - rspamd_lua_new_class (L, KANN_NETWORK_CLASS, NULL); /* TODO: add methods */ + rspamd_lua_new_class (L, KANN_NETWORK_CLASS, rspamd_kann_m); lua_pop (L, 1); /* No need in metatable... */ rspamd_lua_add_preload (L, "rspamd_kann", lua_load_kann); lua_settop (L, 0); @@ -820,3 +847,155 @@ lua_kann_new_kann (lua_State *L) return 1; } + +static int +lua_kann_destroy (lua_State *L) +{ + kann_t *k = lua_check_kann (L, 1); + + kann_delete (k); + + return 0; +} + +static int +lua_kann_save (lua_State *L) +{ + kann_t *k = lua_check_kann (L, 1); + + if (k) { + if (lua_istable (L, 2)) { + lua_getfield (L, 2, "filename"); + + if (lua_isstring (L, -1)) { + const gchar *fname = lua_tostring (L, -1); + FILE *f; + + f = fopen (fname, "w"); + + if (!f) { + lua_pop (L, 1); + + return luaL_error (L, "cannot open %s for writing: %s", + fname, strerror (errno)); + } + + kann_save_fp (f, k); + fclose (f); + + lua_pushboolean (L, true); + } + else { + lua_pop (L, 1); + + return luaL_error (L, "invalid arguments: missing filename"); + } + + lua_pop (L, 1); + } + else { + /* Save to Rspamd text */ +#ifndef HAVE_OPENMEMSTREAM + return luaL_error (L, "no support of saving to memory on your system"); +#endif + FILE *f; + char *buf = NULL; + size_t buflen; + struct rspamd_lua_text *t; + + f = open_memstream (&buf, &buflen); + g_assert (f != NULL); + + kann_save_fp (f, k); + fclose (f); + + t = lua_newuserdata (L, sizeof (*t)); + rspamd_lua_setclass (L, "rspamd{text}", -1); + t->flags = RSPAMD_TEXT_FLAG_OWN; + t->start = (const gchar *)buf; + t->len = buflen; + } + } + else { + return luaL_error (L, "invalid arguments"); + } + + return 1; +} + +static int +lua_kann_load (lua_State *L) +{ + kann_t *k; + FILE *f = NULL; + + if (lua_istable (L, 1)) { + lua_getfield (L, 2, "filename"); + + if (lua_isstring (L, -1)) { + const gchar *fname = lua_tostring (L, -1); + + f = fopen (fname, "rb"); + } + else { + lua_pop (L, 1); + + return luaL_error (L, "invalid arguments: missing filename"); + } + + lua_pop (L, 1); + } + else if (lua_isstring (L, 1)) { + gsize dlen; + const gchar *data; + + data = lua_tolstring (L, 1, &dlen); + +#ifndef HAVE_FMEMOPEN + return luaL_error (L, "no support of loading from memory on your system"); +#endif + f = fmemopen ((void *)data, dlen, "rb"); + } + else if (lua_isuserdata (L, 1)) { + struct rspamd_lua_text *t; + + t = lua_check_text (L, 1); + +#ifndef HAVE_FMEMOPEN + return luaL_error (L, "no support of loading from memory on your system"); +#endif + f = fmemopen ((void *)t->start, t->len, "rb"); + } + + if (f == NULL) { + return luaL_error (L, "invalid arguments or cannot open file"); + } + + k = kann_load_fp (f); + fclose (f); + + if (k == NULL) { + lua_pushnil (L); + } + else { + PUSH_KAN_NETWORK (k); + } + + return 1; +} + +static int +lua_kann_train (lua_State *L) +{ + kann_t *k = lua_check_kann (L, 1); + + g_assert_not_reached (); /* TODO: implement */ +} + +static int +lua_kann_forward (lua_State *L) +{ + kann_t *k = lua_check_kann (L, 1); + + g_assert_not_reached (); /* TODO: implement */ +} \ No newline at end of file diff --git a/src/lua/lua_text.c b/src/lua/lua_text.c index e130490ef..ab52bd32d 100644 --- a/src/lua/lua_text.c +++ b/src/lua/lua_text.c @@ -350,7 +350,12 @@ lua_text_gc (lua_State *L) munmap ((gpointer)t->start, t->len); } else { - g_free ((gpointer)t->start); + if (t->flags & RSPAMD_TEXT_FLAG_SYSMALLOC) { + free ((gpointer) t->start); + } + else { + g_free ((gpointer) t->start); + } } } -- 2.39.5