modules/rostests/kmtests/ntos_mm/ZwAllocateVirtualMemory.c

   1 /*
   2  * PROJECT:         ReactOS kernel-mode tests
   3  * LICENSE:         GPLv2+ - See COPYING in the top level directory
   4  * PURPOSE:         Kernel-Mode Test Suite ZwAllocateVirtualMemory/ZwFreeVirtualMemory
   5  * PROGRAMMER:      Nikolay Borisov <nib9@aber.ac.uk>
   6  */
   7
   8 #include <kmt_test.h>
   9
  10 #define ROUND_DOWN(n,align) (((ULONG_PTR)n) & ~((align) - 1l))
  11 #define DEFAULT_ALLOC_SIZE 200
  12 #define IGNORE -1
  13 #define PAGE_NOPROT 0x0 //MEM_RESERVE has this type of "protection"
  14
  15 /* These are being used in ZwMapViewOfSection as well */
  16 const char TestString[] = "TheLongBrownFoxJumpedTheWhiteRabbitTheLongBrownFoxJumpedTheWhiteRabbitTheLongBrownFoxJumpedTheWhiteRabbitTheLongBrownFoxJumpedTheWhiteRabbitTheLongBrownFoxJumpedTheWhiteRabbitTheLongBrownFoxJumpedThe";
  17 const ULONG TestStringSize = sizeof(TestString);
  18
  19 VOID Test_ZwAllocateVirtualMemory(VOID);
  20
  21 typedef struct _TEST_CONTEXT
  22 {
  23     HANDLE ProcessHandle;
  24     SIZE_T RegionSize;
  25     ULONG AllocationType;
  26     ULONG Protect;
  27     PVOID Bases[1024];
  28     SHORT ThreadId;
  29 } TEST_CONTEXT, *PTEST_CONTEXT;
  30
  31
  32 #define ALLOC_MEMORY_WITH_FREE(ProcessHandle, BaseAddress, ZeroBits, RegionSize, AllocationType, Protect, RetStatus, FreeStatus)   \
  33     do {                                                                                                                   \
  34         PVOID __BaseSave = BaseAddress;                                                                                    \
  35         Status = ZwAllocateVirtualMemory(ProcessHandle, &BaseAddress, ZeroBits, &RegionSize, AllocationType, Protect);     \
  36         ok_eq_hex(Status, RetStatus);                                                                                      \
  37         if (__BaseSave != NULL)                                                                                            \
  38             ok_eq_pointer(BaseAddress, __BaseSave);                                                                        \
  39         else if (!NT_SUCCESS(Status))                                                                                      \
  40             ok_eq_pointer(BaseAddress, NULL);                                                                              \
  41         RegionSize = 0;                                                                                                    \
  42         Status = ZwFreeVirtualMemory(ProcessHandle, &BaseAddress, &RegionSize, MEM_RELEASE);                               \
  43         if (FreeStatus != IGNORE) ok_eq_hex(Status, FreeStatus);                                                           \
  44         BaseAddress = NULL;                                                                                                \
  45         RegionSize = DEFAULT_ALLOC_SIZE;                                                                                   \
  46     } while (0)                                                                                                            \
  47
  48
  49
  50 static
  51 BOOLEAN
  52 CheckBuffer(PVOID Buffer, SIZE_T Size, UCHAR Value)
  53 {
  54     PUCHAR Array = Buffer;
  55     SIZE_T i;
  56
  57     for (i = 0; i < Size; i++)
  58     {
  59         if (Array[i] != Value)
  60         {
  61             trace("Expected %x, found %x at offset %lu\n", Value, Array[i], (ULONG)i);
  62             return FALSE;
  63         }
  64     }
  65     return TRUE;
  66 }
  67
  68 static
  69 SIZE_T
  70 CheckBufferRead(CONST VOID *Source, CONST VOID *Destination, SIZE_T Length, NTSTATUS ExpectedStatus)
  71 {
  72     SIZE_T Match = 0;
  73
  74     KmtStartSeh()
  75         Match = RtlCompareMemory(Source, Destination, Length);
  76     KmtEndSeh(ExpectedStatus);
  77
  78     return Match;
  79 }
  80
  81 static
  82 VOID
  83 CheckBufferReadWrite(PVOID Destination, CONST VOID *Source, SIZE_T Length, NTSTATUS ExpectedStatus)
  84 {
  85     //do a little bit of writing/reading to memory
  86     SIZE_T Match = 0;
  87
  88     KmtStartSeh()
  89         RtlCopyMemory(Destination, Source, Length);
  90     KmtEndSeh(ExpectedStatus);
  91
  92     Match = CheckBufferRead(Source, Destination, Length, ExpectedStatus);
  93     if (ExpectedStatus == STATUS_SUCCESS) ok_eq_int(Match, Length);
  94 }
  95
  96
  97 static
  98 VOID
  99 SimpleErrorChecks(VOID)
 100 {
 101     NTSTATUS Status;
 102     PVOID Base = NULL;
 103     SIZE_T RegionSize = DEFAULT_ALLOC_SIZE;
 104
 105     //HANDLE TESTS
 106     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_SUCCESS, STATUS_SUCCESS);
 107     ALLOC_MEMORY_WITH_FREE(NULL, Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_INVALID_HANDLE, STATUS_INVALID_HANDLE);
 108     ALLOC_MEMORY_WITH_FREE((HANDLE)(ULONG_PTR)0xDEADBEEFDEADBEEFull, Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_INVALID_HANDLE, STATUS_INVALID_HANDLE);
 109
 110     //BASE ADDRESS TESTS
 111     Base = (PVOID)0x00567A20;
 112     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_CONFLICTING_ADDRESSES, STATUS_UNABLE_TO_DELETE_SECTION);
 113
 114     Base = (PVOID) 0x60000000;
 115     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_SUCCESS, STATUS_SUCCESS);
 116
 117     Base = (PVOID)((char *)MmSystemRangeStart + 200);
 118     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_INVALID_PARAMETER_2, STATUS_INVALID_PARAMETER_2);
 119
 120     /* http://jira.reactos.org/browse/CORE-6814 */
 121     RegionSize = 0x1000;
 122     Base = Test_ZwAllocateVirtualMemory;
 123     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, MEM_COMMIT, PAGE_READWRITE, STATUS_INVALID_PARAMETER_2, STATUS_INVALID_PARAMETER_2);
 124
 125     //ZERO BITS TESTS
 126     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 21, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_NO_MEMORY, STATUS_MEMORY_NOT_ALLOCATED);
 127     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 22, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_INVALID_PARAMETER_3, STATUS_MEMORY_NOT_ALLOCATED);
 128     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, -1, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_INVALID_PARAMETER_3, STATUS_MEMORY_NOT_ALLOCATED);
 129     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 3, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_SUCCESS, STATUS_SUCCESS);
 130
 131     //REGION SIZE TESTS
 132     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_SUCCESS, STATUS_SUCCESS);
 133     RegionSize = -1;
 134     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_INVALID_PARAMETER_4, STATUS_MEMORY_NOT_ALLOCATED);
 135     RegionSize = 0;
 136     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_INVALID_PARAMETER_4, STATUS_MEMORY_NOT_ALLOCATED);
 137     RegionSize = 0xFFFFFFFF; // 4 gb  is invalid
 138     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_INVALID_PARAMETER_4, STATUS_MEMORY_NOT_ALLOCATED);
 139
 140     //Allocation type tests
 141     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, MEM_PHYSICAL, PAGE_READWRITE, STATUS_INVALID_PARAMETER_5, STATUS_MEMORY_NOT_ALLOCATED);
 142     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESET), PAGE_READWRITE, STATUS_INVALID_PARAMETER_5, STATUS_MEMORY_NOT_ALLOCATED);
 143     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, 0, PAGE_READWRITE, STATUS_INVALID_PARAMETER_5, STATUS_MEMORY_NOT_ALLOCATED);
 144     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, MEM_TOP_DOWN, PAGE_READWRITE, STATUS_INVALID_PARAMETER_5, STATUS_MEMORY_NOT_ALLOCATED);
 145     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_TOP_DOWN | MEM_RESET), PAGE_READWRITE, STATUS_INVALID_PARAMETER_5, STATUS_MEMORY_NOT_ALLOCATED);
 146     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_TOP_DOWN | MEM_COMMIT), PAGE_READWRITE, STATUS_SUCCESS, STATUS_SUCCESS);
 147     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_PHYSICAL | MEM_RESERVE), PAGE_READWRITE, STATUS_SUCCESS, STATUS_SUCCESS);
 148     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_PHYSICAL | MEM_COMMIT), PAGE_READWRITE, STATUS_INVALID_PARAMETER_5, STATUS_MEMORY_NOT_ALLOCATED);
 149     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_RESET | MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE, STATUS_INVALID_PARAMETER_5, STATUS_MEMORY_NOT_ALLOCATED);
 150     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, -1, PAGE_READWRITE, STATUS_INVALID_PARAMETER_5, STATUS_MEMORY_NOT_ALLOCATED);
 151     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize,  MEM_COMMIT, PAGE_READWRITE, STATUS_SUCCESS, STATUS_SUCCESS);
 152     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize,  MEM_RESERVE, PAGE_READWRITE, STATUS_SUCCESS, STATUS_SUCCESS);
 153     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize,  MEM_RESERVE, PAGE_WRITECOPY, STATUS_INVALID_PAGE_PROTECTION, STATUS_MEMORY_NOT_ALLOCATED);
 154     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize,  MEM_RESERVE, PAGE_EXECUTE_WRITECOPY, STATUS_INVALID_PAGE_PROTECTION, STATUS_MEMORY_NOT_ALLOCATED);
 155
 156     //Memory protection tests
 157     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), 0, STATUS_INVALID_PAGE_PROTECTION, STATUS_MEMORY_NOT_ALLOCATED);
 158     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), -1, STATUS_INVALID_PAGE_PROTECTION, STATUS_MEMORY_NOT_ALLOCATED);
 159     if (!KmtIsCheckedBuild)
 160     {
 161         ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), (PAGE_NOACCESS | PAGE_GUARD), STATUS_INVALID_PAGE_PROTECTION, STATUS_MEMORY_NOT_ALLOCATED);
 162         ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), (PAGE_NOACCESS | PAGE_WRITECOMBINE), STATUS_INVALID_PAGE_PROTECTION, STATUS_MEMORY_NOT_ALLOCATED);
 163     }
 164     ALLOC_MEMORY_WITH_FREE(NtCurrentProcess(), Base, 0, RegionSize, (MEM_COMMIT | MEM_RESERVE), (PAGE_READONLY | PAGE_WRITECOMBINE), STATUS_SUCCESS, STATUS_SUCCESS);
 165 }
 166
 167
 168 static
 169 NTSTATUS
 170 SimpleAllocation(VOID)
 171 {
 172     NTSTATUS Status;
 173     PVOID Base = NULL;
 174     SIZE_T RegionSize = DEFAULT_ALLOC_SIZE;
 175
 176     //////////////////////////////////////////////////////////////////////////
 177     //Normal operation
 178     //////////////////////////////////////////////////////////////////////////
 179     Status = ZwAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &RegionSize, MEM_COMMIT, PAGE_READWRITE);
 180     ok_eq_size(RegionSize, 4096);
 181
 182     //check for the zero-filled pages
 183     ok_bool_true(CheckBuffer(Base, RegionSize, 0), "The buffer is not zero-filled");
 184
 185     CheckBufferReadWrite(Base, TestString, TestStringSize, STATUS_SUCCESS);
 186
 187     // try freeing
 188     RegionSize = 0;
 189     Status = ZwFreeVirtualMemory(NtCurrentProcess(), &Base, &RegionSize, MEM_RELEASE);
 190     ok_eq_hex(Status, STATUS_SUCCESS);
 191     ok_eq_size(RegionSize, PAGE_SIZE);
 192
 193     //////////////////////////////////////////////////////////////////////////
 194     // COMMIT AND RESERVE SCENARIO AND STATE CHANGE
 195     //////////////////////////////////////////////////////////////////////////
 196     //reserve and then commit
 197     Base = NULL;
 198     RegionSize = DEFAULT_ALLOC_SIZE;
 199     Status = ZwAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &RegionSize, MEM_RESERVE, PAGE_READWRITE);
 200     Test_NtQueryVirtualMemory(Base, RegionSize, MEM_RESERVE, PAGE_NOPROT);
 201     CheckBufferReadWrite(Base, TestString, TestStringSize, STATUS_ACCESS_VIOLATION);
 202
 203
 204     Status = ZwAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &RegionSize, MEM_COMMIT, PAGE_READWRITE);
 205     CheckBufferReadWrite(Base, TestString, TestStringSize, STATUS_SUCCESS);
 206     Test_NtQueryVirtualMemory(Base, RegionSize, MEM_COMMIT, PAGE_READWRITE);
 207
 208     RegionSize = 0;
 209     ZwFreeVirtualMemory(NtCurrentProcess(), &Base, &RegionSize, MEM_RELEASE);
 210
 211     //////////////////////////////////////////////////////////////////////////
 212     // TRY READING/WRITING TO INVALID PROTECTION PAGES
 213     //////////////////////////////////////////////////////////////////////////
 214     RegionSize = DEFAULT_ALLOC_SIZE;
 215     Base = NULL;
 216     ZwAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_NOACCESS);
 217
 218     KmtStartSeh()
 219         RtlCopyMemory(Base, TestString, TestStringSize);
 220     KmtEndSeh(STATUS_ACCESS_VIOLATION);
 221
 222     Test_NtQueryVirtualMemory(Base, RegionSize, MEM_COMMIT, PAGE_NOACCESS);
 223     CheckBufferRead(Base, TestString, TestStringSize, STATUS_ACCESS_VIOLATION);
 224
 225     RegionSize = 0;
 226     ZwFreeVirtualMemory(NtCurrentProcess(), &Base, &RegionSize, MEM_RELEASE);
 227
 228     ZwAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READONLY);
 229     KmtStartSeh()
 230         RtlCopyMemory(Base, TestString, TestStringSize);
 231     KmtEndSeh(STATUS_ACCESS_VIOLATION);
 232
 233     Test_NtQueryVirtualMemory(Base, RegionSize, MEM_COMMIT, PAGE_READONLY);
 234
 235     ok_bool_true(CheckBuffer(Base, TestStringSize, 0), "Couldn't read a read-only buffer");
 236
 237     RegionSize = 0;
 238     ZwFreeVirtualMemory(NtCurrentProcess(), &Base, &RegionSize, MEM_RELEASE);
 239
 240     //////////////////////////////////////////////////////////////////////////
 241     // GUARD PAGES
 242     //////////////////////////////////////////////////////////////////////////
 243
 244     RegionSize = 1000;
 245     Base = NULL;
 246     ZwAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &RegionSize, (MEM_COMMIT | MEM_RESERVE), (PAGE_GUARD | PAGE_READWRITE));
 247
 248     Test_NtQueryVirtualMemory(Base, RegionSize, MEM_COMMIT, (PAGE_GUARD | PAGE_READWRITE));
 249     KmtStartSeh()
 250         RtlCopyMemory(Base, TestString, TestStringSize);
 251     KmtEndSeh(STATUS_GUARD_PAGE_VIOLATION);
 252
 253     Test_NtQueryVirtualMemory(Base, RegionSize, MEM_COMMIT, PAGE_READWRITE);
 254
 255     KmtStartSeh()
 256         RtlCopyMemory(Base, TestString, TestStringSize);
 257     KmtEndSeh(STATUS_SUCCESS);
 258
 259     RegionSize = 0;
 260     ZwFreeVirtualMemory(NtCurrentProcess(), &Base, &RegionSize, MEM_RELEASE);
 261
 262     return Status;
 263 }
 264
 265
 266 static
 267 VOID
 268 CustomBaseAllocation(VOID)
 269 {
 270     NTSTATUS Status;
 271     SIZE_T RegionSize = 200;
 272     PVOID Base =  (PVOID) 0x60025000;
 273     PVOID ActualStartingAddress = (PVOID)ROUND_DOWN(Base, MM_ALLOCATION_GRANULARITY); //it is rounded down to the nearest allocation granularity (64k) address
 274     PVOID EndingAddress = (PVOID)(((ULONG_PTR)Base + RegionSize - 1) | (PAGE_SIZE - 1));
 275     SIZE_T ActualSize = BYTES_TO_PAGES((ULONG_PTR)EndingAddress - (ULONG_PTR)ActualStartingAddress) * PAGE_SIZE; //calculates the actual size based on the required pages
 276
 277     // allocate the memory
 278     Status = ZwAllocateVirtualMemory(NtCurrentProcess(), (PVOID *)&Base, 0, &RegionSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE);
 279     ok_eq_hex(Status, STATUS_SUCCESS);
 280     ok_eq_size(RegionSize, ActualSize);
 281     ok_eq_ulong(Base, ActualStartingAddress);
 282     Test_NtQueryVirtualMemory(ActualStartingAddress, ActualSize, MEM_COMMIT, PAGE_READWRITE);
 283
 284     // try freeing
 285     RegionSize = 0;
 286     Status = ZwFreeVirtualMemory(NtCurrentProcess(), (PVOID *)&Base, &RegionSize, MEM_RELEASE);
 287     ok_eq_hex(Status, STATUS_SUCCESS);
 288     ok_eq_ulong(RegionSize, ActualSize);
 289 }
 290
 291
 292 static
 293 NTSTATUS
 294 StressTesting(ULONG AllocationType)
 295 {
 296     NTSTATUS Status = STATUS_SUCCESS;
 297     NTSTATUS ReturnStatus = STATUS_SUCCESS;
 298     static PVOID bases[1024]; //assume we are going to allocate only 5 gigs. static here means the arrays is not allocated on the stack but in the BSS segment of the driver
 299     ULONG Index = 0;
 300     PVOID Base = NULL;
 301     SIZE_T RegionSize = 5 * 1024 * 1024; // 5 megabytes;
 302
 303     RtlZeroMemory(bases, sizeof(bases));
 304
 305     for (Index = 0; Index < RTL_NUMBER_OF(bases) && NT_SUCCESS(Status); Index++)
 306     {
 307         Status = ZwAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &RegionSize, AllocationType, PAGE_READWRITE);
 308
 309         bases[Index] = Base;
 310         if ((Index % 10) == 0)
 311         {
 312             if (AllocationType == MEM_COMMIT && NT_SUCCESS(Status))
 313             {
 314                 CheckBufferReadWrite(Base, TestString, TestStringSize, STATUS_SUCCESS);
 315             }
 316             else
 317             {
 318                 CheckBufferReadWrite(Base, TestString, TestStringSize, STATUS_ACCESS_VIOLATION);
 319             }
 320         }
 321
 322         Base = NULL;
 323     }
 324
 325     trace("Finished reserving. Error code %x. Chunks allocated: %d\n", Status, Index );
 326
 327     ReturnStatus = Status;
 328
 329     //free the allocated memory so that we can continue with the tests
 330     Status = STATUS_SUCCESS;
 331     Index = 0;
 332     while (NT_SUCCESS(Status) && Index < RTL_NUMBER_OF(bases))
 333     {
 334         RegionSize = 0;
 335         Status = ZwFreeVirtualMemory(NtCurrentProcess(), &bases[Index], &RegionSize, MEM_RELEASE);
 336         bases[Index++] = NULL;
 337     }
 338
 339     return ReturnStatus;
 340 }
 341
 342
 343 static
 344 VOID
 345 NTAPI
 346 SystemProcessTestWorker(PVOID StartContext)
 347 {
 348    NTSTATUS Status = STATUS_SUCCESS;
 349    PTEST_CONTEXT Context = (PTEST_CONTEXT)StartContext;
 350    ULONG Index = 0;
 351    PVOID Base = NULL;
 352
 353    PAGED_CODE();
 354
 355    RtlZeroMemory(Context->Bases, sizeof(Context->Bases));
 356
 357    Status = ZwAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &Context->RegionSize, Context->AllocationType, Context->Protect);
 358    ZwFreeVirtualMemory(NtCurrentProcess(), &Base, &Context->RegionSize, MEM_RELEASE);
 359    Base = NULL;
 360
 361     //if the previous allocation has failed there is no need to do the loop
 362     while (NT_SUCCESS(Status) && Index < RTL_NUMBER_OF(Context->Bases))
 363     {
 364         Status = ZwAllocateVirtualMemory(NtCurrentProcess(), &Base, 0, &Context->RegionSize, Context->AllocationType, Context->Protect);
 365
 366         Context->Bases[Index] = Base;
 367         if ((Index % 10) == 0)
 368         {
 369             if (Context->AllocationType == MEM_COMMIT)
 370             {
 371                 CheckBufferReadWrite(Base, TestString, TestStringSize, STATUS_SUCCESS);
 372             }
 373             else
 374             {
 375                 CheckBufferReadWrite(Base, TestString, TestStringSize, STATUS_ACCESS_VIOLATION);
 376             }
 377         }
 378
 379         Base = NULL;
 380         Index++;
 381     }
 382
 383     trace("[SYSTEM THREAD %d]. Error code %x. Chunks allocated: %d\n", Context->ThreadId, Status, Index);
 384
 385     //free the allocated memory so that we can continue with the tests
 386     Status = STATUS_SUCCESS;
 387     Index = 0;
 388     while (NT_SUCCESS(Status) && Index < RTL_NUMBER_OF(Context->Bases))
 389     {
 390         Context->RegionSize = 0;
 391         Status = ZwFreeVirtualMemory(NtCurrentProcess(), &Context->Bases[Index], &Context->RegionSize, MEM_RELEASE);
 392         Context->Bases[Index++] = NULL;
 393     }
 394
 395     PsTerminateSystemThread(Status);
 396 }
 397
 398
 399 static
 400 VOID
 401 KmtInitTestContext(PTEST_CONTEXT Ctx, SHORT ThreadId, ULONG RegionSize, ULONG AllocationType, ULONG Protect)
 402 {
 403     PAGED_CODE();
 404
 405     Ctx->AllocationType = AllocationType;
 406     Ctx->Protect = Protect;
 407     Ctx->RegionSize = RegionSize;
 408     Ctx->ThreadId = ThreadId;
 409 }
 410
 411
 412 static
 413 VOID
 414 SystemProcessTest(VOID)
 415 {
 416     NTSTATUS Status;
 417     HANDLE Thread1 = INVALID_HANDLE_VALUE;
 418     HANDLE Thread2 = INVALID_HANDLE_VALUE;
 419     PVOID ThreadObjects[2] = { NULL };
 420     OBJECT_ATTRIBUTES ObjectAttributes;
 421     PTEST_CONTEXT StartContext1;
 422     PTEST_CONTEXT StartContext2;
 423
 424     PAGED_CODE();
 425
 426     StartContext1 = ExAllocatePoolWithTag(PagedPool, sizeof(TEST_CONTEXT), 'tXTC');
 427     StartContext2 = ExAllocatePoolWithTag(PagedPool, sizeof(TEST_CONTEXT), 'tXTC');
 428     if (StartContext1 == NULL || StartContext2 == NULL)
 429     {
 430         trace("Error allocating space for context structs\n");
 431         goto cleanup;
 432     }
 433
 434     KmtInitTestContext(StartContext1, 1, 1 * 1024 * 1024, MEM_COMMIT, PAGE_READWRITE);
 435     KmtInitTestContext(StartContext2, 2, 3 * 1024 * 1024, MEM_COMMIT, PAGE_READWRITE);
 436     InitializeObjectAttributes(&ObjectAttributes, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);
 437
 438     Status = PsCreateSystemThread(&Thread1, THREAD_ALL_ACCESS, &ObjectAttributes, NULL, NULL, SystemProcessTestWorker, StartContext1);
 439     if (!NT_SUCCESS(Status))
 440     {
 441         trace("Error creating thread1\n");
 442         goto cleanup;
 443     }
 444
 445     Status = ObReferenceObjectByHandle(Thread1, THREAD_ALL_ACCESS, *PsThreadType, KernelMode, &ThreadObjects[0], NULL);
 446     if (!NT_SUCCESS(Status))
 447     {
 448         trace("error referencing thread1\n");
 449         goto cleanup;
 450     }
 451
 452     Status = PsCreateSystemThread(&Thread2, THREAD_ALL_ACCESS, &ObjectAttributes, NULL, NULL, SystemProcessTestWorker, StartContext2);
 453     if (!NT_SUCCESS(Status))
 454     {
 455         trace("Error creating thread2\n");
 456         goto cleanup;
 457     }
 458
 459     Status = ObReferenceObjectByHandle(Thread2, THREAD_ALL_ACCESS, *PsThreadType, KernelMode, &ThreadObjects[1], NULL);
 460     if (!NT_SUCCESS(Status))
 461     {
 462         trace("error referencing thread2\n");
 463         goto cleanup;
 464     }
 465
 466 cleanup:
 467
 468     if (ThreadObjects[0])
 469         Status = KeWaitForSingleObject(ThreadObjects[0], Executive, KernelMode, FALSE, NULL);
 470
 471     if (StartContext1 != NULL)
 472         ExFreePoolWithTag(StartContext1, 'tXTC');
 473
 474     if (ThreadObjects[1])
 475         Status = KeWaitForSingleObject(ThreadObjects[1], Executive, KernelMode, FALSE, NULL);
 476
 477     if (StartContext2 != NULL)
 478         ExFreePoolWithTag(StartContext2, 'tXTC');
 479
 480     if (ThreadObjects[0] != NULL)
 481         ObDereferenceObject(ThreadObjects[0]);
 482
 483     if (ThreadObjects[1] != NULL)
 484         ObDereferenceObject(ThreadObjects[1]);
 485
 486     if (Thread1 != INVALID_HANDLE_VALUE)
 487         ZwClose(Thread1);
 488
 489     if (Thread2 != INVALID_HANDLE_VALUE)
 490         ZwClose(Thread2);
 491 }
 492
 493
 494 START_TEST(ZwAllocateVirtualMemory)
 495 {
 496     NTSTATUS Status;
 497
 498     SimpleErrorChecks();
 499
 500     SimpleAllocation();
 501
 502     CustomBaseAllocation();
 503
 504     Status = StressTesting(MEM_RESERVE);
 505     ok_eq_hex(Status, STATUS_NO_MEMORY);
 506
 507     Status = STATUS_SUCCESS;
 508     Status = StressTesting(MEM_COMMIT);
 509     ok_eq_hex(Status, STATUS_NO_MEMORY);
 510
 511     SystemProcessTest();
 512 }