stratosphere: Fix fs.mitm bis_protect race condition.

stratosphere/fs_mitm/source/fsmitm_boot0storage.cpp

@@ -88,7 +88,7 @@ Result Boot0Storage::Write(void *_buffer, size_t size, u64 offset) {
         if (R_FAILED((rc = ProxyStorage::Write(buffer + diff, size - diff, BctEndOffset)))) {
             return rc;
         }
-        size -= diff;
+        size = diff;
     }
     
     /* Read in the current BCT region. */

stratosphere/fs_mitm/source/fsmitm_main.cpp

@@ -89,6 +89,12 @@ int main(int argc, char **argv)
     Thread hid_initializer_thread = {0};
     consoleDebugInit(debugDevice_SVC);
     
+    /* TODO: What's a good timeout value to use here? */
+    auto server_manager = new FsMitmManager(5);
+
+    /* Create fsp-srv mitm. */
+    AddMitmServerToManager<FsMitmService>(server_manager, "fsp-srv", 61);
+
     if (R_FAILED(threadCreate(&sd_initializer_thread, &Utils::InitializeSdThreadFunc, NULL, 0x4000, 0x15, 0))) {
         /* TODO: Panic. */
     }
@@ -103,12 +109,6 @@ int main(int argc, char **argv)
         /* TODO: Panic. */
     }
                 
-    /* TODO: What's a good timeout value to use here? */
-    auto server_manager = new FsMitmManager(5);
-        
-    /* Create fsp-srv mitm. */
-    AddMitmServerToManager<FsMitmService>(server_manager, "fsp-srv", 61);
-                
     /* Loop forever, servicing our services. */
     server_manager->Process();
     

stratosphere/libstratosphere

@@ -1 +1 @@
-Subproject commit 0bec72ca36084e9780a8c28abd4a0b24c03c18af
+Subproject commit dd31b3d2e2450a458ef9c0d6268f3780fe70e083

stratosphere/pm/source/pm_boot2.cpp

@@ -124,6 +124,21 @@ static void MountSdCard() {
 }
 
 void EmbeddedBoot2::Main() {
+    /* Wait until fs.mitm has installed itself. We want this to happen as early as possible. */
+    bool fs_mitm_installed = false;
+
+    Result rc = smManagerAmsInitialize();
+    if (R_FAILED(rc)) {
+        std::abort();
+    }
+    while (R_FAILED((rc = smManagerAmsHasMitm(&fs_mitm_installed, "fsp-srv"))) || !fs_mitm_installed) {
+        if (R_FAILED(rc)) {
+            std::abort();
+        }
+        svcSleepThread(1000ull);
+    }
+    smManagerAmsExit();
+
     /* psc, bus, pcv is the minimal set of required titles to get SD card. */ 
     /* bus depends on pcie, and pcv depends on settings. */
     /* Launch psc. */

stratosphere/pm/source/pm_main.cpp

@@ -102,7 +102,6 @@ void __appInit(void) {
     rc = smManagerAmsInitialize();
     if (R_SUCCEEDED(rc)) {
         smManagerAmsEndInitialDefers();
-        smManagerAmsExit();
     } else {
         fatalSimple(0xCAFE << 4 | 2);
     }

stratosphere/sm/source/sm_manager_service.cpp

@@ -31,3 +31,7 @@ void ManagerService::AtmosphereEndInitDefers() {
     Registration::EndInitDefers();
 }
 
+void ManagerService::AtmosphereHasMitm(Out<bool> out, SmServiceName service) {
+    out.SetValue(Registration::HasMitm(smEncodeName(service.name)));
+}
+

stratosphere/sm/source/sm_manager_service.hpp

@@ -17,6 +17,7 @@
 #pragma once
 #include <switch.h>
 #include <stratosphere.hpp>
+#include "sm_types.hpp"
 
 enum ManagerServiceCmd {
     Manager_Cmd_RegisterProcess = 0,
@@ -24,6 +25,7 @@ enum ManagerServiceCmd {
     
     
     Manager_Cmd_AtmosphereEndInitDefers = 65000,
+    Manager_Cmd_AtmosphereHasMitm = 65001,
 };
 
 class ManagerService final : public IServiceObject {
@@ -32,11 +34,13 @@ class ManagerService final : public IServiceObject {
         virtual Result RegisterProcess(u64 pid, InBuffer<u8> acid_sac, InBuffer<u8> aci0_sac);
         virtual Result UnregisterProcess(u64 pid);
         virtual void AtmosphereEndInitDefers();
+        virtual void AtmosphereHasMitm(Out<bool> out, SmServiceName service);
     public:
         DEFINE_SERVICE_DISPATCH_TABLE {
             MakeServiceCommandMeta<Manager_Cmd_RegisterProcess, &ManagerService::RegisterProcess>(),
             MakeServiceCommandMeta<Manager_Cmd_UnregisterProcess, &ManagerService::UnregisterProcess>(),
             
             MakeServiceCommandMeta<Manager_Cmd_AtmosphereEndInitDefers, &ManagerService::AtmosphereEndInitDefers>(),
+            MakeServiceCommandMeta<Manager_Cmd_AtmosphereHasMitm, &ManagerService::AtmosphereHasMitm>(),
         };
 };

stratosphere/sm/source/sm_registration.cpp

@@ -180,6 +180,9 @@ bool Registration::IsInitialProcess(u64 pid) {
 
 u64 Registration::GetInitialProcessId() {
     CacheInitialProcessIdLimits();
+    if (IsInitialProcess(1)) {
+        return 1;
+    }
     return g_initial_process_id_low;
 }
 
@@ -219,6 +222,11 @@ bool Registration::HasService(u64 service) {
     return std::any_of(g_service_list.begin(), g_service_list.end(), member_equals_fn(&Service::service_name, service));
 }
 
+bool Registration::HasMitm(u64 service) {
+    Registration::Service *target_service = GetService(service);
+    return target_service != NULL && target_service->mitm_pid != 0;
+}
+
 Result Registration::GetServiceHandle(u64 pid, u64 service, Handle *out) {
     Registration::Service *target_service = GetService(service);
     if (target_service == NULL || ShouldInitDefer(service) || target_service->mitm_waiting_ack) {

stratosphere/sm/source/sm_registration.hpp

@@ -69,6 +69,7 @@ class Registration {
         
         /* Service management. */
         static bool HasService(u64 service);
+        static bool HasMitm(u64 service);
         static Result GetServiceHandle(u64 pid, u64 service, Handle *out);
         static Result GetServiceForPid(u64 pid, u64 service, Handle *out);
         static Result RegisterServiceForPid(u64 pid, u64 service, u64 max_sessions, bool is_light, Handle *out);