[Feature] Add support for dictionary in client compression

4 files changed, 58 insertions, 5 deletions

diff --git a/src/client/rspamc.c b/src/client/rspamc.c
index 92983b491..ae3b3fe82 100644
--- a/src/client/rspamc.c
+++ b/src/client/rspamc.c
@@ -44,6 +44,7 @@ static gchar **http_headers = NULL;
 static gint weight = 0;
 static gint flag = 0;
 static gchar *fuzzy_symbol = NULL;
+static gchar *dictionary = NULL;
 static gint max_requests = 8;
 static gdouble timeout = 10.0;
 static gboolean pass_all;
@@ -139,6 +140,8 @@ static GOptionEntry entries[] =
 	   "Learn the specified fuzzy symbol", NULL },
 	{ "compressed", 'z', 0, G_OPTION_ARG_NONE, &compressed,
 	   "Enable zstd compression", NULL },
+	{ "dictionary", 'D', 0, G_OPTION_ARG_FILENAME, &dictionary,
+	   "Use dictionary to compress data", NULL },
 	{ NULL, 0, 0, G_OPTION_ARG_NONE, NULL, NULL, NULL }
 };
 
@@ -1386,7 +1389,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, compressed, &err);
+				cbdata, compressed, dictionary, &err);
 		}
 		else {
 			rspamd_client_command (conn,
@@ -1396,6 +1399,7 @@ rspamc_process_input (struct event_base *ev_base, struct rspamc_command *cmd,
 				rspamc_client_cb,
 				cbdata,
 				compressed,
+				dictionary,
 				&err);
 		}
 	}
diff --git a/src/client/rspamdclient.c b/src/client/rspamdclient.c
index b71d6f335..ba737a9a6 100644
--- a/src/client/rspamdclient.c
+++ b/src/client/rspamdclient.c
@@ -19,6 +19,7 @@
 #include "libutil/http_private.h"
 #include "unix-std.h"
 #include "contrib/zstd/zstd.h"
+#include "contrib/zstd/zdict.h"
 
 #ifdef HAVE_FETCH_H
 #include <fetch.h>
@@ -200,7 +201,9 @@ gboolean
 rspamd_client_command (struct rspamd_client_connection *conn,
 		const gchar *command, GQueue *attrs,
 		FILE *in, rspamd_client_callback cb,
-		gpointer ud, gboolean compressed, GError **err)
+		gpointer ud, gboolean compressed,
+		const gchar *comp_dictionary,
+		GError **err)
 {
 	struct rspamd_client_request *req;
 	struct rspamd_http_client_header *nh;
@@ -209,6 +212,10 @@ rspamd_client_command (struct rspamd_client_connection *conn,
 	GList *cur;
 	GString *input = NULL;
 	rspamd_fstring_t *body;
+	guint dict_id = 0;
+	gsize dict_len = 0;
+	void *dict = NULL;
+	ZSTD_CCtx *zctx;
 
 	req = g_slice_alloc0 (sizeof (struct rspamd_client_request));
 	req->conn = conn;
@@ -251,9 +258,43 @@ rspamd_client_command (struct rspamd_client_connection *conn,
 			body = rspamd_fstring_new_init (input->str, input->len);
 		}
 		else {
+			if (comp_dictionary) {
+				dict = rspamd_file_xmap (comp_dictionary, PROT_READ, &dict_len);
+
+				if (dict == NULL) {
+					g_set_error (err, RCLIENT_ERROR, errno,
+							"cannot open dictionary %s: %s",
+							comp_dictionary,
+							strerror (errno));
+					g_slice_free1 (sizeof (struct rspamd_client_request), req);
+					g_string_free (input, TRUE);
+
+					return FALSE;
+				}
+
+				dict_id = ZDICT_getDictID (comp_dictionary, dict_len);
+
+				if (dict_id == 0) {
+					g_set_error (err, RCLIENT_ERROR, errno,
+							"cannot open dictionary %s: %s",
+							comp_dictionary,
+							strerror (errno));
+					g_slice_free1 (sizeof (struct rspamd_client_request), req);
+					g_string_free (input, TRUE);
+					munmap (dict, dict_len);
+
+					return FALSE;
+				}
+			}
+
 			body = rspamd_fstring_sized_new (ZSTD_compressBound (input->len));
-			body->len = ZSTD_compress (body->str, body->allocated, input->str,
-					input->len, 1);
+			zctx = ZSTD_createCCtx ();
+			body->len = ZSTD_compress_usingDict (zctx, body->str, body->allocated,
+					input->str, input->len,
+					dict, dict_len,
+					1);
+
+			munmap (dict, dict_len);
 
 			if (ZSTD_isError (body->len)) {
 				g_set_error (err, RCLIENT_ERROR, ferror (
@@ -284,6 +325,13 @@ rspamd_client_command (struct rspamd_client_connection *conn,
 
 	if (compressed) {
 		rspamd_http_message_add_header (req->msg, "Compression", "zstd");
+
+		if (dict_id != 0) {
+			gchar dict_str[32];
+
+			rspamd_snprintf (dict_str, sizeof (dict_str), "%ud", dict_id);
+			rspamd_http_message_add_header (req->msg, "Dictionary", dict_str);
+		}
 	}
 
 	req->msg->url = rspamd_fstring_append (req->msg->url, "/", 1);
diff --git a/src/client/rspamdclient.h b/src/client/rspamdclient.h
index c7c8d274c..5313bb88b 100644
--- a/src/client/rspamdclient.h
+++ b/src/client/rspamdclient.h
@@ -78,6 +78,7 @@ gboolean rspamd_client_command (
 	rspamd_client_callback cb,
 	gpointer ud,
 	gboolean compressed,
+	const gchar *comp_dictionary,
 	GError **err);
 
 /**
diff --git a/src/libserver/task.c b/src/libserver/task.c
index d4b5cb696..95c3375b2 100644
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -521,7 +521,7 @@ rspamd_task_load_message (struct rspamd_task *task,
 
 				zstream = ZSTD_createDStream ();
 				g_assert (zstream != NULL);
-				g_assert (!ZSTD_isError (ZSTD_initDStream_usingDict(zstream,
+				g_assert (!ZSTD_isError (ZSTD_initDStream_usingDict (zstream,
 						task->cfg->libs_ctx->in_dict->dict,
 						task->cfg->libs_ctx->in_dict->size)));
 			}