From 7fdb4317d87601d6ac7d32996cf859473c746823 Mon Sep 17 00:00:00 2001
From: Stenzek <stenzek@gmail.com>
Date: Sun, 25 Aug 2024 17:48:12 +1000
Subject: [PATCH] Host: Add ReadCompressedResourceFile()

---
 src/core/host.cpp | 50 ++++++++++++++++++++++++++++++++++++++++++++++-
 src/util/host.h   |  5 ++++-
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/src/core/host.cpp b/src/core/host.cpp
index 2069f2044..9be5577a1 100644
--- a/src/core/host.cpp
+++ b/src/core/host.cpp
@@ -1,5 +1,5 @@
 // SPDX-FileCopyrightText: 2019-2024 Connor McLaughlin <stenzek@gmail.com>
-// SPDX-License-Identifier: (GPL-3.0 OR CC-BY-NC-ND-4.0)
+// SPDX-License-Identifier: (GPL-3.0 OR PolyForm-Strict-1.0.0)
 
 #include "host.h"
 #include "fullscreen_ui.h"
@@ -18,11 +18,15 @@
 #include "common/error.h"
 #include "common/layered_settings_interface.h"
 #include "common/log.h"
+#include "common/path.h"
 #include "common/string_util.h"
 
 #include "imgui.h"
 
 #include <cstdarg>
+#include <limits>
+#include <zstd.h>
+#include <zstd_errors.h>
 
 Log_SetChannel(Host);
 
@@ -41,6 +45,50 @@ SettingsInterface* Host::GetSettingsInterface()
   return &s_layered_settings_interface;
 }
 
+std::optional<DynamicHeapArray<u8>> Host::ReadCompressedResourceFile(std::string_view filename, bool allow_override)
+{
+  std::optional<DynamicHeapArray<u8>> ret = Host::ReadResourceFile(filename, allow_override);
+  if (ret.has_value())
+  {
+    const std::string_view extension = Path::GetExtension(filename);
+    if (StringUtil::EqualNoCase(extension, "zst"))
+    {
+      // Need to zstd decompress.
+      const unsigned long long decompressed_size = ZSTD_getFrameContentSize(ret->data(), ret->size());
+      if (decompressed_size == ZSTD_CONTENTSIZE_UNKNOWN || decompressed_size == ZSTD_CONTENTSIZE_ERROR ||
+          decompressed_size >= std::numeric_limits<size_t>::max()) [[unlikely]]
+      {
+        ERROR_LOG("Failed to get size of {}.", filename);
+        ret.reset();
+        return ret;
+      }
+
+      DEV_LOG("Decompressing resource {}: {} => {} bytes", filename, ret->size(), decompressed_size);
+
+      DynamicHeapArray<u8> decompressed_data(decompressed_size);
+      const size_t result =
+        ZSTD_decompress(decompressed_data.data(), decompressed_data.size(), ret->data(), ret->size());
+      if (ZSTD_isError(result)) [[unlikely]]
+      {
+        const char* errstr = ZSTD_getErrorString(ZSTD_getErrorCode(result));
+        ERROR_LOG("ZSTD_decompress() for {} failed: {}", filename, errstr ? errstr : "<unknown>");
+        ret.reset();
+        return ret;
+      }
+      if (result < decompressed_size) [[unlikely]]
+      {
+        ERROR_LOG("ZSTD_decompress() for {} only returned {} of {} bytes.", filename, result, decompressed_size);
+        ret.reset();
+        return ret;
+      }
+
+      ret = std::move(decompressed_data);
+    }
+  }
+
+  return ret;
+}
+
 std::string Host::GetBaseStringSettingValue(const char* section, const char* key, const char* default_value /*= ""*/)
 {
   std::unique_lock lock(s_settings_mutex);
diff --git a/src/util/host.h b/src/util/host.h
index 52f5cc6bd..26567a8c2 100644
--- a/src/util/host.h
+++ b/src/util/host.h
@@ -3,8 +3,8 @@
 
 #pragma once
 
-#include "common/types.h"
 #include "common/heap_array.h"
+#include "common/types.h"
 
 #include <ctime>
 #include <optional>
@@ -25,6 +25,9 @@ std::optional<std::string> ReadResourceFileToString(std::string_view filename, b
 /// Returns the modified time of a resource.
 std::optional<std::time_t> GetResourceFileTimestamp(std::string_view filename, bool allow_override);
 
+/// Reads a potentially-compressed file from the resources directory of the application.
+std::optional<DynamicHeapArray<u8>> ReadCompressedResourceFile(std::string_view filename, bool allow_override);
+
 /// Reports a fatal error on the main thread. This does not assume that the main window exists,
 /// unlike ReportErrorAsync(), and will exit the application after the popup is closed.
 void ReportFatalError(std::string_view title, std::string_view message);