sys: FspIrpHook and IRP_MN_NOTIFY_CHANGE_DIRECTORY

src/sys/dirctl.c

@@ -29,6 +29,8 @@ static NTSTATUS FspFsvolQueryDirectory(
     PDEVICE_OBJECT FsvolDeviceObject, PIRP Irp, PIO_STACK_LOCATION IrpSp);
 static NTSTATUS FspFsvolNotifyChangeDirectory(
     PDEVICE_OBJECT FsvolDeviceObject, PIRP Irp, PIO_STACK_LOCATION IrpSp);
+static IO_COMPLETION_ROUTINE FspFsvolNotifyChangeDirectoryCompletion;
+static WORKER_THREAD_ROUTINE FspFsvolNotifyChangeDirectoryCompletionWork;
 static NTSTATUS FspFsvolDirectoryControl(
     PDEVICE_OBJECT FsvolDeviceObject, PIRP Irp, PIO_STACK_LOCATION IrpSp);
 FSP_IOPREP_DISPATCH FspFsvolDirectoryControlPrepare;
@@ -43,6 +45,8 @@ FSP_DRIVER_DISPATCH FspDirectoryControl;
 #pragma alloc_text(PAGE, FspFsvolQueryDirectoryRetry)
 #pragma alloc_text(PAGE, FspFsvolQueryDirectory)
 #pragma alloc_text(PAGE, FspFsvolNotifyChangeDirectory)
+// !#pragma alloc_text(PAGE, FspFsvolNotifyChangeDirectoryCompletion)
+#pragma alloc_text(PAGE, FspFsvolNotifyChangeDirectoryCompletionWork)
 #pragma alloc_text(PAGE, FspFsvolDirectoryControl)
 #pragma alloc_text(PAGE, FspFsvolDirectoryControlPrepare)
 #pragma alloc_text(PAGE, FspFsvolDirectoryControlComplete)
@@ -566,6 +570,12 @@ static NTSTATUS FspFsvolQueryDirectory(
     return Result;
 }
 
+typedef struct
+{
+    PDEVICE_OBJECT FsvolDeviceObject;
+    WORK_QUEUE_ITEM WorkItem;
+} FSP_FSVOL_NOTIFY_CHANGE_DIRECTORY_COMPLETION_CONTEXT;
+
 static NTSTATUS FspFsvolNotifyChangeDirectory(
     PDEVICE_OBJECT FsvolDeviceObject, PIRP Irp, PIO_STACK_LOCATION IrpSp)
 {
@@ -583,6 +593,7 @@ static NTSTATUS FspFsvolNotifyChangeDirectory(
     ULONG CompletionFilter = IrpSp->Parameters.NotifyDirectory.CompletionFilter;
     BOOLEAN WatchTree = BooleanFlagOn(IrpSp->Flags, SL_WATCH_TREE);
     BOOLEAN DeletePending;
+    FSP_FSVOL_NOTIFY_CHANGE_DIRECTORY_COMPLETION_CONTEXT *CompletionContext;
 
     ASSERT(FileNode == FileDesc->FileNode);
 
@@ -600,6 +611,27 @@ static NTSTATUS FspFsvolNotifyChangeDirectory(
     if (DeletePending)
         return STATUS_DELETE_PENDING;
 
+    /*
+     * This IRP will be completed by the FSRTL package and therefore
+     * we will have no chance to do our normal IRP completion processing.
+     * Hook the IRP completion and perform the IRP completion processing
+     * there.
+     */
+
+    CompletionContext = FspAllocNonPaged(sizeof *CompletionContext);
+    if (0 == CompletionContext)
+        return STATUS_INSUFFICIENT_RESOURCES;
+    CompletionContext->FsvolDeviceObject = FsvolDeviceObject;
+    ExInitializeWorkItem(&CompletionContext->WorkItem,
+        FspFsvolNotifyChangeDirectoryCompletionWork, CompletionContext);
+
+    Result = FspIrpHook(Irp, FspFsvolNotifyChangeDirectoryCompletion, CompletionContext);
+    if (!NT_SUCCESS(Result))
+    {
+        FspFree(CompletionContext);
+        return Result;
+    }
+
     FspFileNodeAcquireExclusive(FileNode, Main);
 
     Result = FspNotifyChangeDirectory(
@@ -613,12 +645,37 @@ static NTSTATUS FspFsvolNotifyChangeDirectory(
 
     FspFileNodeRelease(FileNode, Main);
 
-    if (NT_SUCCESS(Result))
+    if (!NT_SUCCESS(Result))
-        Result = STATUS_PENDING;
+    {
-
+        FspIrpHookReset(Irp);
+        FspFree(CompletionContext);
         return Result;
     }
 
+    return STATUS_PENDING;
+}
+
+static NTSTATUS FspFsvolNotifyChangeDirectoryCompletion(
+    PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context)
+{
+    // !PAGED_CODE();
+
+    FSP_FSVOL_NOTIFY_CHANGE_DIRECTORY_COMPLETION_CONTEXT *CompletionContext =
+        FspIrpHookContext(Context);
+    ExQueueWorkItem(&CompletionContext->WorkItem, DelayedWorkQueue);
+
+    return FspIrpHookNext(DeviceObject, Irp, Context);
+}
+
+static VOID FspFsvolNotifyChangeDirectoryCompletionWork(PVOID Context)
+{
+    PAGED_CODE();
+
+    FSP_FSVOL_NOTIFY_CHANGE_DIRECTORY_COMPLETION_CONTEXT *CompletionContext = Context;
+    FspDeviceDereference(CompletionContext->FsvolDeviceObject);
+    FspFree(CompletionContext);
+}
+
 static NTSTATUS FspFsvolDirectoryControl(
     PDEVICE_OBJECT FsvolDeviceObject, PIRP Irp, PIO_STACK_LOCATION IrpSp)
 {

src/sys/driver.h

@@ -484,6 +484,12 @@ NTSTATUS FspSafeMdlCreate(PMDL UserMdl, LOCK_OPERATION Operation, FSP_SAFE_MDL *
 VOID FspSafeMdlCopyBack(FSP_SAFE_MDL *SafeMdl);
 VOID FspSafeMdlDelete(FSP_SAFE_MDL *SafeMdl);
 
+/* utility: hook IRP completion */
+NTSTATUS FspIrpHook(PIRP Irp, PIO_COMPLETION_ROUTINE CompletionRoutine, PVOID OwnContext);
+VOID FspIrpHookReset(PIRP Irp);
+PVOID FspIrpHookContext(PVOID Context);
+NTSTATUS FspIrpHookNext(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context);
+
 /* IRP context */
 #define FspIrpTimestampInfinity         ((ULONG)-1L)
 #define FspIrpTimestamp(Irp)            \

src/sys/util.c

@@ -68,6 +68,10 @@ BOOLEAN FspSafeMdlCheck(PMDL Mdl);
 NTSTATUS FspSafeMdlCreate(PMDL UserMdl, LOCK_OPERATION Operation, FSP_SAFE_MDL **PSafeMdl);
 VOID FspSafeMdlCopyBack(FSP_SAFE_MDL *SafeMdl);
 VOID FspSafeMdlDelete(FSP_SAFE_MDL *SafeMdl);
+NTSTATUS FspIrpHook(PIRP Irp, PIO_COMPLETION_ROUTINE CompletionRoutine, PVOID OwnContext);
+VOID FspIrpHookReset(PIRP Irp);
+PVOID FspIrpHookContext(PVOID Context);
+NTSTATUS FspIrpHookNext(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context);
 
 #ifdef ALLOC_PRAGMA
 #pragma alloc_text(PAGE, FspUnicodePathIsValid)
@@ -99,6 +103,10 @@ VOID FspSafeMdlDelete(FSP_SAFE_MDL *SafeMdl);
 #pragma alloc_text(PAGE, FspSafeMdlCreate)
 #pragma alloc_text(PAGE, FspSafeMdlCopyBack)
 #pragma alloc_text(PAGE, FspSafeMdlDelete)
+#pragma alloc_text(PAGE, FspIrpHook)
+#pragma alloc_text(PAGE, FspIrpHookReset)
+// !#pragma alloc_text(PAGE, FspIrpHookContext)
+// !#pragma alloc_text(PAGE, FspIrpHookNext)
 #endif
 
 static const LONG Delays[] =
@@ -957,3 +965,98 @@ VOID FspSafeMdlDelete(FSP_SAFE_MDL *SafeMdl)
         IoFreeMdl(SafeMdl->Mdl);
     FspFree(SafeMdl);
 }
+
+typedef struct
+{
+    PIO_COMPLETION_ROUTINE CompletionRoutine;
+    PVOID Context;
+    ULONG Control;
+    PVOID OwnContext;
+} FSP_IRP_HOOK_CONTEXT;
+
+NTSTATUS FspIrpHook(PIRP Irp, PIO_COMPLETION_ROUTINE CompletionRoutine, PVOID OwnContext)
+{
+    PAGED_CODE();
+
+    PIO_STACK_LOCATION IrpSp = IoGetCurrentIrpStackLocation(Irp);
+    FSP_IRP_HOOK_CONTEXT *HookContext = 0;
+
+    if (0 != IrpSp->CompletionRoutine || 0 != OwnContext)
+    {
+        HookContext = FspAllocNonPaged(sizeof *HookContext);
+        if (0 == HookContext)
+            return STATUS_INSUFFICIENT_RESOURCES;
+
+        HookContext->CompletionRoutine = IrpSp->CompletionRoutine;
+        HookContext->Context = IrpSp->Context;
+        HookContext->Control = FlagOn(IrpSp->Control,
+            SL_INVOKE_ON_SUCCESS | SL_INVOKE_ON_ERROR | SL_INVOKE_ON_CANCEL);
+        HookContext->OwnContext = OwnContext;
+    }
+
+    IrpSp->CompletionRoutine = CompletionRoutine;
+    IrpSp->Context = HookContext;
+    SetFlag(IrpSp->Control, SL_INVOKE_ON_SUCCESS | SL_INVOKE_ON_ERROR | SL_INVOKE_ON_CANCEL);
+
+    return STATUS_SUCCESS;
+}
+
+VOID FspIrpHookReset(PIRP Irp)
+{
+    PAGED_CODE();
+
+    PIO_STACK_LOCATION IrpSp = IoGetCurrentIrpStackLocation(Irp);
+    FSP_IRP_HOOK_CONTEXT *HookContext = IrpSp->Context;
+
+    if (0 != HookContext)
+    {
+        IrpSp->CompletionRoutine = HookContext->CompletionRoutine;
+        IrpSp->Context = HookContext->Context;
+        ClearFlag(IrpSp->Control, SL_INVOKE_ON_SUCCESS | SL_INVOKE_ON_ERROR | SL_INVOKE_ON_CANCEL);
+        SetFlag(IrpSp->Control, HookContext->Control);
+
+        FspFree(HookContext);
+    }
+    else
+    {
+        IrpSp->CompletionRoutine = 0;
+        IrpSp->Context = 0;
+        ClearFlag(IrpSp->Control, SL_INVOKE_ON_SUCCESS | SL_INVOKE_ON_ERROR | SL_INVOKE_ON_CANCEL);
+    }
+}
+
+PVOID FspIrpHookContext(PVOID Context)
+{
+    // !PAGED_CODE();
+
+    FSP_IRP_HOOK_CONTEXT *HookContext = Context;
+    return 0 != HookContext ? HookContext->OwnContext : 0;
+}
+
+NTSTATUS FspIrpHookNext(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context)
+{
+    // !PAGED_CODE();
+
+    FSP_IRP_HOOK_CONTEXT *HookContext = Context;
+    NTSTATUS Result;
+
+    if (0 != HookContext && 0 != HookContext->CompletionRoutine && (
+        (NT_SUCCESS(Irp->IoStatus.Status) && FlagOn(HookContext->Control, SL_INVOKE_ON_SUCCESS)) ||
+        (!NT_SUCCESS(Irp->IoStatus.Status) && FlagOn(HookContext->Control, SL_INVOKE_ON_ERROR)) ||
+        (Irp->Cancel && FlagOn(HookContext->Control, SL_INVOKE_ON_CANCEL))))
+    {
+        Result = HookContext->CompletionRoutine(DeviceObject, Irp, HookContext->Context);
+    }
+    else
+    {
+        if (Irp->PendingReturned && Irp->CurrentLocation <= Irp->StackCount)
+            IoMarkIrpPending(Irp);
+
+        Result = STATUS_SUCCESS;
+    }
+
+    if (0 != HookContext)
+        FspFree(HookContext);
+
+    return Result;
+}