Greatly improve mitm session acquire semantics.

include/stratosphere/mitm/imitmserviceobject.hpp

@@ -36,6 +36,11 @@ class IMitmServiceObject : public IServiceObject {
             return this->process_id;
         }
         
+        void SetPidTid(u64 pid, u64 tid) {
+            this->process_id = pid;
+            this->title_id = tid;
+        }
+        
         static bool ShouldMitm(u64 pid, u64 tid);
                 
     protected:

include/stratosphere/mitm/mitm_server.hpp

@@ -89,13 +89,15 @@ class MitmServer : public IWaitable {
                 fatalSimple(rc);
             }
             
-            if (R_FAILED(smMitMGetService(forward_service.get(), mitm_name))) {
+            u64 client_pid;
+            
+            if (R_FAILED(smMitMAcknowledgeSession(forward_service.get(), &client_pid, mitm_name))) {
                 /* TODO: Panic. */
             }
             
             smMitMExit();
             
-            this->GetSessionManager()->AddWaitable(new MitmSession(session_h, forward_service, std::make_shared<T>(forward_service)));
+            this->GetSessionManager()->AddWaitable(new MitmSession(session_h, client_pid, forward_service, std::make_shared<T>(forward_service)));
             return 0;
         }
 

include/stratosphere/mitm/mitm_session.hpp

@@ -32,14 +32,19 @@ class MitmSession final : public ServiceSession {
         void (*service_post_process_handler)(IMitmServiceObject *, IpcResponseContext *);
         
         /* For cleanup usage. */
+        u64 client_pid;
         u32 num_fwd_copy_hnds = 0;
         Handle fwd_copy_hnds[8];
     public:
         template<typename T>
-        MitmSession(Handle s_h, std::shared_ptr<Service> fs, std::shared_ptr<T> srv) : ServiceSession(s_h) {
+        MitmSession(Handle s_h, u64 pid, std::shared_ptr<Service> fs, std::shared_ptr<T> srv) : ServiceSession(s_h), client_pid(pid) {
             this->forward_service = std::move(fs);
             this->obj_holder = std::move(ServiceObjectHolder(std::move(srv)));
             
+            u64 tid = 0;
+            MitmQueryUtils::GetAssociatedTidForPid(client_pid, &tid);
+            this->obj_holder.GetServiceObjectUnsafe<IMitmServiceObject>()->SetPidTid(client_pid, tid);
+            
             this->service_post_process_handler = T::PostProcess;
             
             size_t pbs;
@@ -51,7 +56,7 @@ class MitmSession final : public ServiceSession {
             this->control_holder = std::move(ServiceObjectHolder(std::move(std::make_shared<IMitmHipcControlService>(this))));
         }
         
-        MitmSession(Handle s_h, std::shared_ptr<Service> fs, ServiceObjectHolder &&h, void (*pph)(IMitmServiceObject *, IpcResponseContext *)) : ServiceSession(s_h) {
+        MitmSession(Handle s_h, u64 pid, std::shared_ptr<Service> fs, ServiceObjectHolder &&h, void (*pph)(IMitmServiceObject *, IpcResponseContext *)) : ServiceSession(s_h), client_pid(pid) {
             this->session_handle = s_h;
             this->forward_service = std::move(fs);
             this->obj_holder = std::move(h);
@@ -277,7 +282,7 @@ class MitmSession final : public ServiceSession {
                     out_h.SetValue(client_h);
                     
                     if (id == serviceGetObjectId(this->session->forward_service.get())) {
-                        this->session->GetSessionManager()->AddWaitable(new MitmSession(server_h, this->session->forward_service, std::move(object->Clone()), this->session->service_post_process_handler));
+                        this->session->GetSessionManager()->AddWaitable(new MitmSession(server_h, this->session->client_pid, this->session->forward_service, std::move(object->Clone()), this->session->service_post_process_handler));
                     } else {
                         this->session->GetSessionManager()->AddSession(server_h, std::move(object->Clone()));
                     }
@@ -291,7 +296,7 @@ class MitmSession final : public ServiceSession {
                         std::abort();
                     }
                     
-                    this->session->GetSessionManager()->AddWaitable(new MitmSession(server_h, this->session->forward_service, std::move(this->session->obj_holder.Clone()), this->session->service_post_process_handler));
+                    this->session->GetSessionManager()->AddWaitable(new MitmSession(server_h, this->session->client_pid, this->session->forward_service, std::move(this->session->obj_holder.Clone()), this->session->service_post_process_handler));
                     out_h.SetValue(client_h);
                 }
                 

include/stratosphere/mitm/sm_mitm.h

@@ -16,6 +16,7 @@ void smMitMExit(void);
 Result smMitMGetService(Service* service_out, const char *name);
 Result smMitMInstall(Handle *handle_out, Handle *query_out, const char *name);
 Result smMitMUninstall(const char *name);
+Result smMitMAcknowledgeSession(Service *srv_out, u64 *pid_out, const char *name);
 
 Result smMitMIsRegistered(const char *name);
 

source/sm_mitm.c

@@ -188,3 +188,43 @@ Result smMitMUninstall(const char *name) {
 
     return rc;
 }
+
+Result smMitMAcknowledgeSession(Service *srv_out, u64 *pid_out, const char *name) {
+    IpcCommand c;
+    ipcInitialize(&c);
+
+    struct {
+        u64 magic;
+        u64 cmd_id;
+        u64 service_name;
+        u64 reserved;
+    } *raw;
+
+    raw = ipcPrepareHeader(&c, sizeof(*raw));
+
+    raw->magic = SFCI_MAGIC;
+    raw->cmd_id = 65003;
+    raw->service_name = smEncodeName(name);
+
+    Result rc = ipcDispatch(g_smMitmHandle);
+
+    if (R_SUCCEEDED(rc)) {
+        IpcParsedCommand r;
+        ipcParse(&r);
+
+        struct {
+            u64 magic;
+            u64 result;
+            u64 pid;
+        } *resp = r.Raw;
+
+        rc = resp->result;
+        if (R_SUCCEEDED(rc)) {
+            *pid_out = resp->pid;
+            serviceCreate(srv_out, r.Handles[0]);
+        }
+    }
+
+    return rc;
+
+}