From a30a1f43a50bd0d8d368cd857297409869fb5bb5 Mon Sep 17 00:00:00 2001 From: kvojacheck <93647470+kvojacheck@users.noreply.github.com> Date: Mon, 8 Nov 2021 10:52:52 +0100 Subject: [PATCH 1/3] Cleanup warnings --- UsbDk/RedirectorStrategy.cpp | 2 +- UsbDk/UsbTarget.h | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/UsbDk/RedirectorStrategy.cpp b/UsbDk/RedirectorStrategy.cpp index f8755f3..02edfb7 100644 --- a/UsbDk/RedirectorStrategy.cpp +++ b/UsbDk/RedirectorStrategy.cpp @@ -83,7 +83,7 @@ NTSTATUS CUsbDkRedirectorStrategy::Create(CUsbDkFilterDevice *Owner) return status; } -using USBDK_REDIRECTOR_REQUEST_CONTEXT = struct : public USBDK_TARGET_REQUEST_CONTEXT +struct USBDK_REDIRECTOR_REQUEST_CONTEXT : public USBDK_TARGET_REQUEST_CONTEXT { bool PreprocessingDone; diff --git a/UsbDk/UsbTarget.h b/UsbDk/UsbTarget.h index 9c2a3fa..eec896a 100644 --- a/UsbDk/UsbTarget.h +++ b/UsbDk/UsbTarget.h @@ -28,7 +28,7 @@ #include "Urb.h" #include "WdfRequest.h" -using USBDK_TARGET_REQUEST_CONTEXT = struct : public WDF_REQUEST_CONTEXT +struct USBDK_TARGET_REQUEST_CONTEXT : public WDF_REQUEST_CONTEXT { ULONG64 RequestId; }; From a58993cb35c89f6915ed05125e7300d5503f3eaa Mon Sep 17 00:00:00 2001 From: kvojacheck <93647470+kvojacheck@users.noreply.github.com> Date: Mon, 8 Nov 2021 11:01:53 +0100 Subject: [PATCH 2/3] Add ability to pass process ID to OnClose() Refactor CUsbDkFilterStrategy::OnClose() in a way that the process id can be passed from the beginning of the call-chain. This is just a technical change needed for the next commit. --- UsbDk/ControlDevice.cpp | 7 +++---- UsbDk/ControlDevice.h | 4 ++-- UsbDk/FilterDevice.cpp | 3 ++- UsbDk/FilterStrategy.h | 3 ++- UsbDk/RedirectorStrategy.cpp | 4 ++-- UsbDk/RedirectorStrategy.h | 2 +- 6 files changed, 12 insertions(+), 11 deletions(-) diff --git a/UsbDk/ControlDevice.cpp b/UsbDk/ControlDevice.cpp index a244a31..ab2b7de 100644 --- a/UsbDk/ControlDevice.cpp +++ b/UsbDk/ControlDevice.cpp @@ -1025,9 +1025,9 @@ NTSTATUS CUsbDkControlDevice::AddRedirectionToSet(const USB_DK_DEVICE_ID &Device return STATUS_SUCCESS; } -NTSTATUS CUsbDkControlDevice::RemoveRedirect(const USB_DK_DEVICE_ID &DeviceId) +NTSTATUS CUsbDkControlDevice::RemoveRedirect(const USB_DK_DEVICE_ID &DeviceId, ULONG pid) { - if (NotifyRedirectorRemovalStarted(DeviceId)) + if (NotifyRedirectorRemovalStarted(DeviceId, pid)) { auto res = ResetUsbDevice(DeviceId, false); if (NT_SUCCESS(res)) @@ -1079,9 +1079,8 @@ bool CUsbDkControlDevice::NotifyRedirectorAttached(CRegText *DeviceID, CRegText return m_Redirections.ModifyOne(&ID, [RedirectorDevice](CUsbDkRedirection *R){ R->NotifyRedirectorCreated(RedirectorDevice); }); } -bool CUsbDkControlDevice::NotifyRedirectorRemovalStarted(const USB_DK_DEVICE_ID &ID) +bool CUsbDkControlDevice::NotifyRedirectorRemovalStarted(const USB_DK_DEVICE_ID &ID, ULONG pid) { - ULONG pid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId(); return m_Redirections.ModifyOne(&ID, [](CUsbDkRedirection *R){ R->NotifyRedirectionRemovalStarted(); }, pid); } diff --git a/UsbDk/ControlDevice.h b/UsbDk/ControlDevice.h index 9223036..1ed72a7 100644 --- a/UsbDk/ControlDevice.h +++ b/UsbDk/ControlDevice.h @@ -274,7 +274,7 @@ class CUsbDkControlDevice : private CWdfControlDevice, public CAllocatableOnClose(); + ULONG pid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId(); + Strategy(Device)->OnClose(pid); }, WDF_NO_EVENT_CALLBACK); diff --git a/UsbDk/FilterStrategy.h b/UsbDk/FilterStrategy.h index 82a664c..fdab0b0 100644 --- a/UsbDk/FilterStrategy.h +++ b/UsbDk/FilterStrategy.h @@ -62,7 +62,8 @@ class CUsbDkFilterStrategy CUsbDkControlDevice* GetControlDevice() { return m_ControlDevice; } - virtual void OnClose(){} + virtual void OnClose(ULONG pid) + { UNREFERENCED_PARAMETER(pid); } protected: CUsbDkFilterDevice *m_Owner = nullptr; diff --git a/UsbDk/RedirectorStrategy.cpp b/UsbDk/RedirectorStrategy.cpp index 02edfb7..262db43 100644 --- a/UsbDk/RedirectorStrategy.cpp +++ b/UsbDk/RedirectorStrategy.cpp @@ -626,13 +626,13 @@ size_t CUsbDkRedirectorStrategy::GetRequestContextSize() return sizeof(USBDK_REDIRECTOR_REQUEST_CONTEXT); } -void CUsbDkRedirectorStrategy::OnClose() +void CUsbDkRedirectorStrategy::OnClose(ULONG pid) { USB_DK_DEVICE_ID ID; UsbDkFillIDStruct(&ID, *m_DeviceID->begin(), *m_InstanceID->begin()); TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_REDIRECTOR, "%!FUNC!"); - auto status = m_ControlDevice->RemoveRedirect(ID); + auto status = m_ControlDevice->RemoveRedirect(ID, pid); if (!NT_SUCCESS(status)) { TraceEvents(TRACE_LEVEL_ERROR, TRACE_REDIRECTOR, "%!FUNC! RemoveRedirect failed: %!STATUS!", status); diff --git a/UsbDk/RedirectorStrategy.h b/UsbDk/RedirectorStrategy.h index 31d07c0..d78c05d 100644 --- a/UsbDk/RedirectorStrategy.h +++ b/UsbDk/RedirectorStrategy.h @@ -75,7 +75,7 @@ class CUsbDkRedirectorStrategy : public CUsbDkHiderStrategy size_t InputBufferLength, ULONG IoControlCode) override; - virtual void OnClose() override; + virtual void OnClose(ULONG pid) override; void SetDeviceID(CRegText *DevID) { m_DeviceID = DevID; } From bc6cae1561a7a7ad460793347e09f9f287115dcf Mon Sep 17 00:00:00 2001 From: kvojacheck <93647470+kvojacheck@users.noreply.github.com> Date: Mon, 8 Nov 2021 12:01:45 +0100 Subject: [PATCH 3/3] Stop redirection upon closing the last reference Add reference counting, since the previously used 'PID-based' method might not work reliably: for example, if the program doing redirection is forcefully closed from the task manager, the USB redirection could remain stuck. After this commit, we are always stopping redirection if the reference count drops to zero. --- UsbDk/FilterDevice.cpp | 24 ++++++++++++++++++++++-- UsbDk/FilterDevice.h | 3 +++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/UsbDk/FilterDevice.cpp b/UsbDk/FilterDevice.cpp index a2e1ceb..7bbabed 100644 --- a/UsbDk/FilterDevice.cpp +++ b/UsbDk/FilterDevice.cpp @@ -86,12 +86,32 @@ NTSTATUS CUsbDkFilterDeviceInit::Configure(ULONG InstanceNumber) SetFileEventCallbacks([](_In_ WDFDEVICE Device, _In_ WDFREQUEST Request, _In_ WDFFILEOBJECT FileObject) { UNREFERENCED_PARAMETER(FileObject); - UsbDkFilterGetContext(Device)->UsbDkFilter->OnFileCreate(Request); + auto filter = UsbDkFilterGetContext(Device)->UsbDkFilter; + filter->m_open_count.AddRef(); + filter->OnFileCreate(Request); }, [](_In_ WDFFILEOBJECT FileObject) { WDFDEVICE Device = WdfFileObjectGetDevice(FileObject); - ULONG pid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId(); + auto filter = UsbDkFilterGetContext(Device)->UsbDkFilter; + ULONG pid = 0; // zero means always match + + // Check PID only if there are multiple open references to the file. + // If this was the last reference, always close the redirection. + // + // This callback function might run in a different process-context + // than the initiator process, therefore the 'current process ID' + // isn't always the ID of the 'owning' process. + // + // In the worst case, the USB redirection will be kept until the last + // open file handle to the device is closed. + // + // From KMDF 1.21, there's a new method that should give us the expected ID: + // WdfFileObjectGetInitiatorProcessId(FileObject) + + if (filter->m_open_count.Release()) { + pid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId(); + } Strategy(Device)->OnClose(pid); }, WDF_NO_EVENT_CALLBACK); diff --git a/UsbDk/FilterDevice.h b/UsbDk/FilterDevice.h index 76b38ca..1de5597 100644 --- a/UsbDk/FilterDevice.h +++ b/UsbDk/FilterDevice.h @@ -176,6 +176,9 @@ class CUsbDkFilterDevice : public CWdfDevice, { m_SerialNumber = Number; } void OnFileCreate(WDFREQUEST Request); + + CWdmRefCounter m_open_count; + private: ~CUsbDkFilterDevice() {