ControlDevice: Split dynamic and persistent hide rules lists

Required to properly support removal of persistent hide rule

Signed-off-by: Dmitry Fleytman <[email protected]>

Author: Dmitry Fleytman

UsbDk/ControlDevice.cpp

@@ -238,7 +238,7 @@ bool CUsbDkControlDevice::ShouldHide(const USB_DEVICE_DESCRIPTOR &DevDescriptor)
 {
     auto Hide = false;
 
-    const_cast<HideRulesSet*>(&m_HideRules)->ForEach([&DevDescriptor, &Hide](CUsbDkHideRule *Entry) -> bool
+    const auto &HideVisitor = [&DevDescriptor, &Hide](CUsbDkHideRule *Entry) -> bool
     {
         if (Entry->Match(DevDescriptor))
         {
@@ -247,7 +247,10 @@ bool CUsbDkControlDevice::ShouldHide(const USB_DEVICE_DESCRIPTOR &DevDescriptor)
         }
 
         return true;
-    });
+    };
+
+    const_cast<HideRulesSet*>(&m_HideRules)->ForEach(HideVisitor);
+    const_cast<HideRulesSet*>(&m_PersistentHideRules)->ForEach(HideVisitor);
 
     return Hide;
 }
@@ -424,7 +427,7 @@ NTSTATUS CUsbDkControlDevice::Register()
     if (NT_SUCCESS(status))
     {
         FinishInitializing();
-        LoadPersistentHideRules();
+        ReloadPersistentHideRules();
     }
 
     return status;
@@ -637,7 +640,7 @@ NTSTATUS CUsbDkControlDevice::AddRedirect(const USB_DK_DEVICE_ID &DeviceId, PHAN
     return STATUS_SUCCESS;
 }
 
-NTSTATUS CUsbDkControlDevice::AddHideRule(const USB_DK_HIDE_RULE &UsbDkRule)
+NTSTATUS CUsbDkControlDevice::AddHideRuleToSet(const USB_DK_HIDE_RULE &UsbDkRule, HideRulesSet &Set)
 {
     TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_CONTROLDEVICE, "%!FUNC! entry");
 
@@ -656,14 +659,14 @@ NTSTATUS CUsbDkControlDevice::AddHideRule(const USB_DK_HIDE_RULE &UsbDkRule)
         return STATUS_INSUFFICIENT_RESOURCES;
     }
 
-    if(!m_HideRules.Add(NewRule))
+    if(!Set.Add(NewRule))
     {
         TraceEvents(TRACE_LEVEL_ERROR, TRACE_CONTROLDEVICE, "%!FUNC! failed. Hide rule already present.");
         return STATUS_OBJECT_NAME_COLLISION;
     }
 
     TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_CONTROLDEVICE, "%!FUNC! Current hide rules:");
-    m_HideRules.Dump();
+    Set.Dump();
 
     NewRule.detach();
     return STATUS_SUCCESS;
@@ -672,9 +675,7 @@ NTSTATUS CUsbDkControlDevice::AddHideRule(const USB_DK_HIDE_RULE &UsbDkRule)
 void CUsbDkControlDevice::ClearHideRules()
 {
     m_HideRules.Clear();
-    TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_CONTROLDEVICE, "%!FUNC! All hide rules dropped.");
-    TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_CONTROLDEVICE, "%!FUNC! Reloading persistent hide rules.");
-    LoadPersistentHideRules();
+    TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_CONTROLDEVICE, "%!FUNC! All dynamic hide rules dropped.");
 }
 
 class CHideRulesRegKey final : public CRegKey
@@ -825,8 +826,10 @@ class CRegHideRule final : private CRegKey
     }
 };
 
-NTSTATUS CUsbDkControlDevice::LoadPersistentHideRules()
+NTSTATUS CUsbDkControlDevice::ReloadPersistentHideRules()
 {
+    m_PersistentHideRules.Clear();
+
     CHideRulesRegKey RulesKey;
     auto status = RulesKey.Open();
     if (NT_SUCCESS(status))
@@ -839,7 +842,7 @@ NTSTATUS CUsbDkControlDevice::LoadPersistentHideRules()
             if (NT_SUCCESS(Rule.Open(RulesKey, *Name)) &&
                 NT_SUCCESS(Rule.Read(ParsedRule)))
             {
-                AddHideRule(ParsedRule);
+                AddPersistentHideRule(ParsedRule);
             }
         });
     }

UsbDk/ControlDevice.h

@@ -218,12 +218,17 @@ class CUsbDkControlDevice : private CWdfControlDevice, public CAllocatable<NonPa
 
     ULONG CountDevices();
     NTSTATUS RescanRegistry()
-    { return LoadPersistentHideRules(); }
+    { return ReloadPersistentHideRules(); }
 
     bool EnumerateDevices(USB_DK_DEVICE_INFO *outBuff, size_t numberAllocatedDevices, size_t &numberExistingDevices);
     NTSTATUS ResetUsbDevice(const USB_DK_DEVICE_ID &DeviceId);
     NTSTATUS AddRedirect(const USB_DK_DEVICE_ID &DeviceId, PHANDLE ObjectHandle);
-    NTSTATUS AddHideRule(const USB_DK_HIDE_RULE &UsbDkRule);
+
+    NTSTATUS AddHideRule(const USB_DK_HIDE_RULE &UsbDkRule)
+    { return AddHideRuleToSet(UsbDkRule, m_HideRules); }
+    NTSTATUS AddPersistentHideRule(const USB_DK_HIDE_RULE &UsbDkRule)
+    { return AddHideRuleToSet(UsbDkRule, m_PersistentHideRules); }
+
     void ClearHideRules();
 
     NTSTATUS RemoveRedirect(const USB_DK_DEVICE_ID &DeviceId);
@@ -261,7 +266,7 @@ class CUsbDkControlDevice : private CWdfControlDevice, public CAllocatable<NonPa
     bool WaitForDetachment(const USB_DK_DEVICE_ID &ID);
 
 private:
-    NTSTATUS LoadPersistentHideRules();
+    NTSTATUS ReloadPersistentHideRules();
 
     CObjHolder<CUsbDkControlDeviceQueue> m_DeviceQueue;
     static CRefCountingHolder<CUsbDkControlDevice> *m_UsbDkControlDevice;
@@ -275,6 +280,9 @@ class CUsbDkControlDevice : private CWdfControlDevice, public CAllocatable<NonPa
 
     typedef CWdmSet<CUsbDkHideRule, CLockedAccess, CNonCountingObject> HideRulesSet;
     HideRulesSet m_HideRules;
+    HideRulesSet m_PersistentHideRules;
+
+    NTSTATUS AddHideRuleToSet(const USB_DK_HIDE_RULE &UsbDkRule, HideRulesSet &Set);
 
     template <typename TPredicate, typename TFunctor>
     bool UsbDevicesForEachIf(TPredicate Predicate, TFunctor Functor)