From 0e8042b2568401df26a056d8d46e72e1cfb76d70 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Fri, 9 Oct 2020 12:31:24 +0100
Subject: [PATCH] [Feature] Support ping in milter mode

Requested by: @andryyy
---
 src/libserver/milter.c          | 88 ++++++++++++++++++++++++++++++++-
 src/libserver/milter_internal.h |  1 +
 2 files changed, 87 insertions(+), 2 deletions(-)

diff --git a/src/libserver/milter.c b/src/libserver/milter.c
index 8cb7ed2f0..4fb5fda44 100644
--- a/src/libserver/milter.c
+++ b/src/libserver/milter.c
@@ -236,6 +236,48 @@ rspamd_milter_on_protocol_error (struct rspamd_milter_session *session,
 	priv->err_cb (priv->fd, session, priv->ud, err);
 	REF_RELEASE (session);
 	g_error_free (err);
+
+	rspamd_milter_plan_io (session, priv, EV_WRITE);
+}
+
+static void
+rspamd_milter_on_protocol_ping (struct rspamd_milter_session *session,
+								 struct rspamd_milter_private *priv)
+{
+	GError *err = NULL;
+	static const gchar reply[] = "HTTP/1.1 200 OK\r\n"
+								 "Connection: close\r\n"
+								 "Server: rspamd/2.7 (milter mode)\r\n"
+								 "Content-Length: 6\r\n"
+								 "Content-Type: text/plain\r\n"
+								 "\r\n"
+								 "pong\r\n";
+
+	if (write (priv->fd, reply, sizeof (reply)) == -1) {
+		gint serrno = errno;
+		msg_err_milter ("cannot write pong reply: %s", strerror (serrno));
+		g_set_error (&err, rspamd_milter_quark (), serrno, "ping command IO error: %s",
+				strerror (serrno));
+		priv->state = RSPAMD_MILTER_WANNA_DIE;
+		REF_RETAIN (session);
+		priv->err_cb (priv->fd, session, priv->ud, err);
+		REF_RELEASE (session);
+		g_error_free (err);
+	}
+	else {
+		priv->state = RSPAMD_MILTER_PONG_AND_DIE;
+		rspamd_milter_plan_io (session, priv, EV_WRITE);
+	}
+}
+
+static gint
+rspamd_milter_http_on_url (http_parser * parser, const gchar *at, size_t length)
+{
+	GString *url = (GString *)parser->data;
+
+	g_string_append_len (url, at, length);
+
+	return 0;
 }
 
 static void
@@ -851,8 +893,40 @@ rspamd_milter_consume_input (struct rspamd_milter_session *session,
 				/* Check if we have HTTP input instead of milter */
 				if (priv->parser.buf->len > sizeof ("GET") &&
 						memcmp (priv->parser.buf->str, "GET", 3) == 0) {
-					err = g_error_new (rspamd_milter_quark (), EINVAL,
-							"HTTP GET request is not supported in milter mode");
+					struct http_parser http_parser;
+					struct http_parser_settings http_callbacks;
+					GString *url = g_string_new (NULL);
+
+					/* Hack, hack, hack */
+					/*
+					 * This code is assumed to read `/ping` command and
+					 * handle it to monitor port's availability since
+					 * milter protocol is stupid and does not allow to do that
+					 * This code also assumes that HTTP request can be read
+					 * as as single data chunk which is not true in some cases
+					 * In general, don't use it for anything but ping checks
+					 */
+					memset (&http_callbacks, 0, sizeof (http_callbacks));
+					http_parser_init (&http_parser, HTTP_REQUEST);
+					http_parser.data = url;
+					http_callbacks.on_url = rspamd_milter_http_on_url;
+					http_parser_execute (&http_parser, &http_callbacks,
+							priv->parser.buf->str, priv->parser.buf->len);
+
+					if (url->len == sizeof ("/ping") - 1 &&
+						rspamd_lc_cmp (url->str, "/ping", url->len) == 0) {
+						rspamd_milter_on_protocol_ping (session, priv);
+						g_string_free (url, TRUE);
+
+						return TRUE;
+					}
+					else {
+						err = g_error_new (rspamd_milter_quark (), EINVAL,
+								"HTTP GET request is not supported in milter mode, url: %s",
+								url->str);
+					}
+
+					g_string_free (url, TRUE);
 				}
 				else if (priv->parser.buf->len > sizeof ("POST") &&
 					 memcmp (priv->parser.buf->str, "POST", 4) == 0) {
@@ -1091,6 +1165,16 @@ rspamd_milter_handle_session (struct rspamd_milter_session *session,
 		REF_RELEASE (session);
 		return FALSE;
 		break;
+	case RSPAMD_MILTER_PONG_AND_DIE:
+		err = g_error_new (rspamd_milter_quark (), 0,
+				"ping command");
+		REF_RETAIN (session);
+		priv->err_cb (priv->fd, session, priv->ud, err);
+		REF_RELEASE (session);
+		g_error_free (err);
+		REF_RELEASE (session);
+		return FALSE;
+		break;
 	}
 
 	return TRUE;
diff --git a/src/libserver/milter_internal.h b/src/libserver/milter_internal.h
index ffca57101..693869c6f 100644
--- a/src/libserver/milter_internal.h
+++ b/src/libserver/milter_internal.h
@@ -57,6 +57,7 @@ enum rspamd_milter_io_state {
 	RSPAMD_MILTER_WRITE_REPLY,
 	RSPAMD_MILTER_WANNA_DIE,
 	RSPAMD_MILTER_WRITE_AND_DIE,
+	RSPAMD_MILTER_PONG_AND_DIE,
 };
 
 KHASH_INIT (milter_headers_hash_t, char *, GArray *, true,
-- 
2.39.5