[Feature] Lua_task: Allow to load data into the existing task

1 files changed, 72 insertions, 21 deletions

diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index b368ad4e6..355680881 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -1233,6 +1233,8 @@ static const struct luaL_reg tasklib_f[] = {
 	{NULL, NULL}};
 
 static const struct luaL_reg tasklib_m[] = {
+	LUA_INTERFACE_DEF(task, load_from_file),
+	LUA_INTERFACE_DEF(task, load_from_string),
 	LUA_INTERFACE_DEF(task, get_message),
 	LUA_INTERFACE_DEF(task, set_message),
 	LUA_INTERFACE_DEF(task, destroy),
@@ -1724,20 +1726,32 @@ lua_task_load_from_file(lua_State *L)
 {
 	LUA_TRACE_POINT;
 	struct rspamd_task *task = NULL, **ptask;
-	const char *fname = luaL_checkstring(L, 1), *err = NULL;
+	const char *fname, *err = NULL;
 	struct rspamd_config *cfg = NULL;
-	gboolean res = FALSE;
+	gboolean res = FALSE, new_task = FALSE;
 	gpointer map;
 	gsize sz;
 
+	if (lua_type(L, 1) == LUA_TSTRING) {
+		fname = luaL_checkstring(L, 1);
+		new_task = TRUE;
+	}
+	else {
+		/* Method */
+		task = lua_check_task(L, 1);
+		fname = luaL_checkstring(L, 2);
+	}
+
 	if (fname) {
 
-		if (lua_type(L, 2) == LUA_TUSERDATA) {
-			gpointer p;
-			p = rspamd_lua_check_udata_maybe(L, 2, rspamd_config_classname);
+		if (!task) {
+			if (lua_type(L, 2) == LUA_TUSERDATA) {
+				gpointer p;
+				p = rspamd_lua_check_udata_maybe(L, 2, rspamd_config_classname);
 
-			if (p) {
-				cfg = *(struct rspamd_config **) p;
+				if (p) {
+					cfg = *(struct rspamd_config **) p;
+				}
 			}
 		}
 
@@ -1763,11 +1777,17 @@ lua_task_load_from_file(lua_State *L)
 				}
 			}
 
-			task = rspamd_task_new(NULL, cfg, NULL, NULL, NULL, FALSE);
+			if (!task) {
+				task = rspamd_task_new(NULL, cfg, NULL, NULL, NULL, FALSE);
+			}
+
 			task->msg.begin = data->str;
 			task->msg.len = data->len;
 			rspamd_mempool_add_destructor(task->task_pool,
 										  lua_task_free_dtor, data->str);
+			if (data->len > 0) {
+				task->flags &= ~RSPAMD_TASK_FLAG_EMPTY;
+			}
 			res = TRUE;
 			g_string_free(data, FALSE); /* Buffer is still valid */
 		}
@@ -1778,9 +1798,16 @@ lua_task_load_from_file(lua_State *L)
 				err = strerror(errno);
 			}
 			else {
-				task = rspamd_task_new(NULL, cfg, NULL, NULL, NULL, FALSE);
+				if (!task) {
+					task = rspamd_task_new(NULL, cfg, NULL, NULL, NULL, FALSE);
+				}
+
 				task->msg.begin = map;
 				task->msg.len = sz;
+
+				if (sz > 0) {
+					task->flags &= ~RSPAMD_TASK_FLAG_EMPTY;
+				}
 				rspamd_mempool_add_destructor(task->task_pool,
 											  lua_task_unmap_dtor, task);
 				res = TRUE;
@@ -1793,21 +1820,26 @@ lua_task_load_from_file(lua_State *L)
 
 	lua_pushboolean(L, res);
 
-	if (res) {
+	if (res && new_task) {
 		ptask = lua_newuserdata(L, sizeof(*ptask));
 		*ptask = task;
 		rspamd_lua_setclass(L, rspamd_task_classname, -1);
+
+		return 2;
 	}
-	else {
+	else if (!res) {
 		if (err) {
 			lua_pushstring(L, err);
 		}
 		else {
 			lua_pushnil(L);
 		}
+		return 2;
+	}
+	else {
+		/* No new task */
+		return 1;
 	}
-
-	return 2;
 }
 
 static int
@@ -1816,14 +1848,23 @@ lua_task_load_from_string(lua_State *L)
 	LUA_TRACE_POINT;
 	struct rspamd_task *task = NULL, **ptask;
 	const char *str_message;
-	gsize message_len;
+	gsize message_len = 0;
 	struct rspamd_config *cfg = NULL;
+	bool new_task = false;
 
-	str_message = luaL_checklstring(L, 1, &message_len);
+	if (lua_type(L, 1) == LUA_TSTRING) {
+		str_message = luaL_checklstring(L, 1, &message_len);
+		new_task = true;
+	}
+	else {
+		/* Method */
+		task = lua_check_task(L, 1);
+		str_message = luaL_checklstring(L, 2, &message_len);
+	}
 
 	if (str_message) {
 
-		if (lua_type(L, 2) == LUA_TUSERDATA) {
+		if (!task && lua_type(L, 2) == LUA_TUSERDATA) {
 			gpointer p;
 			p = rspamd_lua_check_udata_maybe(L, 2, rspamd_config_classname);
 
@@ -1832,10 +1873,15 @@ lua_task_load_from_string(lua_State *L)
 			}
 		}
 
-		task = rspamd_task_new(NULL, cfg, NULL, NULL, NULL, FALSE);
+		if (!task) {
+			task = rspamd_task_new(NULL, cfg, NULL, NULL, NULL, FALSE);
+		}
 		task->msg.begin = g_malloc(message_len);
 		memcpy((char *) task->msg.begin, str_message, message_len);
 		task->msg.len = message_len;
+		if (message_len > 0) {
+			task->flags &= ~RSPAMD_TASK_FLAG_EMPTY;
+		}
 		rspamd_mempool_add_destructor(task->task_pool, lua_task_free_dtor,
 									  (gpointer) task->msg.begin);
 	}
@@ -1845,11 +1891,16 @@ lua_task_load_from_string(lua_State *L)
 
 	lua_pushboolean(L, true);
 
-	ptask = lua_newuserdata(L, sizeof(*ptask));
-	*ptask = task;
-	rspamd_lua_setclass(L, rspamd_task_classname, -1);
+	if (new_task) {
+		ptask = lua_newuserdata(L, sizeof(*ptask));
+		*ptask = task;
+		rspamd_lua_setclass(L, rspamd_task_classname, -1);
 
-	return 2;
+		return 2;
+	}
+	else {
+		return 1;
+	}
 }
 
 static int