#include <stdio.h>

#define STATUS_ACCESS_VIOLATION          ((NTSTATUS)0xC0000005L)    // winnt

#define OVERFLOW_GUARD   1

#define BUFFER_GUARD OVERFLOW_GUARD

//#define SIMULATE_MEMORY_LACK  1
#define NUMBERS_OF_MEMORY_REQUESTS_TO_FAIL_AT 2

#define ALLOC_SIZE 10000

struct __allocation {
	void* mem;
	int len;
	char guard_type;
} __allocs [ALLOC_SIZE], __allocs_bck [ALLOC_SIZE];

void DebugMessage(char* s, ...);

// exception handler
LONG __stdcall MyUnhandledExceptionFilter(EXCEPTION_POINTERS* ep)
{
	if( ep->ExceptionRecord->ExceptionCode == STATUS_ACCESS_VIOLATION )
	{
		for( int i=0; i<sizeof(__allocs)/sizeof(*__allocs); i++ )
		{
			long addr = ep->ExceptionRecord->ExceptionInformation[1];
			long dist;

			if( 
				(addr >= (long)__allocs[i].mem + __allocs[i].len) &&
				(addr <= (long)__allocs[i].mem + __allocs[i].len + 10) &&
				__allocs[i].guard_type == OVERFLOW_GUARD
				)
			{
				dist = addr - (long)__allocs[i].mem - __allocs[i].len + 1;
				DebugMessage(
					"Buffer overflow detected!\n"
					"Buffer starts at 0x%08x\n"
					"Buffer size is %d byte(s) [0x%x]\n"
					"Memory referenced: 0x%08x\n"
					"Distance between end of buffer: %d byte(s) [0x%x]\n"
					"%s"
					,
					__allocs[i].mem,
					__allocs[i].len,
					__allocs[i].len,
					addr,
					dist,
					dist,
					(dist==1) ?
						"\nDistance equals to one byte.\n"
						"It's very common case of overflow called one-byte buffer oveflow.\n"
						"Possible it is a result of string allocation error (maybe you've forgotten to allocate space for terminating NULL-character?)"
						: ""
					);
			}
		}
	}
	return EXCEPTION_CONTINUE_SEARCH;
		
}

// Clear allocation log __allocs
void EnterMemoryCheckRegion()
{
	void* p = __allocs;
	int l = sizeof(__allocs) / 4;

	__asm
	{
		mov edi, p
		xor eax, eax
		mov ecx, l
		rep stosd
	}
}

// Check for memory leak. Should be called at the end of guarded region of code
void LeaveMemoryCheckRegion()
{
	for( int i=0; i<sizeof(__allocs)/sizeof(*__allocs); i++ )
	{
		if( __allocs[i].mem && __allocs[i].len != -1 )
		{
			DebugMessage(
				"Possible memory leak detected. "
				"Buffer at address 0x%08x is not freed at the end of the memory-check region."
				,
				__allocs[i].mem
				);
		}
	}
}

// Show debug message
void DebugMessage(char* s, ...)
{
	char t[ 1024 ];
	va_list va;
	va_start( va, s );
	vsprintf( t, s, va );

	char msg[ 1024 ];
	sprintf(msg,
		"Memory Debugging Message:\n"
		"\n"
		"%s\n"
		"\n"
		"If you want to abort execution press 'Abort'\n"
		"If you want to debug application press 'Retry'\n"
		"If you want to continue execution press 'Ignore'\n"
		,
		t
		);

	OutputDebugString( msg );

	switch( MessageBox( NULL, msg, "Memory Debugging Message", MB_ABORTRETRYIGNORE|MB_ICONERROR ) )
	{
	case IDRETRY:
		__asm int 3;
		break;

	case IDABORT:
		ExitProcess(0);
		
	};
}

// new[] operator handler for the overflow-guard protection mode
// Allocates buffer at the end of the page, next page will be marked as invalid.
// Any access behind the end of the buffer will be failed
void* AllocateOverflowGuardedBuffer(int size)
{
	SYSTEM_INFO si = {0};
	GetSystemInfo( &si );
	/*if( size > (signed)si.dwPageSize ) return NULL;*/

	void* mem = VirtualAlloc( NULL, /*si.dwPageSize*2*/ size + si.dwPageSize, MEM_RESERVE, PAGE_NOACCESS );
	if( !mem ) {
		OutputDebugStringEx(__FUNCTION__" : ERROR : VirtualAlloc(... %d ...) fails! : LasrErr = %08X", size + si.dwPageSize, GetLastError());
		return NULL;
	}
	mem = VirtualAlloc( mem, /*si.dwPageSize*/ size, MEM_COMMIT, PAGE_READWRITE );

	DWORD dwSmth = size % si.dwPageSize ? size % si.dwPageSize : si.dwPageSize;
	DWORD ps = /*si.dwPageSize - size*/ si.dwPageSize - dwSmth;
	__asm
	{
		mov al, 0xfd
		mov ecx, ps
		mov edi, mem
		rep stosb
	}

	mem = (LPVOID)( (DWORD)mem + /*si.dwPageSize - size*/si.dwPageSize - dwSmth );

	int i;
	for( i=0; i<sizeof(__allocs)/sizeof(*__allocs); i++ )
		if( !__allocs[i].mem )
		{
			__allocs[i].guard_type = OVERFLOW_GUARD;
			__allocs[i].len = size;
			return (__allocs[i].mem = mem);
		}
	
	DWORD dwWholeRegionSize = (size - dwSmth) + si.dwPageSize;

	VirtualFree( mem, /*si.dwPageSize*/ dwWholeRegionSize,  MEM_DECOMMIT );
	VirtualFree( mem, /*si.dwPageSize*2*/ dwWholeRegionSize + si.dwPageSize, MEM_RELEASE );

	OutputDebugStringEx(__FUNCTION__" : ERROR : RESULT == NULL : i == %d", i);

	return NULL;
}

// delete[] operator handler for the overflow-guard protection mode
void FreeOverflowGuardedBuffer(void* mem)
{
	SYSTEM_INFO si = {0};
	GetSystemInfo( &si );
	
	for( int i=0; i<sizeof(__allocs)/sizeof(*__allocs); i++ )
		if( __allocs[i].mem == mem )
		{
			if( __allocs[i].len == -1 )
			{
				DebugMessage(
					"Attempt to free already freed buffer 0x%08x.\n"
					"Possible it's a result of incorrect destructor call."
					,
					mem
					);
				return;
			}

			DWORD dwSmth = __allocs[i].len % si.dwPageSize ? __allocs[i].len % si.dwPageSize : si.dwPageSize;
			long ps = /*si.dwPageSize - __allocs[i].len*/ si.dwPageSize - dwSmth;
			LPVOID addr = (LPVOID)( (DWORD)mem + /*__allocs[i].len - si.dwPageSize*/ dwSmth - si.dwPageSize );

			int j;
			for( j=0; j < ps; j++ )
				if( ((unsigned char*)addr)[j] != 0xfd )
					break;
			ps -= j;
			
			if( ps )
			{
				DebugMessage(
					"Buffer underflow detected while overflow-guard protection at address 0x%08x.\n"
					"It's highly recommended to test program in underflow-guard mode to detect faulting instruction."
					,
					mem
					);
			}

			DWORD dwWholeRegionSize = (__allocs[i].len - dwSmth) + si.dwPageSize;

			VirtualFree( mem, /*si.dwPageSize*/ dwWholeRegionSize,  MEM_DECOMMIT );
			VirtualFree( mem, /*si.dwPageSize*2*/ dwWholeRegionSize + si.dwPageSize, MEM_RELEASE );

			__allocs[i].len = -1;
			__allocs[i].mem = NULL;
			return;
		}
}