usbclerk: add temporary driver install per session

-add message type USB_CLERK_DRIVER_SESSION_INSTALL, bump version to 3
-handle multiple pipe connections concurrently with thread for each one
-keep pipe open, and on pipe close - cleanup the session-specific devices
-add test for temporary driver install and multiple devices

v2: fixed most of Uri's comments:
-remove the device from the list upon USB_CLERK_DRIVER_REMOVE
 e.g. prevent the following scenario:
 client1 install+uninstall, client2 install, client1 disconnect
-limit the number of concurrent connected clients to 32
-according to the request type, set device auto_remove flag
 for driver uninstall upon disconnect

FIXME: prevent uninstall for a device used by other clients

rhbz#845216

3 files changed, 132 insertions, 63 deletions

diff --git a/usbclerk.cpp b/usbclerk.cpp
index 9ce12b8..f03dd7d 100644
--- a/usbclerk.cpp
+++ b/usbclerk.cpp
@@ -4,6 +4,7 @@
 #include <stdio.h>
 #include <string.h>
 #include <tchar.h>
+#include <list>
 #include "usbclerk.h"
 #include "usbredirfilter.h"
 #include "libwdi.h"
@@ -18,6 +19,7 @@
 #define USB_CLERK_LOG_PATH          TEXT("%susbclerk.log")
 #define USB_CLERK_PIPE_TIMEOUT      10000
 #define USB_CLERK_PIPE_BUF_SIZE     1024
+#define USB_CLERK_PIPE_MAX_CLIENTS  32
 #define USB_DRIVER_PATH             "%S\\wdi_usb_driver"
 #define USB_DRIVER_INFNAME_LEN      64
 #define USB_DRIVER_INSTALL_RETRIES  10
@@ -26,6 +28,14 @@
 #define MAX_DEVICE_HCID_LEN         1024
 #define MAX_DEVICE_FILTER_LEN       1024
 
+typedef struct USBDev {
+    UINT16 vid;
+    UINT16 pid;
+    bool auto_remove;
+} USBDev;
+
+typedef std::list<USBDev> USBDevs;
+
 class USBClerk {
 public:
     static USBClerk* get();
@@ -37,7 +47,7 @@ public:
 private:
     USBClerk();
     bool execute();
-    bool dispatch_message(CHAR *buffer, DWORD bytes, USBClerkReply *reply);
+    bool dispatch_message(CHAR *buffer, DWORD bytes, USBClerkReply *reply, USBDevs *devs);
     bool install_winusb_driver(int vid, int pid);
     bool remove_winusb_driver(int vid, int pid);
     bool uninstall_inf(HDEVINFO devs, PSP_DEVINFO_DATA dev_info);
@@ -51,6 +61,7 @@ private:
     bool dev_filter_check(int vid, int pid, bool *has_winusb);
     static DWORD WINAPI control_handler(DWORD control, DWORD event_type,
                                         LPVOID event_data, LPVOID context);
+    static DWORD WINAPI pipe_thread(LPVOID param);
     static VOID WINAPI main(DWORD argc, TCHAR * argv[]);
 
 private:
@@ -60,7 +71,6 @@ private:
     struct usbredirfilter_rule *_filter_rules;
     int _filter_count;
     char _wdi_path[MAX_PATH];
-    HANDLE _pipe;
     bool _running;
     VDLog* _log;
 };
@@ -276,11 +286,9 @@ bool USBClerk::execute()
 {
     SECURITY_ATTRIBUTES sec_attr;
     SECURITY_DESCRIPTOR* sec_desr;
-    USBClerkReply reply = {{USB_CLERK_MAGIC, USB_CLERK_VERSION,
-        USB_CLERK_REPLY, sizeof(USBClerkReply)}};
     CHAR filter_str[MAX_DEVICE_FILTER_LEN];
-    CHAR buffer[USB_CLERK_PIPE_BUF_SIZE];
-    DWORD bytes;
+    HANDLE pipe, thread;
+    DWORD tid;
     HKEY hkey;
     LONG ret;
 
@@ -296,14 +304,6 @@ bool USBClerk::execute()
     sec_attr.nLength = sizeof(sec_attr);
     sec_attr.bInheritHandle = TRUE;
     sec_attr.lpSecurityDescriptor = sec_desr;
-    _pipe = CreateNamedPipe(USB_CLERK_PIPE_NAME, PIPE_ACCESS_DUPLEX |
-                            FILE_FLAG_FIRST_PIPE_INSTANCE,
-                            PIPE_TYPE_MESSAGE | PIPE_READMODE_MESSAGE | PIPE_WAIT, 1,
-                            USB_CLERK_PIPE_BUF_SIZE, USB_CLERK_PIPE_BUF_SIZE, 0, &sec_attr);
-    if (_pipe == INVALID_HANDLE_VALUE) {
-        vd_printf("CreatePipe() failed: %u", GetLastError());
-        return false;
-    }
 
     /* Read filter rules from registry */
     ret = RegOpenKeyEx(HKEY_LOCAL_MACHINE, L"Software\\USBClerk", 0, KEY_READ, &hkey);
@@ -323,34 +323,62 @@ bool USBClerk::execute()
         RegCloseKey(hkey);
     }
     while (_running) {
-        if (!ConnectNamedPipe(_pipe, NULL) && GetLastError() != ERROR_PIPE_CONNECTED) {
-            vd_printf("ConnectNamedPipe() failed: %u", GetLastError());
+        pipe = CreateNamedPipe(USB_CLERK_PIPE_NAME, PIPE_ACCESS_DUPLEX,
+                               PIPE_TYPE_MESSAGE | PIPE_READMODE_MESSAGE | PIPE_WAIT,
+                               USB_CLERK_PIPE_MAX_CLIENTS, USB_CLERK_PIPE_BUF_SIZE,
+                               USB_CLERK_PIPE_BUF_SIZE, 0, &sec_attr);
+        if (pipe == INVALID_HANDLE_VALUE) {
+            vd_printf("CreatePipe() failed: %u", GetLastError());
             break;
         }
-        if (!ReadFile(_pipe, &buffer, sizeof(buffer), &bytes, NULL)) {
-            vd_printf("ReadFile() failed: %d", GetLastError());
-            goto disconnect;
-        }
-        if (!dispatch_message(buffer, bytes, &reply)) {
-            goto disconnect;
+        if (!ConnectNamedPipe(pipe, NULL) && GetLastError() != ERROR_PIPE_CONNECTED) {
+            vd_printf("ConnectNamedPipe() failed: %u", GetLastError());
+            CloseHandle(pipe);
+            break;
         }
-        if (!WriteFile(_pipe, &reply, sizeof(reply), &bytes, NULL)) {
-            vd_printf("WriteFile() failed: %d", GetLastError());
-            goto disconnect;
+        thread = CreateThread(NULL, 0, pipe_thread, (LPVOID)pipe, 0, &tid);
+        if (thread == NULL) {
+            vd_printf("CreateThread() failed: %u", GetLastError());
+            break;
         }
-        FlushFileBuffers(_pipe);
-disconnect:
-        DisconnectNamedPipe(_pipe);
+        CloseHandle(thread);
     }
     free(_filter_rules);
-    CloseHandle(_pipe);
     return true;
 }
 
-bool USBClerk::dispatch_message(CHAR *buffer, DWORD bytes, USBClerkReply *reply)
+DWORD WINAPI USBClerk::pipe_thread(LPVOID param)
+{
+    USBClerkReply reply = {{USB_CLERK_MAGIC, USB_CLERK_VERSION,
+        USB_CLERK_REPLY, sizeof(USBClerkReply)}};
+    CHAR buffer[USB_CLERK_PIPE_BUF_SIZE];
+    HANDLE pipe = (HANDLE)param;
+    USBClerk* usbclerk = get();
+    USBDevs devs;
+    DWORD bytes;
+
+    while (usbclerk->_running) {
+        if (!ReadFile(pipe, &buffer, sizeof(buffer), &bytes, NULL) ||
+            !usbclerk->dispatch_message(buffer, bytes, &reply, &devs) ||
+            !WriteFile(pipe, &reply, sizeof(reply), &bytes, NULL)) {
+            break;
+        }
+        FlushFileBuffers(pipe);
+    }
+    DisconnectNamedPipe(pipe);
+    CloseHandle(pipe);
+    for (USBDevs::iterator dev = devs.begin(); dev != devs.end(); dev++) {
+        if (dev->auto_remove) {
+            usbclerk->remove_winusb_driver(dev->vid, dev->pid);
+        }
+    }
+    return 0;
+}
+
+bool USBClerk::dispatch_message(CHAR *buffer, DWORD bytes, USBClerkReply *reply, USBDevs *devs)
 {
     USBClerkHeader *hdr = (USBClerkHeader *)buffer;
-    USBClerkDriverOp *dev;
+    USBClerkDriverOp *op;
 
     if (hdr->magic != USB_CLERK_MAGIC) {
         vd_printf("Bad message received, magic %u", hdr->magic);
@@ -360,15 +388,28 @@ bool USBClerk::dispatch_message(CHAR *buffer, DWORD bytes, USBClerkReply *reply)
         vd_printf("Wrong mesage size %u type %u", hdr->size, hdr->type);
         return false;
     }
-    dev = (USBClerkDriverOp *)buffer;
+    op = (USBClerkDriverOp *)buffer;
     switch (hdr->type) {
+    case USB_CLERK_DRIVER_SESSION_INSTALL:
     case USB_CLERK_DRIVER_INSTALL:
-        vd_printf("Installing winusb driver for %04x:%04x", dev->vid, dev->pid);
-        reply->status = install_winusb_driver(dev->vid, dev->pid);
+        vd_printf("Installing winusb driver for %04x:%04x", op->vid, op->pid);
+        reply->status = install_winusb_driver(op->vid, op->pid);
+        if (reply->status) {
+            USBDev dev = {op->vid, op->pid, hdr->type == USB_CLERK_DRIVER_SESSION_INSTALL};
+            devs->push_back(dev);
+        }
         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);        
+        // FIXME: check device is not used by another client
+        vd_printf("Removing winusb driver for %04x:%04x", op->vid, op->pid);
+        reply->status = remove_winusb_driver(op->vid, op->pid);
+        // remove device from list to prevent another driver removal in pipe disconnect
+        for (USBDevs::iterator d = devs->begin(); d != devs->end(); d++) {
+            if (d->vid == op->vid && d->pid == op->pid) {
+                devs->erase(d);
+                break;
+            }
+        }
         break;
     default:
         vd_printf("Unknown message received, type %u", hdr->type);
diff --git a/usbclerk.h b/usbclerk.h
index 5b1e3cf..24da3b4 100644
--- a/usbclerk.h
+++ b/usbclerk.h
@@ -5,7 +5,7 @@
 
 #define USB_CLERK_PIPE_NAME     TEXT("\\\\.\\pipe\\usbclerkpipe")
 #define USB_CLERK_MAGIC         0xDADA
-#define USB_CLERK_VERSION       0x0002
+#define USB_CLERK_VERSION       0x0003
 
 typedef struct USBClerkHeader {
     UINT16 magic;
@@ -18,6 +18,7 @@ enum {
     USB_CLERK_DRIVER_INSTALL = 1,
     USB_CLERK_DRIVER_REMOVE,
     USB_CLERK_REPLY,
+    USB_CLERK_DRIVER_SESSION_INSTALL,
     USB_CLERK_END_MESSAGE,
 };
 
diff --git a/usbclerktest.cpp b/usbclerktest.cpp
index 712fe96..efaf3df 100644
--- a/usbclerktest.cpp
+++ b/usbclerktest.cpp
@@ -1,19 +1,35 @@
 #include <stdio.h>
+#include <conio.h>

 #include <tchar.h>
 #include "usbclerk.h"
 
 int _tmain(int argc, TCHAR* argv[], TCHAR* envp[])
 {
     HANDLE pipe;
-    USBClerkDriverOp dev = {{USB_CLERK_MAGIC, USB_CLERK_VERSION, 0, sizeof(USBClerkDriverOp)}};
+    USBClerkDriverOp dev = {{USB_CLERK_MAGIC, USB_CLERK_VERSION,
+        USB_CLERK_DRIVER_INSTALL, sizeof(USBClerkDriverOp)}};
     USBClerkReply reply;
     DWORD pipe_mode;
     DWORD bytes = 0;
-    bool do_remove = false;
+    bool err = false;
+    int i, devs = 0;
 
-    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");
+    for (i = 1; i < argc && !err; i++) {
+        if (wcscmp(argv[i], L"/t") == 0) {
+            dev.hdr.type = USB_CLERK_DRIVER_SESSION_INSTALL;
+        } else if (wcscmp(argv[i], L"/u") == 0) {
+            dev.hdr.type = USB_CLERK_DRIVER_REMOVE;
+        } else if (swscanf_s(argv[i], L"%hx:%hx", &dev.vid, &dev.pid) == 2) {
+            devs++;
+        } else {
+            err = true;
+        }
+    }
+    if (argc < 2 || err || devs < argc - 2) {
+        printf("Usage: usbclerktest [/t][/u] vid:pid [vid1:pid1...]\n"
+               "default - install driver for device vid:pid (in hex)\n"
+               "/t - temporary install until session terminated\n"
+               "/u - uninstall driver\n");
         return 1;
     }
     pipe = CreateFile(USB_CLERK_PIPE_NAME, GENERIC_READ | GENERIC_WRITE,
@@ -27,29 +43,40 @@ int _tmain(int argc, TCHAR* argv[], TCHAR* envp[])
         printf("SetNamedPipeHandleState() failed: %d\n", GetLastError());
         return 1;
     }
-    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;
+
+    for (i = 1; i < argc && !err; i++) {
+        if (swscanf_s(argv[i], L"%hx:%hx", &dev.vid, &dev.pid) < 2) continue;
+        switch (dev.hdr.type) {
+        case USB_CLERK_DRIVER_SESSION_INSTALL:
+        case USB_CLERK_DRIVER_INSTALL:
+            printf("Signing & installing %04x:%04x...", dev.vid, dev.pid);
+            break;
+        case USB_CLERK_DRIVER_REMOVE:
+            printf("Removing %04x:%04x...", dev.vid, dev.pid);
+            break;
+        }
+        if (!TransactNamedPipe(pipe, &dev, sizeof(dev), &reply, sizeof(reply), &bytes, NULL)) {
+            printf("TransactNamedPipe() failed: %d\n", GetLastError());
+            CloseHandle(pipe);
+            return 1;
+        }
+        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\n",
+                   reply.hdr.magic, reply.hdr.type, reply.hdr.size);
+            return 1;
+        }
+        if (reply.status) {
+            printf("Completed successfully\n");
+        } else {
+            printf("Failed\n");
+        }
     }
-    if (!TransactNamedPipe(pipe, &dev, sizeof(dev), &reply, sizeof(reply), &bytes, NULL)) {
-        printf("TransactNamedPipe() failed: %d\n", GetLastError());
-        CloseHandle(pipe);
-        return 1;
+
+    if (dev.hdr.type == USB_CLERK_DRIVER_SESSION_INSTALL) {
+        printf("Hit any key to terminate session\n");
+        _getch();
     }
     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\n",
-               reply.hdr.magic, reply.hdr.type, reply.hdr.size);
-        return 1;
-    }
-    if (reply.status) {
-        printf("Completed successfully\n");
-    } else {
-        printf("Failed\n");
-    }
     return 0;
 }