[RDBSS][RXCE] Implement IRP cancellation

CORE-15441

sdk/include/ddk/rxcontx.h

@@ -517,6 +517,8 @@ RxReinitializeContext(
 }
 #endif
 
+extern FAST_MUTEX RxContextPerFileSerializationMutex;
+
 VOID
 NTAPI
 RxResumeBlockedOperations_Serially(
@@ -527,4 +529,19 @@ VOID
 RxResumeBlockedOperations_ALL(
     _Inout_ PRX_CONTEXT RxContext);
 
+#if (_WIN32_WINNT >= 0x0600)
+VOID
+RxCancelBlockingOperation(
+    _Inout_ PRX_CONTEXT RxContext,
+    _In_ PIRP Irp);
+#else
+VOID
+RxCancelBlockingOperation(
+    _Inout_ PRX_CONTEXT RxContext);
+#endif
+
+VOID
+RxRemoveOperationFromBlockingQueue(
+    _Inout_ PRX_CONTEXT RxContext);
+
 #endif

sdk/lib/drivers/rdbsslib/rdbss.c

@@ -361,6 +361,11 @@ NTSTATUS
 RxNotifyChangeDirectory(
     PRX_CONTEXT RxContext);
 
+VOID
+NTAPI
+RxpCancelRoutine(
+    PVOID Context);
+
 NTSTATUS
 RxpQueryInfoMiniRdr(
     PRX_CONTEXT RxContext,
@@ -529,7 +534,6 @@ BOOLEAN DisableFlushOnCleanup = FALSE;
 ULONG ReadAheadGranularity = 1 << PAGE_SHIFT;
 LIST_ENTRY RxActiveContexts;
 NPAGED_LOOKASIDE_LIST RxContextLookasideList;
-FAST_MUTEX RxContextPerFileSerializationMutex;
 RDBSS_DATA RxData;
 FCB RxDeviceFCB;
 BOOLEAN RxLoudLowIoOpsEnabled = FALSE;
@@ -1128,13 +1132,126 @@ RxCancelNotifyChangeDirectoryRequestsForVNetRoot(
     return Status;
 }
 
+/*
+ * @implemented
+ */
+BOOLEAN
+RxCancelOperationInOverflowQueue(
+    PRX_CONTEXT RxContext)
+{
+    KIRQL OldIrql;
+    BOOLEAN OperationToCancel;
+
+    /* By default, nothing cancelled */
+    OperationToCancel = FALSE;
+
+    /* Acquire the overflow spinlock */
+    KeAcquireSpinLock(&RxFileSystemDeviceObject->OverflowQueueSpinLock, &OldIrql);
+
+    /* Is our context in any queue? */
+    if (BooleanFlagOn(RxContext->Flags, (RX_CONTEXT_FLAG_FSP_DELAYED_OVERFLOW_QUEUE | RX_CONTEXT_FLAG_FSP_CRITICAL_OVERFLOW_QUEUE)))
+    {
+        /* Make sure flag is consistent with facts... */
+        if (RxContext->OverflowListEntry.Flink != NULL)
+        {
+            /* Remove it from the list */
+            RemoveEntryList(&RxContext->OverflowListEntry);
+            RxContext->OverflowListEntry.Flink == NULL;
+
+            /* Decrement appropriate count */
+            if (BooleanFlagOn(RxContext->Flags, RX_CONTEXT_FLAG_FSP_CRITICAL_OVERFLOW_QUEUE))
+            {
+                --RxFileSystemDeviceObject->OverflowQueueCount[CriticalWorkQueue];
+            }
+            else
+            {
+                --RxFileSystemDeviceObject->OverflowQueueCount[DelayedWorkQueue];
+            }
+
+            /* Clear the flag */
+            ClearFlag(RxContext->Flags, ~(RX_CONTEXT_FLAG_FSP_DELAYED_OVERFLOW_QUEUE | RX_CONTEXT_FLAG_FSP_CRITICAL_OVERFLOW_QUEUE));
+
+            /* Time to cancel! */
+            OperationToCancel = TRUE;
+        }
+    }
+
+    KeReleaseSpinLock(&RxFileSystemDeviceObject->OverflowQueueSpinLock, OldIrql);
+
+    /* We have something to cancel & complete */
+    if (OperationToCancel)
+    {
+        RxRemoveOperationFromBlockingQueue(RxContext);
+        RxCompleteRequest(RxContext, STATUS_CANCELLED);
+    }
+
+    return OperationToCancel;
+}
+
+/*
+ * @implemented
+ */
 VOID
 NTAPI
 RxCancelRoutine(
     PDEVICE_OBJECT DeviceObject,
     PIRP Irp)
 {
-    UNIMPLEMENTED;
+    KIRQL OldIrql;
+    PLIST_ENTRY Entry;
+    PRX_CONTEXT RxContext;
+
+    /* Lock our contexts list */
+    KeAcquireSpinLock(&RxStrucSupSpinLock, &OldIrql);
+
+    /* Now, find a context that matches the cancelled IRP */
+    Entry = RxActiveContexts.Flink;
+    while (Entry != &RxActiveContexts)
+    {
+        RxContext = CONTAINING_RECORD(Entry, RX_CONTEXT, ContextListEntry);
+        Entry = Entry->Flink;
+
+        /* Found! */
+        if (RxContext->CurrentIrp == Irp)
+        {
+            break;
+        }
+    }
+
+    /* If we reached the end of the list, we didn't find any context, so zero the buffer
+     * If the context is already under cancellation, forget about it too
+     */
+    if (Entry == &RxActiveContexts || BooleanFlagOn(RxContext->Flags, RX_CONTEXT_FLAG_CANCELLED))
+    {
+        RxContext = NULL;
+    }
+    else
+    {
+        /* Otherwise, reference it and mark it cancelled */
+        SetFlag(RxContext->Flags, RX_CONTEXT_FLAG_CANCELLED);
+        InterlockedIncrement((volatile long *)&RxContext->ReferenceCount);
+    }
+
+    /* Done with the contexts list */
+    KeReleaseSpinLock(&RxStrucSupSpinLock, OldIrql);
+
+    /* And done with the cancellation, we'll do it now */
+    IoReleaseCancelSpinLock(Irp->CancelIrql);
+
+    /* If we have a context to cancel */
+    if (RxContext != NULL)
+    {
+        /* We cannot executed at dispatch, so queue a deferred cancel */
+        if (KeGetCurrentIrql() >= DISPATCH_LEVEL)
+        {
+            RxDispatchToWorkerThread(RxFileSystemDeviceObject, CriticalWorkQueue, RxpCancelRoutine, RxContext);
+        }
+        /* Cancel now! */
+        else
+        {
+            RxpCancelRoutine(RxContext);
+        }
+    }
 }
 
 /*
@@ -7491,6 +7608,35 @@ RxNotifyChangeDirectory(
     return Status;
 }
 
+/*
+ * @implemented
+ */
+VOID
+NTAPI
+RxpCancelRoutine(
+    PVOID Context)
+{
+    PRX_CONTEXT RxContext;
+
+    PAGED_CODE();
+
+    RxContext = Context;
+
+    /* First, notify mini-rdr about cancellation */
+    if (RxContext->MRxCancelRoutine != NULL)
+    {
+        RxContext->MRxCancelRoutine(RxContext);
+    }
+    /* If we didn't find in overflow queue, try in blocking operations */
+    else if (!RxCancelOperationInOverflowQueue(RxContext))
+    {
+        RxCancelBlockingOperation(RxContext);
+    }
+
+    /* And delete the context */
+    RxDereferenceAndDeleteRxContext_Real(RxContext);
+}
+
 NTSTATUS
 RxPostStackOverflowRead (
     IN PRX_CONTEXT RxContext)

sdk/lib/drivers/rxce/rxce.c

@@ -143,6 +143,7 @@ LIST_ENTRY RxRecurrentWorkItemsList;
 KDPC RxTimerDpc;
 KTIMER RxTimer;
 ULONG RxTimerTickCount;
+FAST_MUTEX RxContextPerFileSerializationMutex;
 #if DBG
 BOOLEAN DumpDispatchRoutine = TRUE;
 #else
@@ -715,6 +716,66 @@ RxBootstrapWorkerThreadDispatcher(
     RxpWorkerThreadDispatcher(RxWorkQueue, NULL);
 }
 
+/*
+ * @implemented
+ */
+VOID
+RxCancelBlockingOperation(
+    IN OUT PRX_CONTEXT RxContext)
+{
+    PFOBX Fobx;
+    BOOLEAN PostRequest;
+    C_ASSERT(FIELD_OFFSET(RX_CONTEXT, IoStatusBlock.Status) == 100);
+
+    PAGED_CODE();
+
+    Fobx = (PFOBX)RxContext->pFobx;
+    PostRequest = FALSE;
+
+    /* Acquire the pipe mutex */
+    ExAcquireFastMutex(&RxContextPerFileSerializationMutex);
+
+    /* If that's a blocking pipe operation which is not the CCB one, then handle it */
+    if (BooleanFlagOn(RxContext->FlagsForLowIo, RXCONTEXT_FLAG4LOWIO_PIPE_SYNC_OPERATION) &&
+        RxContext->RxContextSerializationQLinks.Flink != NULL &&
+        RxContext != CONTAINING_RECORD(&Fobx->Specific.NamedPipe.ReadSerializationQueue, RX_CONTEXT, RxContextSerializationQLinks) &&
+        RxContext != CONTAINING_RECORD(&Fobx->Specific.NamedPipe.WriteSerializationQueue, RX_CONTEXT, RxContextSerializationQLinks))
+    {
+        /* Clear it! */
+        ClearFlag(RxContext->FlagsForLowIo, RXCONTEXT_FLAG4LOWIO_PIPE_SYNC_OPERATION);
+
+        /* Drop it off the list */
+        RemoveEntryList(&RxContext->RxContextSerializationQLinks);
+        RxContext->RxContextSerializationQLinks.Flink = NULL;
+        RxContext->RxContextSerializationQLinks.Blink = NULL;
+
+        /* Set we've been cancelled */
+        RxContext->IoStatusBlock.Status = STATUS_CANCELLED;
+
+        /*
+         * If it's async, we'll post completion, otherwise, we signal to waiters
+         * it's being cancelled
+         */
+        if (BooleanFlagOn(RxContext->Flags, RX_CONTEXT_FLAG_ASYNC_OPERATION))
+        {
+            PostRequest = TRUE;
+        }
+        else
+        {
+            RxSignalSynchronousWaiter(RxContext);
+        }
+    }
+
+    /* Done */
+    ExReleaseFastMutex(&RxContextPerFileSerializationMutex);
+
+    /* Post if async */
+    if (PostRequest)
+    {
+        RxFsdPostRequest(RxContext);
+    }
+}
+
 /*
  * @implemented
  */
@@ -7562,6 +7623,32 @@ RxRemoveNameNetFcb(
 #endif
 }
 
+/*
+ * @implemented
+ */
+VOID
+RxRemoveOperationFromBlockingQueue(
+    IN OUT PRX_CONTEXT RxContext)
+{
+    /* Acquire the pipe mutex */
+    ExAcquireFastMutex(&RxContextPerFileSerializationMutex);
+
+    /* Is that a blocking serial operation? */
+    if (BooleanFlagOn(RxContext->FlagsForLowIo, RXCONTEXT_FLAG4LOWIO_PIPE_SYNC_OPERATION))
+    {
+        /* Clear it! */
+        ClearFlag(RxContext->FlagsForLowIo, RXCONTEXT_FLAG4LOWIO_PIPE_SYNC_OPERATION);
+
+        /* Drop it off the list */
+        RemoveEntryList(&RxContext->RxContextSerializationQLinks);
+        RxContext->RxContextSerializationQLinks.Flink = NULL;
+        RxContext->RxContextSerializationQLinks.Blink = NULL;
+    }
+
+    /* Done */
+    ExReleaseFastMutex(&RxContextPerFileSerializationMutex);
+}
+
 /*
  * @implemented
  */