123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622 |
- /*-
- * 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, int pos)
- {
- void *ud = rspamd_lua_check_udata(L, pos, rspamd_zstd_compress_classname);
- 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, int pos)
- {
- void *ud = rspamd_lua_check_udata(L, pos, rspamd_zstd_decompress_classname);
- 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;
- }
-
- int lua_compress_zstd_compress(lua_State *L)
- {
- LUA_TRACE_POINT;
- struct rspamd_lua_text *t = NULL, *res;
- gsize sz, r;
- int 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_classname, -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;
- }
-
- int 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;
- char *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_classname, -1);
- res->len = zout.pos;
-
- return 2;
- }
-
- int 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;
- int rc;
- unsigned char *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 = (unsigned char *) 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_classname, -1);
-
- p = (unsigned char *) 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 = (unsigned char *) res->start + remain;
- remain = sz - remain;
- }
- }
-
- inflateEnd(&strm);
- res->len = strm.total_out;
-
- return 1;
- }
-
- int lua_compress_zlib_compress(lua_State *L)
- {
- LUA_TRACE_POINT;
- struct rspamd_lua_text *t = NULL, *res;
- gsize sz;
- z_stream strm;
- int rc, comp_level = Z_DEFAULT_COMPRESSION;
- unsigned char *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 = (unsigned char *) 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_classname, -1);
-
- p = (unsigned char *) 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 = (unsigned char *) 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 int
- 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_classname, -1);
- return 1;
- }
-
- static int
- 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 int
- 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 int
- 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 int
- 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 int
- 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_classname, -1);
- return 1;
- }
-
- static int
- 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 int
- 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_classname, zstd_compress_lib_m);
- rspamd_lua_new_class(L, rspamd_zstd_decompress_classname, zstd_decompress_lib_m);
- lua_pop(L, 2);
-
- rspamd_lua_add_preload(L, "rspamd_zstd", lua_load_zstd);
- }
|