fs.mitm: add OpenDataStorageWithProgramIndex support (theoretically, closes #1250)

libraries/libstratosphere/include/stratosphere/fssrv.hpp

@@ -18,6 +18,7 @@
 #include <stratosphere/fssrv/sf/fssrv_sf_path.hpp>
 #include <stratosphere/fssrv/sf/fssrv_sf_ifile.hpp>
 #include <stratosphere/fssrv/sf/fssrv_sf_i_event_notifier.hpp>
+#include <stratosphere/fssrv/impl/fssrv_impl_program_index_map_info_manager.hpp>
 #include <stratosphere/fssrv/fssrv_path_normalizer.hpp>
 #include <stratosphere/fssrv/fssrv_nca_crypto_configuration.hpp>
 #include <stratosphere/fssrv/fssrv_memory_resource_from_standard_allocator.hpp>

libraries/libstratosphere/include/stratosphere/fssrv/impl/fssrv_impl_program_index_map_info_manager.hpp

@@ -0,0 +1,159 @@
+/*
+ * Copyright (c) 2018-2020 Atmosphère-NX
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+#pragma once
+#include <vapours.hpp>
+#include <stratosphere/ncm/ncm_ids.hpp>
+#include <stratosphere/fs/impl/fs_newable.hpp>
+#include <stratosphere/fs/fs_program_index_map_info.hpp>
+
+namespace ams::fssrv::impl {
+
+    struct ProgramIndexMapInfoEntry : public ::ams::util::IntrusiveListBaseNode<ProgramIndexMapInfoEntry>, public ::ams::fs::impl::Newable {
+        ncm::ProgramId program_id;
+        ncm::ProgramId base_program_id;
+        u8 program_index;
+    };
+
+    class ProgramIndexMapInfoManager {
+        NON_COPYABLE(ProgramIndexMapInfoManager);
+        NON_MOVEABLE(ProgramIndexMapInfoManager);
+        private:
+            using ProgramIndexMapInfoList = util::IntrusiveListBaseTraits<ProgramIndexMapInfoEntry>::ListType;
+        private:
+            ProgramIndexMapInfoList m_list;
+            mutable os::SdkMutex m_mutex;
+        public:
+            constexpr ProgramIndexMapInfoManager() : m_list(), m_mutex() { /* ... */ }
+
+            void Clear() {
+                /* Acquire exclusive access to the map. */
+                std::scoped_lock lk(m_mutex);
+
+                /* Actually clear. */
+                this->ClearImpl();
+            }
+
+            size_t GetProgramCount() const {
+                /* Acquire exclusive access to the map. */
+                std::scoped_lock lk(m_mutex);
+
+                /* Get the size. */
+                return m_list.size();
+            }
+
+            std::optional<fs::ProgramIndexMapInfo> Get(const ncm::ProgramId &program_id) const {
+                /* Acquire exclusive access to the map. */
+                std::scoped_lock lk(m_mutex);
+
+                /* Get the entry from the map. */
+                return this->GetImpl([&] (const ProgramIndexMapInfoEntry &entry) {
+                    return entry.program_id == program_id;
+                });
+            }
+
+            ncm::ProgramId GetProgramId(const ncm::ProgramId &program_id, u8 program_index) const {
+                /* Acquire exclusive access to the map. */
+                std::scoped_lock lk(m_mutex);
+
+                /* Get the program info for the desired program id. */
+                const auto base_info = this->GetImpl([&] (const ProgramIndexMapInfoEntry &entry) {
+                    return entry.program_id == program_id;
+                });
+
+                /* Check that an entry exists for the program id. */
+                if (!base_info.has_value()) {
+                    return ncm::InvalidProgramId;
+                }
+
+                /* Get a program info which matches the same base program with the desired index. */
+                const auto target_info = this->GetImpl([&] (const ProgramIndexMapInfoEntry &entry) {
+                    return entry.base_program_id == base_info->base_program_id && entry.program_index == program_index;
+                });
+
+                /* Return the desired program id. */
+                if (target_info.has_value()) {
+                    return target_info->program_id;
+                } else {
+                    return ncm::InvalidProgramId;
+                }
+            }
+
+            Result Reset(const fs::ProgramIndexMapInfo *infos, int count) {
+                /* Acquire exclusive access to the map. */
+                std::scoped_lock lk(m_mutex);
+
+                /* Clear the map, and ensure we remain clear if we fail after this point. */
+                this->ClearImpl();
+                auto clear_guard = SCOPE_GUARD { this->ClearImpl(); };
+
+                /* Add each info to the list. */
+                for (int i = 0; i < count; ++i) {
+                    /* Allocate new entry. */
+                    auto *entry = new ProgramIndexMapInfoEntry;
+                    R_UNLESS(entry != nullptr, fs::ResultAllocationFailureInNew());
+
+                    /* Copy over the info. */
+                    entry->program_id      = infos[i].program_id;
+                    entry->base_program_id = infos[i].base_program_id;
+                    entry->program_index   = infos[i].program_index;
+
+                    /* Add to the list. */
+                    m_list.push_back(*entry);
+                }
+
+                /* We successfully imported the map. */
+                clear_guard.Cancel();
+                return ResultSuccess();
+            }
+        private:
+            void ClearImpl() {
+                /* Delete all entries. */
+                while (!m_list.empty()) {
+                    /* Get the first entry. */
+                    ProgramIndexMapInfoEntry *front = std::addressof(*m_list.begin());
+
+                    /* Erase it from the list. */
+                    m_list.erase(m_list.iterator_to(*front));
+
+                    /* Delete the entry. */
+                    delete front;
+                }
+            }
+
+            template<typename F>
+            std::optional<fs::ProgramIndexMapInfo> GetImpl(F f) const {
+                /* Try to find an entry matching the predicate. */
+                std::optional<fs::ProgramIndexMapInfo> match = std::nullopt;
+
+                for (const auto &entry : m_list) {
+                    /* If the predicate matches, we want to return the relevant info. */
+                    if (f(entry)) {
+                        match.emplace();
+
+                        match->program_id      = entry.program_id;
+                        match->base_program_id = entry.base_program_id;
+                        match->program_index   = entry.program_index;
+
+                        break;
+                    }
+                }
+
+                return match;
+            }
+    };
+
+
+}

stratosphere/ams_mitm/source/fs_mitm/fs_mitm_service.cpp

@@ -42,6 +42,8 @@ namespace ams::mitm::fs {
         constinit bool g_detected_boot0_kind = false;
         constinit bool g_is_boot0_custom_public_key = false;
 
+        constinit fssrv::impl::ProgramIndexMapInfoManager g_program_index_map_info_manager;
+
         bool IsBoot0CustomPublicKey(::FsStorage &storage) {
             if (AMS_UNLIKELY(!g_detected_boot0_kind)) {
                 std::scoped_lock lk(g_boot0_detect_lock);
@@ -410,4 +412,65 @@ namespace ams::mitm::fs {
         return ResultSuccess();
     }
 
+    Result FsMitmService::OpenDataStorageWithProgramIndex(sf::Out<sf::SharedPointer<ams::fssrv::sf::IStorage>> out, u8 program_index) {
+        /* Only mitm if we should override contents for the current process. */
+        R_UNLESS(this->client_info.override_status.IsProgramSpecific(), sm::mitm::ResultShouldForwardToSession());
+
+        /* Get the relevant program id. */
+        const ncm::ProgramId program_id = g_program_index_map_info_manager.GetProgramId(this->client_info.program_id, program_index);
+
+        /* If we don't know about the program or don't have content, forward. */
+        R_UNLESS(program_id != ncm::InvalidProgramId,     sm::mitm::ResultShouldForwardToSession());
+        R_UNLESS(mitm::fs::HasSdRomfsContent(program_id), sm::mitm::ResultShouldForwardToSession());
+
+        /* Try to open the process romfs. */
+        FsStorage data_storage;
+        R_TRY(fsOpenDataStorageWithProgramIndexFwd(this->forward_service.get(), &data_storage, program_index));
+        const sf::cmif::DomainObjectId target_object_id{serviceGetObjectId(&data_storage.s)};
+
+        /* Get a scoped lock. */
+        std::scoped_lock lk(g_data_storage_lock);
+
+        /* Try to get a storage from the cache. */
+        {
+            std::shared_ptr<fs::IStorage> cached_storage = GetStorageCacheEntry(program_id);
+            if (cached_storage != nullptr) {
+                out.SetValue(MakeSharedStorage(cached_storage), target_object_id);
+                return ResultSuccess();
+            }
+        }
+
+        /* Make a new layered romfs, and cache to storage. */
+        {
+            std::shared_ptr<fs::IStorage> new_storage = nullptr;
+
+            /* Create the layered storage. */
+            FsFile data_file;
+            if (R_SUCCEEDED(OpenAtmosphereSdFile(&data_file, program_id, "romfs.bin", OpenMode_Read))) {
+                auto layered_storage = std::make_shared<LayeredRomfsStorage>(std::make_unique<ReadOnlyStorageAdapter>(new RemoteStorage(data_storage)), std::make_unique<ReadOnlyStorageAdapter>(new FileStorage(new RemoteFile(data_file))), program_id);
+                layered_storage->BeginInitialize();
+                new_storage = std::move(layered_storage);
+            } else {
+                auto layered_storage = std::make_shared<LayeredRomfsStorage>(std::make_unique<ReadOnlyStorageAdapter>(new RemoteStorage(data_storage)), nullptr, program_id);
+                layered_storage->BeginInitialize();
+                new_storage = std::move(layered_storage);
+            }
+
+            SetStorageCacheEntry(program_id, &new_storage);
+            out.SetValue(MakeSharedStorage(new_storage), target_object_id);
+        }
+
+        return ResultSuccess();
+    }
+
+    Result FsMitmService::RegisterProgramIndexMapInfo(const sf::InBuffer &info_buffer, s32 info_count) {
+        /* Try to register with FS. */
+        R_TRY(fsRegisterProgramIndexMapInfoFwd(this->forward_service.get(), info_buffer.GetPointer(), info_buffer.GetSize(), info_count));
+
+        /* Register with ourselves. */
+        R_ABORT_UNLESS(g_program_index_map_info_manager.Reset(reinterpret_cast<const fs::ProgramIndexMapInfo *>(info_buffer.GetPointer()), info_count));
+
+        return ResultSuccess();
+    }
+
 }

stratosphere/ams_mitm/source/fs_mitm/fs_mitm_service.hpp

@@ -24,7 +24,10 @@
     AMS_SF_METHOD_INFO(C, H,  51, Result, OpenSaveDataFileSystem,          (sf::Out<sf::SharedPointer<ams::fssrv::sf::IFileSystem>> out, u8 space_id, const ams::fs::SaveDataAttribute &attribute),                     (out, space_id, attribute))                                    \
     AMS_SF_METHOD_INFO(C, H,  12, Result, OpenBisStorage,                  (sf::Out<sf::SharedPointer<ams::fssrv::sf::IStorage>> out, u32 bis_partition_id),                                                            (out, bis_partition_id))                                       \
     AMS_SF_METHOD_INFO(C, H, 200, Result, OpenDataStorageByCurrentProcess, (sf::Out<sf::SharedPointer<ams::fssrv::sf::IStorage>> out),                                                                                  (out))                                                         \
-    AMS_SF_METHOD_INFO(C, H, 202, Result, OpenDataStorageByDataId,         (sf::Out<sf::SharedPointer<ams::fssrv::sf::IStorage>> out, ncm::DataId data_id, u8 storage_id),                                              (out, data_id, storage_id))
+    AMS_SF_METHOD_INFO(C, H, 202, Result, OpenDataStorageByDataId,         (sf::Out<sf::SharedPointer<ams::fssrv::sf::IStorage>> out, ncm::DataId data_id, u8 storage_id),                                              (out, data_id, storage_id))                                    \
+    AMS_SF_METHOD_INFO(C, H, 205, Result, OpenDataStorageWithProgramIndex, (sf::Out<sf::SharedPointer<ams::fssrv::sf::IStorage>> out, u8 program_index),                                                                (out, program_index),                      hos::Version_7_0_0) \
+    AMS_SF_METHOD_INFO(C, H, 810, Result, RegisterProgramIndexMapInfo,     (const sf::InBuffer &info_buffer, s32 info_count),                                                                                           (info_buffer, info_count),                 hos::Version_7_0_0)
+
 
 AMS_SF_DEFINE_MITM_INTERFACE(ams::mitm::fs, IFsMitmInterface, AMS_FS_MITM_INTERFACE_INFO)
 
@@ -55,6 +58,11 @@ namespace ams::mitm::fs {
                     return true;
                 }
 
+                /* We want to mitm am, to intercept program info map registration. */
+                if (program_id == ncm::SystemProgramId::Am) {
+                    return true;
+                }
+
                 return false;
             }
 
@@ -80,6 +88,8 @@ namespace ams::mitm::fs {
             Result OpenBisStorage(sf::Out<sf::SharedPointer<ams::fssrv::sf::IStorage>> out, u32 bis_partition_id);
             Result OpenDataStorageByCurrentProcess(sf::Out<sf::SharedPointer<ams::fssrv::sf::IStorage>> out);
             Result OpenDataStorageByDataId(sf::Out<sf::SharedPointer<ams::fssrv::sf::IStorage>> out, ncm::DataId data_id, u8 storage_id);
+            Result OpenDataStorageWithProgramIndex(sf::Out<sf::SharedPointer<ams::fssrv::sf::IStorage>> out, u8 program_index);
+            Result RegisterProgramIndexMapInfo(const sf::InBuffer &info_buffer, s32 info_count);
     };
     static_assert(IsIFsMitmInterface<FsMitmService>);
 

stratosphere/ams_mitm/source/fs_mitm/fs_shim.c

@@ -51,6 +51,20 @@ Result fsOpenDataStorageByDataIdFwd(Service* s, FsStorage* out, u64 data_id, Ncm
     );
 }
 
+Result fsOpenDataStorageWithProgramIndexFwd(Service* s, FsStorage* out, u8 program_index) {
+    return serviceDispatchIn(s, 205, program_index,
+        .out_num_objects = 1,
+        .out_objects = &out->s,
+    );
+}
+
+Result fsRegisterProgramIndexMapInfoFwd(Service* s, const void *buf, size_t buf_size, s32 count) {
+    return serviceDispatchIn(s, 810, count,
+        .buffer_attrs = { SfBufferAttr_HipcMapAlias | SfBufferAttr_In },
+        .buffers = { { buf, buf_size } },
+    );
+}
+
 Result fsOpenSaveDataFileSystemFwd(Service* s, FsFileSystem* out, FsSaveDataSpaceId save_data_space_id, const FsSaveDataAttribute *attr) {
     const struct {
         u8 save_data_space_id;

stratosphere/ams_mitm/source/fs_mitm/fs_shim.h

@@ -16,6 +16,9 @@ Result fsOpenSdCardFileSystemFwd(Service* s, FsFileSystem* out);
 Result fsOpenBisStorageFwd(Service* s, FsStorage* out, FsBisPartitionId partition_id);
 Result fsOpenDataStorageByCurrentProcessFwd(Service* s, FsStorage* out);
 Result fsOpenDataStorageByDataIdFwd(Service* s, FsStorage* out, u64 data_id, NcmStorageId storage_id);
+Result fsOpenDataStorageWithProgramIndexFwd(Service* s, FsStorage* out, u8 program_index);
+
+Result fsRegisterProgramIndexMapInfoFwd(Service* s, const void *buf, size_t buf_size, s32 count);
 
 Result fsOpenSaveDataFileSystemFwd(Service* s, FsFileSystem* out, FsSaveDataSpaceId save_data_space_id, const FsSaveDataAttribute *attr);