usbclerk: support driver removal

2 files changed, 91 insertions, 14 deletions

diff --git a/usbclerk.cpp b/usbclerk.cpp
index 27435e3..5ef4160 100644
--- a/usbclerk.cpp
+++ b/usbclerk.cpp
@@ -1,4 +1,6 @@
 #include <windows.h>
+#include <setupapi.h>
+#include <cfgmgr32.h>

 #include <stdio.h>
 #include <string.h>
 #include <tchar.h>
@@ -33,6 +35,9 @@ private:
     bool execute();
     bool dispatch_message(CHAR *buffer, DWORD bytes, USBClerkReply *reply);
     bool install_winusb_driver(int vid, int pid);
+    bool remove_winusb_driver(int vid, int pid);
+    bool remove_dev(HDEVINFO devs, PSP_DEVINFO_DATA dev_info);

+    bool rescan();

     static DWORD WINAPI control_handler(DWORD control, DWORD event_type,
                                         LPVOID event_data, LPVOID context);
     static VOID WINAPI main(DWORD argc, TCHAR * argv[]);
@@ -320,16 +325,21 @@ bool USBClerk::dispatch_message(CHAR *buffer, DWORD bytes, USBClerkReply *reply)
     switch (hdr->type) {
     case USB_CLERK_DRIVER_INSTALL:
         vd_printf("Installing winusb driver for %04x:%04x", dev->vid, dev->pid);
-        if (reply->status = install_winusb_driver(dev->vid, dev->pid)) {
-            vd_printf("winusb driver install succeed");
-        } else {
-            vd_printf("winusb driver install failed");
-        }
+        reply->status = install_winusb_driver(dev->vid, dev->pid);
+        break;
+    case USB_CLERK_DRIVER_REMOVE:
+        vd_printf("Removing winusb driver for %04x:%04x", dev->vid, dev->pid);
+        reply->status = remove_winusb_driver(dev->vid, dev->pid);        
         break;
     default:
         vd_printf("Unknown message received, type %u", hdr->type);
         return false;
     }
+    if (reply->status) {
+        vd_printf("Completed successfully");
+    } else {
+        vd_printf("Failed");
+    }
     return true;
 }
 
@@ -417,6 +427,67 @@ cleanup:
     return installed;
 }
 
+bool USBClerk::remove_winusb_driver(int vid, int pid)
+{
+    HDEVINFO devs;
+    DWORD dev_index;

+    SP_DEVINFO_DATA dev_info;

+    SP_DEVINFO_LIST_DETAIL_DATA dev_info_list_detail;

+    TCHAR dev_prefix[MAX_DEVICE_ID_LEN];
+    bool ret = false;

+
+    devs = SetupDiGetClassDevsEx(NULL, NULL, NULL, DIGCF_ALLCLASSES, NULL, NULL, NULL);

+    if (devs == INVALID_HANDLE_VALUE) {
+        vd_printf("SetupDiGetClassDevsEx failed: %u\n", GetLastError());
+        return false;
+    }
+

+    dev_info_list_detail.cbSize = sizeof(dev_info_list_detail);

+    if (!SetupDiGetDeviceInfoListDetail(devs, &dev_info_list_detail)) {

+        vd_printf("SetupDiGetDeviceInfoListDetail failed: %u\n", GetLastError());
+        SetupDiDestroyDeviceInfoList(devs);
+        return false;
+    }

+
+    swprintf(dev_prefix, MAX_DEVICE_ID_LEN, L"USB\\VID_%04x&PID_%04x\\", vid, pid);
+    dev_info.cbSize = sizeof(dev_info);

+    for (dev_index = 0; SetupDiEnumDeviceInfo(devs, dev_index, &dev_info); dev_index++) {

+        TCHAR devID[MAX_DEVICE_ID_LEN];

+

+        if (CM_Get_Device_ID_Ex(dev_info.DevInst, devID, MAX_DEVICE_ID_LEN, 0,

+                                dev_info_list_detail.RemoteMachineHandle) == CR_SUCCESS &&

+                                wcsstr(devID, dev_prefix)) {

+            vd_printf("Removing %S\n", devID);

+            ret = remove_dev(devs, &dev_info);

+            break;

+        }

+    }

+    SetupDiDestroyDeviceInfoList(devs);
+    ret = ret && rescan();

+    return ret;
+}
+
+bool USBClerk::remove_dev(HDEVINFO devs, PSP_DEVINFO_DATA dev_info)

+{

+    SP_REMOVEDEVICE_PARAMS rmd_params;

+

+    rmd_params.ClassInstallHeader.cbSize = sizeof(SP_CLASSINSTALL_HEADER);

+    rmd_params.ClassInstallHeader.InstallFunction = DIF_REMOVE;

+    rmd_params.Scope = DI_REMOVEDEVICE_GLOBAL;

+    rmd_params.HwProfile = 0;

+    return (SetupDiSetClassInstallParams(devs, dev_info,

+                &rmd_params.ClassInstallHeader, sizeof(rmd_params)) &&

+            SetupDiCallClassInstaller(DIF_REMOVE, devs, dev_info));

+}
+
+bool USBClerk::rescan()

+{

+    DEVINST dev_root;

+

+    return (CM_Locate_DevNode_Ex(&dev_root, NULL, CM_LOCATE_DEVNODE_NORMAL, NULL) == CR_SUCCESS &&

+            CM_Reenumerate_DevNode_Ex(dev_root, 0, NULL) == CR_SUCCESS);

+}
+
 int _tmain(int argc, TCHAR* argv[], TCHAR* envp[])
 {
     bool success = false;
diff --git a/usbclerktest.cpp b/usbclerktest.cpp
index 0c2189a..712fe96 100644
--- a/usbclerktest.cpp
+++ b/usbclerktest.cpp
@@ -5,15 +5,15 @@
 int _tmain(int argc, TCHAR* argv[], TCHAR* envp[])
 {
     HANDLE pipe;
-    USBClerkDriverInstall dev = {{USB_CLERK_MAGIC, USB_CLERK_VERSION, 
-        USB_CLERK_DRIVER_INSTALL, sizeof(USBClerkDriverInstall)}};
+    USBClerkDriverOp dev = {{USB_CLERK_MAGIC, USB_CLERK_VERSION, 0, sizeof(USBClerkDriverOp)}};
     USBClerkReply reply;
     DWORD pipe_mode;
     DWORD bytes = 0;
+    bool do_remove = false;
 
-    if (argc < 3 || !swscanf_s(argv[1], L"%hx", &dev.vid) ||
-                    !swscanf_s(argv[2], L"%hx", &dev.pid)) {
-        printf("Use: usbclerktest VID PID\n");
+    if (argc < 2 || swscanf_s(argv[argc - 1], L"%hx:%hx", &dev.vid, &dev.pid) < 2 ||
+                               (argc == 3 && !(do_remove = !wcscmp(argv[1], L"/u")))) {
+        printf("Usage: usbclerktest [/u] vid:pid\n/u - uninstall driver\nvid:pid in hex\n");
         return 1;
     }
     pipe = CreateFile(USB_CLERK_PIPE_NAME, GENERIC_READ | GENERIC_WRITE,
@@ -27,7 +27,13 @@ int _tmain(int argc, TCHAR* argv[], TCHAR* envp[])
         printf("SetNamedPipeHandleState() failed: %d\n", GetLastError());
         return 1;
     }
-    printf("Signing & installing %04x:%04x\n", dev.vid, dev.pid);
+    if (do_remove) {
+        printf("Removing %04x:%04x\n", dev.vid, dev.pid);
+        dev.hdr.type = USB_CLERK_DRIVER_REMOVE;
+    } else {
+        printf("Signing & installing %04x:%04x\n", dev.vid, dev.pid);
+        dev.hdr.type = USB_CLERK_DRIVER_INSTALL;
+    }
     if (!TransactNamedPipe(pipe, &dev, sizeof(dev), &reply, sizeof(reply), &bytes, NULL)) {
         printf("TransactNamedPipe() failed: %d\n", GetLastError());
         CloseHandle(pipe);
@@ -36,14 +42,14 @@ int _tmain(int argc, TCHAR* argv[], TCHAR* envp[])
     CloseHandle(pipe);
     if (reply.hdr.magic != USB_CLERK_MAGIC || reply.hdr.type != USB_CLERK_REPLY ||
             reply.hdr.size != sizeof(USBClerkReply)) {
-        printf("Unknown message received, magic 0x%x type %u size %u",
+        printf("Unknown message received, magic 0x%x type %u size %u\n",
                reply.hdr.magic, reply.hdr.type, reply.hdr.size);
         return 1;
     }
     if (reply.status) {
-        printf("winusb driver install succeed");
+        printf("Completed successfully\n");
     } else {
-        printf("winusb driver install failed");
+        printf("Failed\n");
     }
     return 0;
 }