mitm: Improve session acquire semantics.

stratosphere/libstratosphere

@@ -1 +1 @@
-Subproject commit bcd80ab445258c20d968aad1c083fd8cb0937bee
+Subproject commit 8fcac73ab2b91afe184abe48c887ea1cb9bf2076

stratosphere/sm/source/sm_registration.cpp

@@ -221,13 +221,11 @@ bool Registration::HasService(u64 service) {
 
 Result Registration::GetServiceHandle(u64 pid, u64 service, Handle *out) {
     Registration::Service *target_service = GetService(service);
-    if (target_service == NULL || ShouldInitDefer(service)) {
+    if (target_service == NULL || ShouldInitDefer(service) || target_service->mitm_waiting_ack) {
         /* Note: This defers the result until later. */
         return RESULT_DEFER_SESSION;
     }
     
-    /* */
-    
     *out = 0;
     Result rc;
     if (target_service->mitm_pid == 0 || target_service->mitm_pid == pid) {
@@ -255,7 +253,17 @@ Result Registration::GetServiceHandle(u64 pid, u64 service, Handle *out) {
             rc = resp->result;
             if (R_SUCCEEDED(rc)) {
                 if (resp->should_mitm) {
-                    rc = svcConnectToPort(out, target_service->mitm_port_h);
+                    rc = svcConnectToPort(&target_service->mitm_fwd_sess_h, target_service->port_h);
+                    if (R_SUCCEEDED(rc)) {   
+                        rc = svcConnectToPort(out, target_service->mitm_port_h);
+                        if (R_SUCCEEDED(rc)) {
+                            target_service->mitm_waiting_ack_pid = pid;
+                            target_service->mitm_waiting_ack = true;
+                        } else {
+                            svcCloseHandle(target_service->mitm_fwd_sess_h);
+                            target_service->mitm_fwd_sess_h = 0;
+                        }
+                    }
                 } else {
                     rc = svcConnectToPort(out, target_service->port_h);
                 }
@@ -497,6 +505,35 @@ Result Registration::UninstallMitmForPid(u64 pid, u64 service) {
     return 0;
 }
 
+Result Registration::AcknowledgeMitmSessionForPid(u64 pid, u64 service, Handle *out, u64 *out_pid) {
+    if (!service) {
+        return 0xC15;
+    }
+    
+    u64 service_name_len = GetServiceNameLength(service);
+    
+    /* If the service has bytes after a null terminator, that's no good. */
+    if (service_name_len != 8 && (service >> (8 * service_name_len))) {
+        return 0xC15;
+    }
+    
+    Registration::Service *target_service = GetService(service);
+    if (target_service == NULL) {
+        return 0xE15;
+    }
+
+    if ((!IsInitialProcess(pid) && target_service->mitm_pid != pid) || !target_service->mitm_waiting_ack) {
+        return 0x1015;
+    }
+    
+    *out = target_service->mitm_fwd_sess_h;
+    *out_pid = target_service->mitm_waiting_ack_pid;
+    target_service->mitm_fwd_sess_h = 0;
+    target_service->mitm_waiting_ack_pid = 0;
+    target_service->mitm_waiting_ack = false;
+    return 0;
+}
+
 Result Registration::AssociatePidTidForMitm(u64 pid, u64 tid) {
     for (auto &service : g_service_list) {
         if (service.mitm_pid) {

stratosphere/sm/source/sm_registration.hpp

@@ -43,6 +43,10 @@ class Registration {
             u64 mitm_pid;
             Handle mitm_port_h;
             Handle mitm_query_h;
+            
+            bool mitm_waiting_ack;
+            u64 mitm_waiting_ack_pid;
+            Handle mitm_fwd_sess_h;
         };
         
         /* Utilities. */
@@ -74,5 +78,6 @@ class Registration {
         /* Extension. */
         static Result InstallMitmForPid(u64 pid, u64 service, Handle *out, Handle *query_out);
         static Result UninstallMitmForPid(u64 pid, u64 service);
+        static Result AcknowledgeMitmSessionForPid(u64 pid, u64 service, Handle *out, u64 *out_pid);
         static Result AssociatePidTidForMitm(u64 pid, u64 tid);
 };

stratosphere/sm/source/sm_user_service.cpp

@@ -98,6 +98,20 @@ Result UserService::AtmosphereUninstallMitm(SmServiceName service) {
     return rc;
 }
 
+Result UserService::AtmosphereAcknowledgeMitmSession(Out<u64> client_pid, Out<MovedHandle> fwd_h, SmServiceName service) {
+    Result rc = 0x415;
+    Handle out_fwd_h = 0;
+    if (this->has_initialized) {
+        rc = Registration::AcknowledgeMitmSessionForPid(this->pid, smEncodeName(service.name), &out_fwd_h, client_pid.GetPointer());
+    }
+    
+    if (R_SUCCEEDED(rc)) {
+        fwd_h.SetValue(out_fwd_h);
+    }
+    
+    return rc;
+}
+
 Result UserService::AtmosphereAssociatePidTidForMitm(u64 pid, u64 tid) {
     Result rc = 0x415;
     if (this->has_initialized) {

stratosphere/sm/source/sm_user_service.hpp

@@ -27,7 +27,8 @@ enum UserServiceCmd {
     
     User_Cmd_AtmosphereInstallMitm = 65000,
     User_Cmd_AtmosphereUninstallMitm = 65001,
-    User_Cmd_AtmosphereAssociatePidTidForMitm = 65002
+    User_Cmd_AtmosphereAssociatePidTidForMitm = 65002,
+    User_Cmd_AtmosphereAcknowledgeMitmSession = 65003,
 };
 
 class UserService final : public IServiceObject {
@@ -45,6 +46,7 @@ class UserService final : public IServiceObject {
         virtual Result AtmosphereInstallMitm(Out<MovedHandle> srv_h, Out<MovedHandle> qry_h, SmServiceName service);
         virtual Result AtmosphereUninstallMitm(SmServiceName service);
         virtual Result AtmosphereAssociatePidTidForMitm(u64 pid, u64 tid);
+        virtual Result AtmosphereAcknowledgeMitmSession(Out<u64> client_pid, Out<MovedHandle> fwd_h, SmServiceName service);
     public:
         DEFINE_SERVICE_DISPATCH_TABLE {
             MakeServiceCommandMeta<User_Cmd_Initialize, &UserService::Initialize>(),
@@ -56,6 +58,7 @@ class UserService final : public IServiceObject {
             MakeServiceCommandMeta<User_Cmd_AtmosphereInstallMitm, &UserService::AtmosphereInstallMitm>(),
             MakeServiceCommandMeta<User_Cmd_AtmosphereUninstallMitm, &UserService::AtmosphereUninstallMitm>(),
             MakeServiceCommandMeta<User_Cmd_AtmosphereAssociatePidTidForMitm, &UserService::AtmosphereAssociatePidTidForMitm>(),
+            MakeServiceCommandMeta<User_Cmd_AtmosphereAcknowledgeMitmSession, &UserService::AtmosphereAcknowledgeMitmSession>(),
 #endif
         };
 };