From b1233056b76aa803123cc304d323887ad5fdfbae Mon Sep 17 00:00:00 2001
From: 20kdc <asdd2808@gmail.com>
Date: Wed, 28 Sep 2022 14:06:14 +0100
Subject: Add zstd compression support (#12515)

---
 src/script/lua_api/l_util.cpp | 57 +++++++++++++++++++++++++++++++++++++++----
 1 file changed, 52 insertions(+), 5 deletions(-)

(limited to 'src/script/lua_api/l_util.cpp')

diff --git a/src/script/lua_api/l_util.cpp b/src/script/lua_api/l_util.cpp
index f602aed99..bf2bb9137 100644
--- a/src/script/lua_api/l_util.cpp
+++ b/src/script/lua_api/l_util.cpp
@@ -26,6 +26,7 @@ with this program; if not, write to the Free Software Foundation, Inc.,
 #include "cpp_api/s_async.h"
 #include "serialization.h"
 #include <json/json.h>
+#include <zstd.h>
 #include "cpp_api/s_security.h"
 #include "porting.h"
 #include "convert_json.h"
@@ -278,6 +279,34 @@ int ModApiUtil::l_get_user_path(lua_State *L)
 	return 1;
 }
 
+enum LuaCompressMethod
+{
+	LUA_COMPRESS_METHOD_DEFLATE,
+	LUA_COMPRESS_METHOD_ZSTD,
+};
+
+static const struct EnumString es_LuaCompressMethod[] =
+{
+	{LUA_COMPRESS_METHOD_DEFLATE, "deflate"},
+	{LUA_COMPRESS_METHOD_ZSTD, "zstd"},
+	{0, nullptr},
+};
+
+static LuaCompressMethod get_compress_method(lua_State *L, int index)
+{
+	if (lua_isnoneornil(L, index))
+		return LUA_COMPRESS_METHOD_DEFLATE;
+	const char *name = luaL_checkstring(L, index);
+	int value;
+	if (!string_to_enum(es_LuaCompressMethod, value, name)) {
+		// Pretend it's deflate if we don't know, for compatibility reasons.
+		log_deprecated(L, "Unknown compression method \"" + std::string(name)
+			+ "\", defaulting to \"deflate\". You should pass a valid value.");
+		return LUA_COMPRESS_METHOD_DEFLATE;
+	}
+	return (LuaCompressMethod) value;
+}
+
 // compress(data, method, level)
 int ModApiUtil::l_compress(lua_State *L)
 {
@@ -286,12 +315,23 @@ int ModApiUtil::l_compress(lua_State *L)
 	size_t size;
 	const char *data = luaL_checklstring(L, 1, &size);
 
-	int level = -1;
-	if (!lua_isnoneornil(L, 3))
-		level = readParam<int>(L, 3);
+	LuaCompressMethod method = get_compress_method(L, 2);
 
 	std::ostringstream os(std::ios_base::binary);
-	compressZlib(reinterpret_cast<const u8 *>(data), size, os, level);
+
+	if (method == LUA_COMPRESS_METHOD_DEFLATE) {
+		int level = -1;
+		if (!lua_isnoneornil(L, 3))
+			level = readParam<int>(L, 3);
+
+		compressZlib(reinterpret_cast<const u8 *>(data), size, os, level);
+	} else if (method == LUA_COMPRESS_METHOD_ZSTD) {
+		int level = ZSTD_CLEVEL_DEFAULT;
+		if (!lua_isnoneornil(L, 3))
+			level = readParam<int>(L, 3);
+
+		compressZstd(reinterpret_cast<const u8 *>(data), size, os, level);
+	}
 
 	std::string out = os.str();
 
@@ -307,9 +347,16 @@ int ModApiUtil::l_decompress(lua_State *L)
 	size_t size;
 	const char *data = luaL_checklstring(L, 1, &size);
 
+	LuaCompressMethod method = get_compress_method(L, 2);
+
 	std::istringstream is(std::string(data, size), std::ios_base::binary);
 	std::ostringstream os(std::ios_base::binary);
-	decompressZlib(is, os);
+
+	if (method == LUA_COMPRESS_METHOD_DEFLATE) {
+		decompressZlib(is, os);
+	} else if (method == LUA_COMPRESS_METHOD_ZSTD) {
+		decompressZstd(is, os);
+	}
 
 	std::string out = os.str();
 
-- 
cgit v1.2.3