From a750e55f75be93e1df74af126af7b1ce6e290275 Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Sat, 14 Sep 2019 13:37:31 -0700
Subject: [PATCH] tmp-hid-mitm: less aggressively enable controllers

This fixes a crash in SSBU and possibly other games.
---
 .../source/hid_mitm/hid_mitm_service.cpp      | 43 ++++++++++++++++++-
 .../source/hid_mitm/hid_mitm_service.hpp      |  4 ++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/stratosphere/ams_mitm/source/hid_mitm/hid_mitm_service.cpp b/stratosphere/ams_mitm/source/hid_mitm/hid_mitm_service.cpp
index f1c122c85..55be2f534 100644
--- a/stratosphere/ams_mitm/source/hid_mitm/hid_mitm_service.cpp
+++ b/stratosphere/ams_mitm/source/hid_mitm/hid_mitm_service.cpp
@@ -14,17 +14,56 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
+#include <unordered_map>
 #include <switch.h>
 #include <stratosphere.hpp>
 
 #include "hid_shim.h"
 #include "hid_mitm_service.hpp"
 
+namespace {
+
+    std::unordered_map<u64, u64> g_hbl_map;
+
+    bool ShouldSetSystemExt(u64 process_id) {
+        return g_hbl_map.find(process_id) != g_hbl_map.end();
+    }
+
+    void AddSession(u64 process_id) {
+        if (g_hbl_map.find(process_id) != g_hbl_map.end()) {
+            g_hbl_map[process_id]++;
+        } else {
+            g_hbl_map[process_id] = 0;
+        }
+    }
+
+    void RemoveSession(u64 process_id) {
+        if ((--g_hbl_map[process_id]) == 0) {
+            g_hbl_map.erase(process_id);
+        }
+    }
+
+}
+
+HidMitmService::~HidMitmService() {
+    if (this->should_set_system_ext) {
+        RemoveSession(this->process_id);
+    }
+}
+
 void HidMitmService::PostProcess(IMitmServiceObject *obj, IpcResponseContext *ctx) {
     /* Nothing to do here */
 }
 
 Result HidMitmService::SetSupportedNpadStyleSet(u64 applet_resource_user_id, u32 style_set, PidDescriptor pid_desc) {
-    const HidControllerType new_style_set = static_cast<HidControllerType>(style_set | TYPE_SYSTEM | TYPE_SYSTEM_EXT);
-    return hidSetSupportedNpadStyleSetFwd(this->forward_service.get(), this->process_id, applet_resource_user_id, new_style_set);;
+    if (!this->should_set_system_ext && (style_set & (TYPE_SYSTEM | TYPE_SYSTEM_EXT))) {
+        /* Guaranteed: hbmenu 3.1.1 will cause this to be true. */
+        /* This prevents setting this for non-HBL. */
+        this->should_set_system_ext = true;
+        AddSession(this->process_id);
+    }
+    if (ShouldSetSystemExt(this->process_id)) {
+        style_set |= TYPE_SYSTEM | TYPE_SYSTEM_EXT;
+    }
+    return hidSetSupportedNpadStyleSetFwd(this->forward_service.get(), this->process_id, applet_resource_user_id, static_cast<HidControllerType>(style_set));
 }
diff --git a/stratosphere/ams_mitm/source/hid_mitm/hid_mitm_service.hpp b/stratosphere/ams_mitm/source/hid_mitm/hid_mitm_service.hpp
index 33bc075d7..22d09f9a3 100644
--- a/stratosphere/ams_mitm/source/hid_mitm/hid_mitm_service.hpp
+++ b/stratosphere/ams_mitm/source/hid_mitm/hid_mitm_service.hpp
@@ -25,11 +25,15 @@ class HidMitmService : public IMitmServiceObject {
         enum class CommandId {
             SetSupportedNpadStyleSet = 100,
         };
+    private:
+        bool should_set_system_ext = false;
     public:
         HidMitmService(std::shared_ptr<Service> s, u64 pid, sts::ncm::TitleId tid) : IMitmServiceObject(s, pid, tid) {
             /* ... */
         }
 
+        ~HidMitmService();
+
         static bool ShouldMitm(u64 pid, sts::ncm::TitleId tid) {
             /* TODO: Consider removing in Atmosphere 0.10.0/1.0.0. */
             /* We will mitm: