[Feature] Preserve decompression context between tasks

diff --git a/src/libserver/task.c b/src/libserver/task.c
index 95c3375b2ec44bc381d6cee9103633f6a751d097..950564b2da119df79a223e386985211847199e84 100644 (file)
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -494,6 +494,14 @@ rspamd_task_load_message (struct rspamd_task *task,
                        gsize outlen, r;
                        gulong dict_id;
 
+                       if (!rspamd_libs_reset_decompression (task->cfg->libs_ctx)) {
+                               g_set_error (&task->err, rspamd_task_quark(),
+                                               RSPAMD_PROTOCOL_ERROR,
+                                               "Cannot decompress, decompressor init failed");
+
+                               return FALSE;
+                       }
+
                        tok = rspamd_task_get_request_header (task, "dictionary");
 
                        if (tok != NULL) {
@@ -518,19 +526,10 @@ rspamd_task_load_message (struct rspamd_task *task,
 
                                        return FALSE;
                                }
-
-                               zstream = ZSTD_createDStream ();
-                               g_assert (zstream != NULL);
-                               g_assert (!ZSTD_isError (ZSTD_initDStream_usingDict (zstream,
-                                               task->cfg->libs_ctx->in_dict->dict,
-                                               task->cfg->libs_ctx->in_dict->size)));
-                       }
-                       else {
-                               zstream = ZSTD_createDStream ();
-                               g_assert (zstream != NULL);
-                               g_assert (!ZSTD_isError (ZSTD_initDStream (zstream)));
                        }
 
+                       zstream = task->cfg->libs_ctx->in_zstream;
+
                        zin.pos = 0;
                        zin.src = start;
                        zin.size = len;
@@ -548,9 +547,9 @@ rspamd_task_load_message (struct rspamd_task *task,
                                r = ZSTD_decompressStream (zstream, &zout, &zin);
 
                                if (ZSTD_isError (r)) {
-                                       g_set_error (&task->err, rspamd_task_quark(), RSPAMD_PROTOCOL_ERROR,
-                                                       "Decompression error");
-                                       ZSTD_freeDStream (zstream);
+                                       g_set_error (&task->err, rspamd_task_quark(),
+                                                       RSPAMD_PROTOCOL_ERROR,
+                                                       "Decompression error: %s", ZSTD_getErrorName (r));
 
                                        return FALSE;
                                }
@@ -562,7 +561,6 @@ rspamd_task_load_message (struct rspamd_task *task,
                                }
                        }
 
-                       ZSTD_freeDStream (zstream);
                        rspamd_mempool_add_destructor (task->task_pool, g_free, zout.dst);
                        task->msg.begin = zout.dst;
                        task->msg.len = zout.pos;

diff --git a/src/libutil/util.c b/src/libutil/util.c
index fbc06c039e39c81772464667a8097c8faa42bb27..65756693c71298cbab39f9f621c9634378264e77 100644 (file)
--- a/src/libutil/util.c
+++ b/src/libutil/util.c
@@ -23,6 +23,7 @@
 #include "ottery.h"
 #include "cryptobox.h"
 #include "libutil/map.h"
+#include "contrib/zstd/zstd.h"
 #include "contrib/zstd/zdict.h"
 
 #ifdef HAVE_OPENSSL
@@ -2181,7 +2182,78 @@ rspamd_config_libs (struct rspamd_external_libs_ctx *ctx,
                                                cfg->zstd_output_dictionary);
                        }
                }
+
+               /* Init decompression */
+               ctx->in_zstream = ZSTD_createDStream ();
+               rspamd_libs_reset_decompression (ctx);
+
+               /* Init compression */
+               ctx->out_zstream = ZSTD_createCStream ();
+               rspamd_libs_reset_compression (ctx);
+       }
+}
+
+gboolean
+rspamd_libs_reset_decompression (struct rspamd_external_libs_ctx *ctx)
+{
+       gsize r;
+
+       if (ctx->in_zstream == NULL) {
+               msg_err ("cannot create decompression stream");
+               return FALSE;
+       }
+       else {
+               if (ctx->in_dict) {
+                       r = ZSTD_initDStream_usingDict (ctx->in_zstream,
+                                       ctx->in_dict->dict, ctx->in_dict->size);
+               }
+               else {
+                       r = ZSTD_initDStream (ctx->in_zstream);
+               }
+
+               if (ZSTD_isError (r)) {
+                       msg_err ("cannot init decompression stream: %s",
+                                       ZSTD_getErrorName (r));
+                       ZSTD_freeDStream (ctx->in_zstream);
+                       ctx->in_zstream = NULL;
+
+                       return FALSE;
+               }
        }
+
+       return TRUE;
+}
+
+gboolean
+rspamd_libs_reset_compression (struct rspamd_external_libs_ctx *ctx)
+{
+       gsize r;
+
+       if (ctx->out_zstream == NULL) {
+               msg_err ("cannot create compression stream");
+
+               return FALSE;
+       }
+       else {
+               if (ctx->out_dict) {
+                       r = ZSTD_initCStream_usingDict (ctx->out_zstream,
+                                       ctx->out_dict->dict, ctx->out_dict->size, 1);
+               }
+               else {
+                       r = ZSTD_initCStream (ctx->out_zstream, 1);
+               }
+
+               if (ZSTD_isError (r)) {
+                       msg_err ("cannot init compression stream: %s",
+                                       ZSTD_getErrorName (r));
+                       ZSTD_freeCStream (ctx->out_zstream);
+                       ctx->out_zstream = NULL;
+
+                       return FALSE;
+               }
+       }
+
+       return TRUE;
 }
 
 void
@@ -2203,6 +2275,8 @@ rspamd_deinit_libs (struct rspamd_external_libs_ctx *ctx)
                rspamd_inet_library_destroy ();
                rspamd_free_zstd_dictionary (ctx->in_dict);
                rspamd_free_zstd_dictionary (ctx->out_dict);
+               ZSTD_freeCStream (ctx->out_zstream);
+               ZSTD_freeDStream (ctx->in_zstream);
                g_slice_free1 (sizeof (*ctx), ctx);
        }
 }

diff --git a/src/libutil/util.h b/src/libutil/util.h
index 76a02d19820d54dcab7cc1e3b5da6e431bc67be9..564d00b12979f95c92a74ca9816f2fa1dc140816 100644 (file)
--- a/src/libutil/util.h
+++ b/src/libutil/util.h
@@ -393,6 +393,17 @@ struct rspamd_external_libs_ctx* rspamd_init_libs (void);
 void rspamd_config_libs (struct rspamd_external_libs_ctx *ctx,
                struct rspamd_config *cfg);
 
+/**
+ * Reset and initialize decompressor
+ * @param ctx
+ */
+gboolean rspamd_libs_reset_decompression (struct rspamd_external_libs_ctx *ctx);
+/**
+ * Reset and initialize compressor
+ * @param ctx
+ */
+gboolean rspamd_libs_reset_compression (struct rspamd_external_libs_ctx *ctx);
+
 /**
  * Destroy external libraries context
  */

diff --git a/src/rspamd.h b/src/rspamd.h
index 545681bf57921f47e9d18ae689e1546f63186e8f..bb39f38196c415cd52ae5187722c01c08c4fab13 100644 (file)
--- a/src/rspamd.h
+++ b/src/rspamd.h
@@ -300,11 +300,15 @@ struct controller_session {
 };
 
 struct zstd_dictionary {
-       const void *dict;
+       void *dict;
        gsize size;
        guint id;
 };
 
+
+struct ZSTD_CStream_s;
+struct ZSTD_DStream_s;
+
 struct rspamd_external_libs_ctx {
        magic_t libmagic;
        radix_compressed_t **local_addrs;
@@ -313,6 +317,8 @@ struct rspamd_external_libs_ctx {
        SSL_CTX *ssl_ctx;
        struct zstd_dictionary *in_dict;
        struct zstd_dictionary *out_dict;
+       struct ZSTD_CStream_s *out_zstream;
+       struct ZSTD_DStream_s *in_zstream;
        ref_entry_t ref;
 };