From a4bf95a2117d0cbdcbcf9054d72c712aedaf1118 Mon Sep 17 00:00:00 2001
From: Dmitry Fleytman <dfleytma@redhat.com>
Date: Sun, 19 Feb 2017 13:55:24 +0200
Subject: [PATCH] ControlDevice: Generate serial numbers for UsbDk devices

Windows recognizes USB devices by PID/VID/SN combination,
for each new USB device plugged in it generates a new
cached driver information entry in the registry.

When UsbDk hides or redirects a device it creates a virtual
device with UsbDk VID/PID and a generated serial number.

On one hand, this serial number should be unique because
there should be no multiple devices with the same PID/VID/SN
combination at any given moment of time. On the other hand,
UsbDk should re-use serial numbers of unplugged devices,
because we don't want Windows to cache one more driver
information entry each time we redirect of hide a device.

This patch introduces a simple but efficient mechanism
for serial number generation for virtual UsbDk devices.

Signed-off-by: Dmitry Fleytman <dfleytma@redhat.com>
---
 UsbDk/ControlDevice.cpp | 43 +++++++++++++++++++++++++++++++++++++++++
 UsbDk/ControlDevice.h   |  6 ++++++
 2 files changed, 49 insertions(+)

diff --git a/UsbDk/ControlDevice.cpp b/UsbDk/ControlDevice.cpp
index 00c38b4..203f073 100644
--- a/UsbDk/ControlDevice.cpp
+++ b/UsbDk/ControlDevice.cpp
@@ -271,6 +271,49 @@ ULONG CUsbDkControlDevice::CountDevices()
     return numberDevices;
 }
 
+void CUsbDkControlDevice::RegisterHiddenDevice(CUsbDkFilterDevice &FilterDevice)
+{
+    // Windows recognizes USB devices by PID/VID/SN combination,
+    // for each new USB device plugged in it generates a new
+    // cached driver information entry in the registry.
+    //
+    // When UsbDk hides or redirects a device it creates a virtual
+    // device with UsbDk VID / PID and a generated serial number.
+    //
+    // On one hand, this serial number should be unique because
+    // there should be no multiple devices with the same PID/VID/SN
+    // combination at any given moment of time. On the other hand,
+    // UsbDk should re-use serial numbers of unplugged devices,
+    // because we don't want Windows to cache one more driver
+    // information entry each time we redirect of hide a device.
+    //
+    // This function implements a simple but efficient mechanism
+    // for serial number generation for virtual UsbDk devices
+
+    CLockedContext<CWdmSpinLock> Ctx(m_HiddenDevicesLock);
+
+    bool NoSuchSN = false;
+
+    for (ULONG MinSN = 0; !NoSuchSN; MinSN++)
+    {
+        NoSuchSN = m_HiddenDevices.ForEach([MinSN](CUsbDkFilterDevice *Filter)
+                   { return (Filter->GetSerialNumber() != MinSN); });
+
+        if (NoSuchSN)
+        {
+            FilterDevice.SetSerialNumber(MinSN);
+        }
+    }
+
+    m_HiddenDevices.PushBack(&FilterDevice);
+}
+
+void CUsbDkControlDevice::UnregisterHiddenDevice(CUsbDkFilterDevice &FilterDevice)
+{
+    CLockedContext<CWdmSpinLock> Ctx(m_HiddenDevicesLock);
+    m_HiddenDevices.Remove(&FilterDevice);
+}
+
 bool CUsbDkControlDevice::ShouldHide(const USB_DEVICE_DESCRIPTOR &DevDescriptor) const
 {
     auto Hide = false;
diff --git a/UsbDk/ControlDevice.h b/UsbDk/ControlDevice.h
index 28842fc..759479c 100644
--- a/UsbDk/ControlDevice.h
+++ b/UsbDk/ControlDevice.h
@@ -246,6 +246,9 @@ public:
     void UnregisterFilter(CUsbDkFilterDevice &FilterDevice)
     { m_FilterDevices.Remove(&FilterDevice); }
 
+    void RegisterHiddenDevice(CUsbDkFilterDevice &FilterDevice);
+    void UnregisterHiddenDevice(CUsbDkFilterDevice &FilterDevice);
+
     ULONG CountDevices();
     NTSTATUS RescanRegistry()
     { return ReloadPersistentHideRules(); }
@@ -308,6 +311,9 @@ private:
 
     CWdmList<CUsbDkFilterDevice, CLockedAccess, CNonCountingObject, CRefCountingDeleter> m_FilterDevices;
 
+    CWdmList<CUsbDkFilterDevice, CRawAccess, CNonCountingObject, CRefCountingDeleter> m_HiddenDevices;
+    CWdmSpinLock m_HiddenDevicesLock;
+
     typedef CWdmSet<CUsbDkRedirection, CLockedAccess, CNonCountingObject, CRefCountingDeleter> RedirectionsSet;
     RedirectionsSet m_Redirections;