From 105fd6d4c9108f8a0c9a4c582a2fd5f17c023eea Mon Sep 17 00:00:00 2001
From: Lion Kortlepel <development@kortlepel.com>
Date: Sat, 22 Jun 2024 22:47:31 +0200
Subject: [PATCH] rewrite compression and decompression to limit at 30 MB

---
 include/Common.h | 54 +++---------------------------------------------
 src/Common.cpp   | 51 ++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 53 insertions(+), 52 deletions(-)

diff --git a/include/Common.h b/include/Common.h
index 18fe84a..c35b325 100644
--- a/include/Common.h
+++ b/include/Common.h
@@ -29,6 +29,7 @@
 #include <memory>
 #include <mutex>
 #include <shared_mutex>
+#include <span>
 #include <sstream>
 #include <unordered_map>
 #include <zlib.h>
@@ -289,57 +290,8 @@ void RegisterThread(const std::string& str);
 
 void LogChatMessage(const std::string& name, int id, const std::string& msg);
 
-#define Biggest 30000
-
-template <typename T>
-inline T Comp(const T& Data) {
-    std::array<char, Biggest> C {};
-    // obsolete
-    C.fill(0);
-    z_stream defstream;
-    defstream.zalloc = nullptr;
-    defstream.zfree = nullptr;
-    defstream.opaque = nullptr;
-    defstream.avail_in = uInt(Data.size());
-    defstream.next_in = const_cast<Bytef*>(reinterpret_cast<const Bytef*>(&Data[0]));
-    defstream.avail_out = Biggest;
-    defstream.next_out = reinterpret_cast<Bytef*>(C.data());
-    deflateInit(&defstream, Z_BEST_COMPRESSION);
-    deflate(&defstream, Z_SYNC_FLUSH);
-    deflate(&defstream, Z_FINISH);
-    deflateEnd(&defstream);
-    size_t TotalOut = defstream.total_out;
-    T Ret;
-    Ret.resize(TotalOut);
-    std::fill(Ret.begin(), Ret.end(), 0);
-    std::copy_n(C.begin(), TotalOut, Ret.begin());
-    return Ret;
-}
-
-template <typename T>
-inline T DeComp(const T& Compressed) {
-    std::array<char, Biggest> C {};
-    // not needed
-    C.fill(0);
-    z_stream infstream;
-    infstream.zalloc = nullptr;
-    infstream.zfree = nullptr;
-    infstream.opaque = nullptr;
-    infstream.avail_in = Biggest;
-    infstream.next_in = const_cast<Bytef*>(reinterpret_cast<const Bytef*>(&Compressed[0]));
-    infstream.avail_out = Biggest;
-    infstream.next_out = const_cast<Bytef*>(reinterpret_cast<const Bytef*>(C.data()));
-    inflateInit(&infstream);
-    inflate(&infstream, Z_SYNC_FLUSH);
-    inflate(&infstream, Z_FINISH);
-    inflateEnd(&infstream);
-    size_t TotalOut = infstream.total_out;
-    T Ret;
-    Ret.resize(TotalOut);
-    std::fill(Ret.begin(), Ret.end(), 0);
-    std::copy_n(C.begin(), TotalOut, Ret.begin());
-    return Ret;
-}
+std::vector<uint8_t> Comp(std::span<const uint8_t> input);
+std::vector<uint8_t> DeComp(std::span<const uint8_t> input);
 
 std::string GetPlatformAgnosticErrorString();
 #define S_DSN SU_RAW
diff --git a/src/Common.cpp b/src/Common.cpp
index e2b0719..52aea27 100644
--- a/src/Common.cpp
+++ b/src/Common.cpp
@@ -22,13 +22,14 @@
 #include "TConsole.h"
 #include <array>
 #include <charconv>
+#include <chrono>
 #include <fmt/core.h>
+#include <fstream>
 #include <iostream>
 #include <map>
 #include <regex>
 #include <sstream>
 #include <thread>
-#include <chrono>
 
 #include "Compat.h"
 #include "CustomAssert.h"
@@ -386,3 +387,51 @@ void SplitString(const std::string& str, const char delim, std::vector<std::stri
     }
 }
 
+std::vector<uint8_t> DeComp(std::span<const uint8_t> input) {
+    beammp_debugf("got {} bytes of input data", input.size());
+
+    std::vector<uint8_t> output_buffer(std::min<size_t>(input.size() * 5, 15 * 1024 * 1024));
+
+    uLongf output_size = output_buffer.size();
+
+    while (true) {
+        int res = uncompress(
+            reinterpret_cast<Bytef*>(output_buffer.data()),
+            &output_size,
+            reinterpret_cast<const Bytef*>(input.data()),
+            static_cast<uLongf>(input.size()));
+        if (res == Z_BUF_ERROR) {
+            if (output_buffer.size() > 30 * 1024 * 1024) {
+                throw std::runtime_error("decompressed packet size of 30 MB exceeded");
+            }
+            beammp_warn("zlib uncompress() failed, trying with 2x buffer size of " + std::to_string(output_buffer.size() * 2));
+            output_buffer.resize(output_buffer.size() * 2);
+            output_size = output_buffer.size();
+        } else if (res != Z_OK) {
+            beammp_error("zlib uncompress() failed: " + std::to_string(res));
+            throw std::runtime_error("zlib uncompress() failed");
+        } else if (res == Z_OK) {
+            break;
+        }
+    }
+    output_buffer.resize(output_size);
+    return output_buffer;
+}
+
+std::vector<uint8_t> Comp(std::span<const uint8_t> input) {
+    auto max_size = compressBound(input.size());
+    std::vector<uint8_t> output(max_size);
+    uLongf output_size = output.size();
+    int res = compress(
+        reinterpret_cast<Bytef*>(output.data()),
+        &output_size,
+        reinterpret_cast<const Bytef*>(input.data()),
+        static_cast<uLongf>(input.size()));
+    if (res != Z_OK) {
+        beammp_error("zlib compress() failed: " + std::to_string(res));
+        throw std::runtime_error("zlib compress() failed");
+    }
+    beammp_debug("zlib compressed " + std::to_string(input.size()) + " B to " + std::to_string(output_size) + " B");
+    output.resize(output_size);
+    return output;
+}