UsbDkHelper/Controller: Update persistent hide rules after modification

Signed-off-by: Dmitry Fleytman <[email protected]>
fengjixuchui · Mar 25, 2015 · 456334a · 456334a
1 parent 2baf64f
commit 456334a
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 43 deletions.
diff --git a/UsbDkController/UsbDkController.cpp b/UsbDkController/UsbDkController.cpp
@@ -53,27 +53,32 @@ static void ShowUsage()
 }
 //-------------------------------------------------------------------------------
 
-static int Controller_InstallDriver()
+static int Controller_AnalyzeInstallResult(InstallResult Res, const tstring &OpName)
 {
-    tcout << TEXT("Installing UsbDk driver") << endl;
-    auto res = UsbDk_InstallDriver();
-    switch (res)
+    switch (Res)
     {
     case InstallSuccess:
-        tcout << TEXT("UsbDk driver installation succeeded") << endl;
+        tcout << OpName << TEXT(" succeeded") << endl;
         return 0;
     case InstallFailure:
-        tcout << TEXT("UsbDk driver installation failed") << endl;
+        tcout << OpName << TEXT(" failed") << endl;
         return 1;
     case InstallSuccessNeedReboot:
-        tcout << TEXT("UsbDk driver installation succeeded but reboot is required in order to make it functional.") << endl;
+        tcout << OpName << TEXT(" succeeded but reboot is required in order to make it functional") << endl;
         return 2;
     default:
-        tcout << TEXT("UsbDk driver installation returned unknown error code") << endl;
+        tcout << OpName << TEXT(" returned unknown error code") << endl;
         assert(0);
         return 3;
     }
 }
+
+static int Controller_InstallDriver()
+{
+    tcout << TEXT("Installing UsbDk driver") << endl;
+    auto res = UsbDk_InstallDriver();
+    return Controller_AnalyzeInstallResult(res, TEXT("UsbDk driver installation"));
+}
 //-------------------------------------------------------------------------------
 
 static int Controller_UninstallDriver()
@@ -234,7 +239,7 @@ static bool Controller_ParseRule(TCHAR *VID, TCHAR *PID, TCHAR *BCD, TCHAR *UsbC
            Controller_ParseBoolRuleField(TEXT("Hide"), Hide, Rule.Hide);
 }
 
-static bool Controller_AddPersistentHideRule(TCHAR *VID, TCHAR *PID, TCHAR *BCD, TCHAR *UsbClass, TCHAR *Hide)
+static int Controller_AddPersistentHideRule(TCHAR *VID, TCHAR *PID, TCHAR *BCD, TCHAR *UsbClass, TCHAR *Hide)
 {
     USB_DK_HIDE_RULE Rule;
 
@@ -244,19 +249,11 @@ static bool Controller_AddPersistentHideRule(TCHAR *VID, TCHAR *PID, TCHAR *BCD,
         return false;
     }
 
-    if (UsbDk_AddPersistentHideRule(&Rule))
-    {
-        tcout << TEXT("Persistent hide rule created succesfully.");
-        return true;
-    }
-    else
-    {
-        tcout << TEXT("Failed to create persistent hide rule.") << endl;
-        return false;
-    }
+    auto res = UsbDk_AddPersistentHideRule(&Rule);
+    return Controller_AnalyzeInstallResult(res, TEXT("Persistent hide rule creation"));
 }
 
-static bool Controller_DeletePersistentHideRule(TCHAR *VID, TCHAR *PID, TCHAR *BCD, TCHAR *UsbClass, TCHAR *Hide)
+static int Controller_DeletePersistentHideRule(TCHAR *VID, TCHAR *PID, TCHAR *BCD, TCHAR *UsbClass, TCHAR *Hide)
 {
     USB_DK_HIDE_RULE Rule;
 
@@ -266,16 +263,8 @@ static bool Controller_DeletePersistentHideRule(TCHAR *VID, TCHAR *PID, TCHAR *B
         return false;
     }
 
-    if (UsbDk_DeletePersistentHideRule(&Rule))
-    {
-        tcout << TEXT("Persistent rule deleted successfully.") << endl;
-        return true;
-    }
-    else
-    {
-        tcout << TEXT("Failed to delete persistent hide rule.") << endl;
-        return false;
-    }
+    auto res = UsbDk_DeletePersistentHideRule(&Rule);
+    return Controller_AnalyzeInstallResult(res, TEXT("Persistent hide rule removal"));
 }
 
 static void Controller_HideDevice(TCHAR *VID, TCHAR *PID, TCHAR *BCD, TCHAR *UsbClass, TCHAR *Hide)
@@ -402,7 +391,7 @@ int __cdecl _tmain(int argc, TCHAR* argv[])
                 ShowUsage();
                 return -4;
             }
-            return Controller_AddPersistentHideRule(argv[2], argv[3], argv[4], argv[5], argv[6]) ? 0 : -100;
+            return Controller_AddPersistentHideRule(argv[2], argv[3], argv[4], argv[5], argv[6]);
         }
         else if (_tcscmp(L"-D", argv[1]) == 0)
         {
@@ -411,7 +400,7 @@ int __cdecl _tmain(int argc, TCHAR* argv[])
                 ShowUsage();
                 return -5;
             }
-            return Controller_DeletePersistentHideRule(argv[2], argv[3], argv[4], argv[5], argv[6]) ? 0 : -100;
+            return Controller_DeletePersistentHideRule(argv[2], argv[3], argv[4], argv[5], argv[6]);
         }
         else
         {

diff --git a/UsbDkHelper/DriverAccess.cpp b/UsbDkHelper/DriverAccess.cpp
@@ -88,6 +88,11 @@ PUSB_CONFIGURATION_DESCRIPTOR UsbDkDriverAccess::GetConfigurationDescriptor(USB_
 }
 //------------------------------------------------------------------------------------------------
 
+void UsbDkDriverAccess::UpdateRegistryParameters()
+{
+    Ioctl(IOCTL_USBDK_UPDATE_REG_PARAMETERS);
+}
+
 HANDLE UsbDkDriverAccess::AddRedirect(USB_DK_DEVICE_ID &DeviceID)
 {
     ULONG64 RedirectorHandle;

diff --git a/UsbDkHelper/DriverAccess.h b/UsbDkHelper/DriverAccess.h
@@ -50,6 +50,7 @@ class UsbDkDriverAccess : public UsbDkDriverFile
 
     void GetDevicesList(PUSB_DK_DEVICE_INFO &DevicesArray, ULONG &NumberDevice);
     PUSB_CONFIGURATION_DESCRIPTOR GetConfigurationDescriptor(USB_DK_CONFIG_DESCRIPTOR_REQUEST &Request, ULONG &Length);
+    void UpdateRegistryParameters();
     static void ReleaseDevicesList(PUSB_DK_DEVICE_INFO DevicesArray);
     static void ReleaseConfigurationDescriptor(PUSB_CONFIGURATION_DESCRIPTOR Descriptor);
 

diff --git a/UsbDkHelper/UsbDkHelper.cpp b/UsbDkHelper/UsbDkHelper.cpp
@@ -318,28 +318,37 @@ void UsbDk_CloseHiderHandle(HANDLE HiderHandle)
 }
 
 static inline
-BOOL ModifyPersistentHideRules(const USB_DK_HIDE_RULE &Rule,
-                               void(CRulesManager::*Modifier)(const USB_DK_HIDE_RULE&))
+InstallResult ModifyPersistentHideRules(const USB_DK_HIDE_RULE &Rule,
+                                        void(CRulesManager::*Modifier)(const USB_DK_HIDE_RULE&))
 {
     try
     {
         CRulesManager Manager;
-        (Manager.*Modifier)(Rule);
-        return TRUE;
+        (CRulesManager().*Modifier)(Rule);
+
+        UsbDkDriverAccess driver;
+        driver.UpdateRegistryParameters();
+
+        return InstallSuccess;
+    }
+    catch (const UsbDkDriverFileException &e)
+    {
+        printExceptionString(e.what());
+        return InstallSuccessNeedReboot;
     }
     catch (const exception &e)
     {
         printExceptionString(e.what());
-        return FALSE;
+        return InstallFailure;
     }
 }
 
-DLL BOOL UsbDk_AddPersistentHideRule(PUSB_DK_HIDE_RULE Rule)
+DLL InstallResult UsbDk_AddPersistentHideRule(PUSB_DK_HIDE_RULE Rule)
 {
     return ModifyPersistentHideRules(*Rule, &CRulesManager::AddRule);
 }
 
-DLL BOOL UsbDk_DeletePersistentHideRule(PUSB_DK_HIDE_RULE Rule)
+DLL InstallResult UsbDk_DeletePersistentHideRule(PUSB_DK_HIDE_RULE Rule)
 {
     return ModifyPersistentHideRules(*Rule, &CRulesManager::DeleteRule);
 }
diff --git a/UsbDkHelper/UsbDkHelperHider.h b/UsbDkHelper/UsbDkHelperHider.h
@@ -117,7 +117,7 @@ extern "C" {
     *    OUT - None
     *
     * @return
-    *  TRUE if function succeeds
+    *  Rule installation status
     *
     * @note
     * 1. Persistent rule stays until explicitly deleted by
@@ -127,7 +127,7 @@ extern "C" {
     *    device re-plug or system reboot.
     *
     */
-    DLL BOOL             UsbDk_AddPersistentHideRule(PUSB_DK_HIDE_RULE Rule);
+    DLL InstallResult    UsbDk_AddPersistentHideRule(PUSB_DK_HIDE_RULE Rule);
 
     /* Delete specific persistent hide rule
     *
@@ -136,15 +136,15 @@ extern "C" {
     *    OUT - None
     *
     * @return
-    *  TRUE if function succeeds
+    *  Rule removal status
     *
     * @note
     * 1. This API requires administrative privileges
     * 2. For already attached devices the rule will be applied after
     *    device re-plug or system reboot.
     *
     */
-    DLL BOOL             UsbDk_DeletePersistentHideRule(PUSB_DK_HIDE_RULE Rule);
+    DLL InstallResult    UsbDk_DeletePersistentHideRule(PUSB_DK_HIDE_RULE Rule);
 #ifdef __cplusplus
 }
 #endif