fs.mitm: WIP LayeredFS impl (NOTE: UNUSABLE ATM)

Also greatly refactors libstratosphere, and does a lot of other things.
There is a lot of code in this one.

stratosphere/fs_mitm/fs_mitm.json

@@ -1,12 +1,12 @@
 {
     "name" : "fs.mitm",
     "title_id" : "0x010041544D530000",
-	"main_thread_stack_size" : "0x8000",
+	"main_thread_stack_size" : "0x20000",
-	"main_thread_priority":	44,
+	"main_thread_priority":	43,
 	"default_cpu_id":	3,
     "process_category" : 1,
     "kernel_capabilities" : {
-        "handle_table_size" : 256,
+        "handle_table_size" : 512,
 		"syscalls":	{
 			"svcSetHeapSize":	"0x01",
 			"svcSetMemoryPermission":	"0x02",

stratosphere/fs_mitm/source/debug.cpp

@@ -1,17 +1,14 @@
 #include <switch.h>
+#include <stratosphere.hpp>
 #include <cstring>
 #include "debug.hpp"
 
-static u64 g_num_logged = 0;
-
-#define MAX_LOGS U64_MAX
-
 void Reboot() {
-    while (1) {
+    /* ... */
-        /* ... */
-    }
 }
 
 void Log(const void *data, int size) {
+    (void)(data);
+    (void)(size);
     /* ... */
 }

stratosphere/fs_mitm/source/fs_istorage.hpp

@@ -22,7 +22,7 @@ class IStorage {
         
         virtual IStorage *Clone() = 0; 
         
-        virtual Result Read(void *buffer, size_t size, u64 offset, u64 *out_read_size) = 0;
+        virtual Result Read(void *buffer, size_t size, u64 offset) = 0;
         virtual Result Write(void *buffer, size_t size, u64 offset) = 0;
         virtual Result Flush() = 0;
         virtual Result SetSize(u64 size) = 0;
@@ -81,10 +81,8 @@ class IStorageInterface : public IServiceObject {
         };
     private:
         /* Actual command API. */
-        virtual std::tuple<Result, u64> read(OutBuffer<u8, BufferType_Type1> buffer, u64 offset, u64 size) final {
+        virtual std::tuple<Result> read(OutBuffer<u8, BufferType_Type1> buffer, u64 offset, u64 size) final {
-            u64 out_size = 0;
+            return {this->base_storage->Read(buffer.buffer, std::min(buffer.num_elements, size), offset)};
-            Result rc = this->base_storage->Read(buffer.buffer, std::min(buffer.num_elements, size), offset , &out_size);
-            return {rc, out_size};
         };
         virtual std::tuple<Result> write(InBuffer<u8, BufferType_Type1> buffer, u64 offset, u64 size) final {
            return {this->base_storage->Write(buffer.buffer, std::min(buffer.num_elements, size), offset)};
@@ -109,8 +107,8 @@ class IStorageInterface : public IServiceObject {
 };
 
 class IROStorage : public IStorage {
-    protected:
+    public:
-        virtual Result Read(void *buffer, size_t size, u64 offset, u64 *out_read_size) = 0;
+        virtual Result Read(void *buffer, size_t size, u64 offset) = 0;
         Result Write(void *buffer, size_t size, u64 offset) final {
             (void)(buffer);
             (void)(offset);

stratosphere/fs_mitm/source/fs_shim.c

@@ -43,6 +43,77 @@ Result ipcCopyFromDomain(Handle session, u32 object_id, Service *out) {
 
 
 /* Missing fsp-srv commands. */
+Result fsOpenDataStorageByCurrentProcessFwd(Service* s, FsStorage* out) {
+    IpcCommand c;
+    ipcInitialize(&c);
+
+    struct {
+        u64 magic;
+        u64 cmd_id;
+    } *raw;
+
+    raw = ipcPrepareHeader(&c, sizeof(*raw));
+
+    raw->magic = SFCI_MAGIC;
+    raw->cmd_id = 200;
+
+    Result rc = serviceIpcDispatch(s);
+
+    if (R_SUCCEEDED(rc)) {
+        IpcParsedCommand r;
+        ipcParse(&r);
+
+        struct {
+            u64 magic;
+            u64 result;
+        } *resp = r.Raw;
+
+        rc = resp->result;
+
+        if (R_SUCCEEDED(rc)) {
+            serviceCreate(&out->s, r.Handles[0]);
+        }
+    }
+
+    return rc;
+}
+
+Result fsOpenDataStorageByCurrentProcessFromDomainFwd(Service* s, u32 *out_object_id) {
+    IpcCommand c;
+    ipcInitialize(&c);
+
+    struct {
+        u64 magic;
+        u64 cmd_id;
+    } *raw;
+
+    raw = ipcPrepareHeaderForDomain(&c, sizeof(*raw), s->object_id);
+
+    raw->magic = SFCI_MAGIC;
+    raw->cmd_id = 200;
+
+    Result rc = serviceIpcDispatch(s);
+
+    if (R_SUCCEEDED(rc)) {
+        IpcParsedCommand r;
+        ipcParseForDomain(&r);
+
+        struct {
+            u64 magic;
+            u64 result;
+            u32 object_id;
+        } *resp = r.Raw;
+
+        rc = resp->result;
+
+        if (R_SUCCEEDED(rc)) {
+            *out_object_id = resp->object_id;
+        }
+    }
+
+    return rc;
+}
+
 Result fsOpenDataStorageByDataId(Service* s, FsStorageId storage_id, u64 data_id, FsStorage* out) {
     IpcCommand c;
     ipcInitialize(&c);

stratosphere/fs_mitm/source/fs_shim.h

@@ -20,6 +20,8 @@ typedef struct {
 Result ipcCopyFromDomain(Handle session, u32 object_id, Service *out);
 
 /* Missing fsp-srv commands. */
+Result fsOpenDataStorageByCurrentProcessFwd(Service* s, FsStorage* out);
+Result fsOpenDataStorageByCurrentProcessFromDomainFwd(Service* s, u32 *out_object_id);
 Result fsOpenDataStorageByDataId(Service* s, FsStorageId storage_id, u64 data_id, FsStorage* out);
 Result fsOpenDataStorageByDataIdFromDomain(Service* s, FsStorageId storage_id, u64 data_id, u32 *out_object_id);
 

stratosphere/fs_mitm/source/fsmitm_layeredrom.cpp

@@ -0,0 +1,135 @@
+#include <switch.h>
+#include <stratosphere.hpp>
+
+#include "fsmitm_layeredrom.hpp"
+#include "fsmitm_utils.hpp"
+#include "debug.hpp"
+
+
+LayeredRomFS::LayeredRomFS(std::shared_ptr<RomInterfaceStorage> s_r, std::shared_ptr<RomFileStorage> f_r, u64 tid) : storage_romfs(s_r), file_romfs(f_r), title_id(tid) {
+    /* Start building the new virtual romfs. */
+    RomFSBuildContext build_ctx(this->title_id);
+    this->p_source_infos = std::shared_ptr<std::vector<RomFSSourceInfo>>(new std::vector<RomFSSourceInfo>(), [](std::vector<RomFSSourceInfo> *to_delete) {
+        for (unsigned int i = 0; i < to_delete->size(); i++) {
+            (*to_delete)[i].Cleanup();
+        }
+        delete to_delete;
+    });
+    if (Utils::IsSdInitialized()) {
+        build_ctx.MergeSdFiles();
+    }
+    if (this->file_romfs) {
+        build_ctx.MergeRomStorage(this->file_romfs.get(), RomFSDataSource_FileRomFS);
+    }
+    if (this->storage_romfs) {
+        build_ctx.MergeRomStorage(this->storage_romfs.get(), RomFSDataSource_BaseRomFS);
+    }
+    build_ctx.Build(this->p_source_infos.get());
+}
+
+LayeredRomFS::~LayeredRomFS() {
+    /* ... */
+}
+
+
+Result LayeredRomFS::Read(void *buffer, size_t size, u64 offset)  {
+    /* Validate size. */
+    u64 virt_size = (*this->p_source_infos)[this->p_source_infos->size() - 1].virtual_offset + (*this->p_source_infos)[this->p_source_infos->size() - 1].size;
+    if (offset >= virt_size) {
+        return 0x2F5A02;
+    }
+    if (virt_size - offset < size) {
+        size = virt_size - offset;
+    }
+    /* Find first source info via binary search. */
+    u32 cur_source_ind = 0;
+    u32 low = 0, high = this->p_source_infos->size() - 1;
+    while (low <= high) {
+        u32 mid = (low + high) / 2;
+        if ((*this->p_source_infos)[mid].virtual_offset > offset) {
+            /* Too high. */
+            high = mid - 1;
+        } else {
+            /* sources[mid].virtual_offset <= offset, invariant */
+            if (mid == this->p_source_infos->size() - 1 || (*this->p_source_infos)[mid + 1].virtual_offset > offset) {
+                /* Success */
+                cur_source_ind = mid;
+                break;
+            }
+            low = mid + 1;
+        }
+    }
+    
+    Result rc;
+    size_t read_so_far = 0;
+    while (read_so_far < size) {
+        RomFSSourceInfo *cur_source = &((*this->p_source_infos)[cur_source_ind]);
+        if (cur_source->virtual_offset + cur_source->size > offset) {
+            u64 cur_read_size = size - read_so_far;
+            if (cur_read_size > cur_source->size - (offset - cur_source->virtual_offset)) {
+                cur_read_size = cur_source->size - (offset - cur_source->virtual_offset);
+            }
+            switch (cur_source->type) {
+                case RomFSDataSource_LooseFile:
+                    {
+                        FsFile file;
+                        if (R_FAILED((rc = Utils::OpenRomFSSdFile(this->title_id, cur_source->loose_source_info.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_Memory:
+                    {
+                        memcpy((void *)((uintptr_t)buffer + read_so_far), cur_source->memory_source_info.data + (offset - cur_source->virtual_offset), cur_read_size);
+                    }
+                    break;
+                case RomFSDataSource_BaseRomFS:
+                    {
+                        if (R_FAILED((rc = this->storage_romfs->Read((void *)((uintptr_t)buffer + read_so_far), cur_read_size, cur_source->base_source_info.offset + (offset - cur_source->virtual_offset))))) {
+                            /* TODO: Can this ever happen? */
+                            /* fatalSimple(rc); */
+                            return rc;
+                        }
+                    }
+                    break;
+                case RomFSDataSource_FileRomFS:
+                    {
+                        if (R_FAILED((rc = this->file_romfs->Read((void *)((uintptr_t)buffer + read_so_far), cur_read_size, cur_source->base_source_info.offset + (offset - cur_source->virtual_offset))))) {
+                            fatalSimple(rc);
+                        }
+                    }
+                    break;
+                default:
+                    fatalSimple(0xF601);
+            }
+            read_so_far += cur_read_size;
+        } else {
+            /* Handle padding explicitly. */
+            cur_source_ind++;
+            /* Zero out the padding we skip, here. */
+            memset((void *)((uintptr_t)buffer + read_so_far), 0, ((*this->p_source_infos)[cur_source_ind]).virtual_offset - (cur_source->virtual_offset + cur_source->size));
+            read_so_far += ((*this->p_source_infos)[cur_source_ind]).virtual_offset - (cur_source->virtual_offset + cur_source->size);
+        }
+    }
+
+    return 0;
+}
+Result LayeredRomFS::GetSize(u64 *out_size)  {
+    *out_size = (*this->p_source_infos)[this->p_source_infos->size() - 1].virtual_offset + (*this->p_source_infos)[this->p_source_infos->size() - 1].size;
+    return 0x0;
+}
+Result LayeredRomFS::OperateRange(u32 operation_type, u64 offset, u64 size, FsRangeInfo *out_range_info) {
+    /* TODO: How should I implement this for a virtual romfs? */
+    if (operation_type == 3) {
+        *out_range_info = {0};
+    }
+    return 0;
+}

stratosphere/fs_mitm/source/fsmitm_layeredrom.hpp

@@ -0,0 +1,31 @@
+#pragma once
+#include <switch.h>
+#include <stratosphere.hpp>
+
+#include "fsmitm_romstorage.hpp"
+#include "fsmitm_romfsbuild.hpp"
+#include "fsmitm_utils.hpp"
+
+
+/* Represents a merged RomFS. */
+class LayeredRomFS : public IROStorage {
+    private:
+        /* Data Sources. */
+        std::shared_ptr<RomInterfaceStorage> storage_romfs;
+        std::shared_ptr<RomFileStorage> file_romfs;
+        /* Information about the merged RomFS. */
+        u64 title_id;
+        std::shared_ptr<std::vector<RomFSSourceInfo>> p_source_infos;
+
+        LayeredRomFS *Clone() override {
+            return new LayeredRomFS(*this);
+        };
+                
+    public:
+        LayeredRomFS(std::shared_ptr<RomInterfaceStorage> s_r, std::shared_ptr<RomFileStorage> f_r, u64 tid);
+        virtual ~LayeredRomFS();
+        
+        Result Read(void *buffer, size_t size, u64 offset) override;
+        Result GetSize(u64 *out_size) override;
+        Result OperateRange(u32 operation_type, u64 offset, u64 size, FsRangeInfo *out_range_info) override;
+};

stratosphere/fs_mitm/source/fsmitm_main.cpp

@@ -12,7 +12,7 @@
 #include "fsmitm_service.hpp"
 #include "fsmitm_worker.hpp"
 
-#include "mitm_service.hpp"
+#include "mitm_query_service.hpp"
 
 extern "C" {
     extern u32 __start__;
@@ -64,11 +64,6 @@ void __appInit(void) {
         fatalSimple(0xCAFE << 4 | 3);
     }
     
-    rc = pminfoInitialize();
-    if (R_FAILED(rc)) {
-        fatalSimple(0xCAFE << 4 | 4);
-    }
-    
     /* Check for exosphere API compatibility. */
     u64 exosphere_cfg;
     if (R_SUCCEEDED(splGetConfig((SplConfigItem)65000, &exosphere_cfg))) {
@@ -96,7 +91,7 @@ int main(int argc, char **argv)
     Thread worker_thread = {0};
     consoleDebugInit(debugDevice_SVC);
     
-    if (R_FAILED(threadCreate(&worker_thread, &FsMitmWorker::Main, NULL, 0x8000, 45, 0))) {
+    if (R_FAILED(threadCreate(&worker_thread, &FsMitMWorker::Main, NULL, 0x20000, 45, 0))) {
         /* TODO: Panic. */
     }
     if (R_FAILED(threadStart(&worker_thread))) {
@@ -104,11 +99,14 @@ int main(int argc, char **argv)
     }
     
     /* TODO: What's a good timeout value to use here? */
-    WaitableManager *server_manager = new WaitableManager(U64_MAX);
+    MultiThreadedWaitableManager *server_manager = new MultiThreadedWaitableManager(1, U64_MAX, 0x20000);
+    //WaitableManager *server_manager = new WaitableManager(U64_MAX);
         
     /* Create fsp-srv mitm. */
-    //server_manager->add_waitable(new MitMServer<FsMitMService>("fsp-srv", 61));
+    ISession<MitMQueryService<FsMitMService>> *fs_query_srv = NULL;
-    server_manager->add_waitable(new MitMServer<GenericMitMService>("fsp-srv", 61));
+    MitMServer<FsMitMService> *fs_srv = new MitMServer<FsMitMService>(&fs_query_srv, "fsp-srv", 61);
+    server_manager->add_waitable(fs_srv);
+    server_manager->add_waitable(fs_query_srv);
             
     /* Loop forever, servicing our services. */
     server_manager->process();

stratosphere/fs_mitm/source/fsmitm_romfsbuild.cpp

@@ -0,0 +1,420 @@
+#include <switch.h>
+#include <string.h>
+#include <stratosphere.hpp>
+#include "fsmitm_utils.hpp"
+#include "fsmitm_romfsbuild.hpp"
+
+#include "debug.hpp"
+
+void RomFSBuildContext::VisitDirectory(FsFileSystem *filesys, RomFSBuildDirectoryContext *parent) {
+    FsDir dir;
+    Result rc;
+    
+    /* Open the current parent directory. */
+    if (R_FAILED((rc = Utils::OpenRomFSDir(filesys, this->title_id, parent->path, &dir)))) {
+        fatalSimple(rc);
+    }
+    
+    u64 read_entries;
+    while (R_SUCCEEDED((rc = fsDirRead(&dir, 0, &read_entries, 1, &this->dir_entry))) && read_entries == 1) {
+        if (this->dir_entry.type == ENTRYTYPE_DIR) {
+            RomFSBuildDirectoryContext *child = new RomFSBuildDirectoryContext({0});
+            /* Set child's path. */
+            strcpy(child->path, parent->path);
+            child->cur_path_ofs = parent->path_len + 1;
+            child->path_len = child->cur_path_ofs + strlen(this->dir_entry.name);
+            if (child->path_len > sizeof(child->path)) {
+                fatalSimple(0xF601);
+            }
+            strcat(child->path + parent->path_len, "/");
+            strcat(child->path + parent->path_len, this->dir_entry.name);
+                                    
+            if (!this->AddDirectory(parent, child, NULL)) {
+                delete child;
+            }
+        } else if (this->dir_entry.type == ENTRYTYPE_FILE) {
+            RomFSBuildFileContext *child = new RomFSBuildFileContext({0});
+            /* Set child's path. */
+            strcpy(child->path, parent->path);
+            child->cur_path_ofs = parent->path_len + 1;
+            child->path_len = child->cur_path_ofs + strlen(this->dir_entry.name);
+            if (child->path_len > sizeof(child->path)) {
+                fatalSimple(0xF601);
+            }
+            strcat(child->path + parent->path_len, "/");
+            strcat(child->path + parent->path_len, this->dir_entry.name);
+            
+            child->source = this->cur_source_type;
+            
+            child->size = this->dir_entry.fileSize;
+                                    
+            if (!this->AddFile(parent, child)) {
+                delete child;
+            }
+        } else {
+            fatalSimple(rc);
+        }
+    }
+    fsDirClose(&dir);
+    
+    RomFSBuildDirectoryContext *cur_child = parent->child;
+    while (cur_child != NULL) {
+        this->VisitDirectory(filesys, cur_child);
+        cur_child = cur_child->sibling;
+    }
+}
+
+void RomFSBuildContext::MergeSdFiles() {
+    FsFileSystem sd_filesystem;
+    FsDir dir;
+    if (!Utils::IsSdInitialized()) {
+        return;
+    }
+    if (R_FAILED((Utils::OpenSdDirForAtmosphere(this->title_id, "/romfs", &dir)))) {
+        return;
+    }
+    fsDirClose(&dir);
+    if (R_FAILED(fsMountSdcard(&sd_filesystem))) {
+        return;
+    }
+    this->cur_source_type = RomFSDataSource_LooseFile;
+    this->VisitDirectory(&sd_filesystem, this->root);
+    fsFsClose(&sd_filesystem);
+}
+
+void RomFSBuildContext::VisitDirectory(RomFSBuildDirectoryContext *parent, u32 parent_offset, void *dir_table, size_t dir_table_size, void *file_table, size_t file_table_size) {
+    RomFSDirectoryEntry *parent_entry = romfs_get_direntry(dir_table, parent_offset);
+    if (parent_entry->file != ROMFS_ENTRY_EMPTY) {
+        RomFSFileEntry *cur_file = romfs_get_fentry(file_table, parent_entry->file);
+        while (cur_file != NULL) {
+            RomFSBuildFileContext *child = new RomFSBuildFileContext({0});
+            /* Set child's path. */
+            strcpy(child->path, parent->path);
+            child->cur_path_ofs = parent->path_len + 1;
+            child->path_len = child->cur_path_ofs + cur_file->name_size;
+            if (child->path_len > sizeof(child->path)) {
+                fatalSimple(0xF601);
+            }
+            strcat(child->path + parent->path_len, "/");
+            strncat(child->path + parent->path_len, cur_file->name, cur_file->name_size);
+            child->size = cur_file->size;
+            
+            child->source = this->cur_source_type;
+            child->orig_offset = cur_file->offset;
+            if (!this->AddFile(parent, child)) {
+                delete child;
+            }
+            if (cur_file->sibling == ROMFS_ENTRY_EMPTY) {
+                cur_file = NULL;
+            } else {
+                cur_file = romfs_get_fentry(file_table, cur_file->sibling);
+            }
+        }
+    }
+    if (parent_entry->child != ROMFS_ENTRY_EMPTY) {
+        RomFSDirectoryEntry *cur_child = romfs_get_direntry(dir_table, parent_entry->child);
+        u32 cur_child_offset = parent_entry->child;
+        while (cur_child != NULL) {
+            RomFSBuildDirectoryContext *child = new RomFSBuildDirectoryContext({0});
+            /* Set child's path. */
+            strcpy(child->path, parent->path);
+            child->cur_path_ofs = parent->path_len + 1;
+            child->path_len = child->cur_path_ofs + cur_child->name_size;
+            if (child->path_len > sizeof(child->path)) {
+                fatalSimple(0xF601);
+            }
+            strcat(child->path + parent->path_len, "/");
+            strncat(child->path + parent->path_len, cur_child->name, cur_child->name_size);
+            
+            RomFSBuildDirectoryContext *real = NULL;
+            if (!this->AddDirectory(parent, child, &real)) {
+                delete child;
+            }
+            if (real == NULL) {
+                fatalSimple(0xF601);
+            }
+            
+            this->VisitDirectory(real, cur_child_offset, dir_table, dir_table_size, file_table, file_table_size);
+            
+            if (cur_child->sibling == ROMFS_ENTRY_EMPTY) {
+                cur_child = NULL;
+            } else {
+                cur_child_offset = cur_child->sibling;
+                cur_child = romfs_get_direntry(dir_table, cur_child->sibling);
+            }
+        }
+    }
+}
+
+void RomFSBuildContext::MergeRomStorage(IROStorage *storage, RomFSDataSource source) {
+    Result rc;
+    RomFSHeader header;
+    if (R_FAILED((rc = storage->Read(&header, sizeof(header), 0)))) {
+        fatalSimple(rc);
+    }
+    if (header.header_size != sizeof(header)) {
+        /* what */
+        return;
+    }
+    
+    /* Read tables. */
+    u8 *dir_table = new u8[header.dir_table_size];
+    u8 *file_table = new u8[header.file_table_size];
+    if (R_FAILED((rc = storage->Read(dir_table, header.dir_table_size, header.dir_table_ofs)))) {
+        fatalSimple(rc);
+    }
+    if (R_FAILED((rc = storage->Read(file_table, header.file_table_size, header.file_table_ofs)))) {
+        fatalSimple(rc);
+    }
+    
+    this->cur_source_type = source;
+    this->VisitDirectory(this->root, 0x0, dir_table, (size_t)header.dir_table_size, file_table, (size_t)header.file_table_size);
+    delete dir_table;
+    delete file_table;
+}
+
+bool RomFSBuildContext::AddDirectory(RomFSBuildDirectoryContext *parent_dir_ctx, RomFSBuildDirectoryContext *dir_ctx, RomFSBuildDirectoryContext **out_dir_ctx) {    
+    /* Ordered insertion on child. */
+    int cmp_val;
+    if (parent_dir_ctx->child == NULL || (cmp_val = strcmp(dir_ctx->path, parent_dir_ctx->child->path)) < 0) {
+        dir_ctx->sibling = parent_dir_ctx->child;
+        parent_dir_ctx->child = dir_ctx;
+    } else if (cmp_val == 0) {
+        /* This directory already exists! */
+        if (out_dir_ctx) {
+            *out_dir_ctx = parent_dir_ctx->child;
+        }
+        return false;
+    } else {
+        RomFSBuildDirectoryContext *child, *prev;
+        prev = parent_dir_ctx->child;
+        child = prev->sibling;
+        while (child != NULL) {
+            cmp_val = strcmp(dir_ctx->path, child->path);
+            if (cmp_val < 0) {
+                break;
+            } else if (cmp_val == 0) {
+                /* This directory already exists! */
+                if (out_dir_ctx) {
+                    *out_dir_ctx = child;
+                }
+                return false;
+            }
+            prev = child;
+            child = child->sibling;
+        }
+        
+        prev->sibling = dir_ctx;
+        dir_ctx->sibling = child;
+    }
+    
+    /* If we got this far, we're definitely adding a new directory. */
+    this->num_dirs++;
+    this->dir_table_size += sizeof(RomFSDirectoryEntry) + ((dir_ctx->path_len - dir_ctx->cur_path_ofs + 3) & ~3);
+    dir_ctx->parent = parent_dir_ctx;
+    
+    /* Ordered insertion on next */
+    RomFSBuildDirectoryContext *cur = parent_dir_ctx->next, *prev = parent_dir_ctx;
+    while (cur != NULL) {
+        if (strcmp(dir_ctx->path, cur->path) < 0) {
+            break;
+        }
+        prev = cur;
+        cur = cur->next;         
+    }
+    prev->next = dir_ctx;
+    dir_ctx->next = cur;
+    if (out_dir_ctx) {
+        *out_dir_ctx = dir_ctx;
+    }
+    return true;
+}
+
+bool RomFSBuildContext::AddFile(RomFSBuildDirectoryContext *parent_dir_ctx, RomFSBuildFileContext *file_ctx) {
+    /* Ordered insertion on sibling */
+    int cmp_val;
+    if (parent_dir_ctx->file == NULL || (cmp_val = strcmp(file_ctx->path, parent_dir_ctx->file->path)) < 0) {
+        file_ctx->sibling = parent_dir_ctx->file;
+        parent_dir_ctx->file = file_ctx;
+    } else if (cmp_val == 0) {
+        /* This file already exists! */
+        return false;
+    } else {
+        RomFSBuildFileContext *child, *prev;
+        prev = parent_dir_ctx->file;
+        child = prev->sibling;
+        while (child != NULL) {
+            cmp_val = strcmp(file_ctx->path, child->path);
+            if (cmp_val < 0) {
+                break;
+            } else if (cmp_val == 0) {
+                /* This directory already exists! */
+                return false;
+            }
+            prev = child;
+            child = child->sibling;
+        }
+        
+        prev->sibling = file_ctx;
+        file_ctx->sibling = child;
+    }
+    
+    /* If we got this far, we're definitely adding a new file. */
+    this->num_files++;
+    this->file_table_size += sizeof(RomFSFileEntry) + ((file_ctx->path_len - file_ctx->cur_path_ofs + 3) & ~3);
+    file_ctx->parent = parent_dir_ctx;
+    
+    /* Ordered insertion on next */
+    if (this->files == NULL || strcmp(file_ctx->path, this->files->path) < 0) {
+        file_ctx->next = this->files;
+        this->files = file_ctx;
+    } else {
+        RomFSBuildFileContext *cur = this->files->next, *prev = this->files;
+        while (cur != NULL) {
+            if (strcmp(file_ctx->path, cur->path) < 0) {
+                break;
+            }
+            prev = cur;
+            cur = cur->next;         
+        }
+        prev->next = file_ctx;
+        file_ctx->next = cur;
+    }
+    return true;
+}
+
+void RomFSBuildContext::Build(std::vector<RomFSSourceInfo> *out_infos) {
+    RomFSBuildFileContext *cur_file;
+    RomFSBuildDirectoryContext *cur_dir;
+    u32 entry_offset;
+    
+    u32 dir_hash_table_entry_count  = romfs_get_hash_table_count(this->num_dirs);
+    u32 file_hash_table_entry_count = romfs_get_hash_table_count(this->num_files);
+    this->dir_hash_table_size = 4 * dir_hash_table_entry_count;
+    this->file_hash_table_size = 4 * file_hash_table_entry_count;
+    
+    /* Assign metadata pointers */
+    RomFSHeader *header = new RomFSHeader({0});
+    u8 *metadata = new u8[this->dir_hash_table_size + this->dir_table_size + this->file_hash_table_size + this->file_table_size];
+    u32 *dir_hash_table = (u32 *)((uintptr_t)metadata);
+    RomFSDirectoryEntry *dir_table = (RomFSDirectoryEntry *)((uintptr_t)dir_hash_table + this->dir_hash_table_size);
+    u32 *file_hash_table = (u32 *)((uintptr_t)dir_table + this->dir_table_size);
+    RomFSFileEntry *file_table = (RomFSFileEntry *)((uintptr_t)file_hash_table + this->file_hash_table_size);
+    
+    /* Clear out hash tables. */
+    for (u32 i = 0; i < dir_hash_table_entry_count; i++) {
+        dir_hash_table[i] = ROMFS_ENTRY_EMPTY;
+    }
+    for (u32 i = 0; i < file_hash_table_entry_count; i++) {
+        file_hash_table[i] = ROMFS_ENTRY_EMPTY;
+    }
+    
+    out_infos->clear();
+    out_infos->push_back(RomFSSourceInfo(0, sizeof(*header), header, RomFSDataSource_Memory));
+        
+    /* Determine file offsets. */
+    cur_file = this->files;
+    entry_offset = 0;
+    while (cur_file != NULL) {
+        this->file_partition_size = (this->file_partition_size + 0xFULL) & ~0xFULL;
+        cur_file->offset = this->file_partition_size;
+        this->file_partition_size += cur_file->size;
+        cur_file->entry_offset = entry_offset;
+        entry_offset += sizeof(RomFSFileEntry) + ((cur_file->path_len - cur_file->cur_path_ofs + 3) & ~3);
+        cur_file = cur_file->next;
+    }
+    
+    /* Determine directory offsets. */
+    cur_dir = this->root;
+    entry_offset = 0;
+    while (cur_dir != NULL) {
+        cur_dir->entry_offset = entry_offset;
+        entry_offset += sizeof(RomFSDirectoryEntry) + ((cur_dir->path_len - cur_dir->cur_path_ofs + 3) & ~3);
+        cur_dir = cur_dir->next;
+    }
+    
+    
+    /* Populate file tables. */
+    cur_file = this->files;
+    while (cur_file != NULL) {
+        RomFSFileEntry *cur_entry = romfs_get_fentry(file_table, cur_file->entry_offset);
+
+        cur_entry->parent = cur_file->parent->entry_offset;
+        cur_entry->sibling = (cur_file->sibling == NULL) ? ROMFS_ENTRY_EMPTY : cur_file->sibling->entry_offset;
+        cur_entry->offset = cur_file->offset;
+        cur_entry->size = cur_file->size;
+        
+        u32 name_size = cur_file->path_len - cur_file->cur_path_ofs;
+        u32 hash = romfs_calc_path_hash(cur_file->parent->entry_offset, (unsigned char *)cur_file->path + cur_file->cur_path_ofs, 0, name_size);
+        cur_entry->hash = file_hash_table[hash % file_hash_table_entry_count];
+        file_hash_table[hash % file_hash_table_entry_count] = cur_file->entry_offset;
+        
+        
+        cur_entry->name_size = name_size;
+        memset(cur_entry->name, 0, (cur_entry->name_size + 3) & ~3);
+        memcpy(cur_entry->name, cur_file->path + cur_file->cur_path_ofs, name_size);
+        
+        
+        switch (cur_file->source) {
+            case RomFSDataSource_BaseRomFS:
+            case RomFSDataSource_FileRomFS:
+                /* Try to compact, if possible. */
+                if (out_infos->back().type == cur_file->source) {
+                    out_infos->back().size = cur_file->offset + ROMFS_FILEPARTITION_OFS + cur_file->size - out_infos->back().virtual_offset;
+                } else {
+                    out_infos->push_back(RomFSSourceInfo(cur_file->offset + ROMFS_FILEPARTITION_OFS, cur_file->size, cur_file->orig_offset + ROMFS_FILEPARTITION_OFS, cur_file->source));
+                }
+                break;
+            case RomFSDataSource_LooseFile: 
+                {
+                    char *path = new char[cur_file->path_len + 1];
+                    strcpy(path, cur_file->path);
+                    out_infos->push_back(RomFSSourceInfo(cur_file->offset + ROMFS_FILEPARTITION_OFS, cur_file->size, path, cur_file->source));
+                }
+                break;
+            default:
+                fatalSimple(0xF601);
+        }
+        
+        RomFSBuildFileContext *temp = cur_file;
+        cur_file = cur_file->next;
+        delete temp;
+    }
+        
+    /* Populate dir tables. */
+    cur_dir = this->root;
+    while (cur_dir != NULL) {
+        RomFSDirectoryEntry *cur_entry = romfs_get_direntry(dir_table, cur_dir->entry_offset);
+        cur_entry->parent = cur_dir == this->root ? 0 : cur_dir->parent->entry_offset;
+        cur_entry->sibling = (cur_dir->sibling == NULL) ? ROMFS_ENTRY_EMPTY : cur_dir->sibling->entry_offset;
+        cur_entry->child = (cur_dir->child == NULL) ? ROMFS_ENTRY_EMPTY : cur_dir->child->entry_offset;
+        cur_entry->file = (cur_dir->file == NULL) ? ROMFS_ENTRY_EMPTY : cur_dir->file->entry_offset;
+        
+        u32 name_size = cur_dir->path_len - cur_dir->cur_path_ofs;
+        u32 hash = romfs_calc_path_hash(cur_dir == this->root ? 0 : cur_dir->parent->entry_offset, (unsigned char *)cur_dir->path + cur_dir->cur_path_ofs, 0, name_size);
+        cur_entry->hash = dir_hash_table[hash % dir_hash_table_entry_count];
+        dir_hash_table[hash % dir_hash_table_entry_count] = cur_dir->entry_offset;
+        
+        cur_entry->name_size = name_size;
+        memset(cur_entry->name, 0, (cur_entry->name_size + 3) & ~3);
+        memcpy(cur_entry->name, cur_dir->path + cur_dir->cur_path_ofs, name_size);
+        
+        RomFSBuildDirectoryContext *temp = cur_dir;
+        cur_dir = cur_dir->next;
+        delete temp;
+    }
+    
+    /* Set header fields. */
+    header->header_size = sizeof(*header);
+    header->file_hash_table_size = this->file_hash_table_size;
+    header->file_table_size = this->file_table_size;
+    header->dir_hash_table_size = this->dir_hash_table_size;
+    header->dir_table_size = this->dir_table_size;
+    header->file_partition_ofs = ROMFS_FILEPARTITION_OFS;
+    header->dir_hash_table_ofs = (header->file_partition_ofs + this->file_partition_size + 3ULL) & ~3ULL;
+    header->dir_table_ofs = header->dir_hash_table_ofs + header->dir_hash_table_size;
+    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;
+    
+    out_infos->push_back(RomFSSourceInfo(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

@@ -0,0 +1,228 @@
+#pragma once
+#include <switch.h>
+
+#include "fsmitm_romstorage.hpp"
+
+#define ROMFS_ENTRY_EMPTY 0xFFFFFFFF
+#define ROMFS_FILEPARTITION_OFS 0x200
+
+/* Types for RomFS Meta construction. */
+enum RomFSDataSource {
+    RomFSDataSource_BaseRomFS,
+    RomFSDataSource_FileRomFS,
+    RomFSDataSource_LooseFile,
+    RomFSDataSource_Memory,
+};
+
+struct RomFSBaseSourceInfo {
+    u64 offset;
+};
+
+struct RomFSFileSourceInfo {
+    u64 offset;
+};
+
+struct RomFSLooseSourceInfo {
+    const char *path;
+};
+
+struct RomFSMemorySourceInfo {
+    const u8 *data;
+};
+
+struct RomFSSourceInfo {
+    u64 virtual_offset;
+    u64 size;
+    union {
+        RomFSBaseSourceInfo base_source_info;
+        RomFSFileSourceInfo file_source_info;
+        RomFSLooseSourceInfo loose_source_info;
+        RomFSMemorySourceInfo memory_source_info;
+    };
+    RomFSDataSource type;
+    
+    RomFSSourceInfo(u64 v_o, u64 s, u64 offset, RomFSDataSource t) : virtual_offset(v_o), size(s), type(t) {
+        switch (this->type) {
+            case RomFSDataSource_BaseRomFS:
+                this->base_source_info.offset = offset;
+                break;
+            case RomFSDataSource_FileRomFS:
+                this->file_source_info.offset = offset;
+                break;
+            case RomFSDataSource_LooseFile:
+            case RomFSDataSource_Memory:
+            default:
+                fatalSimple(0xF601);
+        }
+    }
+    
+    RomFSSourceInfo(u64 v_o, u64 s, const void *arg, RomFSDataSource t) : virtual_offset(v_o), size(s), type(t) {
+        switch (this->type) {
+            case RomFSDataSource_LooseFile:
+                this->loose_source_info.path = (decltype(this->loose_source_info.path))arg;
+                break;
+            case RomFSDataSource_Memory:
+                this->memory_source_info.data = (decltype(this->memory_source_info.data))arg;
+                break;
+            case RomFSDataSource_BaseRomFS:
+            case RomFSDataSource_FileRomFS:
+            default:
+                fatalSimple(0xF601);
+        }
+    }
+    
+    void Cleanup() {
+        switch (this->type) {
+            case RomFSDataSource_BaseRomFS:
+            case RomFSDataSource_FileRomFS:
+                break;
+            case RomFSDataSource_LooseFile:
+                delete this->loose_source_info.path;
+                break;
+            case RomFSDataSource_Memory:
+                delete this->memory_source_info.data;
+                break;
+            default:
+                fatalSimple(0xF601);
+        }
+    }
+    
+    static bool Compare(RomFSSourceInfo *a, RomFSSourceInfo *b) {
+        return (a->virtual_offset < b->virtual_offset);
+    }
+};
+
+/* Types for building a RomFS. */
+struct RomFSHeader {
+    u64 header_size;
+    u64 dir_hash_table_ofs;
+    u64 dir_hash_table_size;
+    u64 dir_table_ofs;
+    u64 dir_table_size;
+    u64 file_hash_table_ofs;
+    u64 file_hash_table_size;
+    u64 file_table_ofs;
+    u64 file_table_size;
+    u64 file_partition_ofs;
+};
+
+static_assert(sizeof(RomFSHeader) == 0x50, "Incorrect RomFS Header definition!");
+
+struct RomFSDirectoryEntry {
+    u32 parent;
+    u32 sibling;
+    u32 child;
+    u32 file;
+    u32 hash;
+    u32 name_size;
+    char name[];
+};
+
+static_assert(sizeof(RomFSDirectoryEntry) == 0x18, "Incorrect RomFSDirectoryEntry definition!");
+
+struct RomFSFileEntry {
+    u32 parent;
+    u32 sibling;
+    u64 offset;
+    u64 size;
+    u32 hash;
+    u32 name_size;
+    char name[];
+};
+
+static_assert(sizeof(RomFSFileEntry) == 0x20, "Incorrect RomFSFileEntry definition!");
+
+struct RomFSBuildFileContext;
+
+struct RomFSBuildDirectoryContext {
+    char path[FS_MAX_PATH];
+    u32 cur_path_ofs;
+    u32 path_len;
+    u32 entry_offset;
+    RomFSBuildDirectoryContext *parent;
+    RomFSBuildDirectoryContext *child;
+    RomFSBuildDirectoryContext *sibling;
+    RomFSBuildFileContext *file;
+    RomFSBuildDirectoryContext *next;
+};
+
+struct RomFSBuildFileContext {
+    char path[FS_MAX_PATH];
+    u32 cur_path_ofs;
+    u32 path_len;
+    u32 entry_offset;
+    u64 offset;
+    u64 size;
+    RomFSBuildDirectoryContext *parent;
+    RomFSBuildFileContext *sibling;
+    RomFSBuildFileContext *next;
+    RomFSDataSource source;
+    u64 orig_offset;
+};
+
+class RomFSBuildContext {
+    private:
+        u64 title_id;
+        RomFSBuildDirectoryContext *root;
+        RomFSBuildFileContext *files;
+        u64 num_dirs;
+        u64 num_files;
+        u64 dir_table_size;
+        u64 file_table_size;
+        u64 dir_hash_table_size;
+        u64 file_hash_table_size;
+        u64 file_partition_size;
+        
+        FsDirectoryEntry dir_entry;
+        RomFSDataSource cur_source_type;
+        
+        void VisitDirectory(FsFileSystem *filesys, RomFSBuildDirectoryContext *parent);
+        void VisitDirectory(RomFSBuildDirectoryContext *parent, u32 parent_offset, void *dir_table, size_t dir_table_size, void *file_table, size_t file_table_size);
+    
+        bool AddDirectory(RomFSBuildDirectoryContext *parent_dir_ctx, RomFSBuildDirectoryContext *dir_ctx, RomFSBuildDirectoryContext **out_dir_ctx);
+        bool AddFile(RomFSBuildDirectoryContext *parent_dir_ctx, RomFSBuildFileContext *file_ctx);
+    public:
+        RomFSBuildContext(u64 tid) : title_id(tid), root(NULL), files(NULL), num_dirs(0), num_files(0), dir_table_size(0), file_table_size(0), dir_hash_table_size(0), file_hash_table_size(0), file_partition_size(0) {
+            this->root = new RomFSBuildDirectoryContext({0}); 
+            this->num_dirs = 1;
+            this->dir_table_size = 0x18;
+        }
+        
+        void MergeSdFiles();
+        void MergeRomStorage(IROStorage *storage, RomFSDataSource source);
+        
+        /* This finalizes the context. */
+        void Build(std::vector<RomFSSourceInfo> *out_infos);
+};
+
+
+static inline RomFSDirectoryEntry *romfs_get_direntry(void *directories, uint32_t offset) {
+    return (RomFSDirectoryEntry *)((uintptr_t)directories + offset);
+}
+
+static inline RomFSFileEntry *romfs_get_fentry(void *files, uint32_t offset) {
+    return (RomFSFileEntry *)((uintptr_t)files + offset);
+}
+
+static inline uint32_t romfs_calc_path_hash(uint32_t parent, const unsigned char *path, uint32_t start, size_t path_len) {
+    uint32_t hash = parent ^ 123456789;
+    for (uint32_t i = 0; i < path_len; i++) {
+        hash = (hash >> 5) | (hash << 27);
+        hash ^= path[start + i];
+    }
+        
+    return hash;
+}
+
+static inline uint32_t romfs_get_hash_table_count(uint32_t num_entries) {
+    if (num_entries < 3) {
+        return 3;
+    } else if (num_entries < 19) {
+        return num_entries | 1;
+    }
+    uint32_t count = num_entries;
+    while (count % 2 == 0 || count % 3 == 0 || count % 5 == 0 || count % 7 == 0 || count % 11 == 0 || count % 13 == 0 || count % 17 == 0) {
+        count++;
+    }
+    return count;
+}

stratosphere/fs_mitm/source/fsmitm_romstorage.hpp

@@ -23,12 +23,12 @@ class RomFileStorage : public IROStorage {
         RomFileStorage *Clone() override {
             return new RomFileStorage(this->base_file);
         };
-    protected:
+    public:
-        Result Read(void *buffer, size_t size, u64 offset, u64 *out_read_size) override {
+        Result Read(void *buffer, size_t size, u64 offset) override {
             size_t out_sz = 0;
             Result rc = fsFileRead(this->base_file, offset, buffer, size, &out_sz);
-            if (R_SUCCEEDED(rc)) {
+            if (R_SUCCEEDED(rc) && out_sz != size && out_sz) {
-                *out_read_size = out_sz;
+                return this->Read((void *)((uintptr_t)buffer + out_sz), size - out_sz, offset + out_sz);
             }
             return rc;
         };
@@ -36,7 +36,7 @@ class RomFileStorage : public IROStorage {
             return fsFileGetSize(this->base_file, out_size);
         };
         Result OperateRange(u32 operation_type, u64 offset, u64 size, FsRangeInfo *out_range_info) override {
-            /* TODO: Merge into libnx. */
+            /* TODO: Merge into libnx? */
             return fsFileOperateRange(this->base_file, operation_type, offset, size, out_range_info);
         };
 };
@@ -60,20 +60,16 @@ class RomInterfaceStorage : public IROStorage {
         RomInterfaceStorage *Clone() override {
             return new RomInterfaceStorage(this->base_storage);
         };
-    protected:
+    public:
-        Result Read(void *buffer, size_t size, u64 offset, u64 *out_read_size) override {
+        Result Read(void *buffer, size_t size, u64 offset) override {
-            Result rc = fsStorageRead(this->base_storage, offset, buffer, size);
+            return fsStorageRead(this->base_storage, offset, buffer, size);
-            if (R_SUCCEEDED(rc)) {
-                *out_read_size = size;
-            }
-            return rc;
         };
         Result GetSize(u64 *out_size) override {
-            /* TODO: Merge into libnx. */
+            /* TODO: Merge into libnx? */
             return fsStorageGetSize(this->base_storage, out_size);
         };
         Result OperateRange(u32 operation_type, u64 offset, u64 size, FsRangeInfo *out_range_info) override {
-            /* TODO: Merge into libnx. */
+            /* TODO: Merge into libnx? */
             return fsStorageOperateRange(this->base_storage, operation_type, offset, size, out_range_info);
         };
 };

stratosphere/fs_mitm/source/fsmitm_service.cpp

@@ -5,21 +5,28 @@
 #include "fsmitm_worker.hpp"
 #include "fsmitm_utils.hpp"
 #include "fsmitm_romstorage.hpp"
+#include "fsmitm_layeredrom.hpp"
 
+#include "mitm_query_service.hpp"
 #include "debug.hpp"
 
 Result FsMitMService::dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) {
     Result rc = 0xF601;
-    
+    if (this->has_initialized) {
-    switch (cmd_id) {
+        switch (cmd_id) {
-        case FspSrv_Cmd_SetCurrentProcess:
+            case FspSrv_Cmd_OpenDataStorageByCurrentProcess:
-            if (!this->has_initialized && r.HasPid) {
+                rc = WrapIpcCommandImpl<&FsMitMService::open_data_storage_by_current_process>(this, r, out_c, pointer_buffer, pointer_buffer_size);
-                this->process_id = r.Pid;
+                break;
+            case FspSrv_Cmd_OpenDataStorageByDataId:
+                rc = WrapIpcCommandImpl<&FsMitMService::open_data_storage_by_data_id>(this, r, out_c, pointer_buffer, pointer_buffer_size);
+                break;
+        }
+    } else {
+        if (cmd_id == FspSrv_Cmd_SetCurrentProcess) {
+            if (r.HasPid) {
+                this->init_pid = r.Pid;
             }
-            break;
+        }
-        /*case FspSrv_Cmd_OpenDataStorageByDataId:
-            rc = WrapIpcCommandImpl<&FsMitMService::open_data_storage_by_data_id>(this, r, out_c, pointer_buffer, pointer_buffer_size);
-            break; */
     }
     return rc;
 }
@@ -41,24 +48,15 @@ void FsMitMService::postprocess(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_
         case FspSrv_Cmd_SetCurrentProcess:
             if (R_SUCCEEDED(rc)) {
                 this->has_initialized = true;
-                rc = pminfoGetTitleId(&this->title_id, this->process_id);
-                if (R_FAILED(rc)) {
-                    if (rc == 0x20F) {
-                        this->title_id = this->process_id;
-                        rc = 0x0;
-                    } else {
-                        fatalSimple(rc);
-                    }
-                }
             }
-            Log(&this->process_id, 8);
+            this->process_id = this->init_pid;
-            Log(&this->title_id, 8);
+            this->title_id = this->process_id;
+            if (R_FAILED(MitMQueryUtils::get_associated_tid_for_pid(this->process_id, &this->title_id))) {
+                /* Log here, if desired. */
+            }
             for (unsigned int i = 0; i < sizeof(backup_tls)/sizeof(u64); i++) {
                 tls[i] = backup_tls[i];
             }
-            if (this->title_id >= 0x0100000000001000) {
-                Reboot();
-            }
             break;
     }
     resp->result = rc;
@@ -69,13 +67,46 @@ Result FsMitMService::handle_deferred() {
     return 0;
 }
 
+/* Add redirection for RomFS to the SD card. */
+std::tuple<Result, OutSession<IStorageInterface>> FsMitMService::open_data_storage_by_current_process() {
+    IPCSession<IStorageInterface> *out_session = NULL;
+    FsStorage data_storage;
+    FsFile data_file;
+    u32 out_domain_id = 0;
+    Result rc;
+    if (this->get_owner() == NULL) {
+        rc = fsOpenDataStorageByCurrentProcessFwd(this->forward_service, &data_storage);
+    } else {
+        rc = fsOpenDataStorageByCurrentProcessFromDomainFwd(this->forward_service, &out_domain_id);
+        if (R_SUCCEEDED(rc)) {
+            rc = ipcCopyFromDomain(this->forward_service->handle, out_domain_id, &data_storage.s);
+        }
+    }
+    Log(armGetTls(), 0x100);
+    if (R_SUCCEEDED(rc)) {
+        /* TODO: Is there a sensible path that ends in ".romfs" we can use?" */
+        if (R_SUCCEEDED(Utils::OpenSdFileForAtmosphere(this->title_id, "romfs.bin", FS_OPEN_READ, &data_file))) {
+            out_session = new IPCSession<IStorageInterface>(std::make_shared<IStorageInterface>(new LayeredRomFS(std::make_shared<RomInterfaceStorage>(data_storage), std::make_shared<RomFileStorage>(data_file), this->title_id)));
+        } else {
+            out_session = new IPCSession<IStorageInterface>(std::make_shared<IStorageInterface>(new LayeredRomFS(std::make_shared<RomInterfaceStorage>(data_storage), nullptr, this->title_id)));
+        }
+        if (this->get_owner() == NULL) {
+            FsMitMWorker::AddWaitable(out_session);
+        }
+    }
+    
+    OutSession out_s = OutSession(out_session);
+    out_s.domain_id = out_domain_id;
+    return {rc, out_s};
+}
+
 /* Add redirection for System Data Archives to the SD card. */
 std::tuple<Result, OutSession<IStorageInterface>> FsMitMService::open_data_storage_by_data_id(u64 sid, u64 data_id) {
     FsStorageId storage_id = (FsStorageId)sid;
     IPCSession<IStorageInterface> *out_session = NULL;
     FsStorage data_storage;
     FsFile data_file;
-    u32 out_domain_id;
+    u32 out_domain_id = 0;
     Result rc;
     if (this->get_owner() == NULL) {
         rc = fsOpenDataStorageByDataId(this->forward_service, storage_id, data_id, &data_storage);
@@ -86,17 +117,14 @@ std::tuple<Result, OutSession<IStorageInterface>> FsMitMService::open_data_stora
         }
     }
     if (R_SUCCEEDED(rc)) {
-        char path[FS_MAX_PATH] = {0};
         /* TODO: Is there a sensible path that ends in ".romfs" we can use?" */
-        snprintf(path, sizeof(path), "/atmosphere/titles/%016lx/romfs.bin", data_id);
+        if (R_SUCCEEDED(Utils::OpenSdFileForAtmosphere(data_id, "romfs.bin", FS_OPEN_READ, &data_file))) {
-        if (R_SUCCEEDED(Utils::OpenSdFile(path, FS_OPEN_READ, &data_file))) {
+            out_session = new IPCSession<IStorageInterface>(std::make_shared<IStorageInterface>(new LayeredRomFS(std::make_shared<RomInterfaceStorage>(data_storage), std::make_shared<RomFileStorage>(data_file), data_id)));
-            fsStorageClose(&data_storage);
-            out_session = new IPCSession<IStorageInterface>(new IStorageInterface(new RomFileStorage(data_file)));
         } else {
-            out_session = new IPCSession<IStorageInterface>(new IStorageInterface(new RomInterfaceStorage(data_storage)));
+            out_session = new IPCSession<IStorageInterface>(std::make_shared<IStorageInterface>(new LayeredRomFS(std::make_shared<RomInterfaceStorage>(data_storage), nullptr, data_id)));
         }
         if (this->get_owner() == NULL) {
-            FsMitmWorker::AddWaitable(out_session);
+            FsMitMWorker::AddWaitable(out_session);
         }
     }
     

stratosphere/fs_mitm/source/fsmitm_service.hpp

@@ -6,19 +6,23 @@
 
 enum FspSrvCmd {
     FspSrv_Cmd_SetCurrentProcess = 1,
+    FspSrv_Cmd_OpenDataStorageByCurrentProcess = 200,
     FspSrv_Cmd_OpenDataStorageByDataId = 202,
 };
 
 class FsMitMService : public IMitMServiceObject {      
     private:
         bool has_initialized;
-        u64 process_id;
+        u64 init_pid;
-        u64 title_id;
     public:
-        FsMitMService(Service *s) : IMitMServiceObject(s), has_initialized(false), process_id(0), title_id(0) {
+        FsMitMService(Service *s) : IMitMServiceObject(s), has_initialized(false), init_pid(0) {
             /* ... */
         }
         
+        static bool should_mitm(u64 pid, u64 tid) {
+            return tid >= 0x0100000000010000ULL;
+        }
+        
         FsMitMService *clone() override {
             auto new_srv = new FsMitMService((Service *)&this->forward_service);
             this->clone_to(new_srv);
@@ -28,8 +32,7 @@ class FsMitMService : public IMitMServiceObject {
         void clone_to(void *o) override {
             FsMitMService *other = (FsMitMService *)o;
             other->has_initialized = has_initialized;
-            other->process_id = process_id;
+            other->init_pid = init_pid;
-            other->title_id = title_id;
         }
         
         virtual Result dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size);
@@ -38,5 +41,6 @@ class FsMitMService : public IMitMServiceObject {
     
     protected:
         /* Overridden commands. */
+        std::tuple<Result, OutSession<IStorageInterface>> open_data_storage_by_current_process();
         std::tuple<Result, OutSession<IStorageInterface>> open_data_storage_by_data_id(u64 storage_id, u64 data_id);
 };

stratosphere/fs_mitm/source/fsmitm_utils.cpp

@@ -3,10 +3,10 @@
 #include <atomic>
 
 #include "sm_mitm.h"
-
+#include "debug.hpp"
 #include "fsmitm_utils.hpp"
 
-static FsFileSystem g_sd_filesystem;
+static FsFileSystem g_sd_filesystem = {0};
 static bool g_has_initialized = false;
 
 static Result EnsureInitialized() {
@@ -19,7 +19,7 @@ static Result EnsureInitialized() {
         Result rc = smMitMUninstall(required_active_services[i]);
         if (rc == 0xE15) {
             return rc;
-        } else if (rc != 0x1015) {
+        } else if (R_FAILED(rc) && rc != 0x1015) {
             return rc;
         }
     }
@@ -31,12 +31,93 @@ static Result EnsureInitialized() {
     return rc;
 }
 
+bool Utils::IsSdInitialized() {
+    return R_SUCCEEDED(EnsureInitialized());
+}
+
 Result Utils::OpenSdFile(const char *fn, int flags, FsFile *out) {
     Result rc;
     if (R_FAILED((rc = EnsureInitialized()))) {
         return rc;
     }
+    
+    return fsFsOpenFile(&g_sd_filesystem, fn, flags, out);
+}
+
+Result Utils::OpenSdFileForAtmosphere(u64 title_id, const char *fn, int flags, FsFile *out) {
+    Result rc;
+    if (R_FAILED((rc = EnsureInitialized()))) {
+        return rc;
+    }
+    
     char path[FS_MAX_PATH];
-    strcpy(path, fn);
+    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);
+    }
     return fsFsOpenFile(&g_sd_filesystem, path, flags, out);
 }
+
+Result Utils::OpenRomFSSdFile(u64 title_id, const char *fn, int flags, FsFile *out) {
+    Result rc;
+    if (R_FAILED((rc = EnsureInitialized()))) {
+        return rc;
+    }
+    
+    return OpenRomFSFile(&g_sd_filesystem, title_id, fn, flags, out);
+}
+
+Result Utils::OpenSdDir(const char *path, FsDir *out) {
+    Result rc;
+    if (R_FAILED((rc = EnsureInitialized()))) {
+        return rc;
+    }
+    
+    return fsFsOpenDirectory(&g_sd_filesystem, path, FS_DIROPEN_DIRECTORY | FS_DIROPEN_FILE, out);
+}
+
+Result Utils::OpenSdDirForAtmosphere(u64 title_id, const char *path, FsDir *out) {
+    Result rc;
+    if (R_FAILED((rc = EnsureInitialized()))) {
+        return rc;
+    }
+    
+    char safe_path[FS_MAX_PATH];
+    if (*path == '/') {
+        snprintf(safe_path, sizeof(safe_path), "/atmosphere/titles/%016lx%s", title_id, path);
+    } else {
+        snprintf(safe_path, sizeof(safe_path), "/atmosphere/titles/%016lx/%s", title_id, path);
+    }
+    return fsFsOpenDirectory(&g_sd_filesystem, safe_path, FS_DIROPEN_DIRECTORY | FS_DIROPEN_FILE, out);
+}
+
+Result Utils::OpenRomFSSdDir(u64 title_id, const char *path, FsDir *out) {
+    Result rc;
+    if (R_FAILED((rc = EnsureInitialized()))) {
+        return rc;
+    }
+    
+    return OpenRomFSDir(&g_sd_filesystem, title_id, path, out);
+}
+
+
+Result Utils::OpenRomFSFile(FsFileSystem *fs, u64 title_id, const char *fn, int flags, FsFile *out) {
+    char path[FS_MAX_PATH];
+    if (*fn == '/') {
+        snprintf(path, sizeof(path), "/atmosphere/titles/%016lx/romfs%s", title_id, fn);
+    } else {
+        snprintf(path, sizeof(path), "/atmosphere/titles/%016lx/romfs/%s", title_id, fn);
+    }
+    return fsFsOpenFile(fs, path, flags, out);
+}
+
+Result Utils::OpenRomFSDir(FsFileSystem *fs, u64 title_id, const char *path, FsDir *out) {
+    char safe_path[FS_MAX_PATH];
+    if (*path == '/') {
+        snprintf(safe_path, sizeof(safe_path), "/atmosphere/titles/%016lx/romfs%s", title_id, path);
+    } else {
+        snprintf(safe_path, sizeof(safe_path), "/atmosphere/titles/%016lx/romfs/%s", title_id, path);
+    }
+    return fsFsOpenDirectory(fs, safe_path, FS_DIROPEN_DIRECTORY | FS_DIROPEN_FILE, out);
+}

stratosphere/fs_mitm/source/fsmitm_utils.hpp

@@ -4,5 +4,15 @@
 
 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);
+        static Result OpenSdDir(const char *path, FsDir *out);
+        static Result OpenSdDirForAtmosphere(u64 title_id, const char *path, FsDir *out);
+        static Result OpenRomFSSdDir(u64 title_id, const char *path, FsDir *out);
+        
+        
+        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);
 };

stratosphere/fs_mitm/source/fsmitm_worker.cpp

@@ -9,13 +9,14 @@ static HosSemaphore g_sema_new_waitable_finish;
 
 static WaitableManager *g_worker_waiter = NULL;
 
-Result FsMitmWorker::AddWaitableCallback(Handle *handles, size_t num_handles, u64 timeout) {
+Result FsMitMWorker::AddWaitableCallback(void *arg, Handle *handles, size_t num_handles, u64 timeout) {
+    (void)arg;
     svcClearEvent(handles[0]);
     g_sema_new_waitable_finish.Signal();
     return 0;
 }
 
-void FsMitmWorker::AddWaitable(IWaitable *waitable) {
+void FsMitMWorker::AddWaitable(IWaitable *waitable) {
     g_worker_waiter->add_waitable(waitable);
     g_new_waitable_mutex.Lock();
     g_new_waitable_event->signal_event();
@@ -23,9 +24,9 @@ void FsMitmWorker::AddWaitable(IWaitable *waitable) {
     g_new_waitable_mutex.Unlock();
 }
 
-void FsMitmWorker::Main(void *arg) {
+void FsMitMWorker::Main(void *arg) {
     /* Initialize waitable event. */
-    g_new_waitable_event = new SystemEvent(&FsMitmWorker::AddWaitableCallback);
+    g_new_waitable_event = new SystemEvent(NULL, &FsMitMWorker::AddWaitableCallback);
 
     /* Make a new waitable manager. */
     g_worker_waiter = new WaitableManager(U64_MAX);

stratosphere/fs_mitm/source/fsmitm_worker.hpp

@@ -2,9 +2,9 @@
 #include <switch.h>
 #include <stratosphere.hpp>
 
-class FsMitmWorker {
+class FsMitMWorker {
     private:
-        static Result AddWaitableCallback(Handle *handles, size_t num_handles, u64 timeout);
+        static Result AddWaitableCallback(void *arg, Handle *handles, size_t num_handles, u64 timeout);
     public:
         static void Main(void *arg);
         static void AddWaitable(IWaitable *waitable);

stratosphere/fs_mitm/source/imitmserviceobject.hpp

@@ -1,16 +1,25 @@
 #pragma once
 #include <switch.h>
+#include <atomic>
 
 #include <stratosphere.hpp>
 
+#include "debug.hpp"
+
 class IMitMServiceObject : public IServiceObject {
     protected:
         Service *forward_service;
+        u64 process_id;
+        u64 title_id;
     public:
-        IMitMServiceObject(Service *s) : forward_service(s) {
+        IMitMServiceObject(Service *s) : forward_service(s), process_id(0), title_id(0) {
             
         }
         
+        static bool should_mitm(u64 pid, u64 tid) {
+            return true;
+        }
+                
         virtual void clone_to(void *o) = 0;
     protected:
         virtual ~IMitMServiceObject() { }

stratosphere/fs_mitm/source/mitm_query_service.cpp

@@ -0,0 +1,28 @@
+#include <switch.h>
+#include <stratosphere.hpp>
+#include "mitm_query_service.hpp"
+
+static std::vector<u64> g_known_pids;
+static std::vector<u64> g_known_tids;
+static HosMutex g_pid_tid_mutex;
+
+Result MitMQueryUtils::get_associated_tid_for_pid(u64 pid, u64 *tid) {
+    Result rc = 0xCAFE;
+    g_pid_tid_mutex.Lock();
+    for (unsigned int i = 0; i < g_known_pids.size(); i++) {
+        if (g_known_pids[i] == pid) {
+            *tid = g_known_tids[i];
+            rc = 0x0;
+            break;
+        }
+    }
+    g_pid_tid_mutex.Unlock();
+    return rc;
+}
+
+void MitMQueryUtils::associate_pid_to_tid(u64 pid, u64 tid) {
+    g_pid_tid_mutex.Lock();
+    g_known_pids.push_back(pid);
+    g_known_tids.push_back(tid);
+    g_pid_tid_mutex.Unlock();
+}

stratosphere/fs_mitm/source/mitm_query_service.hpp

@@ -0,0 +1,59 @@
+#pragma once
+#include <switch.h>
+#include <stratosphere/iserviceobject.hpp>
+
+#include "debug.hpp"
+
+enum MitMQueryServiceCommand {
+    MQS_Cmd_ShouldMitm = 65000,
+    MQS_Cmd_AssociatePidTid = 65001
+};
+
+namespace MitMQueryUtils {
+    Result get_associated_tid_for_pid(u64 pid, u64 *tid);
+
+    void associate_pid_to_tid(u64 pid, u64 tid); 
+}
+
+
+template <typename T>
+class MitMQueryService : public IServiceObject {
+    public:
+        MitMQueryService<T> *clone() override {
+            return new MitMQueryService<T>();
+        }
+        Result dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) override {
+            Log(armGetTls(), 0x100);
+            switch (cmd_id) {
+                case MQS_Cmd_ShouldMitm:
+                    return WrapIpcCommandImpl<&MitMQueryService::should_mitm>(this, r, out_c, pointer_buffer, pointer_buffer_size);
+                case MQS_Cmd_AssociatePidTid:
+                    return WrapIpcCommandImpl<&MitMQueryService::associate_pid_tid>(this, r, out_c, pointer_buffer, pointer_buffer_size);
+                default:
+                    return 0xF601;
+            }
+            if (cmd_id == 65000) {
+                
+            } else {
+                return 0xF601;
+            }
+        }
+        Result handle_deferred() override {
+            /* This service is never deferrable. */
+            return 0;
+        }
+        
+    protected:
+        std::tuple<Result, u64> should_mitm(u64 pid) {
+            u64 should_mitm = 0;
+            u64 tid = 0;
+            if (R_SUCCEEDED(MitMQueryUtils::get_associated_tid_for_pid(pid, &tid))) {
+                should_mitm = T::should_mitm(pid, tid);
+            }
+            return {0, should_mitm};
+        }
+        std::tuple<Result> associate_pid_tid(u64 pid, u64 tid) {
+            MitMQueryUtils::associate_pid_to_tid(pid, tid);
+            return {0x0};
+        }
+};

stratosphere/fs_mitm/source/mitm_server.hpp

@@ -2,6 +2,7 @@
 #include <switch.h>
 #include <stratosphere.hpp>
 
+#include "mitm_query_service.hpp"
 #include "sm_mitm.h"
 #include "mitm_session.hpp"
 
@@ -17,8 +18,9 @@ class MitMServer final : public IServer<T> {
         char mitm_name[9];
     
     public:
-        MitMServer(const char *service_name, unsigned int max_s, bool s_d = false) : IServer<T>(service_name, max_s, s_d) {
+        MitMServer(ISession<MitMQueryService<T>> **out_query_session, const char *service_name, unsigned int max_s, bool s_d = false) : IServer<T>(service_name, max_s, s_d) {
             Handle tmp_hnd;
+            Handle out_query_h;
             Result rc;
             
             if (R_SUCCEEDED((rc = smGetServiceOriginal(&tmp_hnd, smEncodeName(service_name))))) {
@@ -28,9 +30,10 @@ class MitMServer final : public IServer<T> {
             }
             strncpy(mitm_name, service_name, 8);
             mitm_name[8] = '\x00';
-            if (R_FAILED((rc = smMitMInstall(&this->port_handle, mitm_name)))) {
+            if (R_FAILED((rc = smMitMInstall(&this->port_handle, &out_query_h, mitm_name)))) {
                 fatalSimple(rc);           
             }
+            *out_query_session = new ServiceSession<MitMQueryService<T>>(NULL, out_query_h, 0);
         }
         
         virtual ~MitMServer() {

stratosphere/fs_mitm/source/mitm_service.cpp

@@ -1,16 +0,0 @@
-#include <switch.h>
-#include "mitm_service.hpp"
-
-#include "debug.hpp"
-
-Result GenericMitMService::dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) {
-    return 0xF601;
-}
-
-void GenericMitMService::postprocess(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) {
-}
-
-Result GenericMitMService::handle_deferred() {
-    /* This service is never deferrable. */
-    return 0;
-}

stratosphere/fs_mitm/source/mitm_service.hpp

@@ -1,27 +0,0 @@
-#pragma once
-#include <switch.h>
-#include <stratosphere/iserviceobject.hpp>
-#include "imitmserviceobject.hpp"
-#include "fs_istorage.hpp"
-
-
-class GenericMitMService : public IMitMServiceObject {      
-    public:
-        GenericMitMService(Service *s) : IMitMServiceObject(s) {
-            /* ... */
-        }
-        
-        GenericMitMService *clone() override {
-            auto new_srv = new GenericMitMService((Service *)&this->forward_service);
-            this->clone_to(new_srv);
-            return new_srv;
-        }
-        
-        void clone_to(void *o) override {
-            /* ... */
-        }
-        
-        virtual Result dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size);
-        virtual void postprocess(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size);
-        virtual Result handle_deferred();
-};

stratosphere/fs_mitm/source/mitm_session.hpp

@@ -3,11 +3,12 @@
 #include <stratosphere.hpp>
 #include "imitmserviceobject.hpp"
 
+#include "mitm_query_service.hpp"
 #include "mitm_server.hpp"
+#include "fsmitm_worker.hpp"
 
 #include "debug.hpp"
 
-
 template <typename T>
 class MitMServer;
 
@@ -19,9 +20,10 @@ class MitMSession final : public ISession<T> {
     Service forward_service;
     IpcParsedCommand cur_out_r;
     u32 mitm_domain_id;
+    bool got_first_message;
         
     public:
-        MitMSession<T>(MitMServer<T> *s, Handle s_h, Handle c_h, const char *srv) : ISession<T>(s, s_h, c_h, NULL, 0), mitm_domain_id(0) {
+        MitMSession<T>(MitMServer<T> *s, Handle s_h, Handle c_h, const char *srv) : ISession<T>(s, s_h, c_h, NULL, 0), mitm_domain_id(0), got_first_message(false) {
             this->server = s;
             this->server_handle = s_h;
             this->client_handle = c_h;
@@ -31,16 +33,18 @@ class MitMSession final : public ISession<T> {
             if (R_FAILED(ipcQueryPointerBufferSize(forward_service.handle, &this->pointer_buffer_size))) {
                 /* TODO: Panic. */
             }
-            this->service_object = new T(&forward_service);
+            this->service_object = std::make_shared<T>(&forward_service);
             this->pointer_buffer = new char[this->pointer_buffer_size];
         }
-        MitMSession<T>(MitMServer<T> *s, Handle s_h, Handle c_h, Handle f_h, size_t pbs) : ISession<T>(s, s_h, c_h, NULL, 0), mitm_domain_id(0) {
+        MitMSession<T>(MitMServer<T> *s, Handle s_h, Handle c_h, Handle f_h) : ISession<T>(s, s_h, c_h, NULL, 0), mitm_domain_id(0), got_first_message(true) {
             this->server = s;
             this->server_handle = s_h;
             this->client_handle = c_h;
-            this->pointer_buffer_size = pbs;
+            serviceCreate(&this->forward_service, f_h);
-            this->forward_service = {.handle = f_h};
+            if (R_FAILED(ipcQueryPointerBufferSize(forward_service.handle, &this->pointer_buffer_size))) {
-            this->service_object = new T(&forward_service);
+                /* TODO: Panic. */
+            }
+            this->service_object = std::make_shared<T>(&forward_service);
             this->pointer_buffer = new char[this->pointer_buffer_size];
         }
         
@@ -59,12 +63,19 @@ class MitMSession final : public ISession<T> {
             Log(armGetTls(), 0x100);
                         
             u32 *cmdbuf = (u32 *)armGetTls();
+            if (r.CommandType == IpcCommandType_Close) {
+                Reboot();
+            }
+                        
             if (r.CommandType == IpcCommandType_Request || r.CommandType == IpcCommandType_RequestWithContext) {
-                IServiceObject *obj;
+                std::shared_ptr<IServiceObject> obj;
                 if (r.IsDomainMessage) {
-                    obj = this->get_domain_object(r.ThisObjectId);
+                    obj = this->domain->get_domain_object(r.ThisObjectId);
-                    if (obj && r.MessageType == DomainMessageType_Close) {
+                    if (obj != nullptr && r.MessageType == DomainMessageType_Close) {
-                        this->delete_object(r.ThisObjectId);
+                        if (r.ThisObjectId == this->mitm_domain_id) {
+                            Reboot();
+                        }
+                        this->domain->delete_object(r.ThisObjectId);
                         struct {
                             u64 magic;
                             u64 result;
@@ -75,13 +86,12 @@ class MitMSession final : public ISession<T> {
                         o_resp->magic = SFCO_MAGIC;
                         o_resp->result = 0x0;
                         Log(armGetTls(), 0x100);
-                        Reboot();
                         return o_resp->result;
                     }
                 } else {
                     obj = this->service_object;
                 }
-                if (obj) {
+                if (obj != nullptr) {
                     retval = obj->dispatch(r, c, cmd_id, (u8 *)this->pointer_buffer, this->pointer_buffer_size);
                     if (R_SUCCEEDED(retval)) {
                         if (r.IsDomainMessage) { 
@@ -103,7 +113,7 @@ class MitMSession final : public ISession<T> {
                         u64 result;
                     } *resp = (decltype(resp))cur_out_r.Raw;
                     retval = resp->result;
-                    if (false && (cmd_id == IpcCtrl_Cmd_CloneCurrentObject || cmd_id == IpcCtrl_Cmd_CloneCurrentObjectEx)) {
+                    if ((cmd_id == IpcCtrl_Cmd_CloneCurrentObject || cmd_id == IpcCtrl_Cmd_CloneCurrentObjectEx)) {
                         if (R_SUCCEEDED(retval)) {
                             Handle s_h;
                             Handle c_h;
@@ -112,22 +122,18 @@ class MitMSession final : public ISession<T> {
                                 fatalSimple(rc);
                             }
                             
-                            MitMSession<T> *new_sess = new MitMSession<T>((MitMServer<T> *)this->server, s_h, c_h, cur_out_r.Handles[0], this->pointer_buffer_size);
+                            if (cur_out_r.NumHandles != 1) {
-                            this->get_service_object()->clone_to(new_sess->get_service_object());
+                                Reboot();
+                            }
+                            
+                            MitMSession<T> *new_sess = new MitMSession<T>((MitMServer<T> *)this->server, s_h, c_h, cur_out_r.Handles[0]);
+                            new_sess->service_object = this->service_object;
                             if (this->is_domain) {
                                 new_sess->is_domain = true;
+                                new_sess->domain = this->domain;
                                 new_sess->mitm_domain_id = this->mitm_domain_id;
                                 new_sess->forward_service.type = this->forward_service.type;
                                 new_sess->forward_service.object_id = this->forward_service.object_id;
-                                new_sess->set_object(new_sess->get_service_object(), new_sess->mitm_domain_id);
-                                for (unsigned int i = 0; i < DOMAIN_ID_MAX; i++) {
-                                    if (i != new_sess->mitm_domain_id) {
-                                        IServiceObject *obj = this->get_domain_object(i);
-                                        if (obj) {
-                                            new_sess->set_object(obj->clone(), i);
-                                        }
-                                    }
-                                }
                             }
                             this->get_manager()->add_waitable(new_sess);
                             ipcSendHandleMove(&c, c_h);
@@ -180,6 +186,10 @@ class MitMSession final : public ISession<T> {
                 // Reboot();
             }
             
+            if (retval == 0xA08) {
+                Reboot();
+            }
+            
             return retval;
         }
         
@@ -190,15 +200,15 @@ class MitMSession final : public ISession<T> {
                 this->service_object->postprocess(cur_out_r, c, cmd_id, (u8 *)this->pointer_buffer, this->pointer_buffer_size);
             } else if (r.CommandType == IpcCommandType_Control || r.CommandType == IpcCommandType_ControlWithContext) {
                 if (cmd_id == IpcCtrl_Cmd_ConvertCurrentObjectToDomain) {
-                    return;
                     this->is_domain = true;
+                    this->domain = std::make_shared<DomainOwner>();
                     struct {
                         u64 magic;
                         u64 result;
                         u32 domain_id;
                     } *resp = (decltype(resp))cur_out_r.Raw;
                     Result rc;
-                    if (R_FAILED((rc = this->set_object(this->service_object, resp->domain_id)))) {
+                    if (R_FAILED((rc = this->domain->set_object(this->service_object, resp->domain_id)))) {
                         fatalSimple(rc);
                     }
                     this->mitm_domain_id = resp->domain_id;

stratosphere/fs_mitm/source/sm_mitm.c

@@ -101,7 +101,7 @@ Result smMitMGetService(Service* service_out, const char *name_str)
 }
 
 
-Result smMitMInstall(Handle *handle_out, const char *name) {
+Result smMitMInstall(Handle *handle_out, Handle *query_out, const char *name) {
     IpcCommand c;
     ipcInitialize(&c);
 
@@ -132,6 +132,7 @@ Result smMitMInstall(Handle *handle_out, const char *name) {
 
         if (R_SUCCEEDED(rc)) {
             *handle_out = r.Handles[0];
+            *query_out = r.Handles[1];
         }
     }
 

stratosphere/fs_mitm/source/sm_mitm.h

@@ -14,7 +14,7 @@ extern "C" {
 Result smMitMInitialize(void);
 void smMitMExit(void);
 Result smMitMGetService(Service* service_out, const char *name);
-Result smMitMInstall(Handle *handle_out, const char *name);
+Result smMitMInstall(Handle *handle_out, Handle *query_out, const char *name);
 Result smMitMUninstall(const char *name);
 
 #ifdef __cplusplus

stratosphere/libstratosphere/include/stratosphere.hpp

@@ -16,3 +16,4 @@
 #include "stratosphere/hossynch.hpp"
 
 #include "stratosphere/waitablemanager.hpp"
+#include "stratosphere/multithreadedwaitablemanager.hpp"

stratosphere/libstratosphere/include/stratosphere/domainowner.hpp

@@ -1,37 +1,36 @@
 #pragma once
 #include <switch.h>
+#include <memory>
 #include <type_traits>
 
 #include "iserviceobject.hpp"
 
-#define DOMAIN_ID_MAX 0x200
+#define DOMAIN_ID_MAX 0x1000
 
 class IServiceObject;
 
 class DomainOwner {
     private:
-        IServiceObject *domain_objects[DOMAIN_ID_MAX];
+        std::shared_ptr<IServiceObject> domain_objects[DOMAIN_ID_MAX];
     public:
         DomainOwner() {
             for (unsigned int i = 0; i < DOMAIN_ID_MAX; i++) {
-                domain_objects[i] = NULL;
+                domain_objects[i].reset();
             }
         }
         
         virtual ~DomainOwner() {
-            for (unsigned int i = 0; i < DOMAIN_ID_MAX; i++) {
+            /* Shared ptrs should auto delete here. */
-                this->delete_object(i);
-            }
         }
         
-        IServiceObject *get_domain_object(unsigned int i) {
+        std::shared_ptr<IServiceObject> get_domain_object(unsigned int i) {
             if (i < DOMAIN_ID_MAX) {
                 return domain_objects[i];
             }
-            return NULL;
+            return nullptr;
         }
         
-        Result reserve_object(IServiceObject *object, unsigned int *out_i) {
+        Result reserve_object(std::shared_ptr<IServiceObject> object, unsigned int *out_i) {
             for (unsigned int i = 4; i < DOMAIN_ID_MAX; i++) {
                 if (domain_objects[i] == NULL) {
                     domain_objects[i] = object;
@@ -43,7 +42,7 @@ class DomainOwner {
             return 0x1900B;
         }
         
-        Result set_object(IServiceObject *object, unsigned int i) {
+        Result set_object(std::shared_ptr<IServiceObject> object, unsigned int i) {
             if (domain_objects[i] == NULL) {
                 domain_objects[i] = object;
                 object->set_owner(this);
@@ -52,7 +51,7 @@ class DomainOwner {
             return 0x1900B;
         }
         
-        unsigned int get_object_id(IServiceObject *object) {
+        unsigned int get_object_id(std::shared_ptr<IServiceObject> object) {
             for (unsigned int i = 0; i < DOMAIN_ID_MAX; i++) {
                 if (domain_objects[i] == object) {
                     return i;
@@ -63,16 +62,14 @@ class DomainOwner {
         
         void delete_object(unsigned int i) {
             if (domain_objects[i]) {
-                delete domain_objects[i];
+                domain_objects[i].reset();
-                domain_objects[i] = NULL;
             }
         }
         
-        void delete_object(IServiceObject *object) {
+        void delete_object(std::shared_ptr<IServiceObject> object) {
             for (unsigned int i = 0; i < DOMAIN_ID_MAX; i++) {
                 if (domain_objects[i] == object) {
-                    delete domain_objects[i];
+                    domain_objects[i].reset();
-                    domain_objects[i] = NULL;
                     break;
                 }
             }

stratosphere/libstratosphere/include/stratosphere/ievent.hpp

@@ -4,18 +4,20 @@
 
 #include "iwaitable.hpp"
 
-typedef Result (*EventCallback)(Handle *handles, size_t num_handles, u64 timeout);
+typedef Result (*EventCallback)(void *arg, Handle *handles, size_t num_handles, u64 timeout);
 
 class IEvent : public IWaitable {
     protected:
         std::vector<Handle> handles;
         EventCallback callback;
+        void *arg;
     
     public:
-        IEvent(Handle wait_h, EventCallback callback) {
+        IEvent(Handle wait_h, void *a, EventCallback callback) {
             if (wait_h) {
                 this->handles.push_back(wait_h);
             }
+            this->arg = a;
             this->callback = callback;
         }
         
@@ -41,7 +43,7 @@ class IEvent : public IWaitable {
         }
         
         virtual Result handle_signaled(u64 timeout) {
-            return this->callback(this->handles.data(), this->handles.size(), timeout);
+            return this->callback(this->arg, this->handles.data(), this->handles.size(), timeout);
         }
         
         static Result PanicCallback(Handle *handles, size_t num_handles, u64 timeout) {

stratosphere/libstratosphere/include/stratosphere/ipcsession.hpp

@@ -17,13 +17,13 @@ class IPCSession final : public ISession<T> {
             if (R_FAILED((rc = svcCreateSession(&this->server_handle, &this->client_handle, 0, 0)))) {
                 fatalSimple(rc);
             }
-            this->service_object = new T();
+            this->service_object = std::make_shared<T>();
             this->pointer_buffer_size = pbs;
             this->pointer_buffer = new char[this->pointer_buffer_size];
             this->is_domain = false;
         }
         
-        IPCSession<T>(T *so, size_t pbs = 0x400) : ISession<T>(NULL, 0, 0, so, 0) {
+        IPCSession<T>(std::shared_ptr<T> so, size_t pbs = 0x400) : ISession<T>(NULL, 0, 0, so, 0) {
             Result rc;
             if (R_FAILED((rc = svcCreateSession(&this->server_handle, &this->client_handle, 0, 0)))) {
                 fatalSimple(rc);

stratosphere/libstratosphere/include/stratosphere/iserviceobject.hpp

@@ -18,7 +18,6 @@ class IServiceObject {
         DomainOwner *get_owner() { return this->owner; }
         void set_owner(DomainOwner *owner) { this->owner = owner; }
         virtual Result dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) = 0;
-    protected:
         virtual Result handle_deferred() = 0;
 };
 

stratosphere/libstratosphere/include/stratosphere/isession.hpp

@@ -27,10 +27,10 @@ class IServer;
 class IServiceObject;
 
 template <typename T>
-class ISession : public IWaitable, public DomainOwner {
+class ISession : public IWaitable {
     static_assert(std::is_base_of<IServiceObject, T>::value, "Service Objects must derive from IServiceObject");
     protected:
-        T *service_object;
+        std::shared_ptr<T> service_object;
         IServer<T> *server;
         Handle server_handle;
         Handle client_handle;
@@ -38,35 +38,35 @@ class ISession : public IWaitable, public DomainOwner {
         size_t pointer_buffer_size;
         
         bool is_domain;
+        std::shared_ptr<DomainOwner> domain;
         
         
-        IServiceObject *active_object;
+        std::shared_ptr<IServiceObject> active_object;
     
     static_assert(sizeof(pointer_buffer) <= POINTER_BUFFER_SIZE_MAX, "Incorrect Size for PointerBuffer!");
     
     public:
         ISession<T>(IServer<T> *s, Handle s_h, Handle c_h, size_t pbs = 0x400) : server(s), server_handle(s_h), client_handle(c_h), pointer_buffer_size(pbs) {
-            this->service_object = new T();
+            this->service_object = std::make_shared<T>();
             if (this->pointer_buffer_size) {
                 this->pointer_buffer = new char[this->pointer_buffer_size];
             }
             this->is_domain = false;
-            this->active_object = NULL;
+            this->domain.reset();
+            this->active_object.reset();
         }
         
-        ISession<T>(IServer<T> *s, Handle s_h, Handle c_h, T *so, size_t pbs = 0x400) : service_object(so), server(s), server_handle(s_h), client_handle(c_h), pointer_buffer_size(pbs) {
+        ISession<T>(IServer<T> *s, Handle s_h, Handle c_h, std::shared_ptr<T> so, size_t pbs = 0x400) : service_object(so), server(s), server_handle(s_h), client_handle(c_h), pointer_buffer_size(pbs) {
             if (this->pointer_buffer_size) {
                 this->pointer_buffer = new char[this->pointer_buffer_size];
             }
             this->is_domain = false;
-            this->active_object = NULL;
+            this->domain.reset();
+            this->active_object.reset();
         }
 
         ~ISession() override {
             delete this->pointer_buffer;
-            if (this->service_object && !this->is_domain) {
-                //delete this->service_object;
-            }
             if (server_handle) {
                 svcCloseHandle(server_handle);
             }
@@ -86,12 +86,12 @@ class ISession : public IWaitable, public DomainOwner {
             }
         }
         
-        T *get_service_object() { return this->service_object; }
+        std::shared_ptr<T> get_service_object() { return this->service_object; }
         Handle get_server_handle() { return this->server_handle; }
         Handle get_client_handle() { return this->client_handle; }
         
         
-        DomainOwner *get_owner() { return is_domain ? this : NULL; }
+        DomainOwner *get_owner() { return this->is_domain ? this->domain.get() : NULL; }
         
         /* IWaitable */        
         Handle get_handle() override {
@@ -125,7 +125,7 @@ class ISession : public IWaitable, public DomainOwner {
            
             
             if (r.IsDomainMessage && r.MessageType == DomainMessageType_Close) {
-                this->delete_object(this->active_object);
+                this->domain->delete_object(this->active_object);
                 this->active_object = NULL;
                 struct {
                     u64 magic;
@@ -188,7 +188,7 @@ class ISession : public IWaitable, public DomainOwner {
             ipcAddRecvStatic(&c_for_reply, this->pointer_buffer, this->pointer_buffer_size, 0);
             ipcPrepareHeader(&c_for_reply, 0);
             
-            if (R_SUCCEEDED(rc = svcReplyAndReceive(&handle_index, &this->server_handle, 1, 0, timeout))) {
+            if (R_SUCCEEDED(rc = svcReplyAndReceive(&handle_index, &this->server_handle, 1, 0, U64_MAX))) {
                 if (handle_index != 0) {
                     /* TODO: Panic? */
                 }                                
@@ -203,7 +203,7 @@ class ISession : public IWaitable, public DomainOwner {
                         if (!r.IsDomainMessage || r.ThisObjectId >= DOMAIN_ID_MAX) {
                             retval = 0xF601;
                         } else {
-                            this->active_object = this->get_domain_object(r.ThisObjectId);
+                            this->active_object = this->domain->get_domain_object(r.ThisObjectId);
                         }
                     } else {
                         this->active_object = this->service_object;
@@ -218,19 +218,22 @@ class ISession : public IWaitable, public DomainOwner {
                 
                 if (retval == RESULT_DEFER_SESSION) {
                     /* Session defer. */
-                    this->active_object = NULL;
+                    this->active_object.reset();
                     this->set_deferred(true);
                     rc = retval;
                 } else if (retval == 0xF601) {
                     /* Session close. */
-                    this->active_object = NULL;
+                    this->active_object.reset();
                     rc = retval;
                 } else {
                     if (R_SUCCEEDED(retval)) {
                         this->postprocess(r, cmd_id);
                     }
-                    this->active_object = NULL;
+                    this->active_object.reset();
                     rc = svcReplyAndReceive(&handle_index, &this->server_handle, 0, this->server_handle, 0);
+                    if (rc == 0xEA01) {
+                        rc = 0x0;
+                    }
                     this->cleanup();
                 }
             }

stratosphere/libstratosphere/include/stratosphere/multithreadedwaitablemanager.hpp

@@ -0,0 +1,42 @@
+#pragma once
+#include <switch.h>
+#include <vector>
+
+#include "waitablemanager.hpp"
+#include "systemevent.hpp"
+
+class MultiThreadedWaitableManager : public WaitableManager {
+    protected:
+        u32 num_threads;
+        Thread *threads;
+        HosMutex get_waitable_lock;
+        SystemEvent *new_waitable_event;
+    public:
+        MultiThreadedWaitableManager(u32 n, u64 t, u32 ss = 0x8000) : WaitableManager(t), num_threads(n-1) { 
+            u32 prio;
+            u32 cpuid = svcGetCurrentProcessorNumber();
+            Result rc;
+            threads = new Thread[num_threads];
+            if (R_FAILED((rc = svcGetThreadPriority(&prio, CUR_THREAD_HANDLE)))) {
+                fatalSimple(rc);
+            }
+            for (unsigned int i = 0; i < num_threads; i++) {
+                threads[i] = {0};
+                threadCreate(&threads[i], &MultiThreadedWaitableManager::thread_func, this, ss, prio, cpuid);
+            }
+            new_waitable_event = new SystemEvent(this, &MultiThreadedWaitableManager::add_waitable_callback);
+            this->waitables.push_back(new_waitable_event);
+        }
+        ~MultiThreadedWaitableManager() override {
+            /* TODO: Exit the threads? */
+        }
+        
+        IWaitable *get_waitable();
+        
+        void add_waitable(IWaitable *waitable) override;
+        void process() override;
+        void process_until_timeout() override;
+        
+        static Result add_waitable_callback(void *this_ptr, Handle *handles, size_t num_handles, u64 timeout);
+        static void thread_func(void *this_ptr);
+};

stratosphere/libstratosphere/include/stratosphere/systemevent.hpp

@@ -9,7 +9,7 @@
 
 class SystemEvent final : public IEvent {
     public:
-        SystemEvent(EventCallback callback) : IEvent(0, callback) {
+        SystemEvent(void *a, EventCallback callback) : IEvent(0, a, callback) {
             Handle wait_h;
             Handle sig_h;
             if (R_FAILED(svcCreateEvent(&sig_h, &wait_h))) {

stratosphere/libstratosphere/include/stratosphere/waitablemanager.hpp

@@ -9,16 +9,17 @@
 class IWaitable;
 
 class WaitableManager : public WaitableManagerBase {
-    std::vector<IWaitable *> to_add_waitables;
+    protected:
-    std::vector<IWaitable *> waitables;
+        std::vector<IWaitable *> to_add_waitables;
-    u64 timeout;
+        std::vector<IWaitable *> waitables;
-    HosMutex lock;
+        u64 timeout;
-    std::atomic_bool has_new_items;
+        HosMutex lock;
+        std::atomic_bool has_new_items;
     private:
         void process_internal(bool break_on_timeout);
     public:
         WaitableManager(u64 t) : waitables(0), timeout(t), has_new_items(false) { }
-        ~WaitableManager() {
+        ~WaitableManager() override {
             /* This should call the destructor for every waitable. */
             for (auto & waitable : waitables) {
                 delete waitable;
@@ -26,7 +27,7 @@ class WaitableManager : public WaitableManagerBase {
             waitables.clear();
         }
         
-        void add_waitable(IWaitable *waitable);
+        virtual void add_waitable(IWaitable *waitable);
-        void process();
+        virtual void process();
-        void process_until_timeout();
+        virtual void process_until_timeout();
 };

stratosphere/libstratosphere/include/stratosphere/waitablemanagerbase.hpp

@@ -7,6 +7,7 @@ class WaitableManagerBase {
     std::atomic<u64> cur_priority;
     public:
         WaitableManagerBase() : cur_priority(0) { }
+        virtual ~WaitableManagerBase() { }
         
         u64 get_priority() {            
             return std::atomic_fetch_add(&cur_priority, (u64)1);

stratosphere/libstratosphere/source/multithreadedwaitablemanager.cpp

@@ -0,0 +1,110 @@
+#include <switch.h>
+
+#include <algorithm>
+
+#include <stratosphere/multithreadedwaitablemanager.hpp>
+
+void MultiThreadedWaitableManager::process() {
+    Result rc;
+    for (unsigned int i = 0; i < num_threads; i++) {
+        if (R_FAILED((rc = threadStart(&threads[i])))) {
+            fatalSimple(rc);
+        }
+    }
+    MultiThreadedWaitableManager::thread_func(this);
+}
+
+void MultiThreadedWaitableManager::process_until_timeout() {
+    /* TODO: Panic. */
+}
+
+void MultiThreadedWaitableManager::add_waitable(IWaitable *waitable) {
+    this->lock.Lock();
+    this->to_add_waitables.push_back(waitable);
+    waitable->set_manager(this);
+    this->new_waitable_event->signal_event();
+    this->lock.Unlock();
+}
+
+
+IWaitable *MultiThreadedWaitableManager::get_waitable() {
+    std::vector<Handle> handles;
+    
+    int handle_index = 0;
+    Result rc;
+    this->get_waitable_lock.Lock();
+    while (1) {
+        /* Sort waitables by priority. */
+        std::sort(this->waitables.begin(), this->waitables.end(), IWaitable::compare);
+        
+        /* Copy out handles. */
+        handles.resize(this->waitables.size());
+        std::transform(this->waitables.begin(), this->waitables.end(), handles.begin(), [](IWaitable *w) { return w->get_handle(); });
+        
+        rc = svcWaitSynchronization(&handle_index, handles.data(), this->waitables.size(), this->timeout);
+        IWaitable *w = this->waitables[handle_index];
+        if (R_SUCCEEDED(rc)) {
+            for (int i = 0; i < handle_index; i++) {
+                this->waitables[i]->update_priority();
+            }
+            this->waitables.erase(this->waitables.begin() + handle_index);
+        } else if (rc == 0xEA01) {
+            /* Timeout. */
+            for (auto & waitable : this->waitables) {
+                waitable->update_priority();
+            }
+        } else if (rc != 0xF601) {
+            /* TODO: Panic. When can this happen? */
+        } else {
+            for (int i = 0; i < handle_index; i++) {
+                this->waitables[i]->update_priority();
+            }
+            this->waitables.erase(this->waitables.begin() + handle_index);
+            delete w;
+        }
+        
+        /* Do deferred callback for each waitable. */
+        for (auto & waitable : this->waitables) {
+            if (waitable->get_deferred()) {
+                waitable->handle_deferred();
+            }
+        }
+        
+        /* Return waitable. */
+        if (R_SUCCEEDED(rc)) {
+            if (w == this->new_waitable_event) {
+                w->handle_signaled(0);
+                this->waitables.push_back(w);
+            } else {
+                this->get_waitable_lock.Unlock();
+                return w;
+            }
+        }
+    }
+}
+
+Result MultiThreadedWaitableManager::add_waitable_callback(void *arg, Handle *handles, size_t num_handles, u64 timeout) {
+    MultiThreadedWaitableManager *this_ptr = (MultiThreadedWaitableManager *)arg;
+    svcClearEvent(handles[0]);
+    this_ptr->lock.Lock();
+    this_ptr->waitables.insert(this_ptr->waitables.end(), this_ptr->to_add_waitables.begin(), this_ptr->to_add_waitables.end());
+    this_ptr->to_add_waitables.clear();
+    this_ptr->lock.Unlock();
+    return 0;
+}
+
+void MultiThreadedWaitableManager::thread_func(void *t) {
+    MultiThreadedWaitableManager *this_ptr = (MultiThreadedWaitableManager *)t;
+    while (1) {
+        IWaitable *w = this_ptr->get_waitable();
+        if (w) {
+            Result rc = w->handle_signaled(0);
+            if (rc == 0xF601) {
+                /* Close! */
+                delete w;
+            } else {
+                this_ptr->add_waitable(w);
+            }
+        }
+    }
+}

stratosphere/loader/source/ldr_debug_monitor.hpp

@@ -18,6 +18,10 @@ class DebugMonitorService final : public IServiceObject {
             return 0;
         }
         
+        DebugMonitorService *clone() override {
+            return new DebugMonitorService();
+        }
+        
     private:
         /* Actual commands. */
         std::tuple<Result> add_title_to_launch_queue(u64 tid, InPointer<char> args);

stratosphere/loader/source/ldr_process_creation.cpp

@@ -186,6 +186,9 @@ Result ProcessCreation::CreateProcess(Handle *out_process_h, u64 index, char *nc
         }
     }
     
+    /* Send the pid/tid pair to anyone interested in man-in-the-middle-attacking it. */
+    Registration::AssociatePidTidForMitM(index);
+    
     rc = 0;  
 CREATE_PROCESS_END:
     if (R_SUCCEEDED(rc)) {

stratosphere/loader/source/ldr_process_manager.hpp

@@ -35,6 +35,10 @@ class ProcessManagerService final : public IServiceObject {
             return 0;
         }
         
+        ProcessManagerService *clone() override {
+            return new ProcessManagerService();
+        }
+        
     private:
         /* Actual commands. */
         std::tuple<Result, MovedHandle> create_process(u64 flags, u64 index, CopiedHandle reslimit_h);

stratosphere/loader/source/ldr_registration.cpp

@@ -266,3 +266,65 @@ Result Registration::GetNsoInfosForProcessId(Registration::NsoInfo *out, u32 max
 
     return 0;
 }
+
+void Registration::AssociatePidTidForMitM(u64 index) {
+    Registration::Process *target_process = GetProcess(index);
+    if (target_process == NULL) {
+        return;
+    }
+    
+    Handle sm_hnd;
+    Result rc = svcConnectToNamedPort(&sm_hnd, "sm:");
+    if (R_SUCCEEDED(rc)) {
+        /* Initialize. */
+        {
+            IpcCommand c;
+            ipcInitialize(&c);
+            ipcSendPid(&c);
+
+            struct {
+                u64 magic;
+                u64 cmd_id;
+                u64 zero;
+                u64 reserved[2];
+            } *raw = (decltype(raw))ipcPrepareHeader(&c, sizeof(*raw));
+
+            raw->magic = SFCI_MAGIC;
+            raw->cmd_id = 0;
+            raw->zero = 0;
+
+            rc = ipcDispatch(sm_hnd);
+
+            if (R_SUCCEEDED(rc)) {
+                IpcParsedCommand r;
+                ipcParse(&r);
+
+                struct {
+                    u64 magic;
+                    u64 result;
+                } *resp = (decltype(resp))r.Raw;
+
+                rc = resp->result;
+            }
+        }
+        /* Associate. */
+        if (R_SUCCEEDED(rc)) {
+            IpcCommand c;
+            ipcInitialize(&c);
+            struct {
+                u64 magic;
+                u64 cmd_id;
+                u64 process_id;
+                u64 title_id;
+            } *raw = (decltype(raw))ipcPrepareHeader(&c, sizeof(*raw));
+            
+            raw->magic = SFCI_MAGIC;
+            raw->cmd_id = 65002;
+            raw->process_id = target_process->process_id;
+            raw->title_id = target_process->tid_sid.title_id;
+            
+            ipcDispatch(sm_hnd);
+        }
+        svcCloseHandle(sm_hnd);
+    }
+}

stratosphere/loader/source/ldr_registration.hpp

@@ -76,4 +76,7 @@ class Registration {
         static void AddNroToProcess(u64 index, MappedCodeMemory *nro, MappedCodeMemory *bss, u32 text_size, u32 ro_size, u32 rw_size, u8 *build_id);
         static Result RemoveNroInfo(u64 index, Handle process_h, u64 base_address);
         static Result GetNsoInfosForProcessId(NsoInfo *out, u32 max_out, u64 process_id, u32 *num_written);
+        
+        /* Atmosphere MitM Extension. */
+        static void AssociatePidTidForMitM(u64 index);
 };

stratosphere/loader/source/ldr_ro_service.hpp

@@ -30,6 +30,10 @@ class RelocatableObjectsService final : public IServiceObject {
             return 0;
         }
         
+        RelocatableObjectsService *clone() override {
+            return new RelocatableObjectsService(*this);
+        }
+        
     private:
         /* Actual commands. */
         std::tuple<Result, u64> load_nro(PidDescriptor pid_desc, u64 nro_address, u64 nro_size, u64 bss_address, u64 bss_size);

stratosphere/loader/source/ldr_shell.hpp

@@ -15,6 +15,10 @@ class ShellService final : public IServiceObject {
             return 0;
         }
         
+        ShellService *clone() override {
+            return new ShellService();
+        }
+        
     private:
         /* Actual commands. */
         std::tuple<Result> add_title_to_launch_queue(u64 tid, InPointer<char> args);

stratosphere/pm/source/pm_boot_mode.cpp

@@ -1,4 +1,5 @@
 #include <switch.h>
+#include <stratosphere.hpp>
 #include "pm_boot_mode.hpp"
 
 static bool g_is_maintenance_boot = false;

stratosphere/pm/source/pm_boot_mode.hpp

@@ -7,7 +7,7 @@ enum BootModeCmd {
     BootMode_Cmd_SetMaintenanceBoot = 1
 };
 
-class BootModeService final : IServiceObject {
+class BootModeService final : public IServiceObject {
     public:
         Result dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) override;
         Result handle_deferred() override;

stratosphere/pm/source/pm_debug_monitor.cpp

@@ -1,4 +1,5 @@
 #include <switch.h>
+#include <stratosphere.hpp>
 #include "pm_registration.hpp"
 #include "pm_debug_monitor.hpp"
 

stratosphere/pm/source/pm_debug_monitor.hpp

@@ -27,7 +27,7 @@ enum DmntCmd_5X {
     Dmnt_Cmd_5X_AtmosphereGetProcessHandle = 65000
 };
 
-class DebugMonitorService final : IServiceObject {
+class DebugMonitorService final : public IServiceObject {
     public:
         Result dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) override;
         Result handle_deferred() override;

stratosphere/pm/source/pm_info.hpp

@@ -1,12 +1,13 @@
 #pragma once
 #include <switch.h>
+#include <stratosphere.hpp>
 #include <stratosphere/iserviceobject.hpp>
 
 enum InformationCmd {
     Information_Cmd_GetTitleId = 0,
 };
 
-class InformationService final : IServiceObject {
+class InformationService final : public IServiceObject {
     public:
         Result dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) override;
         Result handle_deferred() override;

stratosphere/pm/source/pm_process_wait.hpp

@@ -19,18 +19,6 @@ class ProcessWaiter final : public IWaitable {
         }
         
         /* IWaitable */        
-        unsigned int get_num_waitables() override {
-            return 1;
-        }
-        
-        void get_waitables(IWaitable **dst) override {
-            dst[0] = this;
-        }
-        
-        void delete_child(IWaitable *child) override {
-            /* TODO: Panic, because we can never have any children. */
-        }
-        
         Handle get_handle() override {
             return this->process.handle;
         }

stratosphere/pm/source/pm_shell.cpp

@@ -1,4 +1,5 @@
 #include <switch.h>
+#include <stratosphere.hpp>
 #include "pm_registration.hpp"
 #include "pm_resource_limits.hpp"
 #include "pm_shell.hpp"

stratosphere/pm/source/pm_shell.hpp

@@ -30,7 +30,7 @@ enum ShellCmd_5X {
     Shell_Cmd_5X_BoostSystemMemoryResourceLimit = 7
 };
 
-class ShellService final : IServiceObject {
+class ShellService final : public IServiceObject {
     public:
         Result dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) override;
         Result handle_deferred() override;

stratosphere/sm/source/sm_registration.cpp

@@ -188,7 +188,37 @@ Result Registration::GetServiceHandle(u64 pid, u64 service, Handle *out) {
     if (target_service->mitm_pid == 0 || target_service->mitm_pid == pid) {
         rc = svcConnectToPort(out, target_service->port_h);
     } else {
-        rc = svcConnectToPort(out, target_service->mitm_port_h);
+        IpcCommand c;
+        ipcInitialize(&c);
+        struct {
+            u64 magic;
+            u64 cmd_id;
+            u64 pid;
+        } *info = ((decltype(info))ipcPrepareHeader(&c, sizeof(*info)));
+        info->magic = SFCI_MAGIC;
+        info->cmd_id = 65000;
+        info->pid = pid;
+        rc = ipcDispatch(target_service->mitm_query_h);
+        if (R_SUCCEEDED(rc)) {
+            IpcParsedCommand r;
+            ipcParse(&r);
+            struct {
+                u64 magic;
+                u64 result;
+                u64 should_mitm;
+            } *resp = ((decltype(resp))r.Raw);
+            rc = resp->result;
+            if (R_SUCCEEDED(rc)) {
+                if (resp->should_mitm) {
+                    rc = svcConnectToPort(out, target_service->mitm_port_h);
+                } else {
+                    rc = svcConnectToPort(out, target_service->port_h);
+                }
+            }
+        }
+        if (R_FAILED(rc)) {
+            rc = svcConnectToPort(out, target_service->port_h);
+        }
     }
     if (R_FAILED(rc)) {
         if ((rc & 0x3FFFFF) == 0xE01) {
@@ -343,12 +373,13 @@ Result Registration::UnregisterServiceForPid(u64 pid, u64 service) {
     
     svcCloseHandle(target_service->port_h);
     svcCloseHandle(target_service->mitm_port_h);
+    svcCloseHandle(target_service->mitm_query_h);
     *target_service = (const Registration::Service){0};
     return 0;
 }
 
 
-Result Registration::InstallMitmForPid(u64 pid, u64 service, Handle *out) {
+Result Registration::InstallMitmForPid(u64 pid, u64 service, Handle *out, Handle *query_out) {
     if (!service) {
         return 0xC15;
     }
@@ -387,7 +418,7 @@ Result Registration::InstallMitmForPid(u64 pid, u64 service, Handle *out) {
     u64 x = 0;
     Result rc = svcCreatePort(out, &target_service->mitm_port_h, target_service->max_sessions, target_service->is_light, (char *)&x);
     
-    if (R_SUCCEEDED(rc)) {
+    if (R_SUCCEEDED(rc) && R_SUCCEEDED((rc = svcCreateSession(query_out, &target_service->mitm_query_h, 0, 0)))) {
         target_service->mitm_pid = pid;
     }
     
@@ -416,6 +447,28 @@ Result Registration::UninstallMitmForPid(u64 pid, u64 service) {
     }
     
     svcCloseHandle(target_service->mitm_port_h);
+    svcCloseHandle(target_service->mitm_query_h);
     target_service->mitm_pid = 0;
     return 0;
 }
+
+Result Registration::AssociatePidTidForMitm(u64 pid, u64 tid) {
+    for (auto &service : g_service_list) {
+        if (service.mitm_pid) {
+            IpcCommand c;
+            ipcInitialize(&c);
+            struct {
+                u64 magic;
+                u64 cmd_id;
+                u64 pid;
+                u64 tid;
+            } *info = ((decltype(info))ipcPrepareHeader(&c, sizeof(*info)));
+            info->magic = SFCI_MAGIC;
+            info->cmd_id = 65001;
+            info->pid = pid;
+            info->tid = tid;
+            ipcDispatch(service.mitm_query_h);
+        }
+    }
+    return 0x0;
+}

stratosphere/sm/source/sm_registration.hpp

@@ -26,6 +26,7 @@ class Registration {
             /* Extension. */
             u64 mitm_pid;
             Handle mitm_port_h;
+            Handle mitm_query_h;
         };
         
         /* Utilities. */
@@ -52,6 +53,7 @@ class Registration {
         static Result UnregisterServiceForPid(u64 pid, u64 service);
         
         /* Extension. */
-        static Result InstallMitmForPid(u64 pid, u64 service, Handle *out);
+        static Result InstallMitmForPid(u64 pid, u64 service, Handle *out, Handle *query_out);
         static Result UninstallMitmForPid(u64 pid, u64 service);
+        static Result AssociatePidTidForMitm(u64 pid, u64 tid);
 };

stratosphere/sm/source/sm_user_service.cpp

@@ -25,6 +25,9 @@ Result UserService::dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id,
         case User_Cmd_AtmosphereUninstallMitm:
             rc = WrapIpcCommandImpl<&UserService::uninstall_mitm>(this, r, out_c, pointer_buffer, pointer_buffer_size);
             break;
+        case User_Cmd_AtmosphereAssociatePidTidForMitm:
+            rc = WrapIpcCommandImpl<&UserService::associate_pid_tid_for_mitm>(this, r, out_c, pointer_buffer, pointer_buffer_size);
+            break;
 #endif
         default:
             break;
@@ -95,13 +98,26 @@ std::tuple<Result> UserService::unregister_service(u64 service) {
     return {rc};
 }
 
-std::tuple<Result, MovedHandle> UserService::install_mitm(u64 service) {
+std::tuple<Result, MovedHandle, MovedHandle> UserService::install_mitm(u64 service) {
     Handle service_h = 0;
+    Handle query_h = 0;
     Result rc = 0x415;
     if (this->has_initialized) {
-        rc = Registration::InstallMitmForPid(this->pid, service, &service_h);
+        rc = Registration::InstallMitmForPid(this->pid, service, &service_h, &query_h);
     }
-    return {rc, MovedHandle{service_h}};
+    return {rc, MovedHandle{service_h}, MovedHandle{query_h}};
+}
+
+std::tuple<Result> UserService::associate_pid_tid_for_mitm(u64 pid, u64 tid) {
+    Result rc = 0x415;
+    if (this->has_initialized) {
+        if (Registration::IsInitialProcess(pid)) {
+            rc = 0x1015;
+        } else {
+            rc = Registration::AssociatePidTidForMitm(pid, tid);
+        }
+    }
+    return {rc};
 }
 
 std::tuple<Result> UserService::uninstall_mitm(u64 service) {

stratosphere/sm/source/sm_user_service.hpp

@@ -9,7 +9,8 @@ enum UserServiceCmd {
     User_Cmd_UnregisterService = 3,
     
     User_Cmd_AtmosphereInstallMitm = 65000,
-    User_Cmd_AtmosphereUninstallMitm = 65001
+    User_Cmd_AtmosphereUninstallMitm = 65001,
+    User_Cmd_AtmosphereAssociatePidTidForMitm = 65002
 };
 
 class UserService final : public IServiceObject {
@@ -39,6 +40,7 @@ class UserService final : public IServiceObject {
         std::tuple<Result> unregister_service(u64 service);
         
         /* Atmosphere commands. */
-        std::tuple<Result, MovedHandle> install_mitm(u64 service);
+        std::tuple<Result, MovedHandle, MovedHandle> install_mitm(u64 service);
         std::tuple<Result> uninstall_mitm(u64 service);
+        std::tuple<Result> associate_pid_tid_for_mitm(u64 pid, u64 tid);
 };