[Feature] Add support for input encryption

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index d0512da833d0a3d644c0dd117771c8f3d291c97d..8896e229f18d27ca501e62c7316ddd68cdcb4ea2 100644 (file)
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -140,6 +140,7 @@ TARGET_LINK_LIBRARIES(rspamd-server rspamd-http-parser)
 TARGET_LINK_LIBRARIES(rspamd-server rspamd-cdb)
 TARGET_LINK_LIBRARIES(rspamd-server rspamd-lpeg)
 TARGET_LINK_LIBRARIES(rspamd-server lcbtrie)
+TARGET_LINK_LIBRARIES(rspamd-server rspamd-zstd)
 
 IF (ENABLE_CLANG_PLUGIN MATCHES "ON")
        ADD_DEPENDENCIES(rspamd-server rspamd-clang)

diff --git a/src/libserver/task.c b/src/libserver/task.c
index 70566b6b06e45aab04f2f163b29bda05f2c2410a..6c4547db50c86d0635d0e4f2f53e91b0ec67f358 100644 (file)
--- a/src/libserver/task.c
+++ b/src/libserver/task.c
@@ -24,6 +24,7 @@
 #include "stat_api.h"
 #include "unix-std.h"
 #include "utlist.h"
+#include "contrib/zstd/zstd.h"
 #include <math.h>
 
 /*
@@ -474,8 +475,75 @@ rspamd_task_load_message (struct rspamd_task *task,
 
        /* Plain data */
        debug_task ("got input of length %z", task->msg.len);
-       task->msg.begin = start;
-       task->msg.len = len;
+
+       /* Check compression */
+       tok = rspamd_task_get_request_header (task, "compression");
+
+       if (tok) {
+               /* Need to uncompress */
+               rspamd_ftok_t t;
+
+               t.begin = "zstd";
+               t.len = 4;
+
+               if (rspamd_ftok_casecmp (tok, &t) == 0) {
+                       ZSTD_DStream *zstream;
+                       ZSTD_inBuffer zin;
+                       ZSTD_outBuffer zout;
+                       guchar *out;
+                       gsize outlen, r;
+
+                       zin.pos = 0;
+                       zin.src = start;
+                       zin.size = len;
+
+                       if ((outlen = ZSTD_getDecompressedSize (start, len)) == 0) {
+                               outlen = ZSTD_DStreamOutSize ();
+                       }
+
+                       out = g_malloc (outlen);
+                       zstream = ZSTD_createDStream ();
+                       g_assert (zstream != NULL);
+                       g_assert (!ZSTD_isError (ZSTD_initDStream (zstream)));
+                       zout.dst = out;
+                       zout.pos = 0;
+                       zout.size = outlen;
+
+                       while (zin.pos < zin.size) {
+                               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);
+
+                                       return FALSE;
+                               }
+
+                               if (zout.pos == zout.size) {
+                                       /* We need to extend output buffer */
+                                       zout.size = zout.size * 1.5 + 1.0;
+                                       zout.dst = g_realloc (zout.dst, zout.size);
+                               }
+                       }
+
+                       ZSTD_freeDStream (zstream);
+                       rspamd_mempool_add_destructor (task->task_pool, g_free, zout.dst);
+                       task->msg.begin = zout.dst;
+                       task->msg.len = zout.pos;
+                       task->flags = RSPAMD_TASK_FLAG_COMPRESSED;
+
+               }
+               else {
+                       g_set_error (&task->err, rspamd_task_quark(), RSPAMD_PROTOCOL_ERROR,
+                                       "Invalid compression method");
+                       return FALSE;
+               }
+       }
+       else {
+               task->msg.begin = start;
+               task->msg.len = len;
+       }
 
        if (task->msg.len == 0) {
                task->flags |= RSPAMD_TASK_FLAG_EMPTY;

diff --git a/src/libserver/task.h b/src/libserver/task.h
index 1d75e31dc1b6f34ccd0c9557562125a34f0aa53d..aa1f52e45436d19665049f085ec7c3669a120238 100644 (file)
--- a/src/libserver/task.h
+++ b/src/libserver/task.h
@@ -104,6 +104,7 @@ enum rspamd_task_stage {
 #define RSPAMD_TASK_FLAG_HAS_HAM_TOKENS (1 << 21)
 #define RSPAMD_TASK_FLAG_EMPTY (1 << 22)
 #define RSPAMD_TASK_FLAG_LOCAL_CLIENT (1 << 23)
+#define RSPAMD_TASK_FLAG_COMPRESSED (1 << 24)
 
 #define RSPAMD_TASK_IS_SKIPPED(task) (((task)->flags & RSPAMD_TASK_FLAG_SKIP))
 #define RSPAMD_TASK_IS_JSON(task) (((task)->flags & RSPAMD_TASK_FLAG_JSON))