From f63c9159c8c25de86ccc2cc6e7dc52b91ae49bdd Mon Sep 17 00:00:00 2001
From: Thomas Faber <thomas.faber@reactos.org>
Date: Thu, 28 Jul 2011 09:00:56 +0000
Subject: [PATCH] [KMTESTS] - probe user address in user mode - don't free the
 result buffer before the driver unlocks it. Fixes assertion in ROS for large
 buffers

svn path=/branches/GSoC_2011/KMTestSuite/; revision=52947
---
 kmtests/include/kmt_test.h      |  8 +++++---
 kmtests/kmtest/kmtest.c         | 19 ++++++++++++-------
 kmtests/kmtest_drv/kmtest_drv.c |  9 +++++----
 3 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/kmtests/include/kmt_test.h b/kmtests/include/kmt_test.h
index cc1e815869c..134645a90b3 100644
--- a/kmtests/include/kmt_test.h
+++ b/kmtests/include/kmt_test.h
@@ -173,15 +173,17 @@ BOOLEAN KmtAreInterruptsEnabled(VOID)
 
 INT __cdecl KmtVSNPrintF(PSTR Buffer, SIZE_T BufferMaxLength, PCSTR Format, va_list Arguments) KMT_FORMAT(ms_printf, 3, 0);
 #elif defined KMT_USER_MODE
-static PKMT_RESULTBUFFER KmtAllocateResultBuffer(SIZE_T LogBufferMaxLength)
+static PKMT_RESULTBUFFER KmtAllocateResultBuffer(SIZE_T ResultBufferSize)
 {
-    PKMT_RESULTBUFFER Buffer = HeapAlloc(GetProcessHeap(), 0, FIELD_OFFSET(KMT_RESULTBUFFER, LogBuffer[LogBufferMaxLength]));
+    PKMT_RESULTBUFFER Buffer = HeapAlloc(GetProcessHeap(), 0, ResultBufferSize);
+    if (!Buffer)
+        return NULL;
 
     Buffer->Successes = 0;
     Buffer->Failures = 0;
     Buffer->Skipped = 0;
     Buffer->LogBufferLength = 0;
-    Buffer->LogBufferMaxLength = LogBufferMaxLength;
+    Buffer->LogBufferMaxLength = ResultBufferSize - FIELD_OFFSET(KMT_RESULTBUFFER, LogBuffer);
 
     return Buffer;
 }
diff --git a/kmtests/kmtest/kmtest.c b/kmtests/kmtest/kmtest.c
index 3be6977f9aa..6ab770948ed 100644
--- a/kmtests/kmtest/kmtest.c
+++ b/kmtests/kmtest/kmtest.c
@@ -24,8 +24,7 @@
 #define SERVICE_PATH        L"kmtest_drv.sys"
 #define SERVICE_DESCRIPTION L"ReactOS Kernel-Mode Test Suite Driver"
 
-#define LOGBUFFER_SIZE      16364
-#define RESULTBUFFER_SIZE   FIELD_OFFSET(KMT_RESULTBUFFER, LogBuffer[LOGBUFFER_SIZE])
+#define RESULTBUFFER_SIZE   (1024 * 1024)
 
 typedef enum
 {
@@ -255,9 +254,14 @@ RunTest(
 
     assert(TestName != NULL);
 
-    ResultBuffer = KmtAllocateResultBuffer(LOGBUFFER_SIZE);
-    if (!DeviceIoControl(KmtestHandle, IOCTL_KMTEST_SET_RESULTBUFFER, ResultBuffer, RESULTBUFFER_SIZE, NULL, 0, &BytesRead, NULL))
-        error_goto(Error, cleanup);
+    if (!ResultBuffer)
+    {
+        ResultBuffer = KmtAllocateResultBuffer(RESULTBUFFER_SIZE);
+        if (!ResultBuffer)
+            error_goto(Error, cleanup);
+        if (!DeviceIoControl(KmtestHandle, IOCTL_KMTEST_SET_RESULTBUFFER, ResultBuffer, RESULTBUFFER_SIZE, NULL, 0, &BytesRead, NULL))
+            error_goto(Error, cleanup);
+    }
 
     // check test list
     TestFunction = FindTest(TestName);
@@ -275,8 +279,6 @@ cleanup:
     if (!Error)
         Error = OutputResult(TestName);
 
-    KmtFreeResultBuffer(ResultBuffer);
-
     return Error;
 }
 
@@ -367,6 +369,9 @@ cleanup:
     if (KmtestHandle)
         CloseHandle(KmtestHandle);
 
+    if (ResultBuffer)
+        KmtFreeResultBuffer(ResultBuffer);
+
     KmtCloseService(&KmtestServiceHandle);
 
     if (Error)
diff --git a/kmtests/kmtest_drv/kmtest_drv.c b/kmtests/kmtest_drv/kmtest_drv.c
index 6612567b070..05b5709da3a 100644
--- a/kmtests/kmtest_drv/kmtest_drv.c
+++ b/kmtests/kmtest_drv/kmtest_drv.c
@@ -314,9 +314,10 @@ DriverIoControl(
         {
             PKMT_DEVICE_EXTENSION DeviceExtension = DeviceObject->DeviceExtension;
 
-            DPRINT("DriverIoControl. IOCTL_KMTEST_SET_RESULTBUFFER, inlen=%lu, outlen=%lu\n",
-                     IoStackLocation->Parameters.DeviceIoControl.InputBufferLength,
-                     IoStackLocation->Parameters.DeviceIoControl.OutputBufferLength);
+            DPRINT("DriverIoControl. IOCTL_KMTEST_SET_RESULTBUFFER, buffer=%p, inlen=%lu, outlen=%lu\n",
+                    IoStackLocation->Parameters.DeviceIoControl.Type3InputBuffer,
+                    IoStackLocation->Parameters.DeviceIoControl.InputBufferLength,
+                    IoStackLocation->Parameters.DeviceIoControl.OutputBufferLength);
 
             if (DeviceExtension->Mdl)
             {
@@ -342,7 +343,7 @@ DriverIoControl(
 
             _SEH2_TRY
             {
-                MmProbeAndLockPages(DeviceExtension->Mdl, KernelMode, IoModifyAccess);
+                MmProbeAndLockPages(DeviceExtension->Mdl, UserMode, IoModifyAccess);
             }
             _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
             {
-- 
2.17.1