fs.mitm: Allow fsmitm.flag specification for mitm

stratosphere/fs_mitm/source/fsmitm_main.cpp

@@ -14,6 +14,8 @@
 
 #include "mitm_query_service.hpp"
 
+#include "fsmitm_utils.hpp"
+
 extern "C" {
     extern u32 __start__;
 
@@ -89,6 +91,9 @@ void __appExit(void) {
 int main(int argc, char **argv)
 {
     Thread worker_thread = {0};
+    Thread sd_initializer_thread = {0};
+    consoleDebugInit(debugDevice_SVC);
+    
     consoleDebugInit(debugDevice_SVC);
     
     if (R_FAILED(threadCreate(&worker_thread, &FsMitMWorker::Main, NULL, 0x20000, 45, 0))) {
@@ -98,6 +103,13 @@ int main(int argc, char **argv)
         /* TODO: Panic. */
     }
     
+    if (R_FAILED(threadCreate(&sd_initializer_thread, &Utils::InitializeSdThreadFunc, NULL, 0x4000, 0x15, 0))) {
+        /* TODO: Panic. */
+    }
+    if (R_FAILED(threadStart(&sd_initializer_thread))) {
+        /* TODO: Panic. */
+    }
+    
     /* TODO: What's a good timeout value to use here? */
     auto server_manager = std::make_unique<MultiThreadedWaitableManager>(1, U64_MAX, 0x20000);
     //auto server_manager = std::make_unique<WaitableManager>(U64_MAX);

stratosphere/fs_mitm/source/fsmitm_service.hpp

@@ -22,12 +22,7 @@ class FsMitMService : public IMitMServiceObject {
         }
         
         static bool should_mitm(u64 pid, u64 tid) {
-            if (tid >= 0x0100000000010000ULL) {
-                return true;
-            }
-            bool has_romfs_content;
-            Result rc = Utils::HasSdRomfsContent(tid, &has_romfs_content);
-            return R_SUCCEEDED(rc) && has_romfs_content;
+            return tid >= 0x0100000000010000ULL || Utils::HasSdMitMFlag(tid);
         }
         
         FsMitMService *clone() override {

stratosphere/fs_mitm/source/fsmitm_utils.cpp

@@ -1,53 +1,88 @@
 #include <switch.h>
 #include <stratosphere.hpp>
 #include <atomic>
+#include <algorithm>
 
 #include "sm_mitm.h"
 #include "debug.hpp"
 #include "fsmitm_utils.hpp"
 
 static FsFileSystem g_sd_filesystem = {0};
-static bool g_has_initialized = false;
+static std::vector<u64> g_mitm_flagged_tids;
+static std::atomic_bool g_has_initialized = false;
 
-static Result EnsureInitialized() {
-    if (g_has_initialized) {
-        return 0x0;
+static bool IsHexadecimal(const char *str) {
+    while (*str) {
+        if (('0' <= *str && *str <= '9') || 
+            ('a' <= *str && *str <= 'f') ||
+            ('A' <= *str && *str <= 'F')) {
+                str++;
+            } else {
+                return false;
             }
+    }
+    return true;
+}
 
+void Utils::InitializeSdThreadFunc(void *args) {
+    /* Get required services. */
+    Handle tmp_hnd = 0;
     static const char * const required_active_services[] = {"pcv", "gpio", "pinmux", "psc:c"};
     for (unsigned int i = 0; i < sizeof(required_active_services) / sizeof(required_active_services[0]); i++) {
-        Result rc = smMitMUninstall(required_active_services[i]);
-        if (rc == 0xE15) {
-            return rc;
-        } else if (R_FAILED(rc) && rc != 0x1015) {
-            return rc;
+        if (R_FAILED(smGetServiceOriginal(&tmp_hnd, smEncodeName(required_active_services[i])))) {
+            /* TODO: Panic */
+        } else {
+            svcCloseHandle(tmp_hnd);   
         }
     }
     
-    Result rc = fsMountSdcard(&g_sd_filesystem);
-    if (R_SUCCEEDED(rc)) {
-        g_has_initialized = true;
+    /* Mount SD. */
+    while (R_FAILED(fsMountSdcard(&g_sd_filesystem))) {
+        svcSleepThread(1000ULL);
     }
-    return rc;
+    
+    /* Check for MitM flags. */
+    FsDir titles_dir;
+    if (R_SUCCEEDED(fsFsOpenDirectory(&g_sd_filesystem, "/atmosphere/titles", FS_DIROPEN_DIRECTORY, &titles_dir))) {
+        FsDirectoryEntry dir_entry;
+        FsFile f;
+        u64 read_entries;
+        while (R_SUCCEEDED((fsDirRead(&titles_dir, 0, &read_entries, 1, &dir_entry))) && read_entries == 1) {
+            if (strlen(dir_entry.name) == 0x10 && IsHexadecimal(dir_entry.name)) {
+                u64 title_id = strtoul(dir_entry.name, NULL, 16);
+                char title_path[FS_MAX_PATH] = {0};
+                strcpy(title_path, "sdmc:/atmosphere/titles/");
+                strcat(title_path, dir_entry.name);
+                strcat(title_path, "/fsmitm.flag");
+                if (R_SUCCEEDED(fsFsOpenFile(&g_sd_filesystem, title_path, FS_OPEN_READ, &f))) {
+                    g_mitm_flagged_tids.push_back(title_id);
+                    fsFileClose(&f);
+                }
+            }
+        }
+        fsDirClose(&titles_dir);
+    }
+    
+    g_has_initialized = true;
+    
+    svcExitThread();
 }
 
 bool Utils::IsSdInitialized() {
-    return R_SUCCEEDED(EnsureInitialized());
+    return g_has_initialized;
 }
 
 Result Utils::OpenSdFile(const char *fn, int flags, FsFile *out) {
-    Result rc;
-    if (R_FAILED((rc = EnsureInitialized()))) {
-        return rc;
+    if (!IsSdInitialized()) {
+        return 0xFA202;
     }
     
     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;
+    if (!IsSdInitialized()) {
+        return 0xFA202;
     }
     
     char path[FS_MAX_PATH];
@@ -60,27 +95,24 @@ Result Utils::OpenSdFileForAtmosphere(u64 title_id, const char *fn, int flags, F
 }
 
 Result Utils::OpenRomFSSdFile(u64 title_id, const char *fn, int flags, FsFile *out) {
-    Result rc;
-    if (R_FAILED((rc = EnsureInitialized()))) {
-        return rc;
+    if (!IsSdInitialized()) {
+        return 0xFA202;
     }
     
     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;
+    if (!IsSdInitialized()) {
+        return 0xFA202;
     }
     
     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;
+    if (!IsSdInitialized()) {
+        return 0xFA202;
     }
     
     char safe_path[FS_MAX_PATH];
@@ -93,9 +125,8 @@ Result Utils::OpenSdDirForAtmosphere(u64 title_id, const char *path, FsDir *out)
 }
 
 Result Utils::OpenRomFSSdDir(u64 title_id, const char *path, FsDir *out) {
-    Result rc;
-    if (R_FAILED((rc = EnsureInitialized()))) {
-        return rc;
+    if (!IsSdInitialized()) {
+        return 0xFA202;
     }
     
     return OpenRomFSDir(&g_sd_filesystem, title_id, path, out);
@@ -137,3 +168,10 @@ Result Utils::HasSdRomfsContent(u64 title_id, bool *out) {
     fsDirClose(&dir);
     return rc;
 }
+
+bool Utils::HasSdMitMFlag(u64 tid) {
+    if (IsSdInitialized()) {
+        return std::find(g_mitm_flagged_tids.begin(), g_mitm_flagged_tids.end(), tid) != g_mitm_flagged_tids.end();
+    }
+    return false;
+}

stratosphere/fs_mitm/source/fsmitm_utils.hpp

@@ -17,4 +17,8 @@ class Utils {
         static Result OpenRomFSDir(FsFileSystem *fs, u64 title_id, const char *path, FsDir *out);
         
         static Result HasSdRomfsContent(u64 title_id, bool *out);
+        
+        /* SD card Initialization + MitM detection. */
+        static void InitializeSdThreadFunc(void *args);
+        static bool HasSdMitMFlag(u64 tid);
 };