#include "stdafx.h"
#include "kxMemoryChunk.h"

#include <iostream>
#include <DbgHelp.h>


// Inicializace statickych clenskych promennych
unsigned long int kxMemoryChunk::m_Count = 0;
HMODULE kxMemoryChunk::m_hImagehlpDll = 0;
DWORD kxMemoryChunk::m_ImageType = 0;
HANDLE kxMemoryChunk::m_hProcess = 0;
HANDLE kxMemoryChunk::m_hThread = 0;
kxMemoryChunk* kxMemoryChunk::m_Begin = NULL;
kxMemoryChunk* kxMemoryChunk::m_End = NULL;


tSC   kxMemoryChunk::pSC = NULL;
tSFTA kxMemoryChunk::pSFTA = NULL;
tSGLFA kxMemoryChunk::pSGLFA = NULL;
tSGMB kxMemoryChunk::pSGMB = NULL;
tSGMI kxMemoryChunk::pSGMI = NULL;
tSGO kxMemoryChunk::pSGO = NULL;
tSGSFA kxMemoryChunk::pSGSFA = NULL;
tSI kxMemoryChunk::pSI = NULL;
tSLM kxMemoryChunk::pSLM = NULL;
tSSO kxMemoryChunk::pSSO = NULL;
tSW kxMemoryChunk::pSW = NULL;
tUDSN kxMemoryChunk::pUDSN = NULL;

kxMemoryChunk::kxMemoryChunk(size_t size)
{		
	if(m_Count > 0 && size > 0)
	{
		ZeroMemory(&m_StackFrame,sizeof(m_StackFrame));
		ZeroMemory(&m_Context,sizeof(m_Context));

		m_Data = malloc(size);
		m_Size = size;	
	}
}

// *******************************************************************************************
kxMemoryChunk::~kxMemoryChunk()
{

}

// Pretizeni lokalniho new 
void* kxMemoryChunk::operator new(size_t size)
{	
	if(m_Count == 0) // Musime inicializovat begin a end
	{
		LazyInit();

	}

	//kxMemoryChunk* ptr = (kxMemoryChunk*)malloc(sizeof(kxMemoryChunk));
	kxMemoryChunk* ptr = (kxMemoryChunk*)malloc(size);
	if(ptr)
	{			
		ptr->m_Prev = m_End->m_Prev;        // Linknuti
		ptr->m_Next = m_End;

		m_End->m_Prev->m_Next = ptr;
		m_End->m_Prev = ptr;

		Dump(ptr); // Dumpnuti dat	
	}

	m_Count++;
	return ptr;
}

// Pretizeni lokalniho delete
void kxMemoryChunk::operator delete(void* ptr)  
{
	kxMemoryChunk* chunk = (kxMemoryChunk*)ptr;

	// Unlinknuti
	chunk->m_Prev->m_Next = chunk->m_Next;
	chunk->m_Next->m_Prev = chunk->m_Prev;

	if(m_Count == 1)
	{
		LazyCleanup();

	}

	free(ptr);
	m_Count--;
}


kxMemoryChunk* kxMemoryChunk::GetBegin()
{
	return m_Begin;
}


kxMemoryChunk* kxMemoryChunk::GetEnd()
{
	return m_End;
}


unsigned long int kxMemoryChunk::GetCount()
{
	return m_Count;
}

void kxMemoryChunk::Delete(void* data_ptr)
{
	// Pre-conditions
	if(m_Begin == NULL || m_End == NULL)
	{
		return;
	}

		// Vykonne telo
	kxMemoryChunk* tmp = m_Begin->m_Next;
	while(tmp != m_End)
	{
		if(tmp->m_Data == data_ptr)
		{
			delete tmp;
			return;
		}
		tmp = tmp->m_Next;
	}
}

void kxMemoryChunk::Check()
{
	if(m_Count > 0)
	{
		std::cout << "Error: " << m_Count << " neuvolnenych chunku bude uvolneno!" << std::endl;
		size_t total = 0;
		while(m_Count > 0)
		{
			total += m_Begin->m_Next->m_Size;
			std::cout << "Address: 0x" << (void*)m_Begin->m_Next->m_Data << ", Length = " << m_Begin->m_Next->m_Size << std::endl;
			delete m_Begin->m_Next;
		}
		std::cout << "Totaly " << total << " has been freed!" << std::endl;
	}
}

void kxMemoryChunk::Dump(kxMemoryChunk* chunk)
{
	/*
	StackWalk( \
		m_ImageType, \
		m_hProcess, \
		m_hThread, \
		(LPSTACKFRAME)&chunk->m_StackFrame, \
		&chunk->m_Context, \
		NULL, \
		(PFUNCTION_TABLE_ACCESS_ROUTINE) GetProcAddress( m_hImagehlpDll, "SymFunctionTableAccess" ), \
		(PGET_MODULE_BASE_ROUTINE) GetProcAddress( m_hImagehlpDll, "SymGetModuleBase" ), \
		NULL \
	);
	*/

	
	#define gle (GetLastError())
	#define lenof(a) (sizeof(a) / sizeof((a)[0]))
	#define MAXNAMELEN 1024 // max name length for found symbols
	#define IMGSYMLEN ( sizeof IMAGEHLP_SYMBOL )
	#define TTBUFLEN 65536 // for a temp buffer
	
	DWORD symOptions; // symbol handler settings
	IMAGEHLP_SYMBOL *pSym = (IMAGEHLP_SYMBOL *) malloc( IMGSYMLEN + MAXNAMELEN );
	DWORD offsetFromSymbol;
	char undName[MAXNAMELEN]; // undecorated name
	char undFullName[MAXNAMELEN]; // undecorated name with all shenanigans
	IMAGEHLP_MODULE Module;
	IMAGEHLP_LINE Line;
	std::string symSearchPath;
	
	char *tt = 0, *p;
	tt = (char*)malloc(TTBUFLEN); // this is a _sample_. you can do the error checking yourself.

	// build symbol search path from:
	symSearchPath = "";
	// current directory
	if ( GetCurrentDirectoryA( TTBUFLEN, tt ) )
		symSearchPath += tt + std::string( ";" );
	// dir with executable
	if ( GetModuleFileNameA( 0, tt, TTBUFLEN ) )
	{
		for ( p = tt + strlen( tt ) - 1; p >= tt; -- p )
		{
			// locate the rightmost path separator
			if ( *p == '\\' || *p == '/' || *p == ':' )
				break;
		}
		// if we found one, p is pointing at it; if not, tt only contains
		// an exe name (no path), and p points before its first byte
		if ( p != tt ) // path sep found?
		{
			if ( *p == ':' ) // we leave colons in place
				++ p;
			*p = '\0'; // eliminate the exe name and last path sep
			symSearchPath += tt + std::string( ";" );
		}
	}
	// environment variable _NT_SYMBOL_PATH
	if ( GetEnvironmentVariableA( "_NT_SYMBOL_PATH", tt, TTBUFLEN ) )
		symSearchPath += tt + std::string( ";" );
	// environment variable _NT_ALTERNATE_SYMBOL_PATH
	if ( GetEnvironmentVariableA( "_NT_ALTERNATE_SYMBOL_PATH", tt, TTBUFLEN ) )
		symSearchPath += tt + std::string( ";" );
	// environment variable SYSTEMROOT
	if ( GetEnvironmentVariableA( "SYSTEMROOT", tt, TTBUFLEN ) )
		symSearchPath += tt + std::string( ";" );

	if ( symSearchPath.size() > 0 ) // if we added anything, we have a trailing semicolon
		symSearchPath = symSearchPath.substr( 0, symSearchPath.size() - 1 );

	printf( "symbols path: %s\n", symSearchPath.c_str() );

	// why oh why does SymInitialize() want a writeable string?
	strncpy( tt, symSearchPath.c_str(), TTBUFLEN );
	tt[TTBUFLEN - 1] = '\0'; // if strncpy() overruns, it doesn't add the null terminator

	// init symbol handler stuff (SymInitialize())
	if ( ! pSI( m_hProcess, tt, false ) )
	{
		printf( "SymInitialize(): gle = %lu\n", gle );
		goto cleanup;
	}

	// SymGetOptions()
	symOptions = pSGO();
	symOptions |= SYMOPT_LOAD_LINES;
	symOptions &= ~SYMOPT_UNDNAME;
	pSSO( symOptions ); // SymSetOptions()

	// Enumerate modules and tell imagehlp.dll about them.
	// On NT, this is not necessary, but it won't hurt.
	EnumumAndLoadModuleSymbols( chunk->m_hProcess, GetCurrentProcessId() );

	// init STACKFRAME for first call
	// Notes: AddrModeFlat is just an assumption. I hate VDM debugging.
	// Notes: will have to be #ifdef-ed for Alphas; MIPSes are dead anyway,
	// and good riddance.
	chunk->m_StackFrame.AddrPC.Offset = chunk->m_Context.Eip;
	chunk->m_StackFrame.AddrPC.Mode = AddrModeFlat;
	chunk->m_StackFrame.AddrFrame.Offset = chunk->m_Context.Ebp;
	chunk->m_StackFrame.AddrFrame.Mode = AddrModeFlat;
	
	chunk->m_hProcess = GetCurrentProcess();
	chunk->m_hThread = GetCurrentThread();

	chunk->m_Context.ContextFlags = CONTEXT_FULL;
	GetThreadContext(chunk->m_hThread,&chunk->m_Context);

	memset( pSym, '\0', IMGSYMLEN + MAXNAMELEN );
	pSym->SizeOfStruct = IMGSYMLEN;
	pSym->MaxNameLength = MAXNAMELEN;

	memset( &Line, '\0', sizeof Line );
	Line.SizeOfStruct = sizeof Line;

	memset( &Module, '\0', sizeof Module );
	Module.SizeOfStruct = sizeof Module;

	offsetFromSymbol = 0;

	for ( int frameNum = 0; ; ++ frameNum )
	{
		// get next stack frame (StackWalk(), SymFunctionTableAccess(), SymGetModuleBase())
		// if this returns ERROR_INVALID_ADDRESS (487) or ERROR_NOACCESS (998), you can
		// assume that either you are done, or that the stack is so hosed that the next
		// deeper frame could not be found.
		if ( ! StackWalk(IMAGE_FILE_MACHINE_I386, chunk->m_hProcess, chunk->m_hThread,(LPSTACKFRAME)&chunk->m_StackFrame, &chunk->m_Context, NULL,	(PFUNCTION_TABLE_ACCESS_ROUTINE) GetProcAddress( m_hImagehlpDll, "SymFunctionTableAccess" ), (PGET_MODULE_BASE_ROUTINE) GetProcAddress( m_hImagehlpDll, "SymGetModuleBase" ), NULL ) )
			break;

		// display its contents
		printf( "\n%3d %c%c %08lx %08lx %08lx %08lx ",	frameNum, chunk->m_StackFrame.Far? 'F': '.', chunk->m_StackFrame.Virtual? 'V': '.',	chunk->m_StackFrame.AddrPC.Offset, chunk->m_StackFrame.AddrReturn.Offset, chunk->m_StackFrame.AddrFrame.Offset, chunk->m_StackFrame.AddrStack.Offset );

		if ( chunk->m_StackFrame.AddrPC.Offset == 0 )
		{
			printf( "(-nosymbols- PC == 0)\n" );
		}
		else
		{ // we seem to have a valid PC
			// show procedure info (SymGetSymFromAddr())
			if ( ! pSGSFA( chunk->m_hProcess, chunk->m_StackFrame.AddrPC.Offset, &offsetFromSymbol, pSym ) )
			{
				if ( gle != 487 )
					printf( "SymGetSymFromAddr(): gle = %lu\n", gle );
			}
			else
			{
				// UnDecorateSymbolName()
				pUDSN( pSym->Name, undName, MAXNAMELEN, UNDNAME_NAME_ONLY );
				pUDSN( pSym->Name, undFullName, MAXNAMELEN, UNDNAME_COMPLETE );
				printf( "%s", undName );
				if ( offsetFromSymbol != 0 )
					printf( " %+ld bytes", (long) offsetFromSymbol );
				putchar( '\n' );
				printf( "    Sig:  %s\n", pSym->Name );
				printf( "    Decl: %s\n", undFullName );
			}

			// show line number info, NT5.0-method (SymGetLineFromAddr())
			if ( pSGLFA != NULL )
			{ // yes, we have SymGetLineFromAddr()
				if ( ! pSGLFA(chunk->m_hProcess, chunk->m_StackFrame.AddrPC.Offset, &offsetFromSymbol, &Line ) )
				{
					if ( gle != 487 )
						printf( "SymGetLineFromAddr(): gle = %lu\n", gle );
				}
				else
				{
					printf( "    Line: %s(%lu) %+ld bytes\n",
						Line.FileName, Line.LineNumber, offsetFromSymbol );
				}
			} // yes, we have SymGetLineFromAddr()

			// show module info (SymGetModuleInfo())
			if ( ! pSGMI( chunk->m_hProcess, chunk->m_StackFrame.AddrPC.Offset, &Module ) )
			{
				printf( "SymGetModuleInfo): gle = %lu\n", gle );
			}
			else
			{ // got module info OK
				char ty[80];
				switch ( Module.SymType )
				{
				case SymNone:
					strcpy( ty, "-nosymbols-" );
					break;
				case SymCoff:
					strcpy( ty, "COFF" );
					break;
				case SymCv:
					strcpy( ty, "CV" );
					break;
				case SymPdb:
					strcpy( ty, "PDB" );
					break;
				case SymExport:
					strcpy( ty, "-exported-" );
					break;
				case SymDeferred:
					strcpy( ty, "-deferred-" );
					break;
				case SymSym:
					strcpy( ty, "SYM" );
					break;
				default:
					_snprintf( ty, sizeof ty, "symtype=%ld", (long) Module.SymType );
					break;
				}

				printf( "    Mod:  %s[%s], base: %08lxh\n",
					Module.ModuleName, Module.ImageName, Module.BaseOfImage );
				printf( "    Sym:  type: %s, file: %s\n",
					ty, Module.LoadedImageName );
			} // got module info OK
		} // we seem to have a valid PC

		// no return address means no deeper stackframe
		if ( chunk->m_StackFrame.AddrReturn.Offset == 0 )
		{
			// avoid misunderstandings in the printf() following the loop
			SetLastError( 0 );
			break;
		}

	} // for ( frameNum )

cleanup:
	pSC( m_hProcess );
	free( pSym );
	free(tt);
}



bool kxMemoryChunk::FillModuleList( ModuleList& modules, DWORD pid, HANDLE hProcess )
{
	// try toolhelp32 first
	if ( FillModuleListTH32( modules, pid ) )
		return true;
	// nope? try psapi, then
	return FillModuleListPSAPI( modules, pid, hProcess );
}





bool kxMemoryChunk::FillModuleListTH32( ModuleList& modules, DWORD pid )
{
	// CreateToolhelp32Snapshot()
	typedef HANDLE (__stdcall *tCT32S)( DWORD dwFlags, DWORD th32ProcessID );
	// Module32First()
	typedef BOOL (__stdcall *tM32F)( HANDLE hSnapshot, LPMODULEENTRY32 lpme );
	// Module32Next()
	typedef BOOL (__stdcall *tM32N)( HANDLE hSnapshot, LPMODULEENTRY32 lpme );

	// I think the DLL is called tlhelp32.dll on Win9X, so we try both
	const char *dllname[] = { "kernel32.dll", "tlhelp32.dll" };
	HINSTANCE hToolhelp;
	tCT32S pCT32S;
	tM32F pM32F;
	tM32N pM32N;

	HANDLE hSnap;
	MODULEENTRY32 me = { sizeof me };
	bool keepGoing;
	ModuleEntry e;
	int i;

	for ( i = 0; i < lenof( dllname ); ++ i )
	{
		hToolhelp = LoadLibraryA( dllname[i] );
		if ( hToolhelp == 0 )
			continue;
		pCT32S = (tCT32S) GetProcAddress( hToolhelp, "CreateToolhelp32Snapshot" );
		pM32F = (tM32F) GetProcAddress( hToolhelp, "Module32First" );
		pM32N = (tM32N) GetProcAddress( hToolhelp, "Module32Next" );
		if ( pCT32S != 0 && pM32F != 0 && pM32N != 0 )
			break; // found the functions!
		FreeLibrary( hToolhelp );
		hToolhelp = 0;
	}

	if ( hToolhelp == 0 ) // nothing found?
		return false;

	hSnap = pCT32S( TH32CS_SNAPMODULE, pid );
	if ( hSnap == (HANDLE) -1 )
		return false;

	keepGoing = !!pM32F( hSnap, &me );
	while ( keepGoing )
	{
		// here, we have a filled-in MODULEENTRY32
		printf( "%08lXh %6lu %-15.15s %s\n", me.modBaseAddr, me.modBaseSize, me.szModule, me.szExePath );
		e.imageName = me.szExePath;
		e.moduleName = me.szModule;
		e.baseAddress = (DWORD) me.modBaseAddr;
		e.size = me.modBaseSize;
		modules.push_back( e );
		keepGoing = !!pM32N( hSnap, &me );
	}

	CloseHandle( hSnap );

	FreeLibrary( hToolhelp );

	return modules.size() != 0;
}



// miscellaneous psapi declarations; we cannot #include the header
// because not all systems may have it
typedef struct _MODULEINFO {
	LPVOID lpBaseOfDll;
	DWORD SizeOfImage;
	LPVOID EntryPoint;
} MODULEINFO, *LPMODULEINFO;



bool kxMemoryChunk::FillModuleListPSAPI( ModuleList& modules, DWORD pid, HANDLE hProcess )
{
	// EnumProcessModules()
	typedef BOOL (__stdcall *tEPM)( HANDLE hProcess, HMODULE *lphModule, DWORD cb, LPDWORD lpcbNeeded );
	// GetModuleFileNameEx()
	typedef DWORD (__stdcall *tGMFNE)( HANDLE hProcess, HMODULE hModule, LPSTR lpFilename, DWORD nSize );
	// GetModuleBaseName() -- redundant, as GMFNE() has the same prototype, but who cares?
	typedef DWORD (__stdcall *tGMBN)( HANDLE hProcess, HMODULE hModule, LPSTR lpFilename, DWORD nSize );
	// GetModuleInformation()
	typedef BOOL (__stdcall *tGMI)( HANDLE hProcess, HMODULE hModule, LPMODULEINFO pmi, DWORD nSize );

	HINSTANCE hPsapi;
	tEPM pEPM;
	tGMFNE pGMFNE;
	tGMBN pGMBN;
	tGMI pGMI;

	int i;
	ModuleEntry e;
	DWORD cbNeeded;
	MODULEINFO mi;
	HMODULE *hMods = 0;
	char *tt = 0;

	hPsapi = LoadLibraryA( "psapi.dll" );
	if ( hPsapi == 0 )
		return false;

	modules.clear();

	pEPM = (tEPM) GetProcAddress( hPsapi, "EnumProcessModules" );
	pGMFNE = (tGMFNE) GetProcAddress( hPsapi, "GetModuleFileNameExA" );
	pGMBN = (tGMFNE) GetProcAddress( hPsapi, "GetModuleBaseNameA" );
	pGMI = (tGMI) GetProcAddress( hPsapi, "GetModuleInformation" );
	if ( pEPM == 0 || pGMFNE == 0 || pGMBN == 0 || pGMI == 0 )
	{
		// yuck. Some API is missing.
		FreeLibrary( hPsapi );
		return false;
	}

	hMods = new HMODULE[TTBUFLEN / sizeof HMODULE];
	tt = new char[TTBUFLEN];
	// not that this is a sample. Which means I can get away with
	// not checking for errors, but you cannot. :)

	if ( ! pEPM( hProcess, hMods, TTBUFLEN, &cbNeeded ) )
	{
		printf( "EPM failed, gle = %lu\n", gle );
		goto cleanup;
	}

	if ( cbNeeded > TTBUFLEN )
	{
		printf( "More than %lu module handles. Huh?\n", lenof( hMods ) );
		goto cleanup;
	}

	for ( i = 0; i < cbNeeded / sizeof hMods[0]; ++ i )
	{
		// for each module, get:
		// base address, size
		pGMI( hProcess, hMods[i], &mi, sizeof mi );
		e.baseAddress = (DWORD) mi.lpBaseOfDll;
		e.size = mi.SizeOfImage;
		// image file name
		tt[0] = '\0';
		pGMFNE( hProcess, hMods[i], tt, TTBUFLEN );
		e.imageName = tt;
		// module name
		tt[0] = '\0';
		pGMBN( hProcess, hMods[i], tt, TTBUFLEN );
		e.moduleName = tt;
		printf( "%08lXh %6lu %-15.15s %s\n", e.baseAddress,
			e.size, e.moduleName.c_str(), e.imageName.c_str() );

		modules.push_back( e );
	}

cleanup:
	if ( hPsapi )
		FreeLibrary( hPsapi );
	delete [] tt;
	delete [] hMods;

	return modules.size() != 0;
}


void kxMemoryChunk::EnumumAndLoadModuleSymbols( HANDLE hProcess, DWORD pid )
{
	ModuleList modules;
	ModuleListIter it;
	char *img, *mod;

	// fill in module list
	FillModuleList( modules, pid, hProcess );

	for ( it = modules.begin(); it != modules.end(); ++ it )
	{
		// unfortunately, SymLoadModule() wants writeable strings
		img = (char*)malloc((*it).imageName.size() + 1);
		strcpy( img, (*it).imageName.c_str() );
		mod = (char*)malloc((*it).moduleName.size() + 1);
		strcpy( mod, (*it).moduleName.c_str() );

		if ( pSLM( hProcess, 0, img, mod, (*it).baseAddress, (*it).size ) == 0 )
			printf( "Error %lu loading symbols for \"%s\"\n",
			gle, (*it).moduleName.c_str() );
		else
			printf( "Symbols loaded: \"%s\"\n", (*it).moduleName.c_str() );

		free(img);
		free(mod);
	}
}


void kxMemoryChunk::LazyInit()
{
	m_Begin = (kxMemoryChunk*)malloc(sizeof(kxMemoryChunk));
	m_End = (kxMemoryChunk*)malloc(sizeof(kxMemoryChunk));

	m_Begin->m_Prev = NULL;
	m_Begin->m_Next = m_End;

	m_End->m_Prev = m_Begin;
	m_End->m_Next = NULL;

	m_hImagehlpDll = LoadLibrary( _T("imagehlp.dll") );

	pSC = (tSC) GetProcAddress( m_hImagehlpDll, "SymCleanup" );
	pSFTA = (tSFTA) GetProcAddress( m_hImagehlpDll, "SymFunctionTableAccess" );
	pSGLFA = (tSGLFA) GetProcAddress( m_hImagehlpDll, "SymGetLineFromAddr" );
	pSGMB = (tSGMB) GetProcAddress( m_hImagehlpDll, "SymGetModuleBase" );
	pSGMI = (tSGMI) GetProcAddress( m_hImagehlpDll, "SymGetModuleInfo" );
	pSGO = (tSGO) GetProcAddress( m_hImagehlpDll, "SymGetOptions" );
	pSGSFA = (tSGSFA) GetProcAddress( m_hImagehlpDll, "SymGetSymFromAddr" );
	pSI = (tSI) GetProcAddress( m_hImagehlpDll, "SymInitialize" );
	pSSO = (tSSO) GetProcAddress( m_hImagehlpDll, "SymSetOptions" );
	pSW = (tSW) GetProcAddress( m_hImagehlpDll, "StackWalk" );
	pUDSN = (tUDSN) GetProcAddress( m_hImagehlpDll, "UnDecorateSymbolName" );
	pSLM = (tSLM) GetProcAddress( m_hImagehlpDll, "SymLoadModule" );
}

void kxMemoryChunk::LazyCleanup()
{
	free(m_Begin);
	free(m_End);

	pSC = (tSC) NULL;
	pSFTA = (tSFTA) NULL;
	pSGLFA = (tSGLFA) NULL;
	pSGMB = (tSGMB) NULL;
	pSGMI = (tSGMI) NULL;
	pSGO = (tSGO) NULL;
	pSGSFA = (tSGSFA) NULL;
	pSI = (tSI) NULL;
	pSSO = (tSSO) NULL;
	pSW = (tSW) NULL;
	pUDSN = (tUDSN) NULL;
	pSLM = (tSLM) NULL;

	m_Begin = m_End = NULL;
}


void* operator new(size_t size)
{
	kxMemoryChunk* tmp = new kxMemoryChunk(size);
	if(tmp->m_Data == NULL) // Pokud se nam nepodarilo alokovat data
	{
		delete tmp;
		return NULL;
	}

	return tmp->m_Data;
}

void operator delete(void* ptr)
{
	kxMemoryChunk::Delete(ptr);
}