Handle EntryType::ZstdMulti better

Co-Authored-By: moonshadow565 <7480230+moonshadow565@users.noreply.github.com>
LeagueToolkit · Sep 17, 2023 · cb6885f · cb6885f · nobobk11 · Nov 4, 2023
1 parent f61b92f
commit cb6885f
Show file tree

Hide file tree

Showing 8 changed files with 127 additions and 6 deletions.
diff --git a/cslol-tools/lib/lol/io/buffer.cpp b/cslol-tools/lib/lol/io/buffer.cpp
@@ -8,6 +8,15 @@
 using namespace lol;
 using namespace lol::io;
 
+static std::size_t find_zstd_magic(std::span<char const> src) {
+    static constexpr char const zstd_magic[4] = {0x28, (char)0xB5, 0x2F, (char)0xFD};
+    auto magic = std::search(src.data(), src.data() + src.size(), zstd_magic, zstd_magic + sizeof(zstd_magic));
+    if (magic == src.data() + src.size()) {
+        return 0;
+    }
+    return (std::size_t)(magic - src.data());
+}
+
 Buffer::~Buffer() noexcept = default;
 
 auto Buffer::copy(std::size_t pos, std::size_t count) const -> Bytes {
@@ -24,14 +33,29 @@ auto Buffer::readsome_decompress_zstd(std::size_t pos,
                                       std::size_t count,
                                       void* dst,
                                       std::size_t dst_count) const noexcept -> std::size_t {
+    [[unlikely]] if (pos > size_ || size_ - pos < count) return 0;
+
+    auto frame_start = find_zstd_magic({data_ + pos, count});
+    if (frame_start >= dst_count) {
+        std::memcpy(dst, data_ + pos, dst_count);
+        return dst_count;
+    }
+    if (frame_start) {
+        std::memcpy(dst, data_ + pos, frame_start);
+        pos += frame_start;
+        count -= frame_start;
+        dst = (char*)dst + frame_start;
+        dst_count -= frame_start;
+    }
+
     thread_local auto ctx = std::shared_ptr<ZSTD_DStream>(ZSTD_createDStream(), ZSTD_freeDStream);
     [[unlikely]] if (!ctx.get()) return 0;
     [[unlikely]] if (ZSTD_isError(ZSTD_initDStream(ctx.get()))) return 0;
     [[unlikely]] if (pos > size_ || size_ - pos < count) return 0;
     auto input = ZSTD_inBuffer{data_ + pos, count, 0};
     auto output = ZSTD_outBuffer{dst, dst_count, 0};
     ZSTD_decompressStream(ctx.get(), &output, &input);
-    return output.pos;
+    return frame_start + output.pos;
 }
 
 auto Buffer::readsome(std::size_t pos, void* dst, std::size_t count) const noexcept -> std::size_t {
@@ -148,7 +172,6 @@ auto Buffer::write_decompress_zstd(std::size_t pos, std::size_t count, void cons
                    lol_trace_var("{:#x}", count),
                    lol_trace_var("{:p}", src),
                    lol_trace_var("{:#x}", src_count));
-    auto ctx = std::shared_ptr<libdeflate_decompressor>(libdeflate_alloc_decompressor(), libdeflate_free_decompressor);
     auto const maxendpos = pos + count;
     lol_throw_if(maxendpos < pos);
     lol_throw_if(impl_reserve(maxendpos));
@@ -158,6 +181,29 @@ auto Buffer::write_decompress_zstd(std::size_t pos, std::size_t count, void cons
     size_ = std::max(size_, pos + result);
 }
 
+auto Buffer::write_decompress_zstd_hack(std::size_t pos, std::size_t count, void const* src, std::size_t src_count)
+    -> void {
+    lol_trace_func(lol_trace_var("{:#x}", size_),
+                   lol_trace_var("{:#x}", pos),
+                   lol_trace_var("{:#x}", count),
+                   lol_trace_var("{:p}", src),
+                   lol_trace_var("{:#x}", src_count));
+    auto const maxendpos = pos + count;
+    lol_throw_if(maxendpos < pos);
+    lol_throw_if(impl_reserve(maxendpos));
+    auto frame_start = find_zstd_magic({(char const*)src, src_count});
+    if (frame_start) {
+        lol_throw_if(frame_start > count);
+        this->write(pos, src, frame_start);
+        pos += frame_start;
+        count -= frame_start;
+        src = (char const*)src + frame_start;
+        src_count -= frame_start;
+    }
+    if (count == 0 && src_count == 0) return;
+    write_decompress_zstd(pos, count, src, src_count);
+}
+
 auto Buffer::write_compress_defl(std::size_t pos, void const* src, std::size_t src_count, int level) -> std::size_t {
     lol_trace_func(lol_trace_var("{:#x}", size_),
                    lol_trace_var("{:#x}", pos),

diff --git a/cslol-tools/lib/lol/io/buffer.hpp b/cslol-tools/lib/lol/io/buffer.hpp
@@ -42,6 +42,9 @@ namespace lol::io {
 
         auto write_decompress_zstd(std::size_t pos, std::size_t count, void const* src, std::size_t src_count) -> void;
 
+        auto write_decompress_zstd_hack(std::size_t pos, std::size_t count, void const* src, std::size_t src_count)
+            -> void;
+
         auto write_compress_defl(std::size_t pos, void const* src, std::size_t src_count, int level = 6) -> std::size_t;
 
         auto write_compress_zlib(std::size_t pos, void const* src, std::size_t src_count, int level = 6) -> std::size_t;

diff --git a/cslol-tools/lib/lol/io/bytes.hpp b/cslol-tools/lib/lol/io/bytes.hpp
@@ -72,6 +72,12 @@ namespace lol::io {
             return result;
         }
 
+        auto copy_decompress_zstd_hack(std::size_t size_decompressed) const -> Bytes {
+            auto result = Bytes();
+            result.write_decompress_zstd_hack(0, size_decompressed, data_, size_);
+            return result;
+        }
+
         auto copy_compress_defl(int level = 6) const -> Bytes {
             auto result = Bytes();
             result.write_compress_defl(0, data_, size_, level);

diff --git a/cslol-tools/lib/lol/wad/archive.cpp b/cslol-tools/lib/lol/wad/archive.cpp
@@ -128,6 +128,8 @@ auto Archive::write_to_file(fs::path const& path) const -> void {
             } else {
                 loc = {
                     .type = data.type(),
+                    .subchunk_count = data.subchunk_count(),
+                    .subchunk_index = data.subchunk_index(),
                     .offset = data_cur,
                     .size = data.bytes_size(),
                     .size_decompressed = data.size_decompressed(),
@@ -146,9 +148,9 @@ auto Archive::write_to_file(fs::path const& path) const -> void {
                 .size = (std::uint32_t)loc.size,
                 .size_decompressed = (std::uint32_t)loc.size_decompressed,
                 .type = loc.type,
-                .subchunk_count = 0,
+                .subchunk_count = loc.subchunk_count,
                 .is_duplicate = 0,
-                .subchunk_index = 0,
+                .subchunk_index = loc.subchunk_index,
                 .checksum = loc.checksum,
             });
         };

diff --git a/cslol-tools/lib/lol/wad/entry.cpp b/cslol-tools/lib/lol/wad/entry.cpp
@@ -54,6 +54,22 @@ auto EntryData::from_zstd(io::Bytes bytes, std::size_t decompressed, std::uint64
     return result;
 }
 
+auto EntryData::from_zstd_multi(io::Bytes bytes,
+                                std::size_t decompressed,
+                                std::uint64_t checksum,
+                                std::uint8_t subchunk_count,
+                                std::uint16_t subchunk_index) -> EntryData {
+    auto result = EntryData(std::make_shared<Impl>());
+    result.impl_->type = EntryType::ZstdMulti;
+    result.impl_->subchunk_count = subchunk_count;
+    result.impl_->subchunk_index = subchunk_index;
+    result.impl_->size_decompressed = decompressed;
+    result.impl_->checksum = checksum;
+    result.impl_->bytes = std::move(bytes);
+    result.impl_->compressed = result.impl_;
+    return result;
+}
+
 auto EntryData::from_file(fs::path const& path) -> EntryData { return from_raw(io::Bytes::from_file(path), 0); }
 
 auto EntryData::from_loc(io::Bytes src, EntryLoc const& loc) -> EntryData {
@@ -67,8 +83,13 @@ auto EntryData::from_loc(io::Bytes src, EntryLoc const& loc) -> EntryData {
         case EntryType::Gzip:
             return EntryData::from_gzip(std::move(bytes), loc.size_decompressed, loc.checksum);
         case EntryType::Zstd:
-        case EntryType::ZstdMulti:
             return EntryData::from_zstd(std::move(bytes), loc.size_decompressed, loc.checksum);
+        case EntryType::ZstdMulti:
+            return EntryData::from_zstd_multi(std::move(bytes),
+                                              loc.size_decompressed,
+                                              loc.checksum,
+                                              loc.subchunk_count,
+                                              loc.subchunk_index);
         default:
             lol_throw_msg("Unknown EntryType: {:#x}", (unsigned)loc.type);
     }
@@ -98,7 +119,8 @@ auto EntryData::extension() const -> std::string_view {
             case EntryType::Gzip:
                 result = into_decompressed().extension();
                 break;
-            case EntryType::Zstd: {
+            case EntryType::Zstd:
+            case EntryType::ZstdMulti: {
                 char buffer[128];
                 auto buffer_cap = std::min(size_decompressed(), sizeof(buffer));
                 auto buffer_size = bytes().readsome_decompress_zstd(0, bytes_size(), buffer, buffer_cap);
@@ -134,6 +156,9 @@ auto EntryData::into_compressed() const -> EntryData {
             case EntryType::Zstd:
                 result = impl_;
                 break;
+            case EntryType::ZstdMulti:
+                result = into_decompressed().into_compressed().impl_;
+                break;
             default:
                 lol_throw_msg("Unreachable Type!");
         }
@@ -168,6 +193,13 @@ auto EntryData::into_decompressed() const -> EntryData {
                 result->compressed = impl_;
                 result->decompressed = result;
                 break;
+            case EntryType::ZstdMulti:
+                result = std::make_shared<Impl>();
+                result->type = EntryType::Raw;
+                result->size_decompressed = size_decompressed();
+                result->bytes = bytes().copy_decompress_zstd_hack(size_decompressed());
+                result->decompressed = result;
+                break;
             default:
                 lol_throw_msg("Unreachable Type!");
         }
@@ -198,6 +230,13 @@ auto EntryData::into_optimal() const -> EntryData {
                 result = into_decompressed();
             }
             break;
+        case EntryType::ZstdMulti:
+            if (auto ext = extension(); ext == ".bnk" || ext == ".wpk") {
+                result = into_decompressed();
+            } else {
+                result = into_compressed();
+            }
+            break;
         default:
             lol_throw_msg("Unreachable Type!");
     }

diff --git a/cslol-tools/lib/lol/wad/entry.hpp b/cslol-tools/lib/lol/wad/entry.hpp
@@ -17,6 +17,8 @@ namespace lol::wad {
 
     struct EntryLoc {
         EntryType type;
+        std::uint8_t subchunk_count = {};
+        std::uint16_t subchunk_index = {};
         std::uint64_t offset;
         std::uint64_t size;
         std::uint64_t size_decompressed;
@@ -43,6 +45,12 @@ namespace lol::wad {
 
         static auto from_zstd(io::Bytes data, std::size_t decompressed, std::uint64_t checksum) -> EntryData;
 
+        static auto from_zstd_multi(io::Bytes data,
+                                    std::size_t decompressed,
+                                    std::uint64_t checksum,
+                                    std::uint8_t subchunk_count,
+                                    std::uint16_t subchunk_index) -> EntryData;
+
         static auto from_file(fs::path const& path) -> EntryData;
 
         static auto from_loc(io::Bytes src, EntryLoc const& loc) -> EntryData;
@@ -59,6 +67,10 @@ namespace lol::wad {
 
         auto type() const noexcept -> EntryType { return impl_->type; }
 
+        auto subchunk_count() const noexcept -> std::uint8_t { return impl_->subchunk_count; }
+
+        auto subchunk_index() const noexcept -> std::uint16_t { return impl_->subchunk_index; }
+
         auto bytes() const noexcept -> io::Bytes { return impl_->bytes; }
 
         auto bytes_data() const noexcept -> char const* { return impl_->bytes.data(); }
@@ -81,6 +93,8 @@ namespace lol::wad {
         struct Impl {
             EntryType type = EntryType::Raw;
             bool is_optimal = false;
+            std::uint8_t subchunk_count = {};
+            std::uint16_t subchunk_index = {};
             std::optional<std::string_view> extension = {};
             io::Bytes bytes = {};
             std::size_t size_decompressed = 0;

diff --git a/cslol-tools/lib/lol/wad/toc.cpp b/cslol-tools/lib/lol/wad/toc.cpp
@@ -42,6 +42,8 @@ auto TOC::read(io::Bytes src) noexcept -> char const* {
                 .loc =
                     {
                         .type = entry_raw.type,
+                        .subchunk_count = entry_raw.subchunk_count,
+                        .subchunk_index = entry_raw.subchunk_index,
                         .offset = entry_raw.offset,
                         .size = entry_raw.size,
                         .size_decompressed = entry_raw.size_decompressed,

diff --git a/cslol-tools/src/main_mod_tools.cpp b/cslol-tools/src/main_mod_tools.cpp
@@ -10,6 +10,7 @@
 #include <lol/wad/archive.hpp>
 #include <lol/wad/index.hpp>
 #include <thread>
+#include <unordered_set>
 
 using namespace lol;
 
@@ -192,11 +193,19 @@ static auto mod_mkoverlay(fs::path src, fs::path dst, fs::path game, fs::names m
     lol_throw_if_msg(game_index.mounts.empty(), "Not a valid Game folder");
     game_index.remove_filter(noTFT ? FILTER_TFT : FILTER_NONE);
 
+    auto blocked = std::unordered_set<hash::Xxh64>{};
+    for (auto const& [_, mounted] : game_index.mounts) {
+        auto subchunk_name = fs::path(mounted.relpath).replace_extension(".SubChunkTOC").generic_string();
+        blocked.insert(hash::Xxh64(subchunk_name));
+    }
     auto mod_queue = std::vector<wad::Index>{};
 
     logi("Reading mods");
     for (auto const& mod_name : mods) {
         auto mod_index = wad::Index::from_mod_folder(src / mod_name);
+        for (auto& [path_, mounted] : mod_index.mounts) {
+            std::erase_if(mounted.archive.entries, [&](auto const& kvp) { return blocked.contains(kvp.first); });
+        }
         if (mod_index.mounts.empty()) {
             logw("Empty mod: {}", mod_index.name);
             continue;