///////////////////////////////////////////////////////////
// boundarycheck.cpp, v1.1
//
// Windows Leaks Detector, 2006
// http://sourceforge.net/projects/winleak
// migosh@users.sourceforge.net
//
// COMPILATION:
// Create empty "Win32 Console Application" project, and
// add this file. Tested with VS2005, but should work
// with other compilers too.
//
// REQUIRES: Windows XP/2000
//
// DESCRIPTION:
// Demonstrates 2 approaches to design of "Memory Boundary
// Check" feature, which should verify that the application
// does not access memory outside of dynamically allocated
// blocks. The feature will be implemented in "Windows Leaks
// Detector" project, and assumes that our code will receive
// all calls to HeapAlloc and HeapFree. In current module
// new allocation functions are called directly.
// The boundary check is not intended to verify access to
// statically allocated memory.
//
// Approach 1:
// Allocate bigger memory buffer than requested, and write
// signature block before and after the returned chunk.
// When the memory is released by application - verify that
// signature blocks are unchanged.
// (+) low memory overhead
// (+) no run-time overhead
// (+) checks both beginning and end of block
// (-) does not alert until the memory is released
// (-) checks write operations only
//
// Approach 2:
// Allocate block of N+1 pages, where N is the minimal number
// of full pages to contain the requested buffer size. The
// extra page will be protected by PAGE_GUARD flag, and the
// pointer returned to the user will be calculated so the
// requested buffer will end at the end of last "unguarded"
// page - so next byte after the block will be located in
// the guarded page. Any attempt to access the memory after
// the allocated block will generate STATUS_GUARD_PAGE_VIOLATION
// exception.
// (+) no run-time overhead
// (+) alerts at the time of out-of-bound access, when the
//     full call stack is available
// (+) checks both read and write access
// (-) high memory overhead (may cause run-time overhead),
//     depends on the page size (the default in XP/2000 is 4K).
// (-) checks only overflow at the end of the block (in
//     similar way the beginning of the block can be protected,
//     but not both at same time, unless the size of allocated
//     block is multiple of page size)
// NOTE: when the memory overflow in not acceptable, this
// technique can be applied to part of allocations:
// 1. Only for allocation of buffers over certain size. This
//    way the ratio between requested and allocated memory will
//    be limited. Anyway, we are mostly interested in overflow
//    control when arrays (e.g. string buffers) are allocated.
// 2. Only for cases where the 1st approach indicated an
//    overflow, and we want to find the violating code.
// 
////////////////////////////////////////////////////////////


// Following define is required in order to compile
// exception-related code.
// See MSDN for "AddVectoredExceptionHandler"
#define _WIN32_WINNT 0x0500 

#include <windows.h>
#include <stdio.h>
#include <malloc.h>
#include <list>


//////////////////////////////////////////////////////////
// 
// IMPLMENTATION OF 1st APPROACH
// Signature blocks
//
//////////////////////////////////////////////////////////

// defines the signature to be placed before and after
// the actual allocated block
static const unsigned char SIGNATURE[] = {0xAB, 0x43, 0xDE, 0x72};
static const unsigned int SIZEOF_SIGNATURE = sizeof(SIGNATURE);


///////////////////////////////////////////////
// This function will replace the regular "malloc". In this
// example the "main" function will invoke it directly, but
// in "Windows Leaks Detector" project this code can be placed
// inside the hook function for HeapAlloc
///////////////////////////////////////////////
void * malloc_signature(size_t size)
{
	// The actual allocated block will contain:
	// + the memory block to be returned to the user
	// + 2 signature borders - before and after user's block
	// + additional size_t value, to store the size of the
	//   block which the user requested. Note, that, when
	//   integrated to Leaks Detector, this field will
	//   not be required, as anyway the size is stored as
	//   part of internal structures.
	char * tmp = (char *) malloc(size + (SIZEOF_SIGNATURE * 2) + sizeof(size_t));
	if (tmp == NULL) return NULL;

	// Layout of the memory block will be:
	// [size of data][signature border][user's data][signature border]
	
	// Store the size of user's block
	*reinterpret_cast<size_t*>(tmp) = size;  
	tmp += sizeof(size_t);

	// Store the signature before the block
	memcpy(tmp, SIGNATURE, SIZEOF_SIGNATURE);
	tmp += SIZEOF_SIGNATURE;  // now "tmp" points to beginning of user's data

	// Store the signature after the block
	memcpy(tmp + size, SIGNATURE, SIZEOF_SIGNATURE);

	// Return the pointer to user's data
	return tmp;
}


///////////////////////////////////////////////
// free_signature is called to release blocks,
// allocated by malloc_signature. For Leaks Detector
// project this code will be executed as part of
// hook function for HeapFree, when the block
// will be identified as having signature protection
// (based on internal data for allocation block)
///////////////////////////////////////////////
void free_signature(void * p)
{
	char * orig = static_cast<char*>(p);

	// the pointer returned to the user in "malloc_signature"
	// is shifted by size of (int + signature block)
	// from the beginning of actual block allocated by
	// "malloc" (or HeapAlloc in general), so we need
	// to restore the original pointer.
	orig -= (SIZEOF_SIGNATURE + sizeof(size_t));

	// read the size of user's block
	// in Leaks Detector application this data will be
	// stored as part of info for allocation block
	size_t allocation_size = *(reinterpret_cast<size_t*>(orig));

	// verify that signature blocks were unchanged
	bool bStartOK = (0 == memcmp(orig + sizeof(size_t), SIGNATURE, SIZEOF_SIGNATURE));
	bool bEndOK = (0 == memcmp(orig + sizeof(size_t) + allocation_size + SIZEOF_SIGNATURE, SIGNATURE, SIZEOF_SIGNATURE));

	// print error message if one of the signature
	// blocks (or both) were changed
	if (!bStartOK) {
		if (!bEndOK) {
			printf("Warning: block of size %d is corrupt - before and after\n", allocation_size);
		} else {
			printf("Warning: block of size %d is corrupt - before\n", allocation_size);
		}
	} else if (!bEndOK) {
		printf("Warning: block of size %d is corrupt - after\n", allocation_size);
	}

	// release the original bloch, allocated by "malloc"
	free (orig);
}


//////////////////////////////////////////////////////////
// 
// IMPLMENTATION OF 2nd APPROACH
// Guarded pages
//
//////////////////////////////////////////////////////////


///////////////////////////////////////////////
// This function will be invoked automatically
// by Windows each time an exception occurs,
// even before it reaches frame-based application
// exception handler (such as try-catch in C++).
// Only when the application is being debugged, the
// debugger will get a notification first - which is
// even better, as a live call stack and all relevant
// context will be available.
///////////////////////////////////////////////
LONG WINAPI MyExceptionsHandler (PEXCEPTION_POINTERS ExceptionInfo)
{
	// check the type of exception
	if (ExceptionInfo->ExceptionRecord->ExceptionCode == STATUS_GUARD_PAGE_VIOLATION) {
		printf("Access violation caught - print call stack, activate a breakpoint etc.\n");

		// when STATUS_GUARD_PAGE_VIOLATION exception is raised, the
		// protection of the page is automatically reset to READ/WRITE,
		// so next attempt to access same memory won't fail
		return EXCEPTION_CONTINUE_EXECUTION;
	}

	// unknown exception - find next exceptions' handler
	return EXCEPTION_CONTINUE_SEARCH;
}


// global variable to hold handle to exception handle
PVOID gExHandler = NULL;

// global variable, which will receive system page size
DWORD gPageSize = 0;


///////////////////////////////////////////////
// This function must be called before any
// protected allocation is made.
// It will initialize all global variables, and 
// set the exception handler.
///////////////////////////////////////////////
void protected_init()
{
	// query the System Info in order to get page size
	SYSTEM_INFO si;
	::GetSystemInfo(&si);
	gPageSize = si.dwPageSize;
	printf("Page size is: %ld\n", gPageSize);

	// set "MyExceptionsHandler" function as the first
	// exception handler. In case of memory access violation
	// we want out function to be called, while the
	// call stack remain unchanged.
	// Notes:
	// 1) In order for this code to compile, we had to put
	//    #define _WIN32_WINNT 0x0500 
	//    at the beginning of the file, before the include of
	//    <windows.h> header.
	// 2) There should be a clean-up function, which will
	//    call "RemoveVectoredExceptionHandler" with returned
	//    handler. In current code it is not implemented.
	// 3) If there's a debugger attached to this process, it
	//    will receive the notification before our handler.
	gExHandler = ::AddVectoredExceptionHandler(1 /*call me first*/, MyExceptionsHandler);
	if (gExHandler == NULL)
		printf("Failed to set MyExceptionsHandler\n");
}


///////////////////////////////////////////////
// Similary to "malloc_signature", this function
// will be invoked by "main" instead of a call
// to malloc/new. In "Windows Leaks Detector" project
// this code can be placed inside the hook function
// for HeapAlloc
///////////////////////////////////////////////
void * malloc_protected(size_t size)
{
	// calculate the number of full pages required to cover
	// the requested data size.
	// (size / gPageSize) will give the number of full
	//                    pages contained in "size"
	// (size % gPageSize) will be bigger than 0 in case
	//                    there's a reminder of division
	//                    operation - additional page is required
	unsigned int iTotalPagesForData = ((unsigned long)size / gPageSize) +
		                              ((unsigned long)size % gPageSize == 0) ? 0 : 1;
	
	// allocate virtual memory of size "iTotalPagesForData + 1" pages,
	// to accomodate the requesed buffer + guarded page.
	char * pRegionStart = (char*) ::VirtualAlloc(NULL, (iTotalPagesForData +1 ) * gPageSize, MEM_COMMIT, PAGE_READWRITE);
	if (pRegionStart == NULL) {
		printf("VirtualAlloc returned NULL\n");
		return NULL;
	}

	// set PAGE_GUARD protection on last allocated page. On any
	// read/write access to this page STATUS_GUARD_PAGE_VIOLATION
	// exception will be raised.
	// Notes:
	// 1) VirtualProtection can't be set on part of a page,
	//    but only on whole pages. For this reason we can't
	//    just set protection on a "border" around allocated
	//    block, and have to allocate much more memory
	//    then requested by user.
	// 2) After the exception is raised, the page permissions will
	//    be changed to PAGE_READWRITE, as we specified
	//    (PAGE_GUARD | PAGE_READWRITE) protection.
	//    For this reason PAGE_GUARD flag can't be used alone,
	//    but has to be or'ed with other flag.
	DWORD dwOldProtect;
	if (! ::VirtualProtect(pRegionStart + (iTotalPagesForData * gPageSize), gPageSize, PAGE_GUARD | PAGE_READWRITE, &dwOldProtect)) {
		printf("VirtualProtecd failrd\n");
		return NULL;
	}

	// when we will go "size" bytes back from the guarded page,
	// we will promise that the the pointer returned to user
	// points to a safe buffer of requested size, but the first
	// byte after this buffer is guarded.
	return pRegionStart + (iTotalPagesForData * gPageSize - size);
}


///////////////////////////////////////////////
// Releases memory allocated by "malloc_protected"
///////////////////////////////////////////////
void free_protected(void *p)
{
	// query the address of the page containing "p"
	// Instead, we could keep the original address
	// returned by VirtualAlloc as part of the info
	// of allocation block (in Leaks Detector)
	MEMORY_BASIC_INFORMATION memInfo;
	::VirtualQuery(p, &memInfo, sizeof(MEMORY_BASIC_INFORMATION));

	// release the memory
	if (! ::VirtualFree(memInfo.AllocationBase, 0 /*has to be 0 for MEM_RELEASE*/, MEM_RELEASE))
		printf("Failed to release virtual memory\n");
}


//////////////////////////////////////////////////////////
// 
// DEMONSTRATION
//
//////////////////////////////////////////////////////////


int main()
{
	//////////////////////
	// Signature
	//////////////////////
	char * arr10 = (char *) malloc_signature(sizeof(char) * 10);
	char * arr11 = (char *) malloc_signature(sizeof(char) * 11);
	char * arr12 = (char *) malloc_signature(sizeof(char) * 12);

	*(arr10 - 1) = '\0';		// 10: overflow before the block
	arr10[10] = '\0';			// 10: overflow after the block
	*(arr11 - 1) = '\0';		// 11: overflow before the block
	arr12[12] = '\0';			// 12: overflow after the block

	// error messages will be printed when releasing the buffers
	free_signature (arr10);
	free_signature (arr11);
	free_signature (arr12);


	//////////////////////
	// Guarded page
	//////////////////////

	// must initialize before first use
	protected_init();

	char * parr10 = (char *) malloc_protected(sizeof(char) * 10);

	// valid read/write access
	parr10[5] = '\0';
	char ch = parr10[2];

	// read access outside the array - will raise
	// exception and print log message
	ch = parr10[10];
	
	// no exceptions on next out-of-bounds accesses
	// for same block - GUARD flag was automatially
	// reset after the first exception
	ch = parr10[11];

	free_protected(parr10);

	printf ("Done\n");
	return 0;
}

// EOF