- Revert 44301
[reactos.git] / drivers / network / tcpip / tcpip / dispatch.c
index 2371495..c6f2d2a 100644 (file)
 #include "precomp.h"
 #include <pseh/pseh2.h>
 
+
+NTSTATUS IRPFinish( PIRP Irp, NTSTATUS Status ) {
+    KIRQL OldIrql;
+
+    Irp->IoStatus.Status = Status;
+
+    if( Status == STATUS_PENDING )
+       IoMarkIrpPending( Irp );
+    else {
+        IoAcquireCancelSpinLock(&OldIrql);
+       (void)IoSetCancelRoutine( Irp, NULL );
+        IoReleaseCancelSpinLock(OldIrql);
+
+       IoCompleteRequest( Irp, IO_NETWORK_INCREMENT );
+    }
+
+    return Status;
+}
+
 NTSTATUS DispPrepareIrpForCancel(
     PTRANSPORT_CONTEXT Context,
     PIRP Irp,
@@ -72,7 +91,6 @@ VOID DispDataRequestComplete(
 {
     PIRP Irp;
     PIO_STACK_LOCATION IrpSp;
-    PTRANSPORT_CONTEXT TranContext;
     KIRQL OldIrql;
 
     TI_DbgPrint(DEBUG_IRP, ("Called for irp %x (%x, %d).\n",
@@ -80,7 +98,6 @@ VOID DispDataRequestComplete(
 
     Irp         = Context;
     IrpSp       = IoGetCurrentIrpStackLocation(Irp);
-    TranContext = (PTRANSPORT_CONTEXT)IrpSp->FileObject->FsContext;
 
     IoAcquireCancelSpinLock(&OldIrql);
 
@@ -102,26 +119,6 @@ VOID DispDataRequestComplete(
     TI_DbgPrint(DEBUG_IRP, ("Done Completing IRP\n"));
 }
 
-VOID DispDoDisconnect( PVOID Data ) {
-    PDISCONNECT_TYPE DisType = (PDISCONNECT_TYPE)Data;
-
-    TI_DbgPrint(DEBUG_IRP, ("PostCancel: DoDisconnect\n"));
-    TcpipRecursiveMutexEnter(&TCPLock, TRUE);
-    TCPDisconnect
-       ( DisType->Context,
-         DisType->Type,
-         NULL,
-         NULL,
-         DispDataRequestComplete,
-         DisType->Irp );
-    TcpipRecursiveMutexLeave(&TCPLock);
-    TI_DbgPrint(DEBUG_IRP, ("PostCancel: DoDisconnect done\n"));
-
-    DispDataRequestComplete(DisType->Irp, STATUS_CANCELLED, 0);
-
-    exFreePool(DisType);
-}
-
 VOID NTAPI DispCancelRequest(
     PDEVICE_OBJECT Device,
     PIRP Irp)
@@ -136,7 +133,9 @@ VOID NTAPI DispCancelRequest(
     PTRANSPORT_CONTEXT TranContext;
     PFILE_OBJECT FileObject;
     UCHAR MinorFunction;
-    PDISCONNECT_TYPE DisType;
+    BOOLEAN DequeuedIrp = TRUE;
+
+    IoReleaseCancelSpinLock(Irp->CancelIrql);
 
     TI_DbgPrint(DEBUG_IRP, ("Called.\n"));
 
@@ -159,22 +158,8 @@ VOID NTAPI DispCancelRequest(
     switch(MinorFunction) {
     case TDI_SEND:
     case TDI_RECEIVE:
-        DisType = exAllocatePool(NonPagedPool, sizeof(DISCONNECT_TYPE));
-        if (DisType)
-        {
-           DisType->Type = TDI_DISCONNECT_RELEASE |
-              ((MinorFunction == TDI_RECEIVE) ? TDI_DISCONNECT_ABORT : 0);
-           DisType->Context = TranContext->Handle.ConnectionContext;
-           DisType->Irp = Irp;
-
-           TCPRemoveIRP( TranContext->Handle.ConnectionContext, Irp );
-
-            if (!ChewCreate(DispDoDisconnect, DisType))
-                exFreePool(DisType);
-        }
-
-       IoReleaseCancelSpinLock(Irp->CancelIrql);
-        return;
+       DequeuedIrp = TCPRemoveIRP( TranContext->Handle.ConnectionContext, Irp );
+        break;
 
     case TDI_SEND_DATAGRAM:
         if (FileObject->FsContext2 != (PVOID)TDI_TRANSPORT_ADDRESS_FILE) {
@@ -182,7 +167,7 @@ VOID NTAPI DispCancelRequest(
             break;
         }
 
-        DGRemoveIRP(TranContext->Handle.AddressHandle, Irp);
+        DequeuedIrp = DGRemoveIRP(TranContext->Handle.AddressHandle, Irp);
         break;
 
     case TDI_RECEIVE_DATAGRAM:
@@ -191,20 +176,21 @@ VOID NTAPI DispCancelRequest(
             break;
         }
 
-        DGRemoveIRP(TranContext->Handle.AddressHandle, Irp);
+        DequeuedIrp = DGRemoveIRP(TranContext->Handle.AddressHandle, Irp);
         break;
 
     case TDI_CONNECT:
-        TCPRemoveIRP(TranContext->Handle.ConnectionContext, Irp);
+        DequeuedIrp = TCPRemoveIRP(TranContext->Handle.ConnectionContext, Irp);
         break;
 
     default:
         TI_DbgPrint(MIN_TRACE, ("Unknown IRP. MinorFunction (0x%X).\n", MinorFunction));
+        ASSERT(FALSE);
         break;
     }
 
-    IoReleaseCancelSpinLock(Irp->CancelIrql);
-    IRPFinish(Irp, STATUS_CANCELLED);
+    if (DequeuedIrp)
+       IRPFinish(Irp, STATUS_CANCELLED);
 
     TI_DbgPrint(MAX_TRACE, ("Leaving.\n"));
 }
@@ -224,7 +210,8 @@ VOID NTAPI DispCancelListenRequest(
     PTRANSPORT_CONTEXT TranContext;
     PFILE_OBJECT FileObject;
     PCONNECTION_ENDPOINT Connection;
-    /*NTSTATUS Status = STATUS_SUCCESS;*/
+
+    IoReleaseCancelSpinLock(Irp->CancelIrql);
 
     TI_DbgPrint(DEBUG_IRP, ("Called.\n"));
 
@@ -243,16 +230,12 @@ VOID NTAPI DispCancelListenRequest(
     /* Try canceling the request */
     Connection = (PCONNECTION_ENDPOINT)TranContext->Handle.ConnectionContext;
 
-    TCPRemoveIRP(Connection, Irp);
-
-    TCPAbortListenForSocket(
-           Connection->AddressFile->Listener,
-           Connection );
-
-    IoReleaseCancelSpinLock(Irp->CancelIrql);
-
-    Irp->IoStatus.Information = 0;
-    IRPFinish(Irp, STATUS_CANCELLED);
+    if (TCPAbortListenForSocket(Connection->AddressFile->Listener,
+                                Connection))
+    {
+        Irp->IoStatus.Information = 0;
+        IRPFinish(Irp, STATUS_CANCELLED);
+    }
 
     TI_DbgPrint(MAX_TRACE, ("Leaving.\n"));
 }
@@ -291,6 +274,7 @@ NTSTATUS DispTdiAssociateAddress(
   PFILE_OBJECT FileObject;
   PADDRESS_FILE AddrFile = NULL;
   NTSTATUS Status;
+  KIRQL OldIrql;
 
   TI_DbgPrint(DEBUG_IRP, ("Called.\n"));
 
@@ -310,11 +294,6 @@ NTSTATUS DispTdiAssociateAddress(
     return STATUS_INVALID_PARAMETER;
   }
 
-  if (Connection->AddressFile) {
-    TI_DbgPrint(MID_TRACE, ("An address file is already asscociated.\n"));
-    return STATUS_INVALID_PARAMETER;
-  }
-
   Parameters = (PTDI_REQUEST_KERNEL_ASSOCIATE)&IrpSp->Parameters;
 
   Status = ObReferenceObjectByHandle(
@@ -330,8 +309,18 @@ NTSTATUS DispTdiAssociateAddress(
     return STATUS_INVALID_PARAMETER;
   }
 
+  KeAcquireSpinLock(&Connection->Lock, &OldIrql);
+
+  if (Connection->AddressFile) {
+    ObDereferenceObject(FileObject);
+    KeReleaseSpinLock(&Connection->Lock, OldIrql);
+    TI_DbgPrint(MID_TRACE, ("An address file is already asscociated.\n"));
+    return STATUS_INVALID_PARAMETER;
+  }
+
   if (FileObject->FsContext2 != (PVOID)TDI_TRANSPORT_ADDRESS_FILE) {
     ObDereferenceObject(FileObject);
+    KeReleaseSpinLock(&Connection->Lock, OldIrql);
     TI_DbgPrint(MID_TRACE, ("Bad address file object. Magic (0x%X).\n",
       FileObject->FsContext2));
     return STATUS_INVALID_PARAMETER;
@@ -342,17 +331,21 @@ NTSTATUS DispTdiAssociateAddress(
   TranContext = FileObject->FsContext;
   if (!TranContext) {
     ObDereferenceObject(FileObject);
+    KeReleaseSpinLock(&Connection->Lock, OldIrql);
     TI_DbgPrint(MID_TRACE, ("Bad transport context.\n"));
     return STATUS_INVALID_PARAMETER;
   }
 
   AddrFile = (PADDRESS_FILE)TranContext->Handle.AddressHandle;
   if (!AddrFile) {
+      KeReleaseSpinLock(&Connection->Lock, OldIrql);
       ObDereferenceObject(FileObject);
       TI_DbgPrint(MID_TRACE, ("No address file object.\n"));
       return STATUS_INVALID_PARAMETER;
   }
 
+  KeAcquireSpinLockAtDpcLevel(&AddrFile->Lock);
+
   Connection->AddressFile = AddrFile;
 
   /* Add connection endpoint to the address file */
@@ -361,6 +354,9 @@ NTSTATUS DispTdiAssociateAddress(
   /* FIXME: Maybe do this in DispTdiDisassociateAddress() instead? */
   ObDereferenceObject(FileObject);
 
+  KeReleaseSpinLockFromDpcLevel(&AddrFile->Lock);
+  KeReleaseSpinLock(&Connection->Lock, OldIrql);
+
   return Status;
 }
 
@@ -387,8 +383,6 @@ NTSTATUS DispTdiConnect(
 
   /* Get associated connection endpoint file object. Quit if none exists */
 
-  TcpipRecursiveMutexEnter( &TCPLock, TRUE );
-
   TranContext = IrpSp->FileObject->FsContext;
   if (!TranContext) {
     TI_DbgPrint(MID_TRACE, ("Bad transport context.\n"));
@@ -419,8 +413,6 @@ NTSTATUS DispTdiConnect(
   }
 
 done:
-  TcpipRecursiveMutexLeave( &TCPLock );
-
   if (Status != STATUS_PENDING) {
       DispDataRequestComplete(Irp, Status, 0);
   } else
@@ -445,6 +437,7 @@ NTSTATUS DispTdiDisassociateAddress(
   PCONNECTION_ENDPOINT Connection;
   PTRANSPORT_CONTEXT TranContext;
   PIO_STACK_LOCATION IrpSp;
+  KIRQL OldIrql;
 
   TI_DbgPrint(DEBUG_IRP, ("Called.\n"));
 
@@ -464,17 +457,26 @@ NTSTATUS DispTdiDisassociateAddress(
     return STATUS_INVALID_PARAMETER;
   }
 
+  KeAcquireSpinLock(&Connection->Lock, &OldIrql);
+
   if (!Connection->AddressFile) {
+    KeReleaseSpinLock(&Connection->Lock, OldIrql);
     TI_DbgPrint(MID_TRACE, ("No address file is asscociated.\n"));
     return STATUS_INVALID_PARAMETER;
   }
 
+  KeAcquireSpinLockAtDpcLevel(&Connection->AddressFile->Lock);
+
   /* Remove this connection from the address file */
   Connection->AddressFile->Connection = NULL;
 
+  KeReleaseSpinLockFromDpcLevel(&Connection->AddressFile->Lock);
+
   /* Remove the address file from this connection */
   Connection->AddressFile = NULL;
 
+  KeReleaseSpinLock(&Connection->Lock, OldIrql);
+
   return STATUS_SUCCESS;
 }
 
@@ -500,8 +502,6 @@ NTSTATUS DispTdiDisconnect(
   IrpSp = IoGetCurrentIrpStackLocation(Irp);
   DisReq = (PTDI_REQUEST_KERNEL_DISCONNECT)&IrpSp->Parameters;
 
-  TcpipRecursiveMutexEnter( &TCPLock, TRUE );
-
   /* Get associated connection endpoint file object. Quit if none exists */
 
   TranContext = IrpSp->FileObject->FsContext;
@@ -527,8 +527,6 @@ NTSTATUS DispTdiDisconnect(
       Irp );
 
 done:
-   TcpipRecursiveMutexLeave( &TCPLock );
-
    if (Status != STATUS_PENDING) {
        DispDataRequestComplete(Irp, Status, 0);
    } else
@@ -555,6 +553,7 @@ NTSTATUS DispTdiListen(
   PTRANSPORT_CONTEXT TranContext;
   PIO_STACK_LOCATION IrpSp;
   NTSTATUS Status = STATUS_SUCCESS;
+  KIRQL OldIrql;
 
   TI_DbgPrint(DEBUG_IRP, ("Called.\n"));
 
@@ -562,8 +561,6 @@ NTSTATUS DispTdiListen(
 
   /* Get associated connection endpoint file object. Quit if none exists */
 
-  TcpipRecursiveMutexEnter( &TCPLock, TRUE );
-
   TranContext = IrpSp->FileObject->FsContext;
   if (TranContext == NULL)
     {
@@ -582,15 +579,23 @@ NTSTATUS DispTdiListen(
 
   Parameters = (PTDI_REQUEST_KERNEL)&IrpSp->Parameters;
 
-  TI_DbgPrint(MIN_TRACE, ("Connection->AddressFile: %x\n",
-                         Connection->AddressFile ));
-  ASSERT(Connection->AddressFile);
-
   Status = DispPrepareIrpForCancel
       (TranContext->Handle.ConnectionContext,
        Irp,
        (PDRIVER_CANCEL)DispCancelListenRequest);
 
+  KeAcquireSpinLock(&Connection->Lock, &OldIrql);
+
+  if (Connection->AddressFile == NULL)
+  {
+     TI_DbgPrint(MID_TRACE, ("No associated address file\n"));
+     KeReleaseSpinLock(&Connection->Lock, OldIrql);
+     Status = STATUS_INVALID_PARAMETER;
+     goto done;
+  }
+
+  KeAcquireSpinLockAtDpcLevel(&Connection->AddressFile->Lock);
+
   /* Listening will require us to create a listening socket and store it in
    * the address file.  It will be signalled, and attempt to complete an irp
    * when a new connection arrives. */
@@ -627,9 +632,10 @@ NTSTATUS DispTdiListen(
            Irp );
   }
 
-done:
-  TcpipRecursiveMutexLeave( &TCPLock );
+  KeReleaseSpinLockFromDpcLevel(&Connection->AddressFile->Lock);
+  KeReleaseSpinLock(&Connection->Lock, OldIrql);
 
+done:
   if (Status != STATUS_PENDING) {
       DispDataRequestComplete(Irp, Status, 0);
   } else
@@ -656,19 +662,15 @@ NTSTATUS DispTdiQueryInformation(
   PTDI_REQUEST_KERNEL_QUERY_INFORMATION Parameters;
   PTRANSPORT_CONTEXT TranContext;
   PIO_STACK_LOCATION IrpSp;
-  NTSTATUS Status;
 
   TI_DbgPrint(DEBUG_IRP, ("Called.\n"));
 
   IrpSp = IoGetCurrentIrpStackLocation(Irp);
   Parameters = (PTDI_REQUEST_KERNEL_QUERY_INFORMATION)&IrpSp->Parameters;
 
-  TcpipRecursiveMutexEnter( &TCPLock, TRUE );
-
   TranContext = IrpSp->FileObject->FsContext;
   if (!TranContext) {
     TI_DbgPrint(MID_TRACE, ("Bad transport context.\n"));
-    TcpipRecursiveMutexLeave(&TCPLock);
     return STATUS_INVALID_PARAMETER;
   }
 
@@ -686,7 +688,6 @@ NTSTATUS DispTdiQueryInformation(
             (FIELD_OFFSET(TDI_ADDRESS_INFO, Address.Address[0].Address) +
              sizeof(TDI_ADDRESS_IP))) {
           TI_DbgPrint(MID_TRACE, ("MDL buffer too small.\n"));
-          TcpipRecursiveMutexLeave(&TCPLock);
           return STATUS_BUFFER_TOO_SMALL;
         }
 
@@ -705,23 +706,18 @@ NTSTATUS DispTdiQueryInformation(
                        RtlZeroMemory(
                                &Address->Address[0].Address[0].sin_zero,
                                sizeof(Address->Address[0].Address[0].sin_zero));
-                       TcpipRecursiveMutexLeave(&TCPLock);
                        return STATUS_SUCCESS;
 
           case TDI_CONNECTION_FILE:
             Endpoint =
                                (PCONNECTION_ENDPOINT)TranContext->Handle.ConnectionContext;
-                       TCPGetSockAddress( Endpoint, (PTRANSPORT_ADDRESS)Address, FALSE );
-                       DbgPrint("Returning socket address %x\n", Address->Address[0].Address[0].in_addr);
                        RtlZeroMemory(
                                &Address->Address[0].Address[0].sin_zero,
                                sizeof(Address->Address[0].Address[0].sin_zero));
-                       TcpipRecursiveMutexLeave(&TCPLock);
-                       return STATUS_SUCCESS;
+                       return TCPGetSockAddress( Endpoint, (PTRANSPORT_ADDRESS)Address, FALSE );
 
           default:
             TI_DbgPrint(MIN_TRACE, ("Invalid transport context\n"));
-            TcpipRecursiveMutexLeave(&TCPLock);
             return STATUS_INVALID_PARAMETER;
         }
       }
@@ -736,7 +732,6 @@ NTSTATUS DispTdiQueryInformation(
             (FIELD_OFFSET(TDI_CONNECTION_INFORMATION, RemoteAddress) +
              sizeof(PVOID))) {
           TI_DbgPrint(MID_TRACE, ("MDL buffer too small (ptr).\n"));
-          TcpipRecursiveMutexLeave(&TCPLock);
           return STATUS_BUFFER_TOO_SMALL;
         }
 
@@ -756,24 +751,18 @@ NTSTATUS DispTdiQueryInformation(
 
           default:
             TI_DbgPrint(MIN_TRACE, ("Invalid transport context\n"));
-            TcpipRecursiveMutexLeave(&TCPLock);
             return STATUS_INVALID_PARAMETER;
         }
 
         if (!Endpoint) {
           TI_DbgPrint(MID_TRACE, ("No connection object.\n"));
-          TcpipRecursiveMutexLeave(&TCPLock);
           return STATUS_INVALID_PARAMETER;
         }
 
-        Status = TCPGetSockAddress( Endpoint, AddressInfo->RemoteAddress, TRUE );
-
-        TcpipRecursiveMutexLeave(&TCPLock);
-        return Status;
+        return TCPGetSockAddress( Endpoint, AddressInfo->RemoteAddress, TRUE );
       }
   }
 
-  TcpipRecursiveMutexLeave(&TCPLock);
   return STATUS_NOT_IMPLEMENTED;
 }
 
@@ -799,8 +788,6 @@ NTSTATUS DispTdiReceive(
   IrpSp = IoGetCurrentIrpStackLocation(Irp);
   ReceiveInfo = (PTDI_REQUEST_KERNEL_RECEIVE)&(IrpSp->Parameters);
 
-  TcpipRecursiveMutexEnter( &TCPLock, TRUE );
-
   TranContext = IrpSp->FileObject->FsContext;
   if (TranContext == NULL)
     {
@@ -836,8 +823,6 @@ NTSTATUS DispTdiReceive(
     }
 
 done:
-  TcpipRecursiveMutexLeave( &TCPLock );
-
   if (Status != STATUS_PENDING) {
       DispDataRequestComplete(Irp, Status, BytesReceived);
   } else
@@ -871,8 +856,6 @@ NTSTATUS DispTdiReceiveDatagram(
   IrpSp     = IoGetCurrentIrpStackLocation(Irp);
   DgramInfo = (PTDI_REQUEST_KERNEL_RECEIVEDG)&(IrpSp->Parameters);
 
-  TcpipRecursiveMutexEnter( &TCPLock, TRUE );
-
   TranContext = IrpSp->FileObject->FsContext;
   if (TranContext == NULL)
     {
@@ -914,8 +897,6 @@ NTSTATUS DispTdiReceiveDatagram(
     }
 
 done:
-   TcpipRecursiveMutexLeave( &TCPLock );
-
    if (Status != STATUS_PENDING) {
        DispDataRequestComplete(Irp, Status, BytesReceived);
    } else
@@ -948,8 +929,6 @@ NTSTATUS DispTdiSend(
   IrpSp = IoGetCurrentIrpStackLocation(Irp);
   SendInfo = (PTDI_REQUEST_KERNEL_SEND)&(IrpSp->Parameters);
 
-  TcpipRecursiveMutexEnter( &TCPLock, TRUE );
-
   TranContext = IrpSp->FileObject->FsContext;
   if (TranContext == NULL)
     {
@@ -990,8 +969,6 @@ NTSTATUS DispTdiSend(
     }
 
 done:
-   TcpipRecursiveMutexLeave( &TCPLock );
-
    if (Status != STATUS_PENDING) {
        DispDataRequestComplete(Irp, Status, BytesSent);
    } else
@@ -1024,8 +1001,6 @@ NTSTATUS DispTdiSendDatagram(
     IrpSp       = IoGetCurrentIrpStackLocation(Irp);
     DgramInfo   = (PTDI_REQUEST_KERNEL_SENDDG)&(IrpSp->Parameters);
 
-    TcpipRecursiveMutexEnter( &TCPLock, TRUE );
-
     TranContext = IrpSp->FileObject->FsContext;
     if (TranContext == NULL)
     {
@@ -1067,13 +1042,13 @@ NTSTATUS DispTdiSendDatagram(
                 DataBuffer,
                 BufferSize,
                 &Irp->IoStatus.Information);
-        else
+        else {
             Status = STATUS_UNSUCCESSFUL;
+            ASSERT(FALSE);
+        }
     }
 
 done:
-    TcpipRecursiveMutexLeave( &TCPLock );
-
     if (Status != STATUS_PENDING) {
         DispDataRequestComplete(Irp, Status, Irp->IoStatus.Information);
     } else
@@ -1122,7 +1097,7 @@ NTSTATUS DispTdiSetEventHandler(PIRP Irp)
   Parameters = (PTDI_REQUEST_KERNEL_SET_EVENT)&IrpSp->Parameters;
   Status     = STATUS_SUCCESS;
 
-  TcpipAcquireSpinLock(&AddrFile->Lock, &OldIrql);
+  KeAcquireSpinLock(&AddrFile->Lock, &OldIrql);
 
   /* Set the event handler. if an event handler is associated with
      a specific event, it's flag (RegisteredXxxHandler) is TRUE.
@@ -1243,7 +1218,7 @@ NTSTATUS DispTdiSetEventHandler(PIRP Irp)
     Status = STATUS_INVALID_PARAMETER;
   }
 
-  TcpipReleaseSpinLock(&AddrFile->Lock, OldIrql);
+  KeReleaseSpinLock(&AddrFile->Lock, OldIrql);
 
   return Status;
 }
@@ -1278,11 +1253,10 @@ VOID DispTdiQueryInformationExComplete(
  */
 {
     PTI_QUERY_CONTEXT QueryContext;
-    UINT Count = 0;
 
     QueryContext = (PTI_QUERY_CONTEXT)Context;
     if (NT_SUCCESS(Status)) {
-        Count = CopyBufferToBufferChain(
+        CopyBufferToBufferChain(
             QueryContext->InputMdl,
             FIELD_OFFSET(TCP_REQUEST_QUERY_INFORMATION_EX, Context),
             (PCHAR)&QueryContext->QueryInfo.Context,
@@ -1299,7 +1273,7 @@ VOID DispTdiQueryInformationExComplete(
     QueryContext->Irp->IoStatus.Information = ByteCount;
     QueryContext->Irp->IoStatus.Status      = Status;
 
-    exFreePool(QueryContext);
+    ExFreePoolWithTag(QueryContext, QUERY_CONTEXT_TAG);
 }
 
 
@@ -1362,7 +1336,7 @@ NTSTATUS DispTdiQueryInformationEx(
             IrpSp->Parameters.DeviceIoControl.Type3InputBuffer;
         OutputBuffer = Irp->UserBuffer;
 
-        QueryContext = exAllocatePool(NonPagedPool, sizeof(TI_QUERY_CONTEXT));
+        QueryContext = ExAllocatePoolWithTag(NonPagedPool, sizeof(TI_QUERY_CONTEXT), QUERY_CONTEXT_TAG);
         if (QueryContext) {
            _SEH2_TRY {
                 InputMdl = IoAllocateMdl(InputBuffer,
@@ -1425,7 +1399,7 @@ NTSTATUS DispTdiQueryInformationEx(
                 IoFreeMdl(OutputMdl);
             }
 
-            exFreePool(QueryContext);
+            ExFreePoolWithTag(QueryContext, QUERY_CONTEXT_TAG);
         } else
             Status = STATUS_INSUFFICIENT_RESOURCES;
     } else if( InputBufferLength ==
@@ -1438,7 +1412,7 @@ NTSTATUS DispTdiQueryInformationEx(
 
        Size = 0;
 
-        QueryContext = exAllocatePool(NonPagedPool, sizeof(TI_QUERY_CONTEXT));
+        QueryContext = ExAllocatePoolWithTag(NonPagedPool, sizeof(TI_QUERY_CONTEXT), QUERY_CONTEXT_TAG);
         if (!QueryContext) return STATUS_INSUFFICIENT_RESOURCES;
 
        _SEH2_TRY {
@@ -1458,7 +1432,7 @@ NTSTATUS DispTdiQueryInformationEx(
 
        if( !NT_SUCCESS(Status) || !InputMdl ) {
            if( InputMdl ) IoFreeMdl( InputMdl );
-           exFreePool(QueryContext);
+           ExFreePoolWithTag(QueryContext, QUERY_CONTEXT_TAG);
            return Status;
        }