fs.mitm: Cache IStorageInterfaces, store meta on SD instead of memory.

stratosphere/fs_mitm/source/fsmitm_layeredrom.cpp

@@ -83,6 +83,22 @@ Result LayeredRomFS::Read(void *buffer, size_t size, u64 offset)  {
                 cur_read_size = cur_source->size - (offset - cur_source->virtual_offset);
             }
             switch (cur_source->type) {
+                case RomFSDataSource::MetaData:
+                    {
+                        FsFile file;
+                        if (R_FAILED((rc = Utils::OpenSdFileForAtmosphere(this->title_id, ROMFS_METADATA_FILE_PATH, FS_OPEN_READ, &file)))) {
+                            fatalSimple(rc);
+                        }
+                        size_t out_read;
+                        if (R_FAILED((rc = fsFileRead(&file, (offset - cur_source->virtual_offset), (void *)((uintptr_t)buffer + read_so_far), cur_read_size, &out_read)))) {
+                            fatalSimple(rc);
+                        }
+                        if (out_read != cur_read_size) {
+                            Reboot();
+                        }
+                        fsFileClose(&file);
+                    }
+                    break;
                 case RomFSDataSource::LooseFile:
                     {
                         FsFile file;

stratosphere/fs_mitm/source/fsmitm_romfsbuild.cpp

@@ -399,23 +399,14 @@ void RomFSBuildContext::Build(std::vector<RomFSSourceInfo> *out_infos) {
     header->file_hash_table_ofs = header->dir_table_ofs + header->dir_table_size;
     header->file_table_ofs = header->file_hash_table_ofs + header->file_hash_table_size;
     
-    /* For debugging, uncomment this to get a log of the generated metadata tables. */
+    const size_t metadata_size = this->dir_hash_table_size + this->dir_table_size + this->file_hash_table_size + this->file_table_size;
     
-        {
+    /* Try to save metadata to the SD card, to save on memory space. */
-            FsFileSystem sd_fs;
+    if (R_SUCCEEDED(Utils::SaveSdFileForAtmosphere(this->title_id, ROMFS_METADATA_FILE_PATH, metadata, metadata_size))) {
-            if (R_SUCCEEDED(fsMountSdcard(&sd_fs))) {
+        out_infos->emplace_back(header->dir_hash_table_ofs, metadata_size, RomFSDataSource::MetaData);
-                FsFile f;
+        delete metadata;
-                fsFsCreateFile(&sd_fs, "/METADATALOG.bin", this->dir_hash_table_size + this->dir_table_size + this->file_hash_table_size + this->file_table_size + sizeof(*header), 0);
+    } else {
-                if (R_SUCCEEDED(fsFsOpenFile(&sd_fs, "/METADATALOG.bin", FS_OPEN_READ | FS_OPEN_WRITE, &f))) {
+        out_infos->emplace_back(header->dir_hash_table_ofs, metadata_size, metadata, RomFSDataSource::Memory);
-                    fsFileSetSize(&f, this->dir_hash_table_size + this->dir_table_size + this->file_hash_table_size + this->file_table_size + sizeof(*header));
-                    fsFileWrite(&f, 0, header, sizeof(*header));
-                    fsFileWrite(&f, sizeof(*header), metadata, this->dir_hash_table_size + this->dir_table_size + this->file_hash_table_size + this->file_table_size);
-                    fsFileClose(&f);
-                }
-                fsFsClose(&sd_fs);
-            }
     }
     
-    
-    out_infos->emplace_back(header->dir_hash_table_ofs, this->dir_hash_table_size + this->dir_table_size + this->file_hash_table_size + this->file_table_size, metadata, RomFSDataSource::Memory);
 }

stratosphere/fs_mitm/source/fsmitm_romfsbuild.hpp

@@ -25,11 +25,14 @@
 #define ROMFS_ENTRY_EMPTY 0xFFFFFFFF
 #define ROMFS_FILEPARTITION_OFS 0x200
 
+#define ROMFS_METADATA_FILE_PATH "romfs_metadata.bin"
+
 /* Types for RomFS Meta construction. */
 enum class RomFSDataSource {
     BaseRomFS,
     FileRomFS,
     LooseFile,
+    MetaData,
     Memory,
 };
 
@@ -49,6 +52,10 @@ struct RomFSMemorySourceInfo {
     const u8 *data;
 };
 
+struct RomFSMetaDataSourceInfo {
+
+};
+
 struct RomFSSourceInfo {
     u64 virtual_offset;
     u64 size;
@@ -57,6 +64,7 @@ struct RomFSSourceInfo {
         RomFSFileSourceInfo file_source_info;
         RomFSLooseSourceInfo loose_source_info;
         RomFSMemorySourceInfo memory_source_info;
+        RomFSMemorySourceInfo metadata_source_info;
     };
     RomFSDataSource type;
     
@@ -69,6 +77,7 @@ struct RomFSSourceInfo {
                 this->file_source_info.offset = offset;
                 break;
             case RomFSDataSource::LooseFile:
+            case RomFSDataSource::MetaData:
             case RomFSDataSource::Memory:
             default:
                 fatalSimple(0xF601);
@@ -83,6 +92,20 @@ struct RomFSSourceInfo {
             case RomFSDataSource::Memory:
                 this->memory_source_info.data = (decltype(this->memory_source_info.data))arg;
                 break;
+            case RomFSDataSource::MetaData:
+            case RomFSDataSource::BaseRomFS:
+            case RomFSDataSource::FileRomFS:
+            default:
+                fatalSimple(0xF601);
+        }
+    }
+    
+    RomFSSourceInfo(u64 v_o, u64 s, RomFSDataSource t) : virtual_offset(v_o), size(s), type(t) {
+        switch (this->type) {
+            case RomFSDataSource::MetaData:
+                break;
+            case RomFSDataSource::LooseFile:
+            case RomFSDataSource::Memory:
             case RomFSDataSource::BaseRomFS:
             case RomFSDataSource::FileRomFS:
             default:
@@ -94,6 +117,7 @@ struct RomFSSourceInfo {
         switch (this->type) {
             case RomFSDataSource::BaseRomFS:
             case RomFSDataSource::FileRomFS:
+            case RomFSDataSource::MetaData:
                 break;
             case RomFSDataSource::LooseFile:
                 delete this->loose_source_info.path;

stratosphere/fs_mitm/source/fsmitm_service.cpp

@@ -14,7 +14,12 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
  
+#include <map>
+#include <memory>
+#include <mutex>
+
 #include <switch.h>
+#include <stratosphere.hpp>
 #include "fsmitm_service.hpp"
 #include "fs_shim.h"
 
@@ -24,6 +29,37 @@
 
 #include "debug.hpp"
 
+static HosMutex g_StorageCacheLock;
+static std::unordered_map<u64, std::weak_ptr<IStorageInterface>> g_StorageCache;
+
+static bool StorageCacheGetEntry(u64 title_id, std::shared_ptr<IStorageInterface> *out) {
+    std::scoped_lock<HosMutex> lock(g_StorageCacheLock);
+    if (g_StorageCache.find(title_id) == g_StorageCache.end()) {
+        return false;
+    }
+    
+    auto intf = g_StorageCache[title_id].lock();
+    if (intf != nullptr) {
+        *out = intf;
+        return true;
+    }
+    return false;
+}
+
+static void StorageCacheSetEntry(u64 title_id, std::shared_ptr<IStorageInterface> *ptr) {
+    std::scoped_lock<HosMutex> lock(g_StorageCacheLock);
+    
+    /* Ensure we always use the cached copy if present. */
+    if (g_StorageCache.find(title_id) != g_StorageCache.end()) {
+        auto intf = g_StorageCache[title_id].lock();
+        if (intf != nullptr) {
+            *ptr = intf;
+        }
+    }
+    
+    g_StorageCache[title_id] = *ptr;
+}
+
 void FsMitmService::PostProcess(IMitmServiceObject *obj, IpcResponseContext *ctx) {
     auto this_ptr = static_cast<FsMitmService *>(obj);
     switch ((FspSrvCmd)ctx->cmd_id) {
@@ -49,8 +85,14 @@ Result FsMitmService::OpenDataStorageByCurrentProcess(Out<std::shared_ptr<IStora
     u32 out_domain_id = 0;
     Result rc = 0;
     
+    bool has_cache = StorageCacheGetEntry(this->title_id, &storage);
+    
     ON_SCOPE_EXIT {
         if (R_SUCCEEDED(rc)) {
+            if (!has_cache) {
+                StorageCacheSetEntry(this->title_id, &storage);
+            }
+            
             out_storage.SetValue(std::move(storage));
             if (out_storage.IsDomain()) {
                 out_storage.ChangeObjectId(out_domain_id);
@@ -58,7 +100,8 @@ Result FsMitmService::OpenDataStorageByCurrentProcess(Out<std::shared_ptr<IStora
         }
     };
     
-    if (this->romfs_storage != nullptr) {
+    
+    if (has_cache) {
         if (out_storage.IsDomain()) {
             FsStorage s = {0};
             rc = fsOpenDataStorageByCurrentProcessFwd(this->forward_service.get(), &s);
@@ -68,8 +111,8 @@ Result FsMitmService::OpenDataStorageByCurrentProcess(Out<std::shared_ptr<IStora
         } else {
             rc = 0;
         }
-        if (R_SUCCEEDED(rc)) {
+        if (R_FAILED(rc)) {
-            storage = this->romfs_storage;
+            storage.reset();
         }
     } else {
         FsStorage data_storage;
@@ -86,7 +129,6 @@ Result FsMitmService::OpenDataStorageByCurrentProcess(Out<std::shared_ptr<IStora
                 } else {
                     storage = std::make_shared<IStorageInterface>(new LayeredRomFS(std::make_shared<RomInterfaceStorage>(data_storage), nullptr, this->title_id));
                 }
-                this->romfs_storage = storage;
                 if (out_storage.IsDomain()) {
                     out_domain_id = data_storage.s.object_id;
                 }
@@ -111,8 +153,14 @@ Result FsMitmService::OpenDataStorageByDataId(Out<std::shared_ptr<IStorageInterf
     u32 out_domain_id = 0;
     Result rc = 0;
     
+    bool has_cache = StorageCacheGetEntry(data_id, &storage);
+    
     ON_SCOPE_EXIT {
         if (R_SUCCEEDED(rc)) {
+            if (!has_cache) {
+                StorageCacheSetEntry(data_id, &storage);
+            }
+            
             out_storage.SetValue(std::move(storage));
             if (out_storage.IsDomain()) {
                 out_storage.ChangeObjectId(out_domain_id);
@@ -120,6 +168,20 @@ Result FsMitmService::OpenDataStorageByDataId(Out<std::shared_ptr<IStorageInterf
         }
     };
     
+    if (has_cache) {
+        if (out_storage.IsDomain()) {
+            FsStorage s = {0};
+            rc = fsOpenDataStorageByDataIdFwd(this->forward_service.get(), storage_id, data_id, &s);
+            if (R_SUCCEEDED(rc)) {
+                out_domain_id = s.s.object_id;
+            }
+        } else {
+            rc = 0;
+        }
+        if (R_FAILED(rc)) {
+            storage.reset();
+        }
+    } else {
         rc = fsOpenDataStorageByDataIdFwd(this->forward_service.get(), storage_id, data_id, &data_storage);
 
         if (R_SUCCEEDED(rc)) {
@@ -139,6 +201,7 @@ Result FsMitmService::OpenDataStorageByDataId(Out<std::shared_ptr<IStorageInterf
                 rc = RESULT_FORWARD_TO_SESSION;
             }
         }
+    }
     
     return rc;
 }

stratosphere/fs_mitm/source/fsmitm_service.hpp

@@ -29,7 +29,6 @@ enum FspSrvCmd : u32 {
 class FsMitmService : public IMitmServiceObject {      
     private:
         bool has_initialized = false;
-        std::shared_ptr<IStorageInterface> romfs_storage;
     public:
         FsMitmService(std::shared_ptr<Service> s) : IMitmServiceObject(s) {
             /* ... */

stratosphere/fs_mitm/source/fsmitm_utils.cpp

@@ -223,6 +223,47 @@ bool Utils::HasSdRomfsContent(u64 title_id) {
     return R_SUCCEEDED(fsDirRead(&dir, 0, &read_entries, 1, &dir_entry)) && read_entries == 1;
 }
 
+Result Utils::SaveSdFileForAtmosphere(u64 title_id, const char *fn, void *data, size_t size) {
+    if (!IsSdInitialized()) {
+        return 0xFA202;
+    }
+    
+    Result rc = 0;
+    
+    char path[FS_MAX_PATH];
+    if (*fn == '/') {
+        snprintf(path, sizeof(path), "/atmosphere/titles/%016lx%s", title_id, fn);
+    } else {
+        snprintf(path, sizeof(path), "/atmosphere/titles/%016lx/%s", title_id, fn);
+    }
+    
+    /* Unconditionally create. */
+    FsFile f;
+    fsFsCreateFile(&g_sd_filesystem, path, size, 0);
+    
+    /* Try to open. */
+    rc = fsFsOpenFile(&g_sd_filesystem, path, FS_OPEN_READ | FS_OPEN_WRITE, &f);
+    if (R_FAILED(rc)) {
+        return rc;
+    }
+    
+    /* Always close, if we opened. */
+    ON_SCOPE_EXIT {
+        fsFileClose(&f);
+    };
+    
+    /* Try to make it big enough. */
+    rc = fsFileSetSize(&f, size);
+    if (R_FAILED(rc)) {
+        return rc;
+    }
+    
+    /* Try to write the data. */
+    rc = fsFileWrite(&f, 0, data, size);
+    
+    return rc;
+}
+
 bool Utils::HasSdMitMFlag(u64 tid) {
     if (IsSdInitialized()) {
         return std::find(g_mitm_flagged_tids.begin(), g_mitm_flagged_tids.end(), tid) != g_mitm_flagged_tids.end();

stratosphere/fs_mitm/source/fsmitm_utils.hpp

@@ -21,6 +21,7 @@
 class Utils {
     public:
         static bool IsSdInitialized();
+        
         static Result OpenSdFile(const char *fn, int flags, FsFile *out);
         static Result OpenSdFileForAtmosphere(u64 title_id, const char *fn, int flags, FsFile *out);
         static Result OpenRomFSSdFile(u64 title_id, const char *fn, int flags, FsFile *out);
@@ -31,6 +32,8 @@ class Utils {
         static Result OpenRomFSFile(FsFileSystem *fs, u64 title_id, const char *fn, int flags, FsFile *out);
         static Result OpenRomFSDir(FsFileSystem *fs, u64 title_id, const char *path, FsDir *out);
         
+        static Result SaveSdFileForAtmosphere(u64 title_id, const char *fn, void *data, size_t size);
+        
         static bool HasSdRomfsContent(u64 title_id);
         
         /* SD card Initialization + MitM detection. */