From 88bde5b28ecaea4a937f87eba43ccbc903dd6dc0 Mon Sep 17 00:00:00 2001
From: "Jason A. Donenfeld" <Jason@zx2c4.com>
Date: Wed, 26 Jun 2019 19:25:03 +0000
Subject: [PATCH] Correct and simplify page locking

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>
---
 wintun.c | 40 ++++++++++++++++++----------------------
 1 file changed, 18 insertions(+), 22 deletions(-)

diff --git a/wintun.c b/wintun.c
index c493ae0..45a94b2 100644
--- a/wintun.c
+++ b/wintun.c
@@ -301,35 +301,31 @@ TunMapUbuffer(_Inout_ TUN_MAPPED_UBUFFER *MappedBuffer, _In_ VOID *UserAddress,
             return STATUS_SUCCESS;
         return STATUS_ALREADY_INITIALIZED;
     }
+
+    MappedBuffer->Mdl = IoAllocateMdl(UserAddress, Size, FALSE, FALSE, NULL);
+    if (!MappedBuffer->Mdl)
+        return STATUS_INSUFFICIENT_RESOURCES;
     try
     {
-        ProbeForWrite(UserAddress, Size, 1);
-        ProbeForRead(UserAddress, Size, 1);
-
-        MappedBuffer->Mdl = IoAllocateMdl(UserAddress, Size, FALSE, FALSE, NULL);
-        if (!MappedBuffer->Mdl)
-            return STATUS_INSUFFICIENT_RESOURCES;
-        MmProbeAndLockPages(MappedBuffer->Mdl, KernelMode, IoWriteAccess);
-        MappedBuffer->KernelAddress =
-            MmGetSystemAddressForMdlSafe(MappedBuffer->Mdl, NormalPagePriority | MdlMappingNoExecute);
-        if (!MappedBuffer->KernelAddress)
-        {
-            IoFreeMdl(MappedBuffer->Mdl);
-            MappedBuffer->Mdl = NULL;
-            return STATUS_INSUFFICIENT_RESOURCES;
-        }
-        MappedBuffer->UserAddress = UserAddress;
-        MappedBuffer->Size = Size;
+        MmProbeAndLockPages(MappedBuffer->Mdl, UserMode, IoWriteAccess);
     }
     except(EXCEPTION_EXECUTE_HANDLER)
     {
-        if (MappedBuffer->Mdl)
-        {
-            IoFreeMdl(MappedBuffer->Mdl);
-            MappedBuffer->Mdl = NULL;
-        }
+        IoFreeMdl(MappedBuffer->Mdl);
+        MappedBuffer->Mdl = NULL;
         return STATUS_INVALID_USER_BUFFER;
     }
+    MappedBuffer->KernelAddress =
+        MmGetSystemAddressForMdlSafe(MappedBuffer->Mdl, NormalPagePriority | MdlMappingNoExecute);
+    if (!MappedBuffer->KernelAddress)
+    {
+        MmUnlockPages(MappedBuffer->Mdl);
+        IoFreeMdl(MappedBuffer->Mdl);
+        MappedBuffer->Mdl = NULL;
+        return STATUS_INSUFFICIENT_RESOURCES;
+    }
+    MappedBuffer->UserAddress = UserAddress;
+    MappedBuffer->Size = Size;
     return STATUS_SUCCESS;
 }