support LUS archivemanager refactor (#14)

5 files changed, 191 insertions, 13 deletions

diff --git a/OTRExporter/CMakeLists.txt b/OTRExporter/CMakeLists.txt
index 9d012af..c582148 100644
--- a/OTRExporter/CMakeLists.txt
+++ b/OTRExporter/CMakeLists.txt
@@ -14,6 +14,7 @@ set(Header_Files
     "CutsceneExporter.h"

     "DisplayListExporter.h"

     "Exporter.h"

+    "ExporterArchive.h"

     "Main.h"

     "MtxExporter.h"

     "PathExporter.h"

@@ -40,6 +41,7 @@ set(Source_Files
     "CutsceneExporter.cpp"

     "DisplayListExporter.cpp"

     "Exporter.cpp"

+    "ExporterArchive.cpp"

     "Main.cpp"

     "MtxExporter.cpp"

     "PathExporter.cpp"

@@ -116,8 +118,6 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Windows")
 			">"

 			"_CONSOLE;"

 			"_CRT_SECURE_NO_WARNINGS;"

-			"UNICODE;"

-			"_UNICODE"

 			STORMLIB_NO_AUTO_LINK

 		)

 	elseif("${CMAKE_VS_PLATFORM_NAME}" STREQUAL "x86")

@@ -131,8 +131,6 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Windows")
 			">"

 			"WIN32;"

 			"_CONSOLE;"

-			"UNICODE;"

-			"_UNICODE"

 			STORMLIB_NO_AUTO_LINK

 		)

 	endif()

@@ -148,8 +146,6 @@ if ("${CMAKE_CXX_COMPILER_ID}" MATCHES "GNU|Clang|AppleClang")
 		">"

 		"_CONSOLE;"

 		"_CRT_SECURE_NO_WARNINGS;"

-		"UNICODE;"

-		"_UNICODE"

 	)

 endif()

 ################################################################################

diff --git a/OTRExporter/ExporterArchive.cpp b/OTRExporter/ExporterArchive.cpp
new file mode 100644
index 0000000..6b65b39
--- /dev/null
+++ b/OTRExporter/ExporterArchive.cpp
@@ -0,0 +1,152 @@
+#include "ExporterArchive.h"
+#include "Utils/StringHelper.h"
+#include <StrHash64/StrHash64.h>
+#include <filesystem>
+
+ExporterArchive::ExporterArchive(const std::string& path, bool enableWriting) : mPath(path) {
+    mMpq = nullptr;
+    Load(enableWriting);
+}
+
+ExporterArchive::~ExporterArchive() {
+    Unload();
+}
+
+bool ExporterArchive::Load(bool enableWriting) {
+    HANDLE mpqHandle = NULL;
+
+    bool baseLoaded = false;
+    std::string fullPath = std::filesystem::absolute(mPath).string();
+
+    bool openArchiveSuccess;
+    {
+        const std::lock_guard<std::mutex> lock(mMutex);
+        openArchiveSuccess =
+            SFileOpenArchive(fullPath.c_str(), 0, enableWriting ? 0 : MPQ_OPEN_READ_ONLY, &mpqHandle);
+    }
+    if (openArchiveSuccess) {
+        printf("Opened mpq file ");
+        printf(fullPath.c_str());
+        printf("\n");
+        mMpq = mpqHandle;
+        mPath = fullPath;
+
+        
+        baseLoaded = true;
+    }
+
+    if (!baseLoaded) {
+        printf("No valid OTR file was provided.");
+        return false;
+    }
+
+    return true;
+}
+
+bool ExporterArchive::Unload() {
+    bool success = true;
+
+    bool closeArchiveSuccess;
+    {
+        const std::lock_guard<std::mutex> lock(mMutex);
+        closeArchiveSuccess = SFileCloseArchive(mMpq);
+    }
+    if (!closeArchiveSuccess) {
+        printf("Failed to close mpq\n");
+        success = false;
+    }
+
+    mMpq = nullptr;
+
+    return success;
+}
+
+std::shared_ptr<ExporterArchive> ExporterArchive::CreateArchive(const std::string& archivePath, size_t fileCapacity) {
+    auto archive = std::make_shared<ExporterArchive>(archivePath, true);
+
+    bool success;
+    {
+        const std::lock_guard<std::mutex> lock(archive->mMutex);
+        success = SFileCreateArchive(archivePath.c_str(), MPQ_CREATE_LISTFILE | MPQ_CREATE_ATTRIBUTES | MPQ_CREATE_ARCHIVE_V2,
+                                     fileCapacity, &archive->mMpq);
+    }
+    int32_t error = GetLastError();
+
+    if (success) {
+        return archive;
+    } else {
+        printf("We tried to create an archive, but it has fallen and cannot get up.\n");
+        return nullptr;
+    }
+}
+
+bool ExporterArchive::AddFile(const std::string& filePath, uintptr_t fileData, DWORD fileSize) {
+    HANDLE hFile;
+#ifdef _WIN32
+    SYSTEMTIME sysTime;
+    GetSystemTime(&sysTime);
+    FILETIME t;
+    SystemTimeToFileTime(&sysTime, &t);
+    ULONGLONG theTime = static_cast<uint64_t>(t.dwHighDateTime) << (sizeof(t.dwHighDateTime) * 8) | t.dwLowDateTime;
+#else
+    time_t theTime;
+    time(&theTime);
+#endif
+
+    std::string updatedPath = filePath;
+
+    StringHelper::ReplaceOriginal(updatedPath, "\\", "/");
+
+    bool createFileSuccess;
+    {
+        const std::lock_guard<std::mutex> lock(mMutex);
+        createFileSuccess =
+            SFileCreateFile(mMpq, updatedPath.c_str(), theTime, fileSize, 0, MPQ_FILE_COMPRESS, &hFile);
+    }
+    if (!createFileSuccess) {
+        printf("Failed to create file.\n");
+        return false;
+    }
+
+    bool writeFileSuccess;
+    {
+        const std::lock_guard<std::mutex> lock(mMutex);
+        writeFileSuccess = SFileWriteFile(hFile, (void*)fileData, fileSize, MPQ_COMPRESSION_ZLIB);
+    }
+    if (!writeFileSuccess) {
+        printf("Failed to write.\n");
+        bool closeFileSuccess;
+        {
+            const std::lock_guard<std::mutex> lock(mMutex);
+            closeFileSuccess = SFileCloseFile(hFile);
+        }
+        if (!closeFileSuccess) {
+            printf("Failed to close.\n");
+        }
+        return false;
+    }
+
+    bool finishFileSuccess;
+    {
+        const std::lock_guard<std::mutex> lock(mMutex);
+        finishFileSuccess = SFileFinishFile(hFile);
+    }
+    if (!finishFileSuccess) {
+        printf("Failed to finish file.\n");
+        bool closeFileSuccess;
+        {
+            const std::lock_guard<std::mutex> lock(mMutex);
+            closeFileSuccess = SFileCloseFile(hFile);
+        }
+        if (!closeFileSuccess) {
+            printf("Failed to close after finish failure.\n");
+        }
+        return false;
+    }
+    // SFileFinishFile already frees the handle, so no need to close it again.
+
+    mAddedFiles.push_back(updatedPath);
+    mHashes[CRC64(updatedPath.c_str())] = updatedPath;
+
+    return true;
+}
diff --git a/OTRExporter/ExporterArchive.h b/OTRExporter/ExporterArchive.h
new file mode 100644
index 0000000..fc65444
--- /dev/null
+++ b/OTRExporter/ExporterArchive.h
@@ -0,0 +1,29 @@
+#pragma once
+
+#undef _DLL
+
+#include <string>
+#include <vector>
+#include <unordered_map>
+#include <memory>
+#include <mutex>
+#include <StormLib.h>
+
+class ExporterArchive : public std::enable_shared_from_this<ExporterArchive> {
+  public:
+    ExporterArchive(const std::string& path, bool enableWriting);
+    ~ExporterArchive();
+
+    static std::shared_ptr<ExporterArchive> CreateArchive(const std::string& archivePath, size_t fileCapacity);
+    bool AddFile(const std::string& filePath, uintptr_t fileData, DWORD fileSize);
+
+  private:
+    std::string mPath;
+    HANDLE mMpq;
+    std::mutex mMutex;
+    std::vector<std::string> mAddedFiles;
+    std::unordered_map<uint64_t, std::string> mHashes;
+
+    bool Load(bool enableWriting);
+    bool Unload();
+};
diff --git a/OTRExporter/Main.cpp b/OTRExporter/Main.cpp
index 4fc3c48..f4e2a00 100644
--- a/OTRExporter/Main.cpp
+++ b/OTRExporter/Main.cpp
@@ -31,7 +31,7 @@ std::string customOtrFileName = "";
 std::string customAssetsPath = "";
 std::string portVersionString = "0.0.0";
 
-std::shared_ptr<LUS::Archive> otrArchive;
+std::shared_ptr<ExporterArchive> otrArchive;
 BinaryWriter* fileWriter;
 std::chrono::steady_clock::time_point fileStart, resStart;
 std::map<std::string, std::vector<char>> files;
@@ -53,9 +53,9 @@ static void ExporterParseFileMode(const std::string& buildMode, ZFileMode& fileM
 		printf("BOTR: Generating OTR Archive...\n");
 
 		if (DiskFile::Exists(otrFileName))
-			otrArchive = std::shared_ptr<LUS::Archive>(new LUS::Archive(otrFileName, true));
+			otrArchive = std::shared_ptr<ExporterArchive>(new ExporterArchive(otrFileName, true));
 		else
-			otrArchive = LUS::Archive::CreateArchive(otrFileName, 40000);
+			otrArchive = ExporterArchive::CreateArchive(otrFileName, 40000);
 
 		auto lst = Directory::ListFiles("Extract");
 
@@ -125,7 +125,7 @@ static void ExporterProgramEnd()
 		printf("Created version file.\n");
 
 		printf("Generating OTR Archive...\n");
-		otrArchive = LUS::Archive::CreateArchive(otrFileName, 40000);
+		otrArchive = ExporterArchive::CreateArchive(otrFileName, 40000);
 
 		printf("Adding game version file.\n");
 		otrArchive->AddFile("version", (uintptr_t)versionStream->ToVector().data(), versionStream->GetLength());
@@ -167,7 +167,7 @@ static void ExporterProgramEnd()
 	const auto& lst = Directory::ListFiles(customAssetsPath);
 
 	printf("Generating Custom OTR Archive...\n");
-	std::shared_ptr<LUS::Archive> customOtr = LUS::Archive::CreateArchive(customOtrFileName, 4096);
+	std::shared_ptr<ExporterArchive> customOtr = ExporterArchive::CreateArchive(customOtrFileName, 4096);
 
 	printf("Adding portVersion file.\n");
 	customOtr->AddFile("portVersion", (uintptr_t)portVersionStream->ToVector().data(), portVersionStream->GetLength());
diff --git a/OTRExporter/Main.h b/OTRExporter/Main.h
index 3aa68c7..6bdc77d 100644
--- a/OTRExporter/Main.h
+++ b/OTRExporter/Main.h
@@ -1,8 +1,9 @@
 #pragma once
 
 #include <libultraship/bridge.h>
+#include "ExporterArchive.h"
 
-extern std::shared_ptr<LUS::Archive> otrArchive;
+extern std::shared_ptr<ExporterArchive> otrArchive;
 extern std::map<std::string, std::vector<char>> files;
 
-void AddFile(std::string fName, std::vector<char> data);
\ No newline at end of file
+void AddFile(std::string fName, std::vector<char> data);