Simplify IRP processing after mapping

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

1 files changed, 13 insertions, 52 deletions

diff --git a/wintun.c b/wintun.c
index 20aad05..c9a9212 100644
--- a/wintun.c
+++ b/wintun.c
@@ -263,35 +263,6 @@ TunCsqCompleteCanceledIrp(IO_CSQ *Csq, IRP *Irp)
     TunCompleteRequest(ctx, Irp, STATUS_CANCELLED, IO_NO_INCREMENT);
 }
 
-_IRQL_requires_max_(DISPATCH_LEVEL)
-_Must_inspect_result_
-static NTSTATUS
-TunGetIrpBuffer(_In_ IRP *Irp, _Out_ UCHAR **buffer, _Out_ ULONG *Size)
-{
-    TUN_MAPPED_UBUFFER *ubuffer;
-    IO_STACK_LOCATION *stack = IoGetCurrentIrpStackLocation(Irp);
-    TUN_FILE_CTX *file_ctx = (TUN_FILE_CTX *)stack->FileObject->FsContext;
-
-    switch (stack->MajorFunction)
-    {
-    case IRP_MJ_READ:
-        *Size = stack->Parameters.Read.Length;
-        ubuffer = &file_ctx->ReadBuffer;
-        break;
-    case IRP_MJ_WRITE:
-        *Size = stack->Parameters.Write.Length;
-        ubuffer = &file_ctx->WriteBuffer;
-        break;
-    default:
-        return STATUS_INVALID_PARAMETER;
-    }
-    if (*Size > ubuffer->Size)
-        return STATUS_INVALID_USER_BUFFER;
-    ASSERT(ubuffer->KernelAddress != NULL);
-    *buffer = ubuffer->KernelAddress;
-    return STATUS_SUCCESS;
-}
-
 _IRQL_requires_max_(APC_LEVEL)
 _Must_inspect_result_
 static NTSTATUS
@@ -300,9 +271,9 @@ TunMapUbuffer(_Inout_ TUN_MAPPED_UBUFFER *MappedBuffer, _In_ VOID *UserAddress,
     VOID *current_uaddr = InterlockedGetPointer(&MappedBuffer->UserAddress);
     if (current_uaddr)
     {
-        if (UserAddress == current_uaddr) // TODO: Check ThreadID
-            return STATUS_SUCCESS;
-        return STATUS_ALREADY_INITIALIZED;
+        if (UserAddress != current_uaddr || Size > MappedBuffer->Size) // TODO: Check ThreadID
+            return STATUS_ALREADY_INITIALIZED;
+        return STATUS_SUCCESS;
     }
 
     NTSTATUS status = STATUS_SUCCESS;
@@ -312,7 +283,7 @@ TunMapUbuffer(_Inout_ TUN_MAPPED_UBUFFER *MappedBuffer, _In_ VOID *UserAddress,
     current_uaddr = InterlockedGetPointer(&MappedBuffer->UserAddress);
     if (current_uaddr)
     {
-        if (UserAddress != current_uaddr) // TODO: Check ThreadID
+        if (UserAddress != current_uaddr || Size > MappedBuffer->Size) // TODO: Check ThreadID
             status = STATUS_ALREADY_INITIALIZED;
         goto err_releasemutex;
     }
@@ -398,22 +369,13 @@ _Must_inspect_result_
 static _Return_type_success_(
     return != NULL) IRP *TunRemoveNextIrp(_Inout_ TUN_CTX *Ctx, _Out_ UCHAR **Buffer, _Out_ ULONG *Size)
 {
-    IRP *irp;
-
-retry:
-    irp = IoCsqRemoveNextIrp(&Ctx->Device.ReadQueue.Csq, NULL);
+    IRP *irp = IoCsqRemoveNextIrp(&Ctx->Device.ReadQueue.Csq, NULL);
     if (!irp)
         return NULL;
-
-    NTSTATUS status = TunGetIrpBuffer(irp, Buffer, Size);
-    if (!NT_SUCCESS(status))
-    {
-        TunCompleteRequest(Ctx, irp, status, IO_NO_INCREMENT);
-        goto retry;
-    }
-
+    IO_STACK_LOCATION *stack = IoGetCurrentIrpStackLocation(irp);
+    *Size = stack->Parameters.Read.Length;
     ASSERT(irp->IoStatus.Information <= (ULONG_PTR)*Size);
-
+    *Buffer = ((TUN_FILE_CTX *)stack->FileObject->FsContext)->ReadBuffer.KernelAddress;
     return irp;
 }
 
@@ -807,12 +769,11 @@ TunDispatchWrite(_Inout_ TUN_CTX *Ctx, _Inout_ IRP *Irp)
     if (status = STATUS_FILE_FORCED_CLOSED, !(flags & TUN_FLAGS_PRESENT))
         goto cleanup_ExReleaseSpinLockShared;
 
-    UCHAR *buffer;
-    ULONG size;
-    if (!NT_SUCCESS(status = TunGetIrpBuffer(Irp, &buffer, &size)))
-        goto cleanup_ExReleaseSpinLockShared;
     IO_STACK_LOCATION *stack = IoGetCurrentIrpStackLocation(Irp);
-    MDL *mdl = ((TUN_FILE_CTX *)stack->FileObject->FsContext)->WriteBuffer.Mdl;
+    TUN_MAPPED_UBUFFER *ubuffer = &((TUN_FILE_CTX *)stack->FileObject->FsContext)->WriteBuffer;
+    UCHAR *buffer = ubuffer->KernelAddress;
+    ULONG size = stack->Parameters.Write.Length;
+
     const UCHAR *b = buffer, *b_end = buffer + size;
     typedef enum _ethtypeidx_t
     {
@@ -868,7 +829,7 @@ TunDispatchWrite(_Inout_ TUN_CTX *Ctx, _Inout_ IRP *Irp)
         }
 
         NET_BUFFER_LIST *nbl =
-            NdisAllocateNetBufferAndNetBufferList(Ctx->NBLPool, 0, 0, mdl, (ULONG)(p->Data - buffer), p->Size);
+            NdisAllocateNetBufferAndNetBufferList(Ctx->NBLPool, 0, 0, ubuffer->Mdl, (ULONG)(p->Data - buffer), p->Size);
         if (!nbl)
         {
             status = STATUS_INSUFFICIENT_RESOURCES;