You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
346 lines
8.2 KiB
346 lines
8.2 KiB
#include <nt.h>
|
|
#include <ntrtl.h>
|
|
#include <nturtl.h>
|
|
|
|
#include <stdlib.h>
|
|
|
|
BOOLEAN DebugFlag;
|
|
|
|
PVOID HeapHandle;
|
|
|
|
PVOID
|
|
TestAlloc(
|
|
IN ULONG Size
|
|
)
|
|
{
|
|
PVOID a;
|
|
|
|
if ((a = RtlAllocateHeap( HeapHandle, 0, Size )) == NULL) {
|
|
RtlValidateHeap( HeapHandle, TRUE );
|
|
DbgPrint( "\nUHEAP: RtlAllocateHeap( %lx ) failed\n", Size );
|
|
DbgBreakPoint();
|
|
NtTerminateProcess( NtCurrentProcess(), STATUS_UNSUCCESSFUL );
|
|
}
|
|
|
|
if (DebugFlag) {
|
|
DbgPrint( "\n" );
|
|
DbgPrint( "\nRtlAllocateHeap( %lx ) => %lx\n", Size, a );
|
|
}
|
|
|
|
if (!RtlValidateHeap( HeapHandle, DebugFlag )) {
|
|
NtTerminateProcess( NtCurrentProcess(), STATUS_UNSUCCESSFUL );
|
|
}
|
|
|
|
return( a );
|
|
}
|
|
|
|
|
|
PVOID
|
|
TestFree(
|
|
IN PVOID BaseAddress,
|
|
IN ULONG Size
|
|
)
|
|
{
|
|
PVOID a;
|
|
|
|
if ((a = RtlFreeHeap( HeapHandle, 0, BaseAddress )) != NULL) {
|
|
DbgPrint( "\nUHEAP: RtlFreeHeap( %lx ) failed\n", BaseAddress );
|
|
RtlValidateHeap( HeapHandle, TRUE );
|
|
DbgBreakPoint();
|
|
NtTerminateProcess( NtCurrentProcess(), STATUS_UNSUCCESSFUL );
|
|
}
|
|
|
|
if (DebugFlag) {
|
|
DbgPrint( "\n" );
|
|
DbgPrint( "\nRtlFreeHeap( %lx ) => %lx\n", BaseAddress, a );
|
|
}
|
|
|
|
if (!RtlValidateHeap( HeapHandle, DebugFlag )) {
|
|
NtTerminateProcess( NtCurrentProcess(), STATUS_UNSUCCESSFUL );
|
|
}
|
|
return( a );
|
|
}
|
|
|
|
|
|
BOOLEAN
|
|
TestHeap(
|
|
IN PVOID UserHeapBase,
|
|
IN BOOLEAN Serialize,
|
|
IN BOOLEAN Sparse,
|
|
IN ULONG GrowthThreshold,
|
|
IN ULONG InitialSize
|
|
)
|
|
{
|
|
PVOID a1,a2,a3,a4;
|
|
DWORD Flags;
|
|
|
|
Flags = 0;
|
|
if (!Serialize) {
|
|
Flags |= HEAP_NO_SERIALIZE;
|
|
}
|
|
|
|
if (!Sparse) {
|
|
Flags |= HEAP_GROWABLE;
|
|
}
|
|
|
|
HeapHandle = RtlCreateHeap( Flags,
|
|
UserHeapBase,
|
|
InitialSize,
|
|
0,
|
|
0,
|
|
GrowthThreshold
|
|
);
|
|
if ( HeapHandle == NULL ) {
|
|
DbgPrint( "UHEAP: RtlCreateHeap failed\n" );
|
|
DbgBreakPoint();
|
|
goto exit;
|
|
}
|
|
if (!RtlValidateHeap( HeapHandle, DebugFlag )) {
|
|
NtTerminateProcess( NtCurrentProcess(), STATUS_UNSUCCESSFUL );
|
|
}
|
|
|
|
|
|
//
|
|
// TEST 0:
|
|
// Allocate and free a large chunk of memory so that the following
|
|
// tests are valid.
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #0\n" );
|
|
a1 = TestAlloc( 4096-16 );
|
|
TestFree( a1, 0 );
|
|
|
|
|
|
//
|
|
// TEST 1:
|
|
// Allocate three chunks, deallocate the middle one, and reallocate it.
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #1\n" );
|
|
a1 = TestAlloc( 16 );
|
|
a2 = TestAlloc( 32 );
|
|
a3 = TestAlloc( 112 );
|
|
TestFree( a2, 32 );
|
|
a4 = TestAlloc( 32 );
|
|
|
|
|
|
//
|
|
// TEST 2:
|
|
// Deallocate first chunk and reallocate it.
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #2\n" );
|
|
TestFree( a1, 16 );
|
|
a4 = TestAlloc( 16 );
|
|
|
|
|
|
//
|
|
// TEST 3:
|
|
// Deallocate last chunk and reallocate it.
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #3\n" );
|
|
TestFree( a3, 112 );
|
|
a4 = TestAlloc( 112 );
|
|
|
|
|
|
//
|
|
// TEST 4:
|
|
// Deallocate last chunk and reallocate larger one.
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #4\n" );
|
|
TestFree( a4, 112 );
|
|
a4 = TestAlloc( 112+64 );
|
|
|
|
|
|
//
|
|
// TEST 5:
|
|
// Deallocate first two chunks and reallocate combined one.
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #5\n" );
|
|
TestFree( a1, 16 );
|
|
TestFree( a2, 32 );
|
|
a4 = TestAlloc( 16+32-4 );
|
|
|
|
|
|
//
|
|
// TEST 6:
|
|
// There should be room between blocks 2 and 3 for a small allocation.
|
|
// Make sure zero byte allocations work.
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #6\n" );
|
|
a4 = TestAlloc( 0 );
|
|
|
|
|
|
//
|
|
// TEST 7:
|
|
// Deallocate last two chunks and reallocate one. Address should change.
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #7\n" );
|
|
TestFree( a3, 112+64 );
|
|
TestFree( a4, 0 );
|
|
a3 = TestAlloc( 112 );
|
|
|
|
|
|
//
|
|
// TEST 8:
|
|
// Deallocate everything and make sure it can be reallocated.
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #8\n" );
|
|
TestFree( a1, 16+32-4 );
|
|
TestFree( a3, 112 );
|
|
a2 = TestAlloc( 200 );
|
|
|
|
|
|
//
|
|
// TEST 9:
|
|
// Allocate more than is committed.
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #9\n" );
|
|
a1 = TestAlloc( 100000 );
|
|
TestFree( a2, 200 );
|
|
TestFree( a1, 100000 );
|
|
|
|
|
|
//
|
|
// TEST 10:
|
|
// Allocate more than maximum size of heap
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #10\n" );
|
|
a3 = TestAlloc( 100000 );
|
|
TestFree( a3, 100000 );
|
|
|
|
|
|
//
|
|
// TEST 11:
|
|
// Destroy the heap
|
|
//
|
|
|
|
DbgPrint( "UHEAP: Test #11\n" );
|
|
HeapHandle = RtlDestroyHeap( HeapHandle );
|
|
if ( HeapHandle != NULL ) {
|
|
DbgPrint( "UHEAP: RtlDestroyHeap failed\n" );
|
|
DbgBreakPoint();
|
|
goto exit;
|
|
}
|
|
|
|
return( TRUE );
|
|
|
|
exit:
|
|
if (HeapHandle != NULL) {
|
|
HeapHandle = RtlDestroyHeap( HeapHandle );
|
|
}
|
|
|
|
return( FALSE );
|
|
}
|
|
|
|
|
|
VOID
|
|
Usage( VOID )
|
|
{
|
|
DbgPrint( "Usage: UHEAP [-s ReserveSize] | [-g InitialSize GrowthThreshold]\n" );
|
|
|
|
(VOID)NtTerminateProcess( NtCurrentProcess(), STATUS_UNSUCCESSFUL );
|
|
}
|
|
|
|
NTSTATUS
|
|
main(
|
|
int argc,
|
|
char *argv[],
|
|
char *envp[],
|
|
ULONG DebugParameter OPTIONAL
|
|
)
|
|
{
|
|
NTSTATUS Status;
|
|
PCH s;
|
|
PVOID UserHeapBase = NULL;
|
|
BOOLEAN Serialize = FALSE;
|
|
BOOLEAN Sparse = FALSE;
|
|
ULONG GrowthThreshold = 0;
|
|
ULONG InitialSize = 0x8000;
|
|
|
|
DebugFlag = DebugParameter;
|
|
|
|
DbgPrint( "** Start of User Mode Test of RtlAllocateHeap/RtlFreeHeap **\n" );
|
|
|
|
while (--argc) {
|
|
s = *++argv;
|
|
if (*s == '-') {
|
|
switch( *++s ) {
|
|
case 'x':
|
|
case 'X':
|
|
Serialize = TRUE;
|
|
break;
|
|
|
|
case 's':
|
|
case 'S':
|
|
Sparse = TRUE;
|
|
if (--argc) {
|
|
InitialSize = atoi( *++argv );
|
|
Status = NtAllocateVirtualMemory( NtCurrentProcess(),
|
|
(PVOID *)&UserHeapBase,
|
|
0,
|
|
&InitialSize,
|
|
MEM_RESERVE,
|
|
PAGE_READWRITE
|
|
);
|
|
if (!NT_SUCCESS( Status )) {
|
|
DbgPrint( "UHEAP: Unable to allocate heap - 0x%lx bytes\n",
|
|
InitialSize
|
|
);
|
|
Usage();
|
|
}
|
|
}
|
|
else {
|
|
Usage();
|
|
}
|
|
break;
|
|
|
|
case 'g':
|
|
case 'G':
|
|
if (argc >= 2) {
|
|
argc -= 2;
|
|
InitialSize = atoi( *++argv );
|
|
GrowthThreshold = atoi( *++argv );
|
|
}
|
|
else {
|
|
Usage();
|
|
}
|
|
break;
|
|
|
|
default:
|
|
Usage();
|
|
}
|
|
}
|
|
else {
|
|
Usage();
|
|
}
|
|
}
|
|
|
|
TestHeap( UserHeapBase,
|
|
Serialize,
|
|
Sparse,
|
|
GrowthThreshold,
|
|
InitialSize
|
|
);
|
|
|
|
if (UserHeapBase != NULL) {
|
|
Status = NtFreeVirtualMemory( NtCurrentProcess(),
|
|
(PVOID *)&UserHeapBase,
|
|
&InitialSize,
|
|
MEM_RELEASE
|
|
);
|
|
}
|
|
|
|
DbgPrint( "** End of User Mode Test of RtlAllocateHeap/RtlFreeHeap **\n" );
|
|
|
|
(VOID)NtTerminateProcess( NtCurrentProcess(), STATUS_SUCCESS );
|
|
return STATUS_SUCCESS;
|
|
}
|