From 448c39791528275606f5c5350f6298de06a18462 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Fri, 9 Sep 2016 12:08:35 +0100 Subject: [PATCH] [Feature] Add support for dictionary in client compression --- src/client/rspamc.c | 6 ++++- src/client/rspamdclient.c | 54 ++++++++++++++++++++++++++++++++++++--- src/client/rspamdclient.h | 1 + src/libserver/task.c | 2 +- 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 @@ -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))); } -- 2.39.5