From 0169d7d77cf7a43aac45508ac0d0f68f848ca318 Mon Sep 17 00:00:00 2001
From: misson20000 <xenotoad@xenotoad.net>
Date: Thu, 20 Dec 2018 16:18:10 -0800
Subject: [PATCH] sm: add query registrations extension

---
 stratosphere/sm/source/sm_manager_service.cpp |  7 +++++
 stratosphere/sm/source/sm_manager_service.hpp |  3 +++
 stratosphere/sm/source/sm_registration.cpp    | 26 +++++++++++++++++++
 stratosphere/sm/source/sm_registration.hpp    |  3 +++
 stratosphere/sm/source/sm_types.hpp           | 13 +++++++++-
 5 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/stratosphere/sm/source/sm_manager_service.cpp b/stratosphere/sm/source/sm_manager_service.cpp
index 16b7c3820..6d8815495 100644
--- a/stratosphere/sm/source/sm_manager_service.cpp
+++ b/stratosphere/sm/source/sm_manager_service.cpp
@@ -35,3 +35,10 @@ void ManagerService::AtmosphereHasMitm(Out<bool> out, SmServiceName service) {
     out.SetValue(Registration::HasMitm(smEncodeName(service.name)));
 }
 
+Result ManagerService::AtmosphereQueryRegistrations(u64 offset, OutBuffer<ServiceRecord> records, Out<u64> out_count, Out<u64> record_size) {
+    uint64_t count = records.num_elements;
+    Registration::QueryRegistrations(offset, records.buffer, &count);
+    out_count.SetValue(count);
+    record_size.SetValue(sizeof(ServiceRecord));
+    return 0;
+}
diff --git a/stratosphere/sm/source/sm_manager_service.hpp b/stratosphere/sm/source/sm_manager_service.hpp
index 454be7edb..6c255a06a 100644
--- a/stratosphere/sm/source/sm_manager_service.hpp
+++ b/stratosphere/sm/source/sm_manager_service.hpp
@@ -26,6 +26,7 @@ enum ManagerServiceCmd {
     
     Manager_Cmd_AtmosphereEndInitDefers = 65000,
     Manager_Cmd_AtmosphereHasMitm = 65001,
+    Manager_Cmd_AtmosphereQueryRegistrations = 65002,
 };
 
 class ManagerService final : public IServiceObject {
@@ -35,6 +36,7 @@ class ManagerService final : public IServiceObject {
         virtual Result UnregisterProcess(u64 pid);
         virtual void AtmosphereEndInitDefers();
         virtual void AtmosphereHasMitm(Out<bool> out, SmServiceName service);
+        virtual Result AtmosphereQueryRegistrations(u64 offset, OutBuffer<ServiceRecord> records, Out<u64> out_count, Out<u64> record_size);
     public:
         DEFINE_SERVICE_DISPATCH_TABLE {
             MakeServiceCommandMeta<Manager_Cmd_RegisterProcess, &ManagerService::RegisterProcess>(),
@@ -42,5 +44,6 @@ class ManagerService final : public IServiceObject {
             
             MakeServiceCommandMeta<Manager_Cmd_AtmosphereEndInitDefers, &ManagerService::AtmosphereEndInitDefers>(),
             MakeServiceCommandMeta<Manager_Cmd_AtmosphereHasMitm, &ManagerService::AtmosphereHasMitm>(),
+            MakeServiceCommandMeta<Manager_Cmd_AtmosphereQueryRegistrations, &ManagerService::AtmosphereQueryRegistrations>(),
         };
 };
diff --git a/stratosphere/sm/source/sm_registration.cpp b/stratosphere/sm/source/sm_registration.cpp
index 3401f800f..839f561a5 100644
--- a/stratosphere/sm/source/sm_registration.cpp
+++ b/stratosphere/sm/source/sm_registration.cpp
@@ -562,3 +562,29 @@ Result Registration::AssociatePidTidForMitm(u64 pid, u64 tid) {
     }
     return 0x0;
 }
+
+void Registration::QueryRegistrations(u64 offset, ServiceRecord *out, u64 *count) {
+    u64 space = *count;
+    *count = 0;
+
+    for(auto i = g_service_list.begin() + offset; i < g_service_list.end() && space > 0; i++) {
+        if(i->service_name != 0) {
+            if(offset > 0) {
+                offset--;
+            } else {
+                ServiceRecord *out_record = out++;
+                
+                out_record->service_name = i->service_name;
+                out_record->owner_pid = i->owner_pid;
+                out_record->max_sessions = i->max_sessions;
+                out_record->mitm_pid = i->mitm_pid;
+                out_record->mitm_waiting_ack_pid = i->mitm_waiting_ack_pid;
+                out_record->is_light = i->is_light;
+                out_record->mitm_waiting_ack = i->mitm_waiting_ack;
+                    
+                space--;
+                (*count)++;
+            }
+        }
+    }
+}
diff --git a/stratosphere/sm/source/sm_registration.hpp b/stratosphere/sm/source/sm_registration.hpp
index 8c4d72660..b2a61c8f6 100644
--- a/stratosphere/sm/source/sm_registration.hpp
+++ b/stratosphere/sm/source/sm_registration.hpp
@@ -17,6 +17,8 @@
 #pragma once
 #include <switch.h>
 
+#include "sm_types.hpp"
+
 #define REGISTRATION_LIST_MAX_PROCESS (0x40)
 #define REGISTRATION_LIST_MAX_SERVICE (0x100)
 #define REGISTRATION_MAX_SAC_SIZE (0x200)
@@ -81,4 +83,5 @@ class Registration {
         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);
+        static void QueryRegistrations(u64 offset, ServiceRecord *out, u64 *count);
 };
diff --git a/stratosphere/sm/source/sm_types.hpp b/stratosphere/sm/source/sm_types.hpp
index 6ac9bce92..529fcc560 100644
--- a/stratosphere/sm/source/sm_types.hpp
+++ b/stratosphere/sm/source/sm_types.hpp
@@ -19,4 +19,15 @@ struct SmServiceName {
     char name[sizeof(u64)];
 };
 
-static_assert(__alignof__(SmServiceName) == 1, "SmServiceName definition!");
\ No newline at end of file
+static_assert(__alignof__(SmServiceName) == 1, "SmServiceName definition!");
+
+// for QueryRegistrations extension
+struct ServiceRecord {
+    u64 service_name;
+    u64 owner_pid;
+    u64 max_sessions;
+    u64 mitm_pid;
+    u64 mitm_waiting_ack_pid;
+    bool is_light;
+    bool mitm_waiting_ack;
+};