windows-nt/Source/XPSP1/NT/base/urtl/uheap.c
2020-09-26 16:20:57 +08:00

347 lines
7.8 KiB
C

#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;
}