From f61f5feaf4b2cdcc69fa5cd125186be25bd29982 Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Sun, 20 Jan 2019 22:21:49 -0800
Subject: [PATCH] pm: Launch set.mitm earlier in boot.

---
 stratosphere/pm/source/pm_boot2.cpp | 43 +++++++++++++++++++++++++----
 stratosphere/pm/source/pm_boot2.hpp |  4 +++
 stratosphere/pm/source/pm_main.cpp  |  6 ++--
 3 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/stratosphere/pm/source/pm_boot2.cpp b/stratosphere/pm/source/pm_boot2.cpp
index 5a88f10bb..181fbb44d 100644
--- a/stratosphere/pm/source/pm_boot2.cpp
+++ b/stratosphere/pm/source/pm_boot2.cpp
@@ -26,6 +26,20 @@
 #include "pm_boot2.hpp"
 #include "pm_registration.hpp"
 
+static std::vector<Boot2KnownTitleId> g_boot2_titles;
+
+static void ClearLaunchedTitles() {
+    g_boot2_titles.clear();
+}
+
+static void SetLaunchedTitle(Boot2KnownTitleId title_id) {
+    g_boot2_titles.push_back(title_id);
+}
+
+static bool HasLaunchedTitle(Boot2KnownTitleId title_id) {
+    return std::find(g_boot2_titles.begin(), g_boot2_titles.end(), title_id) != g_boot2_titles.end();
+}
+
 static bool IsHexadecimal(const char *str) {
     while (*str) {
         if (isxdigit(*str)) {
@@ -40,6 +54,11 @@ static bool IsHexadecimal(const char *str) {
 static void LaunchTitle(Boot2KnownTitleId title_id, FsStorageId storage_id, u32 launch_flags, u64 *pid) {
     u64 local_pid;
     
+    /* Don't launch a title twice during boot2. */
+    if (HasLaunchedTitle(title_id)) {
+        return;
+    }
+    
     Result rc = Registration::LaunchProcessByTidSid(Registration::TidSid{(u64)title_id, storage_id}, launch_flags, &local_pid);
     switch (rc) {
         case 0xCE01:
@@ -61,6 +80,10 @@ static void LaunchTitle(Boot2KnownTitleId title_id, FsStorageId storage_id, u32
     if (pid) {
         *pid = local_pid;
     }
+    
+    if (R_SUCCEEDED(rc)) {
+        SetLaunchedTitle(title_id);
+    }
 }
 
 static bool ShouldForceMaintenanceMode() {
@@ -123,25 +146,28 @@ static void MountSdCard() {
     fsdevMountSdmc();
 }
 
-static void WaitForFsMitm() {
-    bool fs_mitm_installed = false;
+static void WaitForMitm(const char *service) {
+    bool 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) {
+    while (R_FAILED((rc = smManagerAmsHasMitm(&mitm_installed, service))) || !mitm_installed) {
         if (R_FAILED(rc)) {
             std::abort();
         }
-        svcSleepThread(1000ull);
+        svcSleepThread(1000000ull);
     }
     smManagerAmsExit();
 }
 
 void EmbeddedBoot2::Main() {
     /* Wait until fs.mitm has installed itself. We want this to happen as early as possible. */
-    WaitForFsMitm();
+    WaitForMitm("fsp-srv");
+    
+    /* Clear titles. */
+    ClearLaunchedTitles();
 
     /* psc, bus, pcv is the minimal set of required titles to get SD card. */ 
     /* bus depends on pcie, and pcv depends on settings. */
@@ -159,6 +185,10 @@ void EmbeddedBoot2::Main() {
     /* At this point, the SD card can be mounted. */
     MountSdCard();
     
+    /* Launch set:mitm, wait for it. */
+    LaunchTitle(Boot2KnownTitleId::ams_set_mitm, FsStorageId_None, 0, NULL);
+    WaitForMitm("set:sys");
+    
     /* Launch usb. */
     LaunchTitle(Boot2KnownTitleId::usb, FsStorageId_NandSystem, 0, NULL);
     /* Launch tma. */
@@ -206,4 +236,7 @@ void EmbeddedBoot2::Main() {
         
     /* We no longer need the SD card. */
     fsdevUnmountAll();
+    
+    /* Clear titles. */
+    ClearLaunchedTitles();
 }
diff --git a/stratosphere/pm/source/pm_boot2.hpp b/stratosphere/pm/source/pm_boot2.hpp
index 44c4e1de1..7a0cd5dec 100644
--- a/stratosphere/pm/source/pm_boot2.hpp
+++ b/stratosphere/pm/source/pm_boot2.hpp
@@ -77,6 +77,10 @@ enum class Boot2KnownTitleId : u64 {
     jpegdec = 0x010000000000003CUL,
     safemode = 0x010000000000003DUL,
     olsc = 0x010000000000003EUL,
+    
+    
+    /* atmosphere extensions */
+    ams_set_mitm = 0x0100000000000032UL,
 };
 
 class EmbeddedBoot2 {
diff --git a/stratosphere/pm/source/pm_main.cpp b/stratosphere/pm/source/pm_main.cpp
index c38e7ec1b..42f72fdd4 100644
--- a/stratosphere/pm/source/pm_main.cpp
+++ b/stratosphere/pm/source/pm_main.cpp
@@ -148,15 +148,15 @@ void __appExit(void) {
 
 int main(int argc, char **argv)
 {
-    Thread process_track_thread = {0};
+    HosThread process_track_thread;
     consoleDebugInit(debugDevice_SVC);
     
     /* Initialize and spawn the Process Tracking thread. */
     Registration::InitializeSystemResources();
-    if (R_FAILED(threadCreate(&process_track_thread, &ProcessTracking::MainLoop, NULL, 0x4000, 0x15, 0))) {
+    if (R_FAILED(process_track_thread.Initialize(&ProcessTracking::MainLoop, NULL, 0x4000, 0x15))) {
         /* TODO: Panic. */
     }
-    if (R_FAILED(threadStart(&process_track_thread))) {
+    if (R_FAILED(process_track_thread.Start())) {
         /* TODO: Panic. */
     }