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 . */ +#include #include #include #include "hid_shim.h" #include "hid_mitm_service.hpp" +namespace { + + std::unordered_map 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(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(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 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: