diff options
| -rw-r--r-- | usbclerk.cpp | 113 | ||||
| -rw-r--r-- | usbclerk.h | 3 | ||||
| -rw-r--r-- | usbclerktest.cpp | 79 |
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); @@ -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; } |