5 files changed, 38 insertions, 21 deletions

diff --git a/src/core/file_sys/fsmitm_romfsbuild.cpp b/src/core/file_sys/fsmitm_romfsbuild.cpp
index 4aa8ae369..47b7526c7 100644
--- a/src/core/file_sys/fsmitm_romfsbuild.cpp
+++ b/src/core/file_sys/fsmitm_romfsbuild.cpp
@@ -124,7 +124,7 @@ static u64 romfs_get_hash_table_count(u64 num_entries) {
     return count;
 }
 
-void RomFSBuildContext::VisitDirectory(VirtualDir root_romfs,
+void RomFSBuildContext::VisitDirectory(VirtualDir root_romfs, VirtualDir ext,
                                        std::shared_ptr<RomFSBuildDirectoryContext> parent) {
     std::vector<std::shared_ptr<RomFSBuildDirectoryContext>> child_dirs;
 
@@ -139,15 +139,15 @@ void RomFSBuildContext::VisitDirectory(VirtualDir root_romfs,
 
     for (const auto& kv : entries) {
         if (kv.second == VfsEntryType::Directory) {
-            if (dir->GetSubdirectory(kv.first + ".stub") != nullptr)
-                continue;
-
             const auto child = std::make_shared<RomFSBuildDirectoryContext>();
             // Set child's path.
             child->cur_path_ofs = parent->path_len + 1;
             child->path_len = child->cur_path_ofs + static_cast<u32>(kv.first.size());
             child->path = parent->path + "/" + kv.first;
 
+            if (ext != nullptr && ext->GetFileRelative(child->path + ".stub") != nullptr)
+                continue;
+
             // Sanity check on path_len
             ASSERT(child->path_len < FS_MAX_PATH);
 
@@ -155,25 +155,28 @@ void RomFSBuildContext::VisitDirectory(VirtualDir root_romfs,
                 child_dirs.push_back(child);
             }
         } else {
-            if (dir->GetFile(kv.first + ".stub") != nullptr)
-                continue;
-
             const auto child = std::make_shared<RomFSBuildFileContext>();
             // Set child's path.
             child->cur_path_ofs = parent->path_len + 1;
             child->path_len = child->cur_path_ofs + static_cast<u32>(kv.first.size());
             child->path = parent->path + "/" + kv.first;
 
+            if (ext != nullptr && ext->GetFileRelative(child->path + ".stub") != nullptr)
+                continue;
+
             // Sanity check on path_len
             ASSERT(child->path_len < FS_MAX_PATH);
 
             child->source = root_romfs->GetFileRelative(child->path);
 
-            if (dir->GetFile(kv.first + ".ips") != nullptr) {
-                const auto ips = dir->GetFile(kv.first + ".ips");
-                auto patched = PatchIPS(child->source, ips);
-                if (patched != nullptr)
-                    child->source = std::move(patched);
+            if (ext != nullptr) {
+                const auto ips = ext->GetFileRelative(child->path + ".ips");
+
+                if (ips != nullptr) {
+                    auto patched = PatchIPS(child->source, ips);
+                    if (patched != nullptr)
+                        child->source = std::move(patched);
+                }
             }
 
             child->size = child->source->GetSize();
@@ -183,7 +186,7 @@ void RomFSBuildContext::VisitDirectory(VirtualDir root_romfs,
     }
 
     for (auto& child : child_dirs) {
-        this->VisitDirectory(root_romfs, child);
+        this->VisitDirectory(root_romfs, ext, child);
     }
 }
 
@@ -222,14 +225,15 @@ bool RomFSBuildContext::AddFile(std::shared_ptr<RomFSBuildDirectoryContext> pare
     return true;
 }
 
-RomFSBuildContext::RomFSBuildContext(VirtualDir base_) : base(std::move(base_)) {
+RomFSBuildContext::RomFSBuildContext(VirtualDir base_, VirtualDir ext_)
+    : base(std::move(base_)), ext(std::move(ext_)) {
     root = std::make_shared<RomFSBuildDirectoryContext>();
     root->path = "\0";
     directories.emplace(root->path, root);
     num_dirs = 1;
     dir_table_size = 0x18;
 
-    VisitDirectory(base, root);
+    VisitDirectory(base, ext, root);
 }
 
 RomFSBuildContext::~RomFSBuildContext() = default;
diff --git a/src/core/file_sys/fsmitm_romfsbuild.h b/src/core/file_sys/fsmitm_romfsbuild.h
index b0c3c123b..3d377b0af 100644
--- a/src/core/file_sys/fsmitm_romfsbuild.h
+++ b/src/core/file_sys/fsmitm_romfsbuild.h
@@ -40,7 +40,7 @@ struct RomFSFileEntry;
 
 class RomFSBuildContext {
 public:
-    explicit RomFSBuildContext(VirtualDir base);
+    explicit RomFSBuildContext(VirtualDir base, VirtualDir ext = nullptr);
     ~RomFSBuildContext();
 
     // This finalizes the context.
@@ -48,6 +48,7 @@ public:
 
 private:
     VirtualDir base;
+    VirtualDir ext;
     std::shared_ptr<RomFSBuildDirectoryContext> root;
     std::map<std::string, std::shared_ptr<RomFSBuildDirectoryContext>, std::less<>> directories;
     std::map<std::string, std::shared_ptr<RomFSBuildFileContext>, std::less<>> files;
@@ -59,7 +60,8 @@ private:
     u64 file_hash_table_size = 0;
     u64 file_partition_size = 0;
 
-    void VisitDirectory(VirtualDir filesys, std::shared_ptr<RomFSBuildDirectoryContext> parent);
+    void VisitDirectory(VirtualDir filesys, VirtualDir ext,
+                        std::shared_ptr<RomFSBuildDirectoryContext> parent);
 
     bool AddDirectory(std::shared_ptr<RomFSBuildDirectoryContext> parent_dir_ctx,
                       std::shared_ptr<RomFSBuildDirectoryContext> dir_ctx);
diff --git a/src/core/file_sys/patch_manager.cpp b/src/core/file_sys/patch_manager.cpp
index 539698f6e..76b51fe56 100644
--- a/src/core/file_sys/patch_manager.cpp
+++ b/src/core/file_sys/patch_manager.cpp
@@ -162,11 +162,17 @@ static void ApplyLayeredFS(VirtualFile& romfs, u64 title_id, ContentRecordType t
               [](const VirtualDir& l, const VirtualDir& r) { return l->GetName() < r->GetName(); });
 
     std::vector<VirtualDir> layers;
+    std::vector<VirtualDir> layers_ext;
     layers.reserve(patch_dirs.size() + 1);
+    layers_ext.reserve(patch_dirs.size() + 1);
     for (const auto& subdir : patch_dirs) {
         auto romfs_dir = subdir->GetSubdirectory("romfs");
         if (romfs_dir != nullptr)
             layers.push_back(std::move(romfs_dir));
+
+        auto ext_dir = subdir->GetSubdirectory("romfs_ext");
+        if (ext_dir != nullptr)
+            layers.push_back(std::move(ext_dir));
     }
     layers.push_back(std::move(extracted));
 
@@ -175,7 +181,12 @@ static void ApplyLayeredFS(VirtualFile& romfs, u64 title_id, ContentRecordType t
         return;
     }
 
-    auto packed = CreateRomFS(std::move(layered));
+    auto layered_ext = LayeredVfsDirectory::MakeLayeredDirectory(std::move(layers_ext));
+    if (layered_ext == nullptr) {
+        return;
+    }
+
+    auto packed = CreateRomFS(std::move(layered), std::move(layered_ext));
     if (packed == nullptr) {
         return;
     }
diff --git a/src/core/file_sys/romfs.cpp b/src/core/file_sys/romfs.cpp
index 5910f7046..81e1f66ac 100644
--- a/src/core/file_sys/romfs.cpp
+++ b/src/core/file_sys/romfs.cpp
@@ -129,11 +129,11 @@ VirtualDir ExtractRomFS(VirtualFile file, RomFSExtractionType type) {
     return out;
 }
 
-VirtualFile CreateRomFS(VirtualDir dir) {
+VirtualFile CreateRomFS(VirtualDir dir, VirtualDir ext) {
     if (dir == nullptr)
         return nullptr;
 
-    RomFSBuildContext ctx{dir};
+    RomFSBuildContext ctx{dir, ext};
     return ConcatenatedVfsFile::MakeConcatenatedFile(0, ctx.Build(), dir->GetName());
 }
 
diff --git a/src/core/file_sys/romfs.h b/src/core/file_sys/romfs.h
index ecd1eb725..0ec404731 100644
--- a/src/core/file_sys/romfs.h
+++ b/src/core/file_sys/romfs.h
@@ -44,6 +44,6 @@ VirtualDir ExtractRomFS(VirtualFile file,
 
 // Converts a VFS filesystem into a RomFS binary
 // Returns nullptr on failure
-VirtualFile CreateRomFS(VirtualDir dir);
+VirtualFile CreateRomFS(VirtualDir dir, VirtualDir ext = nullptr);
 
 } // namespace FileSys