[KBM] Refactor KBM's thread safety logic to avoid mutex re-entrancy bugs and improve performance (#6803)

* Unlock mutex before ResetModifierForLowerLevelKeyHandlers method to avoid crash if two instances of KBM are running

* Added alias for Shortcut DWORD variant to clean up code

* Removed mutex usage in single key remap method and added GetSingleKeyRemap

* Added more alias

* Moved to boolean disable remapping

* Added missing ! in condition

* Remove lock statement from bad auto-merge

src/modules/keyboardmanager/common/KeyboardManagerState.cpp

@@ -8,7 +8,7 @@
 
 // Constructor
 KeyboardManagerState::KeyboardManagerState() :
-    uiState(KeyboardManagerUIState::Deactivated), currentUIWindow(nullptr), currentShortcutUI1(nullptr), currentShortcutUI2(nullptr), currentSingleKeyUI(nullptr), detectedRemapKey(NULL)
+    uiState(KeyboardManagerUIState::Deactivated), currentUIWindow(nullptr), currentShortcutUI1(nullptr), currentShortcutUI2(nullptr), currentSingleKeyUI(nullptr), detectedRemapKey(NULL), remappingsEnabled(true)
 {
     configFile_mutex = CreateMutex(
         NULL, // default security descriptor
@@ -104,7 +104,6 @@ void KeyboardManagerState::ResetUIState()
 // Function to clear the OS Level shortcut remapping table
 void KeyboardManagerState::ClearOSLevelShortcuts()
 {
-    std::lock_guard<std::mutex> lock(osLevelShortcutReMap_mutex);
     osLevelShortcutReMap.clear();
     osLevelShortcutReMapSortedKeys.clear();
 }
@@ -112,23 +111,19 @@ void KeyboardManagerState::ClearOSLevelShortcuts()
 // Function to clear the Keys remapping table.
 void KeyboardManagerState::ClearSingleKeyRemaps()
 {
-    std::lock_guard<std::mutex> lock(singleKeyReMap_mutex);
     singleKeyReMap.clear();
 }
 
 // Function to clear the App specific shortcut remapping table
 void KeyboardManagerState::ClearAppSpecificShortcuts()
 {
-    std::lock_guard<std::mutex> lock(appSpecificShortcutReMap_mutex);
     appSpecificShortcutReMap.clear();
     appSpecificShortcutReMapSortedKeys.clear();
 }
 
 // Function to add a new OS level shortcut remapping
-bool KeyboardManagerState::AddOSLevelShortcut(const Shortcut& originalSC, const std::variant<DWORD, Shortcut>& newSC)
+bool KeyboardManagerState::AddOSLevelShortcut(const Shortcut& originalSC, const KeyShortcutUnion& newSC)
 {
-    std::lock_guard<std::mutex> lock(osLevelShortcutReMap_mutex);
-
     // Check if the shortcut is already remapped
     auto it = osLevelShortcutReMap.find(originalSC);
     if (it != osLevelShortcutReMap.end())
@@ -144,10 +139,8 @@ bool KeyboardManagerState::AddOSLevelShortcut(const Shortcut& originalSC, const
 }
 
 // Function to add a new single key to key/shortcut remapping
-bool KeyboardManagerState::AddSingleKeyRemap(const DWORD& originalKey, const std::variant<DWORD, Shortcut>& newRemapKey)
+bool KeyboardManagerState::AddSingleKeyRemap(const DWORD& originalKey, const KeyShortcutUnion& newRemapKey)
 {
-    std::lock_guard<std::mutex> lock(singleKeyReMap_mutex);
-
     // Check if the key is already remapped
     auto it = singleKeyReMap.find(originalKey);
     if (it != singleKeyReMap.end())
@@ -160,10 +153,8 @@ bool KeyboardManagerState::AddSingleKeyRemap(const DWORD& originalKey, const std
 }
 
 // Function to add a new App specific shortcut remapping
-bool KeyboardManagerState::AddAppSpecificShortcut(const std::wstring& app, const Shortcut& originalSC, const std::variant<DWORD, Shortcut>& newSC)
+bool KeyboardManagerState::AddAppSpecificShortcut(const std::wstring& app, const Shortcut& originalSC, const KeyShortcutUnion& newSC)
 {
-    std::lock_guard<std::mutex> lock(appSpecificShortcutReMap_mutex);
-
     // Convert app name to lower case
     std::wstring process_name;
     process_name.resize(app.length());
@@ -191,6 +182,54 @@ bool KeyboardManagerState::AddAppSpecificShortcut(const std::wstring& app, const
     return true;
 }
 
+// Function to get the iterator of a single key remap given the source key. Returns nullopt if it isn't remapped
+std::optional<SingleKeyRemapTable::iterator> KeyboardManagerState::GetSingleKeyRemap(const DWORD& originalKey)
+{
+    auto it = singleKeyReMap.find(originalKey);
+    if (it != singleKeyReMap.end())
+    {
+        return it;
+    }
+
+    return std::nullopt;
+}
+
+bool KeyboardManagerState::CheckShortcutRemapInvoked(const std::optional<std::wstring>& appName)
+{
+    // Assumes appName exists in the app-specific remap table
+    ShortcutRemapTable& currentRemapTable = appName ? appSpecificShortcutReMap[*appName] : osLevelShortcutReMap;
+    for (auto& it : currentRemapTable)
+    {
+        if (it.second.isShortcutInvoked)
+        {
+            return true;
+        }
+    }
+
+    return false;
+}
+
+std::vector<Shortcut>& KeyboardManagerState::GetSortedShortcutRemapVector(const std::optional<std::wstring>& appName)
+{
+    // Assumes appName exists in the app-specific remap table
+    return appName ? appSpecificShortcutReMapSortedKeys[*appName] : osLevelShortcutReMapSortedKeys;
+}
+
+// Function to get the source and target of a shortcut remap given the source shortcut. Returns nullopt if it isn't remapped
+ShortcutRemapTable& KeyboardManagerState::GetShortcutRemapTable(const std::optional<std::wstring>& appName)
+{
+    if (appName)
+    {
+        auto itTable = appSpecificShortcutReMap.find(*appName);
+        if (itTable != appSpecificShortcutReMap.end())
+        {
+            return itTable->second;
+        }
+    }
+
+    return osLevelShortcutReMap;
+}
+
 // Function to set the textblock of the detect shortcut UI so that it can be accessed by the hook
 void KeyboardManagerState::ConfigureDetectShortcutUI(const StackPanel& textBlock1, const StackPanel& textBlock2)
 {
@@ -468,7 +507,6 @@ bool KeyboardManagerState::SaveConfigToFile()
     json::JsonArray inProcessRemapKeysArray;
     json::JsonArray appSpecificRemapShortcutsArray;
     json::JsonArray globalRemapShortcutsArray;
-    std::unique_lock<std::mutex> lockSingleKeyReMap(singleKeyReMap_mutex);
     for (const auto& it : singleKeyReMap)
     {
         json::JsonObject keys;
@@ -488,9 +526,7 @@ bool KeyboardManagerState::SaveConfigToFile()
 
         inProcessRemapKeysArray.Append(keys);
     }
-    lockSingleKeyReMap.unlock();
 
-    std::unique_lock<std::mutex> lockOsLevelShortcutReMap(osLevelShortcutReMap_mutex);
     for (const auto& it : osLevelShortcutReMap)
     {
         json::JsonObject keys;
@@ -510,9 +546,7 @@ bool KeyboardManagerState::SaveConfigToFile()
 
         globalRemapShortcutsArray.Append(keys);
     }
-    lockOsLevelShortcutReMap.unlock();
 
-    std::unique_lock<std::mutex> lockAppSpecificShortcutReMap(appSpecificShortcutReMap_mutex);
     for (const auto& itApp : appSpecificShortcutReMap)
     {
         // Iterate over apps
@@ -538,7 +572,6 @@ bool KeyboardManagerState::SaveConfigToFile()
             appSpecificRemapShortcutsArray.Append(keys);
         }
     }
-    lockAppSpecificShortcutReMap.unlock();
 
     remapShortcuts.SetNamedValue(KeyboardManagerConstants::GlobalRemapShortcutsSettingName, globalRemapShortcutsArray);
     remapShortcuts.SetNamedValue(KeyboardManagerConstants::AppSpecificRemapShortcutsSettingName, appSpecificRemapShortcutsArray);
@@ -596,3 +629,20 @@ std::wstring KeyboardManagerState::GetActivatedApp()
 {
     return activatedAppSpecificShortcutTarget;
 }
+
+bool KeyboardManagerState::AreRemappingsEnabled()
+{
+    return remappingsEnabled;
+}
+
+void KeyboardManagerState::RemappingsDisabledWrapper(std::function<void()> method)
+{
+    // Disable keyboard remappings
+    remappingsEnabled = false;
+
+    // Run the method which requires the remappings to be disabled
+    method();
+
+    // Re-enable the keyboard remappings
+    remappingsEnabled = true;
+}

src/modules/keyboardmanager/common/KeyboardManagerState.h

@@ -20,6 +20,10 @@ namespace winrt::Windows::UI::Xaml::Controls
     struct StackPanel;
 }
 
+using SingleKeyRemapTable = std::unordered_map<DWORD, KeyShortcutUnion>;
+using ShortcutRemapTable = std::map<Shortcut, RemapShortcut>;
+using AppSpecificShortcutRemapTable = std::map<std::wstring, ShortcutRemapTable>;
+
 // Enum type to store different states of the UI
 enum class KeyboardManagerUIState
 {
@@ -84,6 +88,9 @@ private:
     // Stores the activated target application in app-specfic shortcut
     std::wstring activatedAppSpecificShortcutTarget;
 
+    // Thread safe boolean value to check if remappings are currently enabled. This is used to disable remappings while the remap tables are being updated by the UI thread
+    std::atomic_bool remappingsEnabled;
+
     // Display a key by appending a border Control as a child of the panel.
     void AddKeyToLayout(const winrt::Windows::UI::Xaml::Controls::StackPanel& panel, const winrt::hstring& key);
 
@@ -91,22 +98,21 @@ public:
     // The map members and their mutexes are left as public since the maps are used extensively in dllmain.cpp.
     // Maps which store the remappings for each of the features. The bool fields should be initialized to false. They are used to check the current state of the shortcut (i.e is that particular shortcut currently pressed down or not).
     // Stores single key remappings
-    std::unordered_map<DWORD, std::variant<DWORD, Shortcut>> singleKeyReMap;
-    std::mutex singleKeyReMap_mutex;
+    std::unordered_map<DWORD, KeyShortcutUnion> singleKeyReMap;
 
+    /* This feature has not been enabled (code from proof of concept stage)
+    * 
     // Stores keys which need to be changed from toggle behavior to modifier behavior. Eg. Caps Lock
     std::unordered_map<DWORD, bool> singleKeyToggleToMod;
-    std::mutex singleKeyToggleToMod_mutex;
+    */
 
     // Stores the os level shortcut remappings
-    std::map<Shortcut, RemapShortcut> osLevelShortcutReMap;
+    ShortcutRemapTable osLevelShortcutReMap;
     std::vector<Shortcut> osLevelShortcutReMapSortedKeys;
-    std::mutex osLevelShortcutReMap_mutex;
 
     // Stores the app-specific shortcut remappings. Maps application name to the shortcut map
-    std::map<std::wstring, std::map<Shortcut, RemapShortcut>> appSpecificShortcutReMap;
+    AppSpecificShortcutRemapTable appSpecificShortcutReMap;
     std::map<std::wstring, std::vector<Shortcut>> appSpecificShortcutReMapSortedKeys;
-    std::mutex appSpecificShortcutReMap_mutex;
 
     // Stores the keyboard layout
     LayoutMap keyboardMap;
@@ -139,13 +145,23 @@ public:
     void ClearAppSpecificShortcuts();
 
     // Function to add a new single key to key remapping
-    bool AddSingleKeyRemap(const DWORD& originalKey, const std::variant<DWORD, Shortcut>& newRemapKey);
+    bool AddSingleKeyRemap(const DWORD& originalKey, const KeyShortcutUnion& newRemapKey);
 
     // Function to add a new OS level shortcut remapping
-    bool AddOSLevelShortcut(const Shortcut& originalSC, const std::variant<DWORD, Shortcut>& newSC);
+    bool AddOSLevelShortcut(const Shortcut& originalSC, const KeyShortcutUnion& newSC);
 
     // Function to add a new App specific level shortcut remapping
-    bool AddAppSpecificShortcut(const std::wstring& app, const Shortcut& originalSC, const std::variant<DWORD, Shortcut>& newSC);
+    bool AddAppSpecificShortcut(const std::wstring& app, const Shortcut& originalSC, const KeyShortcutUnion& newSC);
+
+    // Function to get the iterator of a single key remap given the source key. Returns nullopt if it isn't remapped
+    std::optional<SingleKeyRemapTable::iterator> GetSingleKeyRemap(const DWORD& originalKey);
+
+    bool CheckShortcutRemapInvoked(const std::optional<std::wstring>& appName);
+
+    std::vector<Shortcut>& GetSortedShortcutRemapVector(const std::optional<std::wstring>& appName);
+
+    // Function to get the source and target of a shortcut remap given the source shortcut. Returns nullopt if it isn't remapped
+    ShortcutRemapTable& GetShortcutRemapTable(const std::optional<std::wstring>& appName);
 
     // Function to set the textblock of the detect shortcut UI so that it can be accessed by the hook
     void ConfigureDetectShortcutUI(const winrt::Windows::UI::Xaml::Controls::StackPanel& textBlock1, const winrt::Windows::UI::Xaml::Controls::StackPanel& textBlock2);
@@ -214,4 +230,8 @@ public:
 
     // Gets the activated target application in app-specfic shortcut
     std::wstring GetActivatedApp();
+
+    bool AreRemappingsEnabled();
+
+    void RemappingsDisabledWrapper(std::function<void()> method);
 };

src/modules/keyboardmanager/common/RemapShortcut.h

@@ -6,11 +6,11 @@
 class RemapShortcut
 {
 public:
-    std::variant<DWORD, Shortcut> targetShortcut;
+    KeyShortcutUnion targetShortcut;
     bool isShortcutInvoked;
     ModifierKey winKeyInvoked;
 
-    RemapShortcut(const std::variant<DWORD, Shortcut>& sc) :
+    RemapShortcut(const KeyShortcutUnion& sc) :
         targetShortcut(sc), isShortcutInvoked(false), winKeyInvoked(ModifierKey::Disabled)
     {
     }

src/modules/keyboardmanager/common/Shortcut.h

@@ -171,6 +171,7 @@ public:
     KeyboardManagerHelper::ErrorType IsShortcutIllegal() const;
 };
 
-using RemapBufferItem = std::vector<std::variant<DWORD, Shortcut>>;
+using KeyShortcutUnion = std::variant<DWORD, Shortcut>;
+using RemapBufferItem = std::vector<KeyShortcutUnion>;
 using RemapBufferRow = std::pair<RemapBufferItem, std::wstring>;
 using RemapBuffer = std::vector<RemapBufferRow>;