123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632 |
- /*-
- * Copyright 2021 Vsevolod Stakhov
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
- #include "lua_common.h"
- #include "unix-std.h"
- #include <zlib.h>
-
- #ifdef SYS_ZSTD
- # include "zstd.h"
- # include "zstd_errors.h"
- #else
- # include "contrib/zstd/zstd.h"
- # include "contrib/zstd/error_public.h"
- #endif
-
- /***
- * @module rspamd_compress
- * This module contains compression/decompression routines (zstd and zlib currently)
- */
-
- /***
- * @function zstd.compress_ctx()
- * Creates new compression ctx
- * @return {compress_ctx} new compress ctx
- */
- LUA_FUNCTION_DEF (zstd, compress_ctx);
-
- /***
- * @function zstd.compress_ctx()
- * Creates new compression ctx
- * @return {compress_ctx} new compress ctx
- */
- LUA_FUNCTION_DEF (zstd, decompress_ctx);
-
- LUA_FUNCTION_DEF (zstd_compress, stream);
- LUA_FUNCTION_DEF (zstd_compress, dtor);
-
- LUA_FUNCTION_DEF (zstd_decompress, stream);
- LUA_FUNCTION_DEF (zstd_decompress, dtor);
-
- static const struct luaL_reg zstd_compress_lib_f[] = {
- LUA_INTERFACE_DEF (zstd, compress_ctx),
- LUA_INTERFACE_DEF (zstd, decompress_ctx),
- {NULL, NULL}
- };
-
- static const struct luaL_reg zstd_compress_lib_m[] = {
- LUA_INTERFACE_DEF (zstd_compress, stream),
- {"__gc", lua_zstd_compress_dtor},
- {NULL, NULL}
- };
-
- static const struct luaL_reg zstd_decompress_lib_m[] = {
- LUA_INTERFACE_DEF (zstd_decompress, stream),
- {"__gc", lua_zstd_decompress_dtor},
- {NULL, NULL}
- };
-
- static ZSTD_CStream *
- lua_check_zstd_compress_ctx (lua_State *L, gint pos)
- {
- void *ud = rspamd_lua_check_udata (L, pos, "rspamd{zstd_compress}");
- luaL_argcheck (L, ud != NULL, pos, "'zstd_compress' expected");
- return ud ? *(ZSTD_CStream **)ud : NULL;
- }
-
- static ZSTD_DStream *
- lua_check_zstd_decompress_ctx (lua_State *L, gint pos)
- {
- void *ud = rspamd_lua_check_udata (L, pos, "rspamd{zstd_decompress}");
- luaL_argcheck (L, ud != NULL, pos, "'zstd_decompress' expected");
- return ud ? *(ZSTD_DStream **)ud : NULL;
- }
-
- int
- lua_zstd_push_error (lua_State *L, int err)
- {
- lua_pushnil (L);
- lua_pushfstring (L, "zstd error %d (%s)", err, ZSTD_getErrorString (err));
-
- return 2;
- }
-
- gint
- lua_compress_zstd_compress (lua_State *L)
- {
- LUA_TRACE_POINT;
- struct rspamd_lua_text *t = NULL, *res;
- gsize sz, r;
- gint comp_level = 1;
-
- t = lua_check_text_or_string (L,1);
-
- if (t == NULL || t->start == NULL) {
- return luaL_error (L, "invalid arguments");
- }
-
- if (lua_type (L, 2) == LUA_TNUMBER) {
- comp_level = lua_tointeger (L, 2);
- }
-
- sz = ZSTD_compressBound (t->len);
-
- if (ZSTD_isError (sz)) {
- msg_err ("cannot compress data: %s", ZSTD_getErrorName (sz));
- lua_pushnil (L);
-
- return 1;
- }
-
- res = lua_newuserdata (L, sizeof (*res));
- res->start = g_malloc (sz);
- res->flags = RSPAMD_TEXT_FLAG_OWN;
- rspamd_lua_setclass (L, "rspamd{text}", -1);
- r = ZSTD_compress ((void *)res->start, sz, t->start, t->len, comp_level);
-
- if (ZSTD_isError (r)) {
- msg_err ("cannot compress data: %s", ZSTD_getErrorName (r));
- lua_pop (L, 1); /* Text will be freed here */
- lua_pushnil (L);
-
- return 1;
- }
-
- res->len = r;
-
- return 1;
- }
-
- gint
- lua_compress_zstd_decompress (lua_State *L)
- {
- LUA_TRACE_POINT;
- struct rspamd_lua_text *t = NULL, *res;
- gsize outlen, r;
- ZSTD_DStream *zstream;
- ZSTD_inBuffer zin;
- ZSTD_outBuffer zout;
- gchar *out;
-
- t = lua_check_text_or_string (L,1);
-
- if (t == NULL || t->start == NULL) {
- return luaL_error (L, "invalid arguments");
- }
-
- zstream = ZSTD_createDStream ();
- ZSTD_initDStream (zstream);
-
- zin.pos = 0;
- zin.src = t->start;
- zin.size = t->len;
-
- if ((outlen = ZSTD_getDecompressedSize (zin.src, zin.size)) == 0) {
- outlen = ZSTD_DStreamOutSize ();
- }
-
- out = g_malloc (outlen);
-
- zout.dst = out;
- zout.pos = 0;
- zout.size = outlen;
-
- while (zin.pos < zin.size) {
- r = ZSTD_decompressStream (zstream, &zout, &zin);
-
- if (ZSTD_isError (r)) {
- msg_err ("cannot decompress data: %s", ZSTD_getErrorName (r));
- ZSTD_freeDStream (zstream);
- g_free (out);
- lua_pushstring (L, ZSTD_getErrorName (r));
- lua_pushnil (L);
-
- return 2;
- }
-
- if (zin.pos < zin.size && zout.pos == zout.size) {
- /* We need to extend output buffer */
- zout.size = zout.size * 2;
- out = g_realloc (zout.dst, zout.size);
- zout.dst = out;
- }
- }
-
- ZSTD_freeDStream (zstream);
- lua_pushnil (L); /* Error */
- res = lua_newuserdata (L, sizeof (*res));
- res->start = out;
- res->flags = RSPAMD_TEXT_FLAG_OWN;
- rspamd_lua_setclass (L, "rspamd{text}", -1);
- res->len = zout.pos;
-
- return 2;
- }
-
- gint
- lua_compress_zlib_decompress (lua_State *L, bool is_gzip)
- {
- LUA_TRACE_POINT;
- struct rspamd_lua_text *t = NULL, *res;
- gsize sz;
- z_stream strm;
- gint rc;
- guchar *p;
- gsize remain;
- gssize size_limit = -1;
-
- int windowBits = is_gzip ? (MAX_WBITS + 16) : (MAX_WBITS);
-
- t = lua_check_text_or_string (L,1);
-
- if (t == NULL || t->start == NULL) {
- return luaL_error (L, "invalid arguments");
- }
-
- if (lua_type (L, 2) == LUA_TNUMBER) {
- size_limit = lua_tointeger (L, 2);
- if (size_limit <= 0) {
- return luaL_error (L, "invalid arguments (size_limit)");
- }
-
- sz = MIN (t->len * 2, size_limit);
- }
- else {
- sz = t->len * 2;
- }
-
- memset (&strm, 0, sizeof (strm));
- /* windowBits +16 to decode gzip, zlib 1.2.0.4+ */
-
- /* Here are dragons to distinguish between raw deflate and zlib */
- if (windowBits == MAX_WBITS && t->len > 0) {
- if ((int)(unsigned char)((t->start[0] << 4)) != 0x80) {
- /* Assume raw deflate */
- windowBits = -windowBits;
- }
- }
-
- rc = inflateInit2 (&strm, windowBits);
-
- if (rc != Z_OK) {
- return luaL_error (L, "cannot init zlib");
- }
-
- strm.avail_in = t->len;
- strm.next_in = (guchar *)t->start;
-
- res = lua_newuserdata (L, sizeof (*res));
- res->start = g_malloc (sz);
- res->flags = RSPAMD_TEXT_FLAG_OWN;
- rspamd_lua_setclass (L, "rspamd{text}", -1);
-
- p = (guchar *)res->start;
- remain = sz;
-
- while (strm.avail_in != 0) {
- strm.avail_out = remain;
- strm.next_out = p;
-
- rc = inflate (&strm, Z_NO_FLUSH);
-
- if (rc != Z_OK && rc != Z_BUF_ERROR) {
- if (rc == Z_STREAM_END) {
- break;
- }
- else {
- msg_err ("cannot decompress data: %s (last error: %s)",
- zError (rc), strm.msg);
- lua_pop (L, 1); /* Text will be freed here */
- lua_pushnil (L);
- inflateEnd (&strm);
-
- return 1;
- }
- }
-
- res->len = strm.total_out;
-
- if (strm.avail_out == 0 && strm.avail_in != 0) {
-
- if (size_limit > 0 || res->len >= G_MAXUINT32 / 2) {
- if (res->len > size_limit || res->len >= G_MAXUINT32 / 2) {
- lua_pop (L, 1); /* Text will be freed here */
- lua_pushnil (L);
- inflateEnd (&strm);
-
- return 1;
- }
- }
-
- /* Need to allocate more */
- remain = res->len;
- res->start = g_realloc ((gpointer)res->start, res->len * 2);
- sz = res->len * 2;
- p = (guchar *)res->start + remain;
- remain = sz - remain;
- }
- }
-
- inflateEnd (&strm);
- res->len = strm.total_out;
-
- return 1;
- }
-
- gint
- lua_compress_zlib_compress (lua_State *L)
- {
- LUA_TRACE_POINT;
- struct rspamd_lua_text *t = NULL, *res;
- gsize sz;
- z_stream strm;
- gint rc, comp_level = Z_DEFAULT_COMPRESSION;
- guchar *p;
- gsize remain;
-
- t = lua_check_text_or_string (L,1);
-
- if (t == NULL || t->start == NULL) {
- return luaL_error (L, "invalid arguments");
- }
-
- if (lua_isnumber (L, 2)) {
- comp_level = lua_tointeger (L, 2);
-
- if (comp_level > Z_BEST_COMPRESSION || comp_level < Z_BEST_SPEED) {
- return luaL_error (L, "invalid arguments: compression level must be between %d and %d",
- Z_BEST_SPEED, Z_BEST_COMPRESSION);
- }
- }
-
-
- memset (&strm, 0, sizeof (strm));
- rc = deflateInit2 (&strm, comp_level, Z_DEFLATED,
- MAX_WBITS + 16, MAX_MEM_LEVEL - 1, Z_DEFAULT_STRATEGY);
-
- if (rc != Z_OK) {
- return luaL_error (L, "cannot init zlib: %s", zError (rc));
- }
-
- sz = deflateBound (&strm, t->len);
-
- strm.avail_in = t->len;
- strm.next_in = (guchar *) t->start;
-
- res = lua_newuserdata (L, sizeof (*res));
- res->start = g_malloc (sz);
- res->flags = RSPAMD_TEXT_FLAG_OWN;
- rspamd_lua_setclass (L, "rspamd{text}", -1);
-
- p = (guchar *) res->start;
- remain = sz;
-
- while (strm.avail_in != 0) {
- strm.avail_out = remain;
- strm.next_out = p;
-
- rc = deflate (&strm, Z_FINISH);
-
- if (rc != Z_OK && rc != Z_BUF_ERROR) {
- if (rc == Z_STREAM_END) {
- break;
- }
- else {
- msg_err ("cannot compress data: %s (last error: %s)",
- zError (rc), strm.msg);
- lua_pop (L, 1); /* Text will be freed here */
- lua_pushnil (L);
- deflateEnd (&strm);
-
- return 1;
- }
- }
-
- res->len = strm.total_out;
-
- if (strm.avail_out == 0 && strm.avail_in != 0) {
- /* Need to allocate more */
- remain = res->len;
- res->start = g_realloc ((gpointer) res->start, strm.avail_in + sz);
- sz = strm.avail_in + sz;
- p = (guchar *) res->start + remain;
- remain = sz - remain;
- }
- }
-
- deflateEnd (&strm);
- res->len = strm.total_out;
-
- return 1;
- }
-
- /* Stream API interface for Zstd: both compression and decompression */
-
- /* Operations allowed by zstd stream methods */
- static const char *const zstd_stream_op[] = {
- "continue",
- "flush",
- "end",
- NULL
- };
-
- static gint
- lua_zstd_compress_ctx (lua_State *L)
- {
- ZSTD_CCtx *ctx, **pctx;
-
- pctx = lua_newuserdata (L, sizeof (*pctx));
- ctx = ZSTD_createCCtx ();
-
- if (!ctx) {
- return luaL_error (L, "context create failed");
- }
-
- *pctx = ctx;
- rspamd_lua_setclass (L, "rspamd{zstd_compress}", -1);
- return 1;
- }
-
- static gint
- lua_zstd_compress_dtor (lua_State *L)
- {
- ZSTD_CCtx *ctx = lua_check_zstd_compress_ctx (L, 1);
-
- if (ctx) {
- ZSTD_freeCCtx (ctx);
- }
-
- return 0;
- }
-
- static gint
- lua_zstd_compress_reset (lua_State *L)
- {
- ZSTD_CCtx *ctx = lua_check_zstd_compress_ctx (L, 1);
-
- if (ctx) {
- ZSTD_CCtx_reset (ctx, ZSTD_reset_session_and_parameters);
- }
- else {
- return luaL_error (L, "invalid arguments");
- }
-
- return 0;
- }
-
- static gint
- lua_zstd_compress_stream (lua_State *L)
- {
- ZSTD_CStream *ctx = lua_check_zstd_compress_ctx (L, 1);
- struct rspamd_lua_text *t = lua_check_text_or_string (L, 2);
- int op = luaL_checkoption (L, 3, zstd_stream_op[0], zstd_stream_op);
- int err = 0;
- ZSTD_inBuffer inb;
- ZSTD_outBuffer onb;
-
- if (ctx && t) {
- gsize dlen = 0;
-
- inb.size = t->len;
- inb.pos = 0;
- inb.src = (const void*)t->start;
-
- onb.pos = 0;
- onb.size = ZSTD_CStreamInSize (); /* Initial guess */
- onb.dst = NULL;
-
- for (;;) {
- if ((onb.dst = g_realloc (onb.dst, onb.size)) == NULL) {
- return lua_zstd_push_error (L, ZSTD_error_memory_allocation);
- }
-
- dlen = onb.size;
-
- int res = ZSTD_compressStream2 (ctx, &onb, &inb, op);
-
- if (res == 0) {
- /* All done */
- break;
- }
-
- if ((err = ZSTD_getErrorCode (res))) {
- break;
- }
-
- onb.size *= 2;
- res += dlen; /* Hint returned by compression routine */
-
- /* Either double the buffer, or use the hint provided */
- if (onb.size < res) {
- onb.size = res;
- }
- }
- }
- else {
- return luaL_error (L, "invalid arguments");
- }
-
- if (err) {
- return lua_zstd_push_error (L, err);
- }
-
- lua_new_text (L, onb.dst, onb.pos, TRUE);
-
- return 1;
- }
-
- static gint
- lua_zstd_decompress_dtor (lua_State *L)
- {
- ZSTD_DStream *ctx = lua_check_zstd_decompress_ctx (L, 1);
-
- if (ctx) {
- ZSTD_freeDStream (ctx);
- }
-
- return 0;
- }
-
-
- static gint
- lua_zstd_decompress_ctx (lua_State *L)
- {
- ZSTD_DStream *ctx, **pctx;
-
- pctx = lua_newuserdata (L, sizeof (*pctx));
- ctx = ZSTD_createDStream ();
-
- if (!ctx) {
- return luaL_error (L, "context create failed");
- }
-
- *pctx = ctx;
- rspamd_lua_setclass (L, "rspamd{zstd_decompress}", -1);
- return 1;
- }
-
- static gint
- lua_zstd_decompress_stream (lua_State *L)
- {
- ZSTD_DStream *ctx = lua_check_zstd_decompress_ctx (L, 1);
- struct rspamd_lua_text *t = lua_check_text_or_string (L, 2);
- int err = 0;
- ZSTD_inBuffer inb;
- ZSTD_outBuffer onb;
-
- if (ctx && t) {
- gsize dlen = 0;
-
- if (t->len == 0) {
- return lua_zstd_push_error (L, ZSTD_error_init_missing);
- }
-
- inb.size = t->len;
- inb.pos = 0;
- inb.src = (const void*)t->start;
-
- onb.pos = 0;
- onb.size = ZSTD_DStreamInSize (); /* Initial guess */
- onb.dst = NULL;
-
- for (;;) {
- if ((onb.dst = g_realloc (onb.dst, onb.size)) == NULL) {
- return lua_zstd_push_error (L, ZSTD_error_memory_allocation);
- }
-
- dlen = onb.size;
-
- int res = ZSTD_decompressStream (ctx, &onb, &inb);
-
- if (res == 0) {
- /* All done */
- break;
- }
-
- if ((err = ZSTD_getErrorCode (res))) {
- break;
- }
-
- onb.size *= 2;
- res += dlen; /* Hint returned by compression routine */
-
- /* Either double the buffer, or use the hint provided */
- if (onb.size < res) {
- onb.size = res;
- }
- }
- }
- else {
- return luaL_error (L, "invalid arguments");
- }
-
- if (err) {
- return lua_zstd_push_error (L, err);
- }
-
- lua_new_text (L, onb.dst, onb.pos, TRUE);
-
- return 1;
- }
-
- static gint
- lua_load_zstd (lua_State * L)
- {
- lua_newtable (L);
- luaL_register (L, NULL, zstd_compress_lib_f);
-
- return 1;
- }
-
- void
- luaopen_compress (lua_State *L)
- {
- rspamd_lua_new_class (L, "rspamd{zstd_compress}", zstd_compress_lib_m);
- rspamd_lua_new_class (L, "rspamd{zstd_decompress}", zstd_decompress_lib_m);
- lua_pop (L, 2);
-
- rspamd_lua_add_preload (L, "rspamd_zstd", lua_load_zstd);
- }
|