From 8c7edd99971655dffd54ca39a2e5e6c351730076 Mon Sep 17 00:00:00 2001
From: Dmitry Fleytman <dfleytma@redhat.com>
Date: Sun, 3 Jul 2016 14:57:02 +0300
Subject: [PATCH] FilterDevice: Extend IRP post processing logic

This patch adds support for IRP post processing
on failure and in all cases.

New logic will be used by following commits to
implement device text patching.

Signed-off-by: Dmitry Fleytman <dfleytma@redhat.com>
---
 UsbDk/FilterDevice.cpp  |  5 +++--
 UsbDk/FilterStrategy.h  | 22 +++++++++++++++++-----
 UsbDk/HiderStrategy.cpp |  8 +++++---
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/UsbDk/FilterDevice.cpp b/UsbDk/FilterDevice.cpp
index 1457cb6..7ce1a7e 100644
--- a/UsbDk/FilterDevice.cpp
+++ b/UsbDk/FilterDevice.cpp
@@ -244,7 +244,7 @@ NTSTATUS CUsbDkHubFilterStrategy::PNPPreProcess(PIRP Irp)
         (BusRelations == irpStack->Parameters.QueryDeviceRelations.Type))
     {
         return PostProcessOnSuccess(Irp,
-                                    [this](PIRP Irp)
+                                    [this](PIRP Irp) -> NTSTATUS
                                     {
                                         CNonPagedDeviceRelations Relations;
                                         auto status = Relations.Create((PDEVICE_RELATIONS)Irp->IoStatus.Information);
@@ -252,7 +252,7 @@ NTSTATUS CUsbDkHubFilterStrategy::PNPPreProcess(PIRP Irp)
                                         if (!NT_SUCCESS(status))
                                         {
                                             TraceEvents(TRACE_LEVEL_ERROR, TRACE_FILTERDEVICE, "%!FUNC! Failed to create device relations object: %!STATUS!", status);
-                                            return;
+                                            return STATUS_SUCCESS;
                                         }
 
                                         TraceEvents(TRACE_LEVEL_ERROR, TRACE_FILTERDEVICE, "%!FUNC! Starting relations array processing:");
@@ -261,6 +261,7 @@ NTSTATUS CUsbDkHubFilterStrategy::PNPPreProcess(PIRP Irp)
                                         DropRemovedDevices(Relations);
                                         AddNewDevices(Relations);
                                         TraceEvents(TRACE_LEVEL_ERROR, TRACE_FILTERDEVICE, "%!FUNC! Finished relations array processing");
+                                        return STATUS_SUCCESS;
                                     });
     }
 
diff --git a/UsbDk/FilterStrategy.h b/UsbDk/FilterStrategy.h
index 849e0fc..55fce1b 100644
--- a/UsbDk/FilterStrategy.h
+++ b/UsbDk/FilterStrategy.h
@@ -70,17 +70,29 @@ protected:
 
     template<typename PostProcessFuncT>
     NTSTATUS PostProcessOnSuccess(PIRP Irp, PostProcessFuncT PostProcessFunc)
+    {
+        return PostProcess(Irp, [PostProcessFunc](PIRP Irp, NTSTATUS Status) -> NTSTATUS
+                                { return NT_SUCCESS(Status) ? PostProcessFunc(Irp) : Status; });
+    }
+
+    template<typename PostProcessFuncT>
+    NTSTATUS PostProcessOnFailure(PIRP Irp, PostProcessFuncT PostProcessFunc)
+    {
+        return PostProcess(Irp, [PostProcessFunc](PIRP Irp, NTSTATUS Status) -> NTSTATUS
+                                { return !NT_SUCCESS(Status) ? PostProcessFunc(Irp, Status) : Status; });
+    }
+
+    template<typename PostProcessFuncT>
+    NTSTATUS PostProcess(PIRP Irp, PostProcessFuncT PostProcessFunc)
     {
         IoCopyCurrentIrpStackLocationToNext(Irp);
 
         auto status = CIrp::ForwardAndWait(Irp, [this, Irp]()
-                                           { return WdfDeviceWdmDispatchPreprocessedIrp(m_Owner->WdfObject(), Irp); });
+                                                { return WdfDeviceWdmDispatchPreprocessedIrp(m_Owner->WdfObject(), Irp); });
 
-        if (NT_SUCCESS(status))
-        {
-            PostProcessFunc(Irp);
-        }
+        status = PostProcessFunc(Irp, status);
 
+        Irp->IoStatus.Status = status;
         IoCompleteRequest(Irp, IO_NO_INCREMENT);
         return status;
     }
diff --git a/UsbDk/HiderStrategy.cpp b/UsbDk/HiderStrategy.cpp
index ac67991..1753279 100644
--- a/UsbDk/HiderStrategy.cpp
+++ b/UsbDk/HiderStrategy.cpp
@@ -105,14 +105,15 @@ NTSTATUS CUsbDkHiderStrategy::PNPPreProcess(PIRP Irp)
     {
     case IRP_MN_QUERY_ID:
         return PostProcessOnSuccess(Irp,
-                                    [this](PIRP Irp)
+                                    [this](PIRP Irp) -> NTSTATUS
                                     {
                                         PatchDeviceID(Irp);
+                                        return STATUS_SUCCESS;
                                     });
 
     case IRP_MN_QUERY_CAPABILITIES:
         return PostProcessOnSuccess(Irp,
-                                    [](PIRP Irp)
+                                    [](PIRP Irp) -> NTSTATUS
                                     {
                                         auto irpStack = IoGetCurrentIrpStackLocation(Irp);
                                         irpStack->Parameters.DeviceCapabilities.Capabilities->RawDeviceOK = 1;
@@ -120,7 +121,8 @@ NTSTATUS CUsbDkHiderStrategy::PNPPreProcess(PIRP Irp)
                                         irpStack->Parameters.DeviceCapabilities.Capabilities->Removable = 0;
                                         irpStack->Parameters.DeviceCapabilities.Capabilities->EjectSupported = 0;
                                         irpStack->Parameters.DeviceCapabilities.Capabilities->SilentInstall = 1;
-                                    });
+                                        return STATUS_SUCCESS;
+        });
     default:
         return CUsbDkNullFilterStrategy::PNPPreProcess(Irp);
     }