From a9ccbf5baf751a715bc2ac2d89503f9f27f30665 Mon Sep 17 00:00:00 2001
From: Yuri Benditovich <yuri.benditovich@daynix.com>
Date: Thu, 2 Jan 2020 14:58:39 +0200
Subject: [PATCH] driver: maintain owner process id field of CUsbDkRedirection

Set owner process id when device redirection successfully
established (i.e. filter device is open and respective user
mode handle created for the process).

Signed-off-by: Yuri Benditovich <yuri.benditovich@daynix.com>
---
 UsbDk/ControlDevice.cpp | 9 +++++++++
 UsbDk/ControlDevice.h   | 3 +++
 2 files changed, 12 insertions(+)

diff --git a/UsbDk/ControlDevice.cpp b/UsbDk/ControlDevice.cpp
index 49097f7..74cd2c4 100644
--- a/UsbDk/ControlDevice.cpp
+++ b/UsbDk/ControlDevice.cpp
@@ -1120,6 +1120,12 @@ void CUsbDkRedirection::Dump(LPCSTR message) const
                 m_DeviceID, m_InstanceID);
 }
 
+bool CUsbDkRedirection::MatchProcess(ULONG pid)
+{
+    TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_WDFDEVICE, "%!FUNC! pid 0x%X, owner 0x%X", pid, m_OwnerPid);
+    return pid == m_OwnerPid;
+}
+
 void CUsbDkRedirection::NotifyRedirectorCreated(CUsbDkFilterDevice *RedirectorDevice)
 {
     TraceEvents(TRACE_LEVEL_ERROR, TRACE_WDFDEVICE, "%!FUNC! Redirector created for");
@@ -1209,6 +1215,9 @@ NTSTATUS CUsbDkRedirection::CreateRedirectorHandle(HANDLE RequestorProcess, PHAN
         status = m_RedirectorDevice->CreateUserModeHandle(RequestorProcess, ObjectHandle);
         if (NT_SUCCESS(status))
         {
+            ULONG pid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId();
+            TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_WDFDEVICE, "%!FUNC! done for process 0x%X", pid);
+            m_OwnerPid = pid;
             return status;
         }
 
diff --git a/UsbDk/ControlDevice.h b/UsbDk/ControlDevice.h
index e612050..7e7cd0f 100644
--- a/UsbDk/ControlDevice.h
+++ b/UsbDk/ControlDevice.h
@@ -190,6 +190,8 @@ public:
     bool IsPreparedForRemove() const
     { return m_RemovalInProgress; }
 
+    bool MatchProcess(ULONG pid);
+
     NTSTATUS WaitForAttachment()
     { return m_RedirectionCreated.Wait(true, -SecondsTo100Nanoseconds(120)); }
 
@@ -216,6 +218,7 @@ private:
     CWdmEvent m_RedirectionCreated;
     CWdmEvent m_RedirectionRemoved;
     CUsbDkFilterDevice *m_RedirectorDevice = nullptr;
+    ULONG m_OwnerPid = 0;
 
     bool m_RemovalInProgress = false;