sys: TopLevelIrp functionality

src/sys/driver.h

@@ -124,16 +124,26 @@ extern __declspec(selectany) int fsp_bp = 1;
 #define FSP_ENTER_MJ(...)               \
     NTSTATUS Result = STATUS_SUCCESS;   \
     PIO_STACK_LOCATION IrpSp = IoGetCurrentIrpStackLocation(Irp);\
-    BOOLEAN fsp_device_deref = FALSE; \
+    BOOLEAN fsp_device_deref = FALSE;   \
+    PIRP fsp_top_level_irp = IoGetTopLevelIrp();\
     FSP_ENTER_(__VA_ARGS__);            \
     do                                  \
     {                                   \
+        if (0 != fsp_top_level_irp)     \
+        {                               \
+            if (FSRTL_MAX_TOP_LEVEL_IRP_FLAG < (UINT_PTR)fsp_top_level_irp &&\
+                IO_TYPE_IRP == fsp_top_level_irp->Type)\
+                FspIrpSetTopFlags(Irp, FspIrpFlags(fsp_top_level_irp));\
+            else                        \
+                FspIrpSetTopFlags(Irp, FspFileNodeAcquireFull);\
+        }                               \
+        IoSetTopLevelIrp(Irp);          \
         if (!FspDeviceReference(IrpSp->DeviceObject))\
         {                               \
             Result = STATUS_CANCELLED;  \
             goto fsp_leave_label;       \
         }                               \
-        fsp_device_deref = TRUE;      \
+        fsp_device_deref = TRUE;        \
     } while (0,0)
 #define FSP_LEAVE_MJ(fmt, ...)          \
     FSP_LEAVE_(                         \
@@ -164,6 +174,7 @@ extern __declspec(selectany) int fsp_bp = 1;
             else                        \
                 FspIopCompleteIrpEx(Irp, Result, fsp_device_deref);\
         }                               \
+        IoSetTopLevelIrp(fsp_top_level_irp);\
     );                                  \
     return Result
 #define FSP_ENTER_IOC(...)              \
@@ -404,8 +415,42 @@ VOID FspQueueDelayedWorkItem(FSP_DELAYED_WORK_ITEM *DelayedWorkItem, LARGE_INTEG
     (*(ULONG *)&(Irp)->Tail.Overlay.DriverContext[0])
 #define FspIrpDictNext(Irp)             \
     (*(PIRP *)&(Irp)->Tail.Overlay.DriverContext[1])
-#define FspIrpRequest(Irp)              \
+static inline
-    (*(FSP_FSCTL_TRANSACT_REQ **)&(Irp)->Tail.Overlay.DriverContext[2])
+FSP_FSCTL_TRANSACT_REQ *FspIrpRequest(PIRP Irp)
+{
+    return (PVOID)((UINT_PTR)Irp->Tail.Overlay.DriverContext[2] & ~0xf);
+}
+static inline
+VOID FspIrpSetRequest(PIRP Irp, FSP_FSCTL_TRANSACT_REQ *Request)
+{
+    ASSERT(0 == ((UINT_PTR)Request & 0xf));
+    ULONG Flags = (ULONG)((UINT_PTR)Irp->Tail.Overlay.DriverContext[2] & 0xf);
+    Irp->Tail.Overlay.DriverContext[2] = (PVOID)((UINT_PTR)Request | Flags);
+}
+static inline
+ULONG FspIrpFlags(PIRP Irp)
+{
+    return (ULONG)((UINT_PTR)Irp->Tail.Overlay.DriverContext[2] & 3);
+}
+static inline
+VOID FspIrpSetFlags(PIRP Irp, ULONG Flags)
+{
+    ASSERT(3 > Flags);
+    FSP_FSCTL_TRANSACT_REQ *Request = (PVOID)((UINT_PTR)Irp->Tail.Overlay.DriverContext[2] & ~3);
+    Irp->Tail.Overlay.DriverContext[2] = (PVOID)((UINT_PTR)Request | Flags);
+}
+static inline
+ULONG FspIrpTopFlags(PIRP Irp)
+{
+    return (ULONG)((UINT_PTR)Irp->Tail.Overlay.DriverContext[2] & 0xc) >> 2;
+}
+static inline
+VOID FspIrpSetTopFlags(PIRP Irp, ULONG Flags)
+{
+    ASSERT(3 > Flags);
+    FSP_FSCTL_TRANSACT_REQ *Request = (PVOID)((UINT_PTR)Irp->Tail.Overlay.DriverContext[2] & ~0xc);
+    Irp->Tail.Overlay.DriverContext[2] = (PVOID)((UINT_PTR)Request | (Flags << 2));
+}
 
 /* I/O queue */
 #define FspIoqTimeout                   ((PIRP)1)

src/sys/file.c

@@ -58,6 +58,18 @@ VOID FspFileDescDelete(FSP_FILE_DESC *FileDesc);
 #pragma alloc_text(PAGE, FspFileDescDelete)
 #endif
 
+#define FSP_FILE_NODE_GET_FLAGS()       \
+    PIRP Irp = IoGetTopLevelIrp();      \
+    BOOLEAN IrpValid = FSRTL_MAX_TOP_LEVEL_IRP_FLAG < (UINT_PTR)Irp && IO_TYPE_IRP == Irp->Type;\
+    if (IrpValid)                       \
+        Flags &= ~FspIrpTopFlags(Irp);
+#define FSP_FILE_NODE_SET_FLAGS()       \
+    if (IrpValid)                       \
+        FspIrpSetFlags(Irp, FspIrpFlags(Irp) | Flags);
+#define FSP_FILE_NODE_CLR_FLAGS()       \
+    if (IrpValid)                       \
+        FspIrpSetFlags(Irp, FspIrpFlags(Irp) & (~Flags & 3));
+
 NTSTATUS FspFileNodeCreate(PDEVICE_OBJECT DeviceObject,
     ULONG ExtraSize, FSP_FILE_NODE **PFileNode)
 {
@@ -128,18 +140,24 @@ VOID FspFileNodeAcquireSharedF(FSP_FILE_NODE *FileNode, ULONG Flags)
 {
     PAGED_CODE();
 
+    FSP_FILE_NODE_GET_FLAGS();
+
     if (Flags & FspFileNodeAcquireMain)
         ExAcquireResourceSharedLite(FileNode->Header.Resource, TRUE);
 
     if (Flags & FspFileNodeAcquirePgio)
         ExAcquireResourceSharedLite(FileNode->Header.PagingIoResource, TRUE);
+
+    FSP_FILE_NODE_SET_FLAGS();
 }
 
 BOOLEAN FspFileNodeTryAcquireSharedF(FSP_FILE_NODE *FileNode, ULONG Flags, BOOLEAN Wait)
 {
     PAGED_CODE();
 
-    BOOLEAN Result = FALSE;
+    FSP_FILE_NODE_GET_FLAGS();
+
+    BOOLEAN Result = TRUE;
 
     if (Flags & FspFileNodeAcquireMain)
     {
@@ -159,6 +177,9 @@ BOOLEAN FspFileNodeTryAcquireSharedF(FSP_FILE_NODE *FileNode, ULONG Flags, BOOLE
         }
     }
 
+    if (Result)
+        FSP_FILE_NODE_SET_FLAGS();
+
     return Result;
 }
 
@@ -166,18 +187,24 @@ VOID FspFileNodeAcquireExclusiveF(FSP_FILE_NODE *FileNode, ULONG Flags)
 {
     PAGED_CODE();
 
+    FSP_FILE_NODE_GET_FLAGS();
+
     if (Flags & FspFileNodeAcquireMain)
         ExAcquireResourceExclusiveLite(FileNode->Header.Resource, TRUE);
 
     if (Flags & FspFileNodeAcquirePgio)
         ExAcquireResourceExclusiveLite(FileNode->Header.PagingIoResource, TRUE);
+
+    FSP_FILE_NODE_SET_FLAGS();
 }
 
 BOOLEAN FspFileNodeTryAcquireExclusiveF(FSP_FILE_NODE *FileNode, ULONG Flags, BOOLEAN Wait)
 {
     PAGED_CODE();
 
-    BOOLEAN Result = FALSE;
+    FSP_FILE_NODE_GET_FLAGS();
+
+    BOOLEAN Result = TRUE;
 
     if (Flags & FspFileNodeAcquireMain)
     {
@@ -197,6 +224,9 @@ BOOLEAN FspFileNodeTryAcquireExclusiveF(FSP_FILE_NODE *FileNode, ULONG Flags, BO
         }
     }
 
+    if (Result)
+        FSP_FILE_NODE_SET_FLAGS();
+
     return Result;
 }
 
@@ -204,6 +234,8 @@ VOID FspFileNodeSetOwnerF(FSP_FILE_NODE *FileNode, ULONG Flags, PVOID Owner)
 {
     PAGED_CODE();
 
+    FSP_FILE_NODE_GET_FLAGS();
+
     Owner = (PVOID)((UINT_PTR)Owner | 3);
 
     if (Flags & FspFileNodeAcquireMain)
@@ -217,17 +249,23 @@ VOID FspFileNodeReleaseF(FSP_FILE_NODE *FileNode, ULONG Flags)
 {
     PAGED_CODE();
 
+    FSP_FILE_NODE_GET_FLAGS();
+
     if (Flags & FspFileNodeAcquirePgio)
         ExReleaseResourceLite(FileNode->Header.PagingIoResource);
 
     if (Flags & FspFileNodeAcquireMain)
         ExReleaseResourceLite(FileNode->Header.Resource);
+
+    FSP_FILE_NODE_CLR_FLAGS();
 }
 
 VOID FspFileNodeReleaseOwnerF(FSP_FILE_NODE *FileNode, ULONG Flags, PVOID Owner)
 {
     PAGED_CODE();
 
+    FSP_FILE_NODE_GET_FLAGS();
+
     Owner = (PVOID)((UINT_PTR)Owner | 3);
 
     if (Flags & FspFileNodeAcquirePgio)
@@ -245,6 +283,8 @@ VOID FspFileNodeReleaseOwnerF(FSP_FILE_NODE *FileNode, ULONG Flags, PVOID Owner)
         else
             ExReleaseResourceForThreadLite(FileNode->Header.Resource, (ERESOURCE_THREAD)Owner);
     }
+
+    FSP_FILE_NODE_CLR_FLAGS();
 }
 
 FSP_FILE_NODE *FspFileNodeOpen(FSP_FILE_NODE *FileNode, PFILE_OBJECT FileObject,

src/sys/iop.c

@@ -38,6 +38,13 @@ NTSTATUS FspIopDispatchComplete(PIRP Irp, const FSP_FSCTL_TRANSACT_RSP *Response
 #pragma alloc_text(PAGE, FspIopDispatchComplete)
 #endif
 
+/* Requests (and RequestHeaders) must be 16-byte aligned, because we use the low 4 bits for flags */
+#if 16 != MEMORY_ALLOCATION_ALIGNMENT
+#define REQ_HEADER_ALIGNMASK            15
+#else
+#define REQ_HEADER_ALIGNMASK            0
+#endif
+
 typedef struct
 {
     FSP_IOP_REQUEST_FINI *RequestFini;
@@ -66,18 +73,22 @@ NTSTATUS FspIopCreateRequestFunnel(
     if (FlagOn(Flags, FspIopRequestMustSucceed))
         RequestHeader = FspAllocatePoolMustSucceed(
             FlagOn(Flags, FspIopRequestNonPaged) ? NonPagedPool : PagedPool,
-            sizeof *RequestHeader + sizeof *Request + ExtraSize,
+            sizeof *RequestHeader + sizeof *Request + ExtraSize + REQ_HEADER_ALIGNMASK,
             FSP_ALLOC_INTERNAL_TAG);
     else
     {
         RequestHeader = ExAllocatePoolWithTag(
             FlagOn(Flags, FspIopRequestNonPaged) ? NonPagedPool : PagedPool,
-            sizeof *RequestHeader + sizeof *Request + ExtraSize,
+            sizeof *RequestHeader + sizeof *Request + ExtraSize + REQ_HEADER_ALIGNMASK,
             FSP_ALLOC_INTERNAL_TAG);
         if (0 == RequestHeader)
             return STATUS_INSUFFICIENT_RESOURCES;
     }
 
+#if 0 != REQ_HEADER_ALIGNMASK
+    RequestHeader = (PVOID)(((UINT_PTR)RequestHeader + REQ_HEADER_ALIGNMASK) & REQ_HEADER_ALIGNMASK);
+#endif
+
     RtlZeroMemory(RequestHeader, sizeof *RequestHeader + sizeof *Request + ExtraSize);
     RequestHeader->RequestFini = RequestFini;
 
@@ -94,7 +105,7 @@ NTSTATUS FspIopCreateRequestFunnel(
     }
 
     if (0 != Irp)
-        FspIrpRequest(Irp) = Request;
+        FspIrpSetRequest(Irp, Request);
     *PRequest = Request;
 
     return STATUS_SUCCESS;
@@ -206,7 +217,7 @@ VOID FspIopCompleteIrpEx(PIRP Irp, NTSTATUS Result, BOOLEAN DeviceDereference)
     if (0 != FspIrpRequest(Irp))
     {
         FspIopDeleteRequest(FspIrpRequest(Irp));
-        FspIrpRequest(Irp) = 0;
+        FspIrpSetRequest(Irp, 0);
     }
 
     /* get the device object out of the IRP before completion */
@@ -227,7 +238,12 @@ VOID FspIopCompleteCanceledIrp(PIRP Irp)
 
     DEBUGLOGIRP(Irp, STATUS_CANCELLED);
 
+    PIRP TopLevelIrp = IoGetTopLevelIrp();
+    IoSetTopLevelIrp(Irp);
+
     FspIopCompleteIrpEx(Irp, STATUS_CANCELLED, TRUE);
+
+    IoSetTopLevelIrp(TopLevelIrp);
 }
 
 BOOLEAN FspIopRetryPrepareIrp(PIRP Irp, NTSTATUS *PResult)

src/sys/volume.c

@@ -534,6 +534,7 @@ NTSTATUS FspVolumeTransact(
     FSP_FSCTL_TRANSACT_REQ *Request, *PendingIrpRequest;
     PIRP ProcessIrp, PendingIrp, RetriedIrp, RepostedIrp;
     LARGE_INTEGER Timeout;
+    PIRP TopLevelIrp = IoGetTopLevelIrp();
 
     /* process any user-mode file system responses */
     RepostedIrp = 0;
@@ -557,6 +558,7 @@ NTSTATUS FspVolumeTransact(
         ASSERT((UINT_PTR)ProcessIrp == (UINT_PTR)Response->Hint);
         ASSERT(FspIrpRequest(ProcessIrp)->Hint == Response->Hint);
 
+        IoSetTopLevelIrp(ProcessIrp);
         Result = FspIopDispatchComplete(ProcessIrp, Response);
         if (STATUS_PENDING == Result)
         {
@@ -579,6 +581,7 @@ NTSTATUS FspVolumeTransact(
         if (0 == RetriedIrp)
             break;
 
+        IoSetTopLevelIrp(RetriedIrp);
         Response = FspIopIrpResponse(RetriedIrp);
         Result = FspIopDispatchComplete(RetriedIrp, Response);
         if (STATUS_PENDING == Result)
@@ -642,6 +645,7 @@ NTSTATUS FspVolumeTransact(
     {
         PendingIrpRequest = FspIrpRequest(PendingIrp);
 
+        IoSetTopLevelIrp(PendingIrp);
         Result = FspIopDispatchPrepare(PendingIrp, PendingIrpRequest);
         if (STATUS_PENDING == Result)
         {
@@ -692,6 +696,7 @@ NTSTATUS FspVolumeTransact(
     Result = STATUS_SUCCESS;
 
 exit:
+    IoSetTopLevelIrp(TopLevelIrp);
     FspDeviceDereference(FsvolDeviceObject);
     return Result;
 }
@@ -737,7 +742,7 @@ NTSTATUS FspVolumeWork(
 
     /* associate the passed Request with our Irp; acquire ownership of the Request */
     Request->Hint = (UINT_PTR)Irp;
-    FspIrpRequest(Irp) = Request;
+    FspIrpSetRequest(Irp, Request);
 
     /*
      * Post the IRP to our Ioq; we do this here instead of at FSP_LEAVE_MJ time,
@@ -747,7 +752,7 @@ NTSTATUS FspVolumeWork(
     if (!FspIoqPostIrpEx(FsvolDeviceExtension->Ioq, Irp, BestEffort, &Result))
     {
         Request->Hint = 0;
-        FspIrpRequest(Irp) = 0;
+        FspIrpSetRequest(Irp, 0);
     }
 
     DEBUGLOG("%s(Irp=%p) = %s",