#include "Gos.pch"
#include "Debug.h"
#include <stdio.h>

extern IMAGEHLP_LINE ImageHlp_pline;

int __stdcall gosReadMemory( void* hProcess, const void * lpBaseAddress, void* lpBuffer, DWORD nSize, DWORD* lpNumberOfBytesRead );

SYMFUNCTIONTABLEACCESSPROC	Debug::_SymFunctionTableAccess;
SYMSTACKWALKPROC			Debug::_StackWalk;
SYMGETMODULEBASEPROC		Debug::_SymGetModuleBase;
SYMGETLINEFROMADDRPROC		Debug::_SymGetLineFromAddr;
SYMGETSYMFROMADDRPROC		Debug::_SymGetSymFromAddr;
SYMLOADMODULEPROC			Debug::_SymLoadModule;
SYMINITIALIZEPROC			Debug::_SymInitialize;
SYMSETOPTIONSPROC			Debug::_SymSetOptions;
SYMUNLOADMODULEPROC			Debug::_SymUnloadModule;

HMODULE Debug::hModImagehlp = 0;
HMODULE Debug::hModMsdbi = 0;
void*	Debug::MainModuleBase = 0;
DWORD	Debug::LastOffset = 0;
bool	Debug::bLoadAllSymbols = 0;

int (*GosDebug::HandleError)( int, TCHAR * ) = Debug::HandleError;
bool *(*GosDebug::Cat)( _TCHAR *pCategoryName ) = Debug::Cat;
void (*GosDebug::EnableCat)( bool *pCat ) = Debug::EnableCat;
void (*GosDebug::DisableCat)( bool *pCat ) = Debug::DisableCat;
void (*GosDebug::Output)( _TCHAR *pFmt,... ) = Debug::Output;
void (*GosDebug::OutputV)( _TCHAR *pFmt, va_list arglist ) = Debug::OutputV;

//
// Buffer for symbol names
//
BYTE Debug::sym[ sizeof(IMAGEHLP_SYMBOL) + 512 ];
PIMAGEHLP_SYMBOL Debug::psym = (PIMAGEHLP_SYMBOL)sym;

DWORD GosDebug::s_ErrorLine;
TCHAR *GosDebug::s_pErrorFileName;

struct SCat
{
	bool	        enabled;
	SCat	        *pNext;
	_TCHAR	        *pName;
};

static SCat *pFirstCat = NULL;


int Debug::HandleError( int Flags, TCHAR *Text )
{
	_asm int 3;
	return 0;
}


bool *Debug::Cat( _TCHAR *pCategoryName )
{
	SCat *pCat;
	for( pCat = pFirstCat; pCat; pCat = pCat->pNext )
	{
		if( _tcscmp( pCategoryName, pCat->pName ) == 0 )
			return (bool *)pCat;
	}
	pCat = (SCat *)malloc( sizeof(SCat) );
	pCat->pNext = pFirstCat;
	pCat->pName = pCategoryName;
	pCat->enabled = false;
	pFirstCat = pCat;
	return (bool *)pCat;
}

void Debug::EnableCat( bool *pCat )
{
	*pCat = true;
}

void Debug::DisableCat( bool *pCat )
{
	*pCat = false;
}


void Debug::Output( _TCHAR *pFmt,... )
{
	static _TCHAR msg[200];
	va_list arglist;
	va_start( arglist, pFmt );
	_vsntprintf( msg, 200, pFmt, arglist );
	va_end( arglist );

	OutputDebugString( msg );
}

void Debug::OutputV( _TCHAR *pFmt, va_list arglist )
{
	static _TCHAR msg[200];
	_vsntprintf( msg, 200, pFmt, arglist );

	OutputDebugString( msg );
}

//
// Find the imagehlp.dll and link to functions used
//
bool Debug::LoadImageHlp()
{
//
// Find the base address of the main program module
//
	hModMsdbi = LoadLibrary( _T("dbghelp.dll") );
	if ( hModMsdbi )
	{
		MEMORY_BASIC_INFORMATION mbi;
		VirtualQuery( &LoadImageHlp, &mbi, sizeof(mbi) );
		MainModuleBase=mbi.AllocationBase;
	//
	// Load the imagehlp library
	//
		hModImagehlp = LoadLibrary( _T("imagehlp.dll") );
		if ( !hModImagehlp )
		{
			return 0;
		}

		_SymInitialize = (SYMINITIALIZEPROC)GetProcAddress( hModImagehlp, "SymInitialize" );
		if ( !_SymInitialize )
			return FALSE;

		_StackWalk = (SYMSTACKWALKPROC)GetProcAddress( hModImagehlp, "StackWalk" );
		if ( !_StackWalk )
			return FALSE;

		_SymFunctionTableAccess = (SYMFUNCTIONTABLEACCESSPROC) GetProcAddress( hModImagehlp, "SymFunctionTableAccess" );
		if ( !_SymFunctionTableAccess )
			return FALSE;

		_SymGetModuleBase=(SYMGETMODULEBASEPROC)GetProcAddress( hModImagehlp, "SymGetModuleBase");
		if ( !_SymGetModuleBase )
			return FALSE;

		_SymGetSymFromAddr=(SYMGETSYMFROMADDRPROC)GetProcAddress( hModImagehlp, "SymGetSymFromAddr" );
		if ( !_SymGetSymFromAddr )
			return FALSE;

		_SymGetLineFromAddr=(SYMGETLINEFROMADDRPROC)GetProcAddress( hModImagehlp, "SymGetLineFromAddr" );
		if ( !_SymGetLineFromAddr )
			return FALSE;

		_SymSetOptions=(SYMSETOPTIONSPROC)GetProcAddress( hModImagehlp, "SymSetOptions" );
		if ( !_SymSetOptions )
			return FALSE;

		_SymLoadModule=(SYMLOADMODULEPROC)GetProcAddress( hModImagehlp, "SymLoadModule" );
		if ( !_SymLoadModule )
			return FALSE;

		_SymUnloadModule=(SYMUNLOADMODULEPROC)GetProcAddress( hModImagehlp, "SymUnloadModule" );
		if ( !_SymUnloadModule )
			return FALSE;

		_SymSetOptions(0);
		_SymSetOptions(SYMOPT_UNDNAME|SYMOPT_LOAD_LINES|SYMOPT_OMAP_FIND_NEAREST);
		char drive[_MAX_PATH];
		char dir[_MAX_DIR];
		char fname[_MAX_FNAME];
		char ext[_MAX_EXT];
		_splitpath( _pgmptr, drive, dir, fname, ext );
		strcat( drive, dir );
		strcat( drive, ";" );
		
		//
		// Open .PDB file
		//
		if( !_SymInitialize( GetCurrentProcess(), drive, bLoadAllSymbols ))
			return false;
		if (!_SymLoadModule( GetCurrentProcess(), NULL, _pgmptr, NULL, NULL, NULL ))
			return false;;

		return true;
	}
	return false;
}

//
// Initialize the stack walking function
//
void Debug::InitStackWalk( STACKFRAME* sf, CONTEXT* Context )
{
	memset( sf, 0, sizeof(STACKFRAME) );
	sf->AddrPC.Offset		= Context->Eip;
	sf->AddrPC.Mode			= AddrModeFlat;
	sf->AddrStack.Offset	= Context->Esp;
	sf->AddrStack.Mode		= AddrModeFlat;
	sf->AddrFrame.Offset	= Context->Ebp;
	sf->AddrFrame.Mode		= AddrModeFlat;
	if( hModMsdbi )
	{
		sf->FuncTableEntry=_SymFunctionTableAccess( GetCurrentProcess(), Context->Eip );
	}
}

//
// Walk the stack, return new address 
//
int Debug::WalkStack( STACKFRAME* sf )
{
    // jcem
	if( hModMsdbi && _StackWalk(IMAGE_FILE_MACHINE_I386,
        GetCurrentProcess(),
        GetCurrentThread(),
        sf,
        NULL,
        (PREAD_PROCESS_MEMORY_ROUTINE)gosReadMemory,
        _SymFunctionTableAccess,
        _SymGetModuleBase,
        NULL))
//DWORD, HANDLE, HANDLE,
//LPSTACKFRAME, LPVOID,
//PREAD_PROCESS_MEMORY_ROUTINE,
//PFUNCTION_TABLE_ACCESS_ROUTINE,
//PGET_MODULE_BASE_ROUTINE,
//PTRANSLATE_ADDRESS_ROUTINE 
	{
		if( sf->AddrPC.Offset < 0x400000 )
			return 0;
		return sf->AddrPC.Offset;
	}
	else
		return 0;
}




//
// Return the "file(line number)" for a given address
//
_TCHAR* Debug::GetLocationFromAddress( IMAGEHLP_LINE* ImageHlp_pline, _TCHAR* Buffer, int Address )
{
	memset( ImageHlp_pline, 0, sizeof(ImageHlp_pline) );
	ImageHlp_pline->SizeOfStruct = sizeof(IMAGEHLP_LINE);
	char bBuffer[MAX_PATH];

	DWORD disp=0;

	if( hModMsdbi )
	{
		for( int off=0; off<1024; off++ )
		{
			if( _SymGetLineFromAddr( GetCurrentProcess(), Address-off, &disp, ImageHlp_pline ) )
			{
				sprintf(bBuffer, "%s(%d)", ImageHlp_pline->FileName, ImageHlp_pline->LineNumber );
				mbstowcs(Buffer, bBuffer, strlen(bBuffer)+1);
				return Buffer;
			}
		}
	}
	return _T("");
}



//
// Return the symbol nearest an address
//
_TCHAR* Debug::GetSymbolFromAddress( _TCHAR* Buffer, int Address )
{
	char bBuffer[1024];
	bBuffer[0] = 0;
	psym->SizeOfStruct = sizeof(IMAGEHLP_SYMBOL);
	psym->Flags = 0;
	psym->MaxNameLength = 512;
	psym->Name[0] = 0;
	LastOffset = 0;
		
	if( hModMsdbi && _SymGetSymFromAddr( GetCurrentProcess(), Address, &LastOffset, psym ) )
	{
		MEMORY_BASIC_INFORMATION mbi;
		_TCHAR tcModuleName[MAX_PATH];
		char ModuleName[MAX_PATH];
//
// Check to see if in current program or another module
//
		if( VirtualQuery( (void*)Address, &mbi, sizeof(mbi) ) )
		{
			if( mbi.AllocationBase && mbi.AllocationBase!=MainModuleBase )
			{
				if( GetModuleFileName( (HMODULE)mbi.AllocationBase, tcModuleName, MAX_PATH ) )
				{
					wcstombs(ModuleName, tcModuleName, wcslen(tcModuleName));
					char drive[_MAX_DRIVE];
					char dir[_MAX_DIR];
					char fname[_MAX_FNAME];
					char ext[_MAX_EXT];
   					_splitpath( ModuleName, drive, dir, fname, ext );
					for( DWORD t0=0; t0<strlen(fname); t0++ )
						if( fname[t0]>='a' && fname[t0]<='z' )
							fname[t0]=(char)(fname[t0]-32);
					char Temp[256];
					sprintf( Temp, "%s! %s",fname,psym->Name );
					strcpy( psym->Name, (char *) Temp );
				}
			}
		}
//
// Return function name
//
		strcat( psym->Name, "()" );	// Add () after function

		if( LastOffset )			// If a displacement, display it after symbol
		{
			char Buffer[30];
			sprintf( Buffer, "+0x%X",LastOffset );
			strcat( psym->Name, Buffer);
		}
		strcpy( bBuffer, psym->Name );
		mbstowcs(Buffer, bBuffer, strlen(bBuffer)+1);
		return Buffer;
	}
	else
	{
		MEMORY_BASIC_INFORMATION mbi;
		_TCHAR ModuleName[MAX_PATH];

		if( !VirtualQuery( (void*)Address, &mbi, sizeof(mbi) ) )
			return _T("");

		DWORD hMod=(DWORD)mbi.AllocationBase;

		if( 0==hMod || 0==GetModuleFileName( (HMODULE)hMod, ModuleName, MAX_PATH ) )
			return _T("");

		PIMAGE_DOS_HEADER pDosHdr=(PIMAGE_DOS_HEADER)hMod;
		PIMAGE_NT_HEADERS pNtHdr=(PIMAGE_NT_HEADERS)(hMod + pDosHdr->e_lfanew);
		PIMAGE_SECTION_HEADER pSection=IMAGE_FIRST_SECTION( pNtHdr );

		DWORD rva=Address-hMod;

		for( int i=0; i < pNtHdr->FileHeader.NumberOfSections; i++, pSection++ )
		{
			DWORD sectionStart = pSection->VirtualAddress;
			DWORD sectionEnd = sectionStart +  pSection->SizeOfRawData > pSection->Misc.VirtualSize ? pSection->SizeOfRawData : pSection->Misc.VirtualSize;

			if( (rva>=sectionStart) && (rva<=sectionEnd) )
			{
				TCHAR drive[_MAX_DRIVE];
				TCHAR dir[_MAX_DIR];
				TCHAR fname[_MAX_FNAME];
				TCHAR ext[_MAX_EXT];
   				_wsplitpath( ModuleName, drive, dir, fname, ext );
				int offset = rva - sectionStart;
				for( DWORD t0=0; t0<wcslen(fname); t0++ )
					if( fname[t0]>='a' && fname[t0]<='z' )
						fname[t0]=(char)(fname[t0]-32);
				swprintf( Buffer, _T("%s! (+0x%X) "),fname,offset );				
				return Buffer;
			}
		}
	}
	return _T("");
}


int __stdcall gosReadMemory( void* hProcess, const void * lpBaseAddress, void* lpBuffer, DWORD nSize, DWORD* lpNumberOfBytesRead )
{
	if( IsBadReadPtr( (void*)lpBaseAddress, nSize ))
		return FALSE;
	
	memcpy( lpBuffer, (void*)lpBaseAddress, nSize );

	if( lpNumberOfBytesRead )
		*lpNumberOfBytesRead=nSize;

	return TRUE;
}

//
// Returns an ASCII string describing the address passed
//
_TCHAR* Debug::DecodeAddress( DWORD Address, bool brief )
{
	if( 0==Address )
		return 0;
	static _TCHAR LineBuffer[256];

	LineBuffer[0]=0;
	
	_TCHAR BufferSym[MAX_PATH];
	_TCHAR BufferLoc[MAX_PATH];
	IMAGEHLP_LINE pline;

	_TCHAR* File=GetSymbolFromAddress( BufferSym, Address );
	_TCHAR* Line=GetLocationFromAddress( &pline, BufferLoc, Address );

	if( 0==Line )
		return 0;
	if (brief)
		swprintf( LineBuffer, _T("%s"), Line );
	else
		swprintf( LineBuffer, _T("%s : %s"),Line,File );
	return LineBuffer;
}

DWORD*	OldEBP;
DWORD*	CurrentEBP;
//
// Walk up the stack (only in debug builds - with stack frames)
//
DWORD Debug::WalkUp()
{
	DWORD result=0;
	#ifdef _DEBUG
		if( CurrentEBP )
		{
			result=CurrentEBP[1];
			CurrentEBP=(DWORD*)CurrentEBP[0];
			if( 0!=((DWORD)CurrentEBP&3) || CurrentEBP<OldEBP || (DWORD)CurrentEBP>0x80000000 )
				CurrentEBP=0;
			OldEBP=CurrentEBP;
		}
	#endif
	return result;
}


//
// Walk up a stack, filling an array with addresses of the routines
//
// NumberOfLevels is the size of the array (in DWORDS)
// IgnoreLevels is the number of levels to ignore (so the WalkStack routine does not show up!)
//
void Debug::WalkStack( DWORD* RoutineAddresses, DWORD NumberOfLevels, DWORD IgnoreLevels )
{
	_asm	mov CurrentEBP,ebp
	_asm	mov OldEBP,esp

	if( 0!=((DWORD)CurrentEBP&3) || CurrentEBP<OldEBP || (DWORD)CurrentEBP>0x80000000 )
		CurrentEBP=0;

	for( DWORD t0=0; t0<IgnoreLevels; t0++ )
		WalkUp();

	for( t0=0; t0<NumberOfLevels; t0++ )
		RoutineAddresses[t0]=WalkUp();
}



