1 files changed, 63 insertions, 12 deletions

diff --git a/src/archive/ZWrapper.cpp b/src/archive/ZWrapper.cpp
index 480dafc..335a3a9 100644
--- a/src/archive/ZWrapper.cpp
+++ b/src/archive/ZWrapper.cpp
@@ -2,6 +2,8 @@
 #include <filesystem>
 #include <iostream>
 #include <fstream>
+#include <cstring>
+#include <stdexcept>
 
 #include "spdlog/spdlog.h"
 #include <Companion.h>
@@ -11,17 +13,40 @@ namespace fs = std::filesystem;
 
 ZWrapper::ZWrapper(const std::string& path) {
     this->mPath = path;
-    this->mZip = new miniz_cpp::zip_file;
+    this->mTempPath = path + ".tmp";
 }
 
-int32_t ZWrapper::CreateArchive() {
-    SPDLOG_INFO("Loaded ZIP (O2R) archive: {}", mPath.c_str());
+ZWrapper::~ZWrapper() {
+    if (this->mZip != nullptr) {
+        if (this->mOpen) {
+            mz_zip_writer_end(reinterpret_cast<mz_zip_archive*>(this->mZip));
+        }
+        delete reinterpret_cast<mz_zip_archive*>(this->mZip);
+        this->mZip = nullptr;
+    }
+    if (!this->mFinalized) {
+        std::error_code ec;
+        fs::remove(this->mTempPath, ec);
+    }
+}
+
+int32_t ZWrapper::CreateArchive(void) {
+    auto* zip = new mz_zip_archive;
+    std::memset(zip, 0, sizeof(mz_zip_archive));
+
+    if (!mz_zip_writer_init_file(zip, this->mTempPath.c_str(), 0)) {
+        delete zip;
+        throw std::runtime_error("[ZWrapper] Could not open '" + this->mTempPath + "' for writing");
+    }
+
+    this->mZip = reinterpret_cast<mz_zip_archive_tag*>(zip);
+    this->mOpen = true;
+    SPDLOG_INFO("Loaded ZIP (O2R) archive: {}", this->mPath.c_str());
     return 0;
 }
 
 bool ZWrapper::AddFile(const std::string& path, std::vector<char> data) {
-    char* fileData = data.data();
-    size_t fileSize = data.size();
+    const size_t fileSize = data.size();
 
     if (Companion::Instance != nullptr && Companion::Instance->IsDebug()) {
         SPDLOG_INFO("Creating debug file: debug/{}", path);
@@ -30,21 +55,47 @@ bool ZWrapper::AddFile(const std::string& path, std::vector<char> data) {
             fs::create_directories(fs::path(dpath).parent_path());
         }
         std::ofstream stream(dpath, std::ios::binary);
-        stream.write(fileData, fileSize);
+        stream.write(data.data(), fileSize);
         stream.close();
     }
 
-    try {
-        this->mZip->writebytes(path, data);
-    } catch (const std::exception& e) {
+    if (this->mZip == nullptr || !this->mOpen) {
+        throw std::runtime_error("[ZWrapper] Cannot add '" + path + "': the archive is not open");
+    }
+
+    std::lock_guard<std::mutex> guard(this->mMutex);
+
+    if (!mz_zip_writer_add_mem(reinterpret_cast<mz_zip_archive*>(this->mZip), path.c_str(), data.data(), fileSize,
+                               MZ_BEST_COMPRESSION)) {
         // miniz's own error says nothing useful, so log which entry blew up and how big.
-        SPDLOG_ERROR("[ZWrapper] Failed to add '{}' ({} bytes): {}", path, fileSize, e.what());
-        throw;
+        SPDLOG_ERROR("[ZWrapper] Failed to add '{}' ({} bytes)", path, fileSize);
+        throw std::runtime_error("[ZWrapper] Failed to add '" + path + "'");
     }
+
     return true;
 }
 
 int32_t ZWrapper::Close(void) {
-    this->mZip->save(this->mPath);
+    if (this->mZip == nullptr || !this->mOpen) {
+        return -1;
+    }
+
+    auto* zip = reinterpret_cast<mz_zip_archive*>(this->mZip);
+    const bool finalized = mz_zip_writer_finalize_archive(zip);
+    mz_zip_writer_end(zip);
+    this->mOpen = false;
+
+    if (!finalized) {
+        throw std::runtime_error("[ZWrapper] Failed to finalize '" + this->mPath + "'");
+    }
+
+    std::error_code ec;
+    fs::rename(this->mTempPath, this->mPath, ec);
+    if (ec) {
+        throw std::runtime_error("[ZWrapper] Could not move '" + this->mTempPath + "' to '" + this->mPath +
+                                 "': " + ec.message());
+    }
+    this->mFinalized = true;
+
     return 0;
 }