[Project] Implement keep-alive support in lua_http

1 files changed, 58 insertions, 24 deletions

diff --git a/src/lua/lua_http.c b/src/lua/lua_http.c
index a8616e82a..4378aea89 100644
--- a/src/lua/lua_http.c
+++ b/src/lua/lua_http.c
@@ -58,6 +58,7 @@ static const struct luaL_reg httplib_m[] = {
 #define RSPAMD_LUA_HTTP_FLAG_TEXT (1 << 0)
 #define RSPAMD_LUA_HTTP_FLAG_NOVERIFY (1 << 1)
 #define RSPAMD_LUA_HTTP_FLAG_RESOLVED (1 << 2)
+#define RSPAMD_LUA_HTTP_FLAG_KEEP_ALIVE (1 << 3)
 
 struct lua_http_cbdata {
 	struct rspamd_http_connection *conn;
@@ -370,27 +371,29 @@ lua_http_make_connection (struct lua_http_cbdata *cbd)
 	int fd;
 
 	rspamd_inet_address_set_port (cbd->addr, cbd->msg->port);
-	fd = rspamd_inet_address_connect (cbd->addr, SOCK_STREAM, TRUE);
 
-	if (fd == -1) {
-		msg_info ("cannot connect to %V", cbd->msg->host);
-		return FALSE;
-	}
-	cbd->fd = fd;
+	if (cbd->flags & RSPAMD_LUA_HTTP_FLAG_KEEP_ALIVE) {
+		cbd->fd = -1; /* FD is owned by keepalive connection */
 
-	if (cbd->cfg) {
-		cbd->conn = rspamd_http_connection_new (
-				NULL,
-				fd,
+		cbd->conn = rspamd_http_connection_new_keepalive (
+				NULL, /* Default context */
 				NULL,
 				lua_http_error_handler,
 				lua_http_finish_handler,
-				RSPAMD_HTTP_CLIENT_SIMPLE,
-				RSPAMD_HTTP_CLIENT);
+				cbd->addr,
+				cbd->host);
 	}
 	else {
+		fd = rspamd_inet_address_connect (cbd->addr, SOCK_STREAM, TRUE);
+
+		if (fd == -1) {
+			msg_info ("cannot connect to %V", cbd->msg->host);
+			return FALSE;
+		}
+
+		cbd->fd = fd;
 		cbd->conn = rspamd_http_connection_new (
-				NULL,
+				NULL, /* Default context */
 				fd,
 				NULL,
 				lua_http_error_handler,
@@ -554,7 +557,6 @@ lua_http_request (lua_State *L)
 	struct rspamd_cryptobox_keypair *local_kp = NULL;
 	const gchar *url, *lua_body;
 	rspamd_fstring_t *body = NULL;
-	gchar *to_resolve;
 	gint cbref = -1;
 	gsize bodylen;
 	gdouble timeout = default_http_timeout;
@@ -722,6 +724,9 @@ lua_http_request (lua_State *L)
 				body = rspamd_fstring_new_init (t->start, t->len);
 			}
 			else {
+				rspamd_http_message_unref (msg);
+				g_free (mime_type);
+
 				return luaL_error (L, "invalid body argument type: %s",
 						lua_typename (L, lua_type (L, -1)));
 			}
@@ -741,17 +746,28 @@ lua_http_request (lua_State *L)
 						body = rspamd_fstring_append (body, t->start, t->len);
 					}
 					else {
+						rspamd_http_message_unref (msg);
+						if (body) {
+							rspamd_fstring_free (body);
+						}
+
 						return luaL_error (L, "invalid body argument: %s",
 								lua_typename (L, lua_type (L, -1)));
 					}
 				}
 				else {
+					rspamd_http_message_unref (msg);
+					if (body) {
+						rspamd_fstring_free (body);
+					}
+
 					return luaL_error (L, "invalid body argument type: %s",
 							lua_typename (L, lua_type (L, -1)));
 				}
 			}
 		}
 		else if (lua_type (L, -1) != LUA_TNONE && lua_type (L, -1) != LUA_TNIL) {
+			rspamd_http_message_unref (msg);
 			return luaL_error (L, "invalid body argument type: %s",
 					lua_typename (L, lua_type (L, -1)));
 		}
@@ -810,6 +826,15 @@ lua_http_request (lua_State *L)
 
 		lua_pop (L, 1);
 
+		lua_pushstring (L, "keepalive");
+		lua_gettable (L, 1);
+
+		if (!!lua_toboolean (L, -1)) {
+			flags |= RSPAMD_LUA_HTTP_FLAG_KEEP_ALIVE;
+		}
+
+		lua_pop (L, 1);
+
 		lua_pushstring (L, "max_size");
 		lua_gettable (L, 1);
 
@@ -870,9 +895,21 @@ lua_http_request (lua_State *L)
 	if (session && rspamd_session_blocked (session)) {
 		lua_pushboolean (L, FALSE);
 
+		g_free (auth);
+		rspamd_http_message_unref (msg);
+		if (body) {
+			rspamd_fstring_free (body);
+		}
+
 		return 1;
 	}
 	if (task == NULL && cfg == NULL) {
+		g_free (auth);
+		rspamd_http_message_unref (msg);
+		if (body) {
+			rspamd_fstring_free (body);
+		}
+
 		return luaL_error (L,
 				"Bad params to rspamd_http:request(): either task or config should be set");
 	}
@@ -930,28 +967,25 @@ lua_http_request (lua_State *L)
 		}
 	}
 	else {
+		if (!cbd->host) {
+			lua_http_maybe_free (cbd);
+
+			return luaL_error (L, "no host has been specified");
+		}
 		if (task == NULL) {
-			to_resolve = g_malloc (msg->host->len + 1);
-			rspamd_strlcpy (to_resolve, msg->host->str, msg->host->len + 1);
 
 			if (!make_dns_request (resolver, session, NULL, lua_http_dns_handler, cbd,
 					RDNS_REQUEST_A,
-					to_resolve)) {
+					cbd->host)) {
 				lua_http_maybe_free (cbd);
 				lua_pushboolean (L, FALSE);
-				g_free (to_resolve);
 
 				return 1;
 			}
-
-
-			g_free (to_resolve);
 		}
 		else {
-			to_resolve = rspamd_mempool_fstrdup (task->task_pool, msg->host);
-
 			if (!make_dns_request_task_forced (task, lua_http_dns_handler, cbd,
-					RDNS_REQUEST_A, to_resolve)) {
+					RDNS_REQUEST_A, cbd->host)) {
 				lua_http_maybe_free (cbd);
 				lua_pushboolean (L, FALSE);