Merge pull request #337 from Sssnowyyy/nx2-review-fixes

nx2: fix pokopia regressions, textures that handle shader computes, and more

3 files changed, 235 insertions, 36 deletions

diff --git a/src/core/file_sys/nup_container.cpp b/src/core/file_sys/nup_container.cpp
index 0e485b44ec..27ac4486ca 100644
--- a/src/core/file_sys/nup_container.cpp
+++ b/src/core/file_sys/nup_container.cpp
@@ -8,11 +8,15 @@
 #include <cstring>
 #include <list>
 #include <map>
+#include <random>
+#include <span>
 #include <unordered_map>
 
+#include <fmt/format.h>
 #include <openssl/evp.h>
 #include <zstd.h>
 
+#include "common/fs/path_util.h"
 #include "common/logging.h"
 #include "common/swap.h"
 #include "core/file_sys/vfs/vfs_offset.h"
@@ -31,6 +35,7 @@ constexpr u32 TocMagic = 0x01434F54; // "TOC\x01" little-endian
 constexpr u8 CodecRaw = 0;
 constexpr u8 CodecZstd = 1;
 constexpr u8 PackageFlagStandaloneUpdate = 0x01;
+constexpr u32 MaxDirDepth = 128;
 
 #pragma pack(push, 1)
 struct Superblock {
@@ -142,8 +147,8 @@ bool Sha256(const u8* data, std::size_t size, std::array<u8, 0x20>& out) {
     return length == out.size();
 }
 
-/// A member file whose extent is zstd-compressed. Decompresses on first read
-/// and caches the result (LRU-evicted) unless the file exceeds the cache cap.
+/// A member file whose extent is zstd-compressed. Decompresses on first read and caches the
+/// result (LRU-evicted); a file over half the cache cap is decompressed once to a temporary file.
 class NupVfsFile final : public VfsFile {
 public:
     NupVfsFile(VirtualFile backing_, u64 extent_offset_, u64 extent_size_, u64 uncompressed_size_,
@@ -151,7 +156,7 @@ public:
                std::size_t cache_index_)
         : backing(std::move(backing_)), extent_offset(extent_offset_), extent_size(extent_size_),
           uncompressed_size(uncompressed_size_), name(std::move(name_)), cache(std::move(cache_)),
-          cache_index(cache_index_), cache_enabled(uncompressed_size <= NupDecompressCache::DecompressCacheCap) {}
+          cache_index(cache_index_), cache_enabled(uncompressed_size <= cache->cap / 2) {}
 
     std::string GetName() const override {
         return name;
@@ -183,6 +188,10 @@ public:
         }
         length = std::min(length, static_cast<std::size_t>(uncompressed_size) - offset);
 
+        if (!cache_enabled) {
+            return ReadSpilled(data, length, offset);
+        }
+
         {
             std::scoped_lock lock{cache->mutex};
             if (const auto* cached = cache->Find(cache_index)) {
@@ -198,9 +207,7 @@ public:
         }
         {
             std::scoped_lock lock{cache->mutex};
-            if (cache_enabled) {
-                cache->Insert(cache_index, scratch);
-            }
+            cache->Insert(cache_index, scratch);
             if (const auto* cached = cache->Find(cache_index)) {
                 std::memcpy(data, cached->data() + offset, length);
                 return length;
@@ -220,6 +227,88 @@ public:
     }
 
 private:
+    std::size_t ReadSpilled(u8* data, std::size_t length, std::size_t offset) const {
+        std::shared_ptr<NupSpillFile> spill;
+        {
+            std::scoped_lock lock{cache->mutex};
+            spill = cache->GetSpill(cache_index);
+        }
+        std::scoped_lock lock{spill->mutex};
+        if (!spill->attempted) {
+            spill->attempted = true;
+            DecompressToFile(*spill);
+        }
+        if (!spill->file.IsOpen() || !spill->file.Seek(static_cast<s64>(offset))) {
+            return 0;
+        }
+        return spill->file.ReadSpan(std::span{data, length});
+    }
+
+    void DecompressToFile(NupSpillFile& spill) const {
+        const auto fail = [&](const char* reason) {
+            LOG_ERROR(Loader, "NUP zstd decompression of {} to a temporary file failed: {}", name,
+                      reason);
+            spill.file.Close();
+            if (!spill.path.empty()) {
+                std::error_code ec;
+                std::filesystem::remove(spill.path, ec);
+            }
+        };
+        if (extent_size > backing->GetSize() || extent_offset > backing->GetSize() - extent_size) {
+            return fail("extent out of range");
+        }
+
+        std::random_device random;
+        std::error_code ec;
+        std::filesystem::create_directories(cache->spill_dir, ec);
+        spill.path =
+            cache->spill_dir / fmt::format("{:08x}{:08x}-{}.tmp", random(), random(), cache_index);
+        Common::FS::IOFile out{spill.path, Common::FS::FileAccessMode::Write};
+        if (!out.IsOpen()) {
+            return fail("cannot create the temporary file");
+        }
+
+        std::unique_ptr<ZSTD_DCtx, decltype(&ZSTD_freeDCtx)> dctx{ZSTD_createDCtx(),
+                                                                  &ZSTD_freeDCtx};
+        std::vector<u8> in_buffer(ZSTD_DStreamInSize());
+        std::vector<u8> out_buffer(ZSTD_DStreamOutSize());
+        u64 read = 0;
+        u64 written = 0;
+        bool frame_done = false;
+        while (!frame_done && read < extent_size) {
+            const auto chunk =
+                static_cast<std::size_t>(std::min<u64>(in_buffer.size(), extent_size - read));
+            if (backing->Read(in_buffer.data(), chunk,
+                              static_cast<std::size_t>(extent_offset + read)) != chunk) {
+                return fail("short read of the compressed data");
+            }
+            read += chunk;
+            ZSTD_inBuffer input{in_buffer.data(), chunk, 0};
+            ZSTD_outBuffer output{};
+            do {
+                output = {out_buffer.data(), out_buffer.size(), 0};
+                const std::size_t result = ZSTD_decompressStream(dctx.get(), &output, &input);
+                if (ZSTD_isError(result) || output.pos > uncompressed_size - written) {
+                    return fail("corrupt or oversized data");
+                }
+                if (out.WriteSpan(std::span<const u8>{out_buffer.data(), output.pos}) !=
+                    output.pos) {
+                    return fail("cannot write the temporary file");
+                }
+                written += output.pos;
+                frame_done = result == 0;
+            } while (!frame_done && (input.pos < input.size || output.pos == output.size));
+        }
+        if (!frame_done || written != uncompressed_size) {
+            return fail("size mismatch");
+        }
+        out.Close();
+        spill.file.Open(spill.path, Common::FS::FileAccessMode::Read);
+        if (!spill.file.IsOpen()) {
+            fail("cannot reopen the temporary file");
+        }
+    }
+
     std::vector<u8> Decompress() const {
         if (extent_size > backing->GetSize() || extent_offset > backing->GetSize() - extent_size) {
             return {};
@@ -251,6 +340,13 @@ private:
 
 } // namespace
 
+NupSpillFile::~NupSpillFile() {
+    file.Close();
+    if (!path.empty()) {
+        std::error_code ec;
+        std::filesystem::remove(path, ec);
+    }
+}
 
 bool NupContainer::IsNupFile(const VirtualFile& file) {
     if (file == nullptr || file->GetSize() < SuperblockSize) {
@@ -261,8 +357,13 @@ bool NupContainer::IsNupFile(const VirtualFile& file) {
     return magic == HeaderMagic;
 }
 
-NupContainer::NupContainer(VirtualFile file_)
+NupContainer::NupContainer(VirtualFile file_, std::size_t memory_cache_cap,
+                           std::filesystem::path spill_dir_)
     : file(std::move(file_)), cache(std::make_shared<NupDecompressCache>()) {
+    cache->cap = memory_cache_cap;
+    cache->spill_dir = spill_dir_.empty()
+                           ? Common::FS::GetEdenPath(Common::FS::EdenPath::CacheDir) / "nup"
+                           : std::move(spill_dir_);
     if (file == nullptr) {
         return;
     }
@@ -293,9 +394,9 @@ bool NupContainer::Parse() {
     toc_size = super.toc_size;
     blob_offset = super.blob_offset;
     file_size = super.file_size;
-    if (file_size != file->GetSize() || toc_offset < SuperblockSize ||
-        blob_offset < toc_offset + toc_size || blob_offset % 0x10 != 0 ||
-        blob_offset > file_size) {
+    if (file_size != file->GetSize() || toc_offset < SuperblockSize || blob_offset > file_size ||
+        toc_offset > blob_offset || toc_size > blob_offset - toc_offset ||
+        blob_offset % 0x10 != 0) {
         status = Loader::ResultStatus::ErrorNUPBadFormat;
         return false;
     }
@@ -325,15 +426,22 @@ bool NupContainer::Parse() {
         return false;
     }
 
+    // The counts come from the file, so bound them by the TOC size before multiplying.
+    if (header.dir_count > toc.size() / DirRecordSize ||
+        header.file_count > toc.size() / FileRecordSize) {
+        status = Loader::ResultStatus::ErrorNUPBadFormat;
+        return false;
+    }
     const u64 needed = TocHeaderSize + static_cast<u64>(header.package_count) * PackageRecordSize +
                        header.dir_count * DirRecordSize + header.file_count * FileRecordSize;
-    const u64 children_offset = needed;
-    const u64 children_bytes = header.pool_offset - needed;
-    if (toc.size() < needed || header.pool_offset < needed ||
-        header.pool_offset + header.pool_size > toc.size() || children_bytes % sizeof(u32) != 0) {
+    if (needed > toc.size() || header.pool_offset < needed || header.pool_offset > toc.size() ||
+        header.pool_size > toc.size() - header.pool_offset ||
+        (header.pool_offset - needed) % sizeof(u32) != 0) {
         status = Loader::ResultStatus::ErrorNUPBadFormat;
         return false;
     }
+    const u64 children_offset = needed;
+    const u64 children_bytes = header.pool_offset - needed;
     const u8* base = toc.data();
 
     std::size_t pos = TocHeaderSize;
@@ -358,8 +466,8 @@ bool NupContainer::Parse() {
         package.root_dir = record.root_dir;
         package.exefs_dir = record.exefs_dir;
         package.romfs_dir = record.romfs_dir;
-        if (record.meta_size > 0 && record.meta_offset != 0 &&
-            record.meta_offset + record.meta_size <= file_size) {
+        if (record.meta_size > 0 && record.meta_offset != 0 && record.meta_offset <= file_size &&
+            record.meta_size <= file_size - record.meta_offset) {
             package.meta_offset = record.meta_offset;
             package.meta_size = record.meta_size;
         }
@@ -371,8 +479,9 @@ bool NupContainer::Parse() {
         DirRecord record{};
         std::memcpy(&record, base + pos, sizeof(record));
         pos += DirRecordSize;
-        if (record.first_child + record.child_count > children_bytes / sizeof(u32) ||
-            record.name_offset + record.name_length > header.pool_size) {
+        if (static_cast<u64>(record.first_child) + record.child_count >
+                children_bytes / sizeof(u32) ||
+            static_cast<u64>(record.name_offset) + record.name_length > header.pool_size) {
             status = Loader::ResultStatus::ErrorNUPBadFormat;
             return false;
         }
@@ -392,9 +501,9 @@ bool NupContainer::Parse() {
             (record.size == 0 && (record.extent_offset != 0 || record.extent_size != 0)) ||
             (record.size != 0 &&
              (record.extent_size == 0 || record.extent_offset < blob_offset ||
-              record.extent_offset % 0x10 != 0 ||
-              record.extent_offset + record.extent_size > file_size)) ||
-            record.name_offset + record.name_length > header.pool_size) {
+              record.extent_offset % 0x10 != 0 || record.extent_offset > file_size ||
+              record.extent_size > file_size - record.extent_offset)) ||
+            static_cast<u64>(record.name_offset) + record.name_length > header.pool_size) {
             status = Loader::ResultStatus::ErrorNUPBadFormat;
             return false;
         }
@@ -417,6 +526,11 @@ bool NupContainer::Parse() {
         }
     }
 
+    if (!ValidateDirTree()) {
+        status = Loader::ResultStatus::ErrorNUPBadFormat;
+        return false;
+    }
+
     for (auto& package : packages) {
         if (!ParseMeta(package)) {
             status = Loader::ResultStatus::ErrorNUPBadFormat;
@@ -426,6 +540,57 @@ bool NupContainer::Parse() {
     return true;
 }
 
+bool NupContainer::ValidateDirTree() const {
+    enum : u8 { Unvisited, InProgress, Done };
+    struct Frame {
+        u32 dir;
+        u32 next;
+    };
+    std::vector<u8> state(dirs.size(), Unvisited);
+    std::vector<u32> height(dirs.size(), 0);
+    std::vector<Frame> stack;
+    for (u32 start = 0; start < dirs.size(); ++start) {
+        if (state[start] != Unvisited) {
+            continue;
+        }
+        state[start] = InProgress;
+        stack.push_back({start, 0});
+        while (!stack.empty()) {
+            const Frame frame = stack.back();
+            const DirNode& node = dirs[frame.dir];
+            if (frame.next == node.child_count) {
+                state[frame.dir] = Done;
+                stack.pop_back();
+                if (!stack.empty()) {
+                    u32& parent_height = height[stack.back().dir];
+                    parent_height = std::max(parent_height, height[frame.dir] + 1);
+                }
+                continue;
+            }
+            ++stack.back().next;
+            const u32 child = children[static_cast<std::size_t>(node.first_child) + frame.next];
+            if (child >= dirs.size()) {
+                continue;
+            }
+            if (state[child] == InProgress) {
+                LOG_ERROR(Loader, "NUP directory tree contains a cycle");
+                return false;
+            }
+            if (state[child] == Done) {
+                height[frame.dir] = std::max(height[frame.dir], height[child] + 1);
+            } else {
+                state[child] = InProgress;
+                stack.push_back({child, 0});
+            }
+            if (stack.size() + height[child] > MaxDirDepth) {
+                LOG_ERROR(Loader, "NUP directory tree is nested too deeply");
+                return false;
+            }
+        }
+    }
+    return true;
+}
+
 bool NupContainer::ParseMeta(PackageInfo& package) {
     if (package.meta_offset == 0) {
         return true;
@@ -547,7 +712,8 @@ VirtualFile NupContainer::OpenFile(std::size_t file_index) {
 
 void NupContainer::BuildDirTree(u32 dir_index, VectorVfsDirectory& out) {
     const DirNode& node = dirs[dir_index];
-    for (u32 slot = node.first_child; slot < node.first_child + node.child_count; ++slot) {
+    const u64 end = static_cast<u64>(node.first_child) + node.child_count;
+    for (u64 slot = node.first_child; slot < end; ++slot) {
         const u32 child = children[slot];
         if (child >= dirs.size()) {
             out.AddFile(OpenFile(child - static_cast<u32>(dirs.size())));
diff --git a/src/core/file_sys/nup_container.h b/src/core/file_sys/nup_container.h
index 118b137a7b..8787717d47 100644
--- a/src/core/file_sys/nup_container.h
+++ b/src/core/file_sys/nup_container.h
@@ -4,6 +4,7 @@
 #pragma once
 
 #include <array>
+#include <filesystem>
 #include <functional>
 #include <list>
 #include <memory>
@@ -14,19 +15,36 @@
 #include <vector>
 
 #include "common/common_types.h"
+#include "common/fs/file.h"
 #include "core/file_sys/vfs/vfs_types.h"
 #include "core/file_sys/vfs/vfs_vector.h"
 #include "core/loader/loader.h"
 
 namespace FileSys {
 
+/// A zstd member too large for the in-memory cache, decompressed once into a temporary file.
+struct NupSpillFile {
+    ~NupSpillFile();
+
+    std::mutex mutex;
+    std::filesystem::path path;
+    Common::FS::IOFile file;
+    bool attempted{};
+};
+
 /// LRU over decompressed zstd extents, shared by the files of one container.
 struct NupDecompressCache {
+    static constexpr std::size_t DefaultCap = 512 * 1024 * 1024;
+
+    using Entry = std::pair<std::size_t, std::vector<u8>>;
+
     std::mutex mutex;
-    std::list<std::pair<std::size_t, std::vector<u8>>> lru;
-    std::unordered_map<std::size_t, std::list<std::pair<std::size_t, std::vector<u8>>>::iterator>
-        index;
+    std::list<Entry> lru;
+    std::unordered_map<std::size_t, std::list<Entry>::iterator> index;
     std::size_t bytes{};
+    std::size_t cap{DefaultCap};
+    std::filesystem::path spill_dir;
+    std::unordered_map<std::size_t, std::shared_ptr<NupSpillFile>> spills;
 
     std::vector<u8>* Find(std::size_t file_index) {
         const auto it = index.find(file_index);
@@ -37,21 +55,31 @@ struct NupDecompressCache {
         return &it->second->second;
     }
 
+    /// Takes ownership of `data` only when it is inserted.
     void Insert(std::size_t file_index, std::vector<u8>& data) {
-        while (bytes + data.size() > DecompressCacheCap && !lru.empty()) {
+        if (index.contains(file_index)) {
+            return;
+        }
+        while (bytes + data.size() > cap && !lru.empty()) {
             bytes -= lru.back().second.size();
             index.erase(lru.back().first);
             lru.pop_back();
         }
-        if (bytes + data.size() > DecompressCacheCap) {
+        if (bytes + data.size() > cap) {
             return;
         }
-        lru.emplace_front(file_index, data);
-        index.emplace(file_index, lru.begin());
         bytes += data.size();
+        lru.emplace_front(file_index, std::move(data));
+        index.emplace(file_index, lru.begin());
     }
 
-    static constexpr std::size_t DecompressCacheCap = 512 * 1024 * 1024;
+    std::shared_ptr<NupSpillFile> GetSpill(std::size_t file_index) {
+        auto& spill = spills[file_index];
+        if (!spill) {
+            spill = std::make_shared<NupSpillFile>();
+        }
+        return spill;
+    }
 };
 
 /**
@@ -101,7 +129,11 @@ public:
     /// Recursively merges overlay into a copy of base; overlay paths win.
     static VirtualDir MergeShadowDirs(const VirtualDir& base, const VirtualDir& overlay);
 
-    explicit NupContainer(VirtualFile file_);
+    /// Zstd members larger than half of `memory_cache_cap` are decompressed to a temporary file
+    /// in `spill_dir` (the cache directory by default) instead of being held in memory.
+    explicit NupContainer(VirtualFile file_,
+                          std::size_t memory_cache_cap = NupDecompressCache::DefaultCap,
+                          std::filesystem::path spill_dir_ = {});
     ~NupContainer();
 
     Loader::ResultStatus GetStatus() const {
@@ -119,7 +151,8 @@ public:
     std::optional<std::size_t> FindBestUpdatePackage() const;
 
     /// Opens a member file by node index; raw extents are zero-copy views into
-    /// the container, zstd extents are lazily decompressed and cached.
+    /// the container, zstd extents are lazily decompressed, cached in memory
+    /// or, when large, in a temporary file.
     VirtualFile OpenFile(std::size_t file_index);
 
     VirtualDir GetExeFSDir(std::size_t package_index);
@@ -151,6 +184,7 @@ private:
 
     bool Parse();
     bool ParseMeta(PackageInfo& package);
+    bool ValidateDirTree() const;
     void BuildDirTree(u32 dir_index, VectorVfsDirectory& out);
     VirtualDir BuildPackageDir(u32 dir_index);
     std::vector<u8> ReadExtentRaw(const FileNode& node) const;
diff --git a/src/core/file_sys/nup_writer.h b/src/core/file_sys/nup_writer.h
index d2423dd1f4..dd7b7e0883 100644
--- a/src/core/file_sys/nup_writer.h
+++ b/src/core/file_sys/nup_writer.h
@@ -25,10 +25,9 @@ struct PackProgress {
 
 enum class PackResult { Success, Cancelled, InvalidSource, IoError };
 
-// NupContainer keeps a zstd file in memory only while it fits its 512 MiB decompression cache;
-// a larger one is decompressed again in full on every read, which makes big compressed RomFS
-// files unusable. Files above this size are therefore stored raw, with headroom below the cap so
-// that one file does not evict everything else from the cache.
+// NupContainer keeps a zstd file in memory only while it fits half of its 512 MiB decompression
+// cache; a larger one is decompressed to a temporary file on first read, which costs time and
+// disk space. Files above this size are therefore stored raw.
 constexpr u64 DefaultRawThreshold = 256ULL * 1024 * 1024;
 
 // Application metadata written to the base package (or the first package when there is no base).