commit 070192c: [Minor] Lua_cdb: Allow to use more data types as cdb keys/values

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2021-10-20 11:08:05 +0100
URL: https://github.com/rspamd/rspamd/commit/070192c87bd1618812ab4577256a2501e6bf88a5

[Minor] Lua_cdb: Allow to use more data types as cdb keys/values

---
 src/lua/lua_cdb.c | 45 ++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/lua/lua_cdb.c b/src/lua/lua_cdb.c
index e0b9d9439..8fe271dd5 100644
--- a/src/lua/lua_cdb.c
+++ b/src/lua/lua_cdb.c
@@ -122,6 +122,45 @@ lua_check_cdb_builder (lua_State * L, int pos)
 	return ud ? ((struct cdb_make *)ud) : NULL;
 }
 
+static const char *
+lua_cdb_get_input (lua_State *L, int pos, gsize *olen)
+{
+	int t = lua_type (L, pos);
+
+	switch (t) {
+	case LUA_TSTRING:
+		return lua_tolstring (L, pos, olen);
+	case LUA_TNUMBER: {
+		static char numbuf[sizeof(lua_Number)];
+		lua_Number n = lua_tonumber(L, pos);
+		memcpy(numbuf, &n, sizeof(numbuf));
+		*olen = sizeof(n);
+		return numbuf;
+	}
+	case LUA_TUSERDATA: {
+		void *p = rspamd_lua_check_udata_maybe (L, pos, "rspamd{text}");
+		if (p) {
+			struct rspamd_lua_text *t = (struct rspamd_lua_text *)p;
+			*olen = t->len;
+			return t->start;
+		}
+
+		p = rspamd_lua_check_udata_maybe (L, pos, "rspamd{int64}");
+		if (p) {
+			static char numbuf[sizeof(gint64)];
+
+			memcpy(numbuf, p, sizeof(numbuf));
+			*olen = sizeof(numbuf);
+			return numbuf;
+		}
+	}
+	default:
+		break;
+	}
+
+	return NULL;
+}
+
 static gint
 lua_cdb_create (lua_State *L)
 {
@@ -200,7 +239,7 @@ lua_cdb_lookup (lua_State *L)
 {
 	struct cdb *cdb = lua_check_cdb (L, 1);
 	gsize klen;
-	const gchar *what = luaL_checklstring(L, 2, &klen);
+	const gchar *what = lua_cdb_get_input(L, 2, &klen);
 
 	if (!cdb || what == NULL) {
 		return lua_error (L);
@@ -275,8 +314,8 @@ lua_cdb_builder_add (lua_State *L)
 {
 	struct cdb_make *cdbm = lua_check_cdb_builder(L, 1);
 	gsize data_sz, key_sz;
-	const char *key = lua_tolstring (L, 2, &key_sz);
-	const char *data = lua_tolstring (L, 3, &data_sz);
+	const char *key = lua_cdb_get_input (L, 2, &key_sz);
+	const char *data = lua_cdb_get_input (L, 3, &data_sz);
 
 	if (cdbm == NULL || key == NULL || data == NULL || cdbm->cdb_fd == -1) {
 		return luaL_error(L, "invalid arguments");