pm/sm: add ability to forward declare mitm'd services (closes #557)

stratosphere/libstratosphere

@@ -1 +1 @@
-Subproject commit b9e53052732699e1f7af8ad061181e3798c3f8fc
+Subproject commit 69dbb69e0b32aa7b977e144fe7a8355868abd04a

stratosphere/pm/source/boot2/boot2_api.cpp

@@ -134,6 +134,10 @@ namespace sts::boot2 {
             return true;
         }
 
+        inline bool IsNewLine(char c) {
+            return c == '\r' || c == '\n';
+        }
+
         void LaunchTitle(u64 *out_process_id, const ncm::TitleLocation &loc, u32 launch_flags) {
             u64 process_id = 0;
 
@@ -201,7 +205,8 @@ namespace sts::boot2 {
             }
         }
 
-        void LaunchFlaggedProgramsFromSdCard() {
+        template<typename F>
+        void IterateOverFlaggedTitlesOnSdCard(F f) {
             /* Validate that the titles directory exists. */
             DIR *titles_dir = opendir("sdmc:/atmosphere/titles");
             if (titles_dir == nullptr) {
@@ -216,21 +221,92 @@ namespace sts::boot2 {
                 if (std::strlen(ent->d_name) == 2 * sizeof(ncm::TitleId) && IsHexadecimal(ent->d_name)) {
                     /* Check if we've already launched the title. */
                     ncm::TitleId title_id{std::strtoul(ent->d_name, nullptr, 16)};
-                    if (pm::info::HasLaunchedTitle(title_id)) {
-                        continue;
-                    }
 
                     /* Check if the title is flagged. */
                     if (!cfg::HasTitleSpecificFlag(title_id, "boot2")) {
                         continue;
                     }
 
-                    /* Actually launch the title. */
+                    /* Call the iteration callback. */
-                    LaunchTitle(nullptr, ncm::TitleLocation::Make(title_id, ncm::StorageId::None), 0);
+                    f(title_id);
                 }
             }
         }
 
+        void DetectAndDeclareFutureMitms() {
+            IterateOverFlaggedTitlesOnSdCard([](ncm::TitleId title_id) {
+                /* When we find a flagged program, check if it has a mitm list. */
+                char mitm_list[0x400];
+                size_t mitm_list_size = 0;
+
+                /* Read the mitm list off the SD card. */
+                {
+                    char path[FS_MAX_PATH];
+                    std::snprintf(mitm_list, sizeof(mitm_list), "sdmc:/atmosphere/titles/%016lx/mitm.lst", static_cast<u64>(title_id));
+                    FILE *f = fopen(path, "rb");
+                    if (f == nullptr) {
+                        return;
+                    }
+                    mitm_list_size = static_cast<size_t>(fread(mitm_list, 1, sizeof(mitm_list), f));
+                    fclose(f);
+                }
+
+                /* Validate read size. */
+                if (mitm_list_size > sizeof(mitm_list) || mitm_list_size == 0) {
+                    return;
+                }
+
+                /* Iterate over the contents of the file. */
+                /* We expect one service name per non-empty line, 1-8 characters. */
+                size_t offset = 0;
+                while (offset < mitm_list_size) {
+                    /* Continue to the start of the next name. */
+                    while (IsNewLine(mitm_list[offset])) {
+                        offset++;
+                    }
+                    if (offset >= mitm_list_size) {
+                        break;
+                    }
+
+                    /* Get the length of the current name. */
+                    size_t name_len;
+                    for (name_len = 0; name_len <= sizeof(sm::ServiceName) && offset + name_len < mitm_list_size; name_len++) {
+                        if (IsNewLine(mitm_list[offset + name_len])) {
+                            break;
+                        }
+                    }
+
+                    /* Allow empty lines. */
+                    if (name_len == 0) {
+                        continue;
+                    }
+
+                    /* Don't allow invalid lines. */
+                    if (name_len > sizeof(sm::ServiceName)) {
+                        std::abort();
+                    }
+
+                    /* Declare the service. */
+                    R_ASSERT(sm::mitm::DeclareFutureMitm(sm::ServiceName::Encode(mitm_list + offset, name_len)));
+
+                    /* Advance to the next line. */
+                    offset += name_len;
+                }
+            });
+        }
+
+        void LaunchFlaggedProgramsOnSdCard() {
+            IterateOverFlaggedTitlesOnSdCard([](ncm::TitleId title_id) {
+                /* Check if we've already launched the title. */
+                if (pm::info::HasLaunchedTitle(title_id)) {
+                    return;
+                }
+
+                /* Launch the title. */
+                LaunchTitle(nullptr, ncm::TitleLocation::Make(title_id, ncm::StorageId::None), 0);
+            });
+        }
+
     }
 
     /* Boot2 API. */
@@ -258,9 +334,13 @@ namespace sts::boot2 {
         if (maintenance) {
             pm::bm::SetMaintenanceBoot();
         }
+
         /* Launch Atmosphere dmnt, using FsStorageId_None to force SD card boot. */
         LaunchTitle(nullptr, ncm::TitleLocation::Make(ncm::TitleId::Dmnt, ncm::StorageId::None), 0);
 
+        /* Check for and forward declare non-atmosphere mitm modules. */
+        DetectAndDeclareFutureMitms();
+
         /* Launch additional programs. */
         if (maintenance) {
             LaunchList(AdditionalMaintenanceLaunchPrograms, NumAdditionalMaintenanceLaunchPrograms);
@@ -272,9 +352,8 @@ namespace sts::boot2 {
             LaunchList(AdditionalLaunchPrograms, NumAdditionalLaunchPrograms);
         }
 
-
         /* Launch user programs off of the SD. */
-        LaunchFlaggedProgramsFromSdCard();
+        LaunchFlaggedProgramsOnSdCard();
 
         /* We no longer need the SD card. */
         fsdevUnmountAll();

stratosphere/sm/source/impl/sm_service_manager.cpp

@@ -27,8 +27,9 @@ namespace sts::sm::impl {
     /* Anonymous namespace for implementation details. */
     namespace {
         /* Constexpr definitions. */
-        static constexpr size_t ProcessCountMax = 0x40;
+        static constexpr size_t ProcessCountMax      = 0x40;
-        static constexpr size_t ServiceCountMax = 0x100;
+        static constexpr size_t ServiceCountMax      = 0x100;
+        static constexpr size_t FutureMitmCountMax   = 0x20;
         static constexpr size_t AccessControlSizeMax = 0x200;
 
         /* Types. */
@@ -182,6 +183,7 @@ namespace sts::sm::impl {
         /* Static members. */
         ProcessInfo g_process_list[ProcessCountMax];
         ServiceInfo g_service_list[ServiceCountMax];
+        ServiceName g_future_mitm_list[FutureMitmCountMax];
         InitialProcessIdLimits g_initial_process_id_limits;
         bool g_ended_initial_defers;
 
@@ -203,6 +205,14 @@ namespace sts::sm::impl {
             return GetProcessInfo(pid) != nullptr;
         }
 
+        bool IsInitialProcess(u64 pid) {
+            return g_initial_process_id_limits.IsInitialProcess(pid);
+        }
+
+        bool IsValidProcessId(u64 pid) {
+            return pid != InvalidProcessId;
+        }
+
         ServiceInfo *GetServiceInfo(ServiceName service_name) {
             for (size_t i = 0; i < ServiceCountMax; i++) {
                 if (g_service_list[i].name == service_name) {
@@ -220,6 +230,53 @@ namespace sts::sm::impl {
             return GetServiceInfo(service) != nullptr;
         }
 
+        bool HasMitm(ServiceName service) {
+            const ServiceInfo *service_info = GetServiceInfo(service);
+            return service_info != nullptr && IsValidProcessId(service_info->mitm_pid);
+        }
+
+        ncm::TitleId GetTitleIdForMitm(u64 pid) {
+            /* Anything that can request a mitm session must have a process info. */
+            const auto process_info = GetProcessInfo(pid);
+            if (process_info == nullptr) {
+                std::abort();
+            }
+            return process_info->tid;
+        }
+
+        bool IsMitmDisallowed(ncm::TitleId title_id) {
+            /* Mitm used on certain titles can prevent the boot process from completing. */
+            /* TODO: Is there a way to do this that's less hardcoded? Needs design thought. */
+            return title_id == ncm::TitleId::Loader || title_id == ncm::TitleId::Boot || title_id == ncm::TitleId::AtmosphereMitm;
+        }
+
+        Result AddFutureMitmDeclaration(ServiceName service) {
+            for (size_t i = 0; i < FutureMitmCountMax; i++) {
+                if (g_future_mitm_list[i] == InvalidServiceName) {
+                    g_future_mitm_list[i] = service;
+                    return ResultSuccess;
+                }
+            }
+            return ResultSmInsufficientServices;
+        }
+
+        bool HasFutureMitmDeclaration(ServiceName service) {
+            for (size_t i = 0; i < FutureMitmCountMax; i++) {
+                if (g_future_mitm_list[i] == service) {
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        void ClearFutureMitmDeclaration(ServiceName service) {
+            for (size_t i = 0; i < FutureMitmCountMax; i++) {
+                if (g_future_mitm_list[i] == service) {
+                    g_future_mitm_list[i] = InvalidServiceName;
+                }
+            }
+        }
+
         void GetServiceInfoRecord(ServiceRecord *out_record, const ServiceInfo *service_info) {
             out_record->service              = service_info->name;
             out_record->owner_pid            = service_info->owner_pid;
@@ -287,23 +344,6 @@ namespace sts::sm::impl {
             return ResultSuccess;
         }
 
-        bool IsInitialProcess(u64 pid) {
-            return g_initial_process_id_limits.IsInitialProcess(pid);
-        }
-
-        bool IsValidProcessId(u64 pid) {
-            return pid != InvalidProcessId;
-        }
-
-        ncm::TitleId GetTitleIdForMitm(u64 pid) {
-            /* Anything that can request a mitm session must have a process info. */
-            const auto process_info = GetProcessInfo(pid);
-            if (process_info == nullptr) {
-                std::abort();
-            }
-            return process_info->tid;
-        }
-
         bool ShouldDeferForInit(ServiceName service) {
             /* Once end has been called, we're done. */
             if (g_ended_initial_defers) {
@@ -315,7 +355,7 @@ namespace sts::sm::impl {
             return service == ServiceName::Encode("fsp-srv");
         }
 
-        Result GetMitmServiceHandleImpl(Handle *out, ServiceInfo *service_info, u64 pid) {
+        Result GetMitmServiceHandleImpl(Handle *out, ServiceInfo *service_info, u64 pid, ncm::TitleId title_id) {
             /* Send command to query if we should mitm. */
             {
                 IpcCommand c;
@@ -329,7 +369,7 @@ namespace sts::sm::impl {
                 info->magic = SFCI_MAGIC;
                 info->cmd_id = 65000;
                 info->pid = pid;
-                info->tid = GetTitleIdForMitm(pid);
+                info->tid = title_id;
                 R_TRY(ipcDispatch(service_info->mitm_query_h.Get()));
             }
 
@@ -372,14 +412,19 @@ namespace sts::sm::impl {
             /* Clear handle output. */
             *out = INVALID_HANDLE;
 
-            /* If not mitm'd or mitm service is requesting, get normal session. */
+            /* Check if we should return a mitm handle. */
-            if (!IsValidProcessId(service_info->mitm_pid) || service_info->mitm_pid == pid) {
+            if (IsValidProcessId(service_info->mitm_pid) && service_info->mitm_pid != pid) {
-                return svcConnectToPort(out, service_info->port_h.Get());
+                /* Get title id, ensure that we're allowed to mitm the given title. */
+                const ncm::TitleId title_id = GetTitleIdForMitm(pid);
+                if (!IsMitmDisallowed(title_id)) {
+                    /* We're mitm'd. Assert, because mitm service host dead is an error state. */
+                    R_ASSERT(GetMitmServiceHandleImpl(out, service_info, pid, title_id));
+                    return ResultSuccess;
+                }
             }
 
-            /* We're mitm'd. Assert, because mitm service host dead is an error state. */
+            /* We're not returning a mitm handle, so just return a normal port handle. */
-            R_ASSERT(GetMitmServiceHandleImpl(out, service_info, pid));
+            return svcConnectToPort(out, service_info->port_h.Get());
-            return ResultSuccess;
         }
 
         Result RegisterServiceImpl(Handle *out, u64 pid, ServiceName service, size_t max_sessions, bool is_light) {
@@ -501,7 +546,7 @@ namespace sts::sm::impl {
 
         /* Get service info. Check to see if we need to defer this until later. */
         ServiceInfo *service_info = GetServiceInfo(service);
-        if (service_info == nullptr || ShouldDeferForInit(service) || service_info->mitm_waiting_ack) {
+        if (service_info == nullptr || ShouldDeferForInit(service) || HasFutureMitmDeclaration(service) || service_info->mitm_waiting_ack) {
             return ResultServiceFrameworkRequestDeferredByUser;
         }
 
@@ -637,6 +682,9 @@ namespace sts::sm::impl {
             *out_query = qry_hnd.Move();
         }
 
+        /* Clear the future declaration, if one exists. */
+        ClearFutureMitmDeclaration(service);
+
         return ResultSuccess;
     }
 
@@ -668,6 +716,30 @@ namespace sts::sm::impl {
         return ResultSuccess;
     }
 
+    Result DeclareFutureMitm(u64 pid, ServiceName service) {
+        /* Validate service name. */
+        R_TRY(ValidateServiceName(service));
+
+        /* Check that the process is registered and allowed to register the service. */
+        if (!IsInitialProcess(pid)) {
+            ProcessInfo *proc = GetProcessInfo(pid);
+            if (proc == nullptr) {
+                return ResultSmInvalidClient;
+            }
+
+            R_TRY(ValidateAccessControl(AccessControlEntry(proc->access_control, proc->access_control_size), service, true, false));
+        }
+
+        /* Check that mitm hasn't already been registered or declared. */
+        if (HasMitm(service) || HasFutureMitmDeclaration(service)) {
+            return ResultSmAlreadyRegistered;
+        }
+
+        /* Try to forward declare it. */
+        R_TRY(AddFutureMitmDeclaration(service));
+        return ResultSuccess;
+    }
+
     Result AcknowledgeMitmSession(u64 *out_pid, ncm::TitleId *out_tid, Handle *out_hnd, u64 pid, ServiceName service) {
         /* Validate service name. */
         R_TRY(ValidateServiceName(service));

stratosphere/sm/source/impl/sm_service_manager.hpp

@@ -38,6 +38,7 @@ namespace sts::sm::impl {
     Result WaitMitm(ServiceName service);
     Result InstallMitm(Handle *out, Handle *out_query, u64 pid, ServiceName service);
     Result UninstallMitm(u64 pid, ServiceName service);
+    Result DeclareFutureMitm(u64 pid, ServiceName service);
     Result AcknowledgeMitmSession(u64 *out_pid, ncm::TitleId *out_tid, Handle *out_hnd, u64 pid, ServiceName service);
 
     /* Dmnt record extensions. */

stratosphere/sm/source/sm_user_service.cpp

@@ -75,6 +75,12 @@ namespace sts::sm {
         return impl::WaitMitm(service);
     }
 
+    Result UserService::AtmosphereDeclareFutureMitm(ServiceName service) {
+        R_TRY(this->EnsureInitialized());
+        return impl::DeclareFutureMitm(this->pid, service);
+    }
+
+
     Result UserService::AtmosphereHasService(Out<bool> out, ServiceName service) {
         R_TRY(this->EnsureInitialized());
         return impl::HasService(out.GetPointer(), service);

stratosphere/sm/source/sm_user_service.hpp

@@ -38,6 +38,7 @@ namespace sts::sm {
                 AtmosphereAcknowledgeMitmSession = 65003,
                 AtmosphereHasMitm                = 65004,
                 AtmosphereWaitMitm               = 65005,
+                AtmosphereDeclareFutureMitm      = 65006,
 
                 AtmosphereHasService             = 65100,
                 AtmosphereWaitService            = 65101,
@@ -60,6 +61,7 @@ namespace sts::sm {
             virtual Result AtmosphereAcknowledgeMitmSession(Out<u64> client_pid, Out<ncm::TitleId> client_tid, Out<MovedHandle> fwd_h, ServiceName service);
             virtual Result AtmosphereHasMitm(Out<bool> out, ServiceName service);
             virtual Result AtmosphereWaitMitm(ServiceName service);
+            virtual Result AtmosphereDeclareFutureMitm(ServiceName service);
 
             virtual Result AtmosphereHasService(Out<bool> out, ServiceName service);
             virtual Result AtmosphereWaitService(ServiceName service);
@@ -75,6 +77,7 @@ namespace sts::sm {
                 MAKE_SERVICE_COMMAND_META(UserService, AtmosphereAcknowledgeMitmSession),
                 MAKE_SERVICE_COMMAND_META(UserService, AtmosphereHasMitm),
                 MAKE_SERVICE_COMMAND_META(UserService, AtmosphereWaitMitm),
+                MAKE_SERVICE_COMMAND_META(UserService, AtmosphereDeclareFutureMitm),
 
                 MAKE_SERVICE_COMMAND_META(UserService, AtmosphereHasService),
                 MAKE_SERVICE_COMMAND_META(UserService, AtmosphereWaitService),