commit 6a16924: [Minor] Lua_tensor: Add deserialisation

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2020-08-06 14:43:18 +0100
URL: https://github.com/rspamd/rspamd/commit/6a1692499fd98e12824c29fa8147ecb050380e6a (HEAD -> master)

[Minor] Lua_tensor: Add deserialisation

---
 src/lua/lua_tensor.c | 56 ++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 54 insertions(+), 2 deletions(-)

diff --git a/src/lua/lua_tensor.c b/src/lua/lua_tensor.c
index 9b85779d7..91fcd763e 100644
--- a/src/lua/lua_tensor.c
+++ b/src/lua/lua_tensor.c
@@ -396,10 +396,62 @@ lua_tensor_mul (lua_State *L)
 static gint
 lua_tensor_load (lua_State *L)
 {
-	struct rspamd_lua_tensor *t = lua_check_tensor (L, 1);
+	const guchar *data;
+	gsize sz;
 
-	if (t) {
+	if (lua_type (L, 1) == LUA_TUSERDATA) {
+		struct rspamd_lua_text *t = lua_check_text (L, 1);
+
+		if (!t) {
+			return luaL_error (L, "invalid argument");
+		}
+
+		data = (const guchar *)t->start;
+		sz = t->len;
+	}
+	else {
+		data = (const guchar *)lua_tolstring (L, 1, &sz);
+	}
+
+	if (sz >= sizeof (gint) * 4) {
+		int ndims, nelts, dims[2];
 
+		memcpy (&ndims, data, sizeof (int));
+		memcpy (&nelts, data + sizeof (int), sizeof (int));
+		memcpy (dims, data + sizeof (int) * 2, sizeof (int) * 2);
+
+		if (sz == nelts * sizeof (rspamd_tensor_num_t) + sizeof (int) * 4) {
+			if (ndims == 1) {
+				if (nelts == dims[0]) {
+					struct rspamd_lua_tensor *t = lua_newtensor (L, ndims, dims, false);
+					memcpy (t->data, data + sizeof (int) * 4, nelts *
+							sizeof (rspamd_tensor_num_t));
+				}
+				else {
+					return luaL_error (L, "invalid argument: bad dims: %d x %d != %d",
+							dims[0], 1, nelts);
+				}
+			}
+			else if (ndims == 2) {
+				if (nelts == dims[0] * dims[1]) {
+					struct rspamd_lua_tensor *t = lua_newtensor (L, ndims, dims, false);
+					memcpy (t->data, data + sizeof (int) * 4, nelts *
+							sizeof (rspamd_tensor_num_t));
+				}
+				else {
+					return luaL_error (L, "invalid argument: bad dims: %d x %d != %d",
+							dims[0], dims[1], nelts);
+				}
+			}
+			else {
+				return luaL_error (L, "invalid argument: bad ndims: %d", ndims);
+			}
+		}
+		else {
+			return luaL_error (L, "invalid size: %d, %d required, %d elts", (int)sz,
+					(int)(nelts * sizeof (rspamd_tensor_num_t) + sizeof (int) * 4),
+					nelts);
+		}
 	}
 	else {
 		return luaL_error (L, "invalid arguments");