[Feature] Add support for dictionary in client compression

diff --git a/src/client/rspamc.c b/src/client/rspamc.c
index 92983b491d139436061f67389289264995a72af8..ae3b3fe82ee80493650bc3036f64617a68ade797 100644 (file)
--- 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 b71d6f335bf4546568581a89020d65f1fc1108e2..ba737a9a6919889a79ce5c64e8ce2d0d3ed0bdad 100644 (file)
--- 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 c7c8d274c1970e9958e75c1515f2e48e23f53f53..5313bb88bdf5c648cd508707d34cfe36f82b4ab5 100644 (file)
--- 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 d4b5cb696a89afa89b9d7ab33fac4b4f8ca566cb..95c3375b2ec44bc381d6cee9103633f6a751d097 100644 (file)
--- 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)));
                        }