From 9a1b912929527accc88b9bdb48c63a5855f64c36 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Fri, 9 Sep 2016 12:39:38 +0100 Subject: [PATCH] [Feature] Preserve decompression context between tasks --- src/libserver/task.c | 28 ++++++++--------- src/libutil/util.c | 74 ++++++++++++++++++++++++++++++++++++++++++++ src/libutil/util.h | 11 +++++++ src/rspamd.h | 8 ++++- 4 files changed, 105 insertions(+), 16 deletions(-) diff --git a/src/libserver/task.c b/src/libserver/task.c index 95c3375b2..950564b2d 100644 --- 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 fbc06c039..65756693c 100644 --- 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 76a02d198..564d00b12 100644 --- 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 545681bf5..bb39f3819 100644 --- 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; }; -- 2.39.5