[Fix] Fix milter connections persistence

3 files changed, 100 insertions, 36 deletions

diff --git a/src/libserver/milter.c b/src/libserver/milter.c
index fab8709ba..b733956ab 100644
--- a/src/libserver/milter.c
+++ b/src/libserver/milter.c
@@ -584,6 +584,8 @@ rspamd_milter_process_command (struct rspamd_milter_session *session,
 		REF_RETAIN (session);
 		priv->fin_cb (priv->fd, session, priv->ud);
 		REF_RELEASE (session);
+
+		return FALSE;
 		break;
 	case RSPAMD_MILTER_CMD_RCPT:
 		msg_debug_milter ("rcpt command");
@@ -1182,6 +1184,7 @@ rspamd_milter_to_http (struct rspamd_milter_session *session)
 
 	if (session->message) {
 		rspamd_http_message_set_body_from_fstring_steal (msg, session->message);
+		session->message = NULL;
 	}
 
 	if (session->hostname && session->hostname->len > 0) {
@@ -1209,4 +1212,17 @@ rspamd_milter_to_http (struct rspamd_milter_session *session)
 	rspamd_milter_macro_http (session, msg);
 
 	return msg;
+}
+
+void *
+rspamd_milter_update_userdata (struct rspamd_milter_session *session,
+		void *ud)
+{
+	struct rspamd_milter_private *priv = session->priv;
+	void *prev_ud;
+
+	prev_ud = priv->ud;
+	priv->ud = ud;
+
+	return prev_ud;
 }
\ No newline at end of file
diff --git a/src/libserver/milter.h b/src/libserver/milter.h
index 963cc182c..3786f3b7b 100644
--- a/src/libserver/milter.h
+++ b/src/libserver/milter.h
@@ -71,6 +71,15 @@ gboolean rspamd_milter_handle_socket (gint fd, const struct timeval *tv,
 		rspamd_milter_error error_cb, void *ud);
 
 /**
+ * Updates userdata for a session, returns previous userdata
+ * @param session
+ * @param ud
+ * @return
+ */
+void * rspamd_milter_update_userdata (struct rspamd_milter_session *session,
+		void *ud);
+
+/**
  * Sets SMTP reply string
  * @param session
  * @param xcode
diff --git a/src/rspamd_proxy.c b/src/rspamd_proxy.c
index 5728419b9..5edd682e5 100644
--- a/src/rspamd_proxy.c
+++ b/src/rspamd_proxy.c
@@ -846,21 +846,17 @@ proxy_session_dtor (struct rspamd_proxy_session *session)
 	if (session->master_conn) {
 		proxy_backend_close_connection (session->master_conn);
 	}
-	else if (session->client_milter_conn) {
-		rspamd_milter_session_unref (session->client_milter_conn);
-	}
 
-	if (session->map && session->map_len) {
-		munmap (session->map, session->map_len);
+	if (session->client_milter_conn) {
+		rspamd_milter_session_unref (session->client_milter_conn);
 	}
-
-	if (session->client_conn) {
+	else if (session->client_conn) {
 		rspamd_http_connection_reset (session->client_conn);
 		rspamd_http_connection_unref (session->client_conn);
 	}
 
-	if (session->client_milter_conn) {
-		rspamd_milter_session_unref (session->client_milter_conn);
+	if (session->map && session->map_len) {
+		munmap (session->map, session->map_len);
 	}
 
 	for (i = 0; i < session->mirror_conns->len; i ++) {
@@ -888,12 +884,46 @@ proxy_session_dtor (struct rspamd_proxy_session *session)
 	g_ptr_array_free (session->mirror_conns, TRUE);
 	rspamd_http_message_shmem_unref (session->shmem_ref);
 	rspamd_http_message_unref (session->client_message);
-	rspamd_inet_address_free (session->client_addr);
-	close (session->client_sock);
-	rspamd_mempool_delete (session->pool);
+
+	if (session->client_addr) {
+		rspamd_inet_address_free (session->client_addr);
+	}
+
+	if (session->client_sock != -1) {
+		close (session->client_sock);
+	}
+
+	if (session->pool) {
+		rspamd_mempool_delete (session->pool);
+	}
+
 	g_slice_free1 (sizeof (*session), session);
 }
 
+static struct rspamd_proxy_session *
+proxy_session_refresh (struct rspamd_proxy_session *session)
+{
+	struct rspamd_proxy_session *nsession;
+
+	nsession = g_slice_alloc0 (sizeof (*nsession));
+	nsession->client_milter_conn = session->client_milter_conn;
+	session->client_milter_conn = NULL;
+	rspamd_milter_update_userdata (nsession->client_milter_conn,
+			nsession);
+	nsession->client_addr = session->client_addr;
+	session->client_addr = NULL;
+	nsession->ctx = session->ctx;
+	nsession->worker = session->worker;
+	nsession->pool = rspamd_mempool_new (rspamd_mempool_suggest_size (), "proxy");
+	nsession->client_sock = session->client_sock;
+	session->client_sock = -1;
+	nsession->mirror_conns = g_ptr_array_sized_new (nsession->ctx->mirrors->len);
+
+	REF_INIT_RETAIN (nsession, proxy_session_dtor);
+
+	return nsession;
+}
+
 static gboolean
 proxy_check_file (struct rspamd_http_message *msg,
 		struct rspamd_proxy_session *session)
@@ -1149,6 +1179,7 @@ proxy_client_write_error (struct rspamd_proxy_session *session, gint code,
 	if (session->client_milter_conn) {
 		rspamd_milter_send_action (session->client_milter_conn,
 				RSPAMD_MILTER_TEMPFAIL);
+		REF_RELEASE (session);
 	}
 	else {
 		reply = rspamd_http_new_message (HTTP_RESPONSE);
@@ -1202,7 +1233,7 @@ proxy_backend_master_finish_handler (struct rspamd_http_connection *conn,
 	struct rspamd_http_message *msg)
 {
 	struct rspamd_proxy_backend_connection *bk_conn = conn->ud;
-	struct rspamd_proxy_session *session;
+	struct rspamd_proxy_session *session, *nsession;
 	rspamd_fstring_t *reply;
 
 	session = bk_conn->s;
@@ -1238,8 +1269,10 @@ proxy_backend_master_finish_handler (struct rspamd_http_connection *conn,
 		/*
 		 * TODO: convert reply to milter reply
 		 */
-		rspamd_milter_send_action (session->client_milter_conn,
+		nsession = proxy_session_refresh (session);
+		rspamd_milter_send_action (nsession->client_milter_conn,
 				RSPAMD_MILTER_ACCEPT);
+		REF_RELEASE (session);
 	}
 	else {
 		rspamd_http_connection_write_message (session->client_conn,
@@ -1254,7 +1287,7 @@ static void
 rspamd_proxy_scan_self_reply (struct rspamd_task *task)
 {
 	struct rspamd_http_message *msg;
-	struct rspamd_proxy_session *session = task->fin_arg;
+	struct rspamd_proxy_session *session = task->fin_arg, *nsession;
 	ucl_object_t *rep;
 	const char *ctype = "application/json";
 
@@ -1289,8 +1322,10 @@ rspamd_proxy_scan_self_reply (struct rspamd_task *task)
 		/*
 		 * TODO: convert reply to milter reply
 		 */
-		rspamd_milter_send_action (session->client_milter_conn,
+		nsession = proxy_session_refresh (session);
+		rspamd_milter_send_action (nsession->client_milter_conn,
 				RSPAMD_MILTER_ACCEPT);
+		REF_RELEASE (session);
 	}
 	else {
 		rspamd_http_connection_reset (session->client_conn);
@@ -1499,9 +1534,16 @@ retry:
 	return TRUE;
 
 err:
-	rspamd_http_connection_steal_msg (session->client_conn);
-	rspamd_http_connection_reset (session->client_conn);
-	proxy_client_write_error (session, 404, "Backend not found");
+	if (session->client_milter_conn) {
+		rspamd_milter_send_action (session->client_milter_conn,
+				RSPAMD_MILTER_TEMPFAIL);
+		REF_RELEASE (session);
+	}
+	else {
+		rspamd_http_connection_steal_msg (session->client_conn);
+		rspamd_http_connection_reset (session->client_conn);
+		proxy_client_write_error (session, 404, "Backend not found");
+	}
 
 	return FALSE;
 }
@@ -1582,29 +1624,26 @@ proxy_milter_finish_handler (gint fd,
 	struct rspamd_proxy_session *session = ud;
 	struct rspamd_http_message *msg;
 
-	if (!session->master_conn) {
+	if (rms->message == NULL || rms->message->len == 0) {
+		msg_info_session ("finished milter connection");
+		proxy_backend_close_connection (session->master_conn);
+		REF_RELEASE (session);
+	}
+	else {
+		if (!session->master_conn) {
+			session->master_conn = rspamd_mempool_alloc0 (session->pool,
+					sizeof (*session->master_conn));
+		}
+
 		session->client_milter_conn = rms;
 		msg = rspamd_milter_to_http (rms);
-		session->master_conn = rspamd_mempool_alloc0 (session->pool,
-				sizeof (*session->master_conn));
+
 		session->master_conn->s = session;
 		session->master_conn->name = "master";
 		session->client_message = msg;
 
-		if (msg->body_buf.len == 0) {
-			msg_info_session ("incomplete master connection");
-			proxy_backend_close_connection (session->master_conn);
-			REF_RELEASE (session);
-		}
-		else {
-			proxy_open_mirror_connections (session);
-			proxy_send_master_message (session);
-		}
-	}
-	else {
-		msg_info_session ("finished master connection");
-		proxy_backend_close_connection (session->master_conn);
-		REF_RELEASE (session);
+		proxy_open_mirror_connections (session);
+		proxy_send_master_message (session);
 	}
 }