From 89ccde78c49743c90782e594c4dad923c2d4cdf7 Mon Sep 17 00:00:00 2001
From: goldendeng <907246009@qq.com>
Date: Mon, 23 Apr 2018 09:28:12 +0800
Subject: [PATCH] use mutex to handle driver file access

---
 UsbDkHelper/DriverAccess.h  |  2 +-
 UsbDkHelper/DriverFile.cpp  | 15 ++++++++++++++-
 UsbDkHelper/DriverFile.h    | 30 +++++++++++++++++++++++++++---
 UsbDkHelper/UsbDkHelper.cpp |  5 ++++-
 4 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/UsbDkHelper/DriverAccess.h b/UsbDkHelper/DriverAccess.h
index 2b407d2..7d444e7 100644
--- a/UsbDkHelper/DriverAccess.h
+++ b/UsbDkHelper/DriverAccess.h
@@ -43,7 +43,7 @@ class UsbDkDriverAccess : public UsbDkDriverFile
 {
 public:
     UsbDkDriverAccess()
-        : UsbDkDriverFile(USBDK_USERMODE_NAME)
+        : UsbDkDriverFile(USBDK_USERMODE_NAME, false, true)
     {}
 
     void GetDevicesList(PUSB_DK_DEVICE_INFO &DevicesArray, ULONG &NumberDevice);
diff --git a/UsbDkHelper/DriverFile.cpp b/UsbDkHelper/DriverFile.cpp
index 3f3991d..868cff1 100644
--- a/UsbDkHelper/DriverFile.cpp
+++ b/UsbDkHelper/DriverFile.cpp
@@ -25,9 +25,15 @@
 #include "DriverFile.h"
 
 
-UsbDkDriverFile::UsbDkDriverFile(LPCTSTR lpFileName, bool bOverlapped)
+static MutexT g_mutex;
+UsbDkDriverFile::UsbDkDriverFile(LPCTSTR lpFileName, bool bOverlapped, bool bLocked)
 {
     m_bOverlapped = bOverlapped;
+    m_locked = bLocked;
+
+    if (bLocked)
+        g_mutex.lock();
+
 
     m_hDriver = CreateFile(lpFileName,
                            GENERIC_READ | GENERIC_WRITE,
@@ -43,6 +49,13 @@ UsbDkDriverFile::UsbDkDriverFile(LPCTSTR lpFileName, bool bOverlapped)
     }
 }
 
+UsbDkDriverFile::~UsbDkDriverFile()
+{
+    CloseHandle(m_hDriver);
+    if (m_locked)
+        g_mutex.unlock();
+}
+
 TransferResult UsbDkDriverFile::Ioctl(DWORD Code,
                             bool ShortBufferOk,
                             LPVOID InBuffer,
diff --git a/UsbDkHelper/DriverFile.h b/UsbDkHelper/DriverFile.h
index fb72450..712dd6b 100644
--- a/UsbDkHelper/DriverFile.h
+++ b/UsbDkHelper/DriverFile.h
@@ -27,6 +27,29 @@
 
 #define DRIVER_FILE_EXCEPTION_STRING TEXT("Driver file operation error. ")
 
+class MutexT {
+public:
+	MutexT() {
+		_usbdk_mutex = OpenMutex(MUTEX_ALL_ACCESS, FALSE, TEXT("usbdk_mutex"));
+		if (_usbdk_mutex == NULL)
+			_usbdk_mutex = CreateMutex(NULL, FALSE, TEXT("usbdk_mutex"));
+	}
+	~MutexT() {
+		CloseHandle(_usbdk_mutex);
+	}
+	void lock() {
+		WaitForSingleObject(_usbdk_mutex, INFINITE);
+	}
+	void unlock() {
+		ReleaseMutex(_usbdk_mutex);
+	}
+private:
+	void* _usbdk_mutex;
+	// no copy
+	MutexT(const MutexT&);
+	void operator=(const MutexT&);
+};
+
 class UsbDkDriverFileException : public UsbDkW32ErrorException
 {
 public:
@@ -39,12 +62,11 @@ class UsbDkDriverFileException : public UsbDkW32ErrorException
 class UsbDkDriverFile
 {
 public:
-    UsbDkDriverFile(LPCTSTR lpFileName, bool bOverlapped = false);
+    UsbDkDriverFile(LPCTSTR lpFileName, bool bOverlapped = false, bool bLocked = false);
     UsbDkDriverFile(HANDLE ObjectHandle, bool bOverlapped = false)
                     : m_bOverlapped(bOverlapped), m_hDriver(ObjectHandle)
     {}
-    virtual ~UsbDkDriverFile()
-    { CloseHandle(m_hDriver); }
+	virtual ~UsbDkDriverFile();
 
     TransferResult Ioctl(DWORD Code,
                bool ShortBufferOk = false,
@@ -70,4 +92,6 @@ class UsbDkDriverFile
 
 private:
     bool m_bOverlapped;
+	bool m_locked;
+	//mutex_t _mutex;
 };
diff --git a/UsbDkHelper/UsbDkHelper.cpp b/UsbDkHelper/UsbDkHelper.cpp
index f9f66d3..c8d2490 100644
--- a/UsbDkHelper/UsbDkHelper.cpp
+++ b/UsbDkHelper/UsbDkHelper.cpp
@@ -173,7 +173,7 @@ BOOL UsbDk_StopRedirect(HANDLE DeviceHandle)
 {
     try
     {
-        UsbDkDriverAccess driverAccess;
+        //UsbDkDriverAccess driverAccess;
         unique_ptr<REDIRECTED_DEVICE_HANDLE> deviceHandle(unpackHandle<REDIRECTED_DEVICE_HANDLE>(DeviceHandle));
         deviceHandle->RedirectorAccess.reset();
         return TRUE;
@@ -289,6 +289,7 @@ HANDLE UsbDk_GetRedirectorSystemHandle(HANDLE DeviceHandle)
 
 HANDLE UsbDk_CreateHiderHandle()
 {
+	OutputDebugStringA("UsbDk_CreateHiderHandle");
     try
     {
         unique_ptr<UsbDkHiderAccess> hiderAccess(new UsbDkHiderAccess);
@@ -372,10 +373,12 @@ InstallResult ModifyPersistentHideRules(const USB_DK_HIDE_RULE &Rule,
 
 DLL InstallResult UsbDk_AddPersistentHideRule(PUSB_DK_HIDE_RULE Rule)
 {
+	OutputDebugStringA("UsbDk_AddPersistentHideRule");
     return ModifyPersistentHideRules(*Rule, &CRulesManager::AddRule);
 }
 
 DLL InstallResult UsbDk_DeletePersistentHideRule(PUSB_DK_HIDE_RULE Rule)
 {
+	OutputDebugStringA("UsbDk_DeletePersistentHideRule");
     return ModifyPersistentHideRules(*Rule, &CRulesManager::DeleteRule);
 }