From d2dc75c14ba5c90da42a85176ed6afc1dfdabe67 Mon Sep 17 00:00:00 2001
From: Dmitry Fleytman <dfleytma@redhat.com>
Date: Thu, 6 Mar 2014 09:31:46 +0200
Subject: [PATCH] UsbDk: Fix CIrp cleanup

Do not hold target device reference in case IRP allocation failure

Signed-off-by: Dmitry Fleytman <dfleytma@redhat.com>
---
 UsbDk/Irp.cpp | 31 +++++++++++++++++++++----------
 UsbDk/Irp.h   |  2 ++
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/UsbDk/Irp.cpp b/UsbDk/Irp.cpp
index 9a26049..0ca72aa 100644
--- a/UsbDk/Irp.cpp
+++ b/UsbDk/Irp.cpp
@@ -5,36 +5,47 @@ CIrp::~CIrp()
 {
     if (m_Irp != nullptr)
     {
-        Destroy();
+        DestroyIrp();
+        ReleaseTarget();
     }
 }
 
-void CIrp::Destroy()
+void CIrp::DestroyIrp()
 {
     ASSERT(m_Irp != nullptr);
-    ASSERT(m_TargetDevice != nullptr);
 
     IoFreeIrp(m_Irp);
     m_Irp = nullptr;
+}
+
+void CIrp::ReleaseTarget()
+{
+    ASSERT(m_TargetDevice != nullptr);
 
     ObDereferenceObject(m_TargetDevice);
     m_TargetDevice = nullptr;
 }
 
+void CIrp::Destroy()
+{
+    DestroyIrp();
+    ReleaseTarget();
+}
+
 NTSTATUS CIrp::Create(PDEVICE_OBJECT TargetDevice)
 {
-    m_TargetDevice = TargetDevice;
-    ObReferenceObject(m_TargetDevice);
+    ASSERT(m_TargetDevice == nullptr);
+    ASSERT(m_Irp == nullptr);
 
     m_Irp = IoAllocateIrp(TargetDevice->StackSize, FALSE);
-
-    if (m_Irp != nullptr)
+    if (m_Irp == nullptr)
     {
-        m_Irp->IoStatus.Status = STATUS_NOT_SUPPORTED;
-        return STATUS_SUCCESS;
+        return STATUS_INSUFFICIENT_RESOURCES;
     }
 
-    return STATUS_INSUFFICIENT_RESOURCES;
+    ObReferenceObject(TargetDevice);
+    m_TargetDevice = TargetDevice;
+    return STATUS_SUCCESS;
 }
 
 NTSTATUS CIrp::SynchronousCompletion(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context)
diff --git a/UsbDk/Irp.h b/UsbDk/Irp.h
index 62c07c0..9582d72 100644
--- a/UsbDk/Irp.h
+++ b/UsbDk/Irp.h
@@ -27,6 +27,8 @@ public:
 
 private:
     static NTSTATUS SynchronousCompletion(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context);
+    void DestroyIrp();
+    void ReleaseTarget();
 
     PIRP m_Irp = nullptr;
     PDEVICE_OBJECT m_TargetDevice = nullptr;