commit 4528be0: [Minor] Lua_task: Accept lua_ip userdata when calling set_ip

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2019-11-30 17:21:45 +0000
URL: https://github.com/rspamd/rspamd/commit/4528be09dbab196bed4626e16e33d9e7a5d2dd5d

[Minor] Lua_task: Accept lua_ip userdata when calling set_ip

---
 src/lua/lua_task.c | 51 +++++++++++++++++++++++++++++++++++----------------
 1 file changed, 35 insertions(+), 16 deletions(-)

diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 0ea1367d5..aff6d8499 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -3938,28 +3938,47 @@ lua_task_set_from_ip (lua_State *L)
 {
 	LUA_TRACE_POINT;
 	struct rspamd_task *task = lua_check_task (L, 1);
-	gsize len;
-	const gchar *ip_str = luaL_checklstring (L, 2, &len);
 	rspamd_inet_addr_t *addr = NULL;
 
-	if (!task || !ip_str) {
-		lua_pushstring (L, "invalid parameters");
-		return lua_error (L);
+	if (!task) {
+		return luaL_error (L, "no task");
 	}
 	else {
-		if (!rspamd_parse_inet_address (&addr,
-				ip_str,
-				len,
-				RSPAMD_INET_ADDRESS_PARSE_DEFAULT)) {
-			msg_warn_task ("cannot get IP from received header: '%s'",
-					ip_str);
-		}
-		else {
-			if (task->from_addr) {
-				rspamd_inet_address_free (task->from_addr);
+		if (lua_type (L, 2) == LUA_TSTRING) {
+			gsize len;
+			const gchar *ip_str = lua_tolstring (L, 2, &len);
+
+			if (!rspamd_parse_inet_address (&addr,
+					ip_str,
+					len,
+					RSPAMD_INET_ADDRESS_PARSE_DEFAULT)) {
+				return luaL_error (L, "invalid IP string: %s", ip_str);
 			}
+			else {
+				if (task->from_addr) {
+					rspamd_inet_address_free (task->from_addr);
+				}
 
-			task->from_addr = addr;
+				task->from_addr = addr;
+			}
+		}
+		else if (lua_type (L, 2) == LUA_TUSERDATA) {
+			struct rspamd_lua_ip *ip = lua_check_ip (L, 2);
+
+			if (ip && ip->addr) {
+				if (task->from_addr) {
+					rspamd_inet_address_free (task->from_addr);
+				}
+
+				task->from_addr = rspamd_inet_address_copy (ip->addr);
+			}
+			else {
+				return luaL_error (L, "invalid IP object");
+			}
+		}
+		else {
+			return luaL_error (L, "invalid IP argument type: %s", lua_typename (L,
+					lua_type (L, 2)));
 		}
 	}