From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Thu, 8 Sep 2016 17:42:42 +0000 (+0100)
Subject: [Feature] Add compression support to rspamd client
X-Git-Tag: 1.4.0~458
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=d763da74c87b5939c2bcfdfdfea71c15245cfd20;p=thirdparty%2Frspamd.git

[Feature] Add compression support to rspamd client
---

diff --git a/src/client/rspamc.c b/src/client/rspamc.c
index 1d2dd39cae..92983b491d 100644
--- a/src/client/rspamc.c
+++ b/src/client/rspamc.c
@@ -57,6 +57,7 @@ static gboolean raw = FALSE;
 static gboolean extended_urls = FALSE;
 static gboolean mime_output = FALSE;
 static gboolean empty_input = FALSE;
+static gboolean compressed = FALSE;
 static gchar *key = NULL;
 static GList *children;
 
@@ -136,6 +137,8 @@ static GOptionEntry entries[] =
 	   "Allow empty input instead of reading from stdin", NULL },
 	{ "fuzzy-symbol", 'S', 0, G_OPTION_ARG_STRING, &fuzzy_symbol,
 	   "Learn the specified fuzzy symbol", NULL },
+	{ "compressed", 'z', 0, G_OPTION_ARG_NONE, &compressed,
+	   "Enable zstd compression", NULL },
 	{ NULL, 0, 0, G_OPTION_ARG_NONE, NULL, NULL, NULL }
 };
 
@@ -1383,7 +1386,7 @@ rspamc_process_input (struct event_base *ev_base, struct rspamc_command *cmd,
 
 		if (cmd->need_input) {
 			rspamd_client_command (conn, cmd->path, attrs, in, rspamc_client_cb,
-				cbdata, &err);
+				cbdata, compressed, &err);
 		}
 		else {
 			rspamd_client_command (conn,
@@ -1392,6 +1395,7 @@ rspamc_process_input (struct event_base *ev_base, struct rspamc_command *cmd,
 				NULL,
 				rspamc_client_cb,
 				cbdata,
+				compressed,
 				&err);
 		}
 	}
diff --git a/src/client/rspamdclient.c b/src/client/rspamdclient.c
index 753f64c746..b71d6f335b 100644
--- a/src/client/rspamdclient.c
+++ b/src/client/rspamdclient.c
@@ -18,6 +18,7 @@
 #include "libutil/http.h"
 #include "libutil/http_private.h"
 #include "unix-std.h"
+#include "contrib/zstd/zstd.h"
 
 #ifdef HAVE_FETCH_H
 #include <fetch.h>
@@ -197,9 +198,9 @@ rspamd_client_init (struct event_base *ev_base, const gchar *name,
 
 gboolean
 rspamd_client_command (struct rspamd_client_connection *conn,
-	const gchar *command, GQueue *attrs,
-	FILE *in, rspamd_client_callback cb,
-	gpointer ud, GError **err)
+		const gchar *command, GQueue *attrs,
+		FILE *in, rspamd_client_callback cb,
+		gpointer ud, gboolean compressed, GError **err)
 {
 	struct rspamd_client_request *req;
 	struct rspamd_http_client_header *nh;
@@ -246,7 +247,25 @@ rspamd_client_command (struct rspamd_client_connection *conn,
 			return FALSE;
 		}
 
-		body = rspamd_fstring_new_init (input->str, input->len);
+		if (!compressed) {
+			body = rspamd_fstring_new_init (input->str, input->len);
+		}
+		else {
+			body = rspamd_fstring_sized_new (ZSTD_compressBound (input->len));
+			body->len = ZSTD_compress (body->str, body->allocated, input->str,
+					input->len, 1);
+
+			if (ZSTD_isError (body->len)) {
+				g_set_error (err, RCLIENT_ERROR, ferror (
+						in), "compression error");
+				g_slice_free1 (sizeof (struct rspamd_client_request), req);
+				g_string_free (input, TRUE);
+				rspamd_fstring_free (body);
+
+				return FALSE;
+			}
+		}
+
 		rspamd_http_message_set_body_from_fstring_steal (req->msg, body);
 		req->input = input;
 	}
@@ -263,13 +282,24 @@ rspamd_client_command (struct rspamd_client_connection *conn,
 		cur = g_list_next (cur);
 	}
 
+	if (compressed) {
+		rspamd_http_message_add_header (req->msg, "Compression", "zstd");
+	}
+
 	req->msg->url = rspamd_fstring_append (req->msg->url, "/", 1);
 	req->msg->url = rspamd_fstring_append (req->msg->url, command, strlen (command));
 
 	conn->req = req;
 
-	rspamd_http_connection_write_message (conn->http_conn, req->msg, NULL,
-		"text/plain", req, conn->fd, &conn->timeout, conn->ev_base);
+	if (compressed) {
+		rspamd_http_connection_write_message (conn->http_conn, req->msg, NULL,
+				"application/x-compressed", req, conn->fd,
+				&conn->timeout, conn->ev_base);
+	}
+	else {
+		rspamd_http_connection_write_message (conn->http_conn, req->msg, NULL,
+				"text/plain", req, conn->fd, &conn->timeout, conn->ev_base);
+	}
 
 	return TRUE;
 }
diff --git a/src/client/rspamdclient.h b/src/client/rspamdclient.h
index 3b602eb3b5..c7c8d274c1 100644
--- a/src/client/rspamdclient.h
+++ b/src/client/rspamdclient.h
@@ -77,6 +77,7 @@ gboolean rspamd_client_command (
 	FILE *in,
 	rspamd_client_callback cb,
 	gpointer ud,
+	gboolean compressed,
 	GError **err);
 
 /**