#include "stdafx.h"

#include "SpBase.h"
#include "SpBaseRoutine.h"
#include <locale.h>

#ifdef _WIN32
#include <imagehlp.h>

#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif // WIN32_LEAN_AND_MEAN

#ifndef _WIN32_WINNT
#define _WIN32_WINNT 0x400
#endif // _WIN32_WINNT

#include <windows.h>
#include <psapi.h>
#include <Tlhelp32.h>
#endif //_WIN32

#include <map>
#include <vector>
using namespace std;


////////////////////////////////////////////////////////////
//DumpStack part
#ifndef _countof
#define _countof(array) (sizeof(array)/sizeof(array[0]))
#endif

const DWORD MODULE_NAME_LEN=64;
const DWORD SYMBOL_NAME_LEN=128;

struct tagSYMBOL_INFO
{
	DWORD dwAddress;
	DWORD dwOffset;
	char szModule[MODULE_NAME_LEN];
	char szSymbol[SYMBOL_NAME_LEN];
};

static tagSYMBOL_INFO ResolveSymbol(HANDLE hProcess, DWORD dwAddress)
{

	union
	{
		CHAR rgchSymbol[sizeof(IMAGEHLP_SYMBOL) + 255];
		IMAGEHLP_SYMBOL  sym;
	};

	char szUndec[256];
	char szWithOffset[256];
	char *pszSymbol = NULL;
	IMAGEHLP_MODULE mi;
	tagSYMBOL_INFO siSymbol;

	memset(&mi,0,sizeof(mi));
	memset(&siSymbol,0,sizeof(siSymbol));
	siSymbol.dwAddress = dwAddress;

	mi.SizeOfStruct = sizeof(IMAGEHLP_MODULE);

	if (!SymGetModuleInfo(hProcess, dwAddress, &mi))
		lstrcpyA(siSymbol.szModule, "<no module>");
	else
	{
		char *pszModule = strchr(mi.ImageName, '\\');
		if (pszModule == NULL)
			pszModule = mi.ImageName;
		else
			pszModule++;

		lstrcpynA(siSymbol.szModule, pszModule, _countof(siSymbol.szModule));
		lstrcatA(siSymbol.szModule, "! ");
	}

	__try
	{
		sym.SizeOfStruct = sizeof(IMAGEHLP_SYMBOL);
		sym.Address = dwAddress;
		sym.MaxNameLength = 255;

		if (SymGetSymFromAddr(hProcess, dwAddress, &(siSymbol.dwOffset), &sym))
		{
			pszSymbol = sym.Name;

			if (UnDecorateSymbolName(sym.Name, szUndec, _countof(szUndec),
				UNDNAME_NO_MS_KEYWORDS | UNDNAME_NO_ACCESS_SPECIFIERS))
			{
				pszSymbol = szUndec;
			}
			else if (SymUnDName(&sym, szUndec, _countof(szUndec)))
			{
				pszSymbol = szUndec;
			}

			if (siSymbol.dwOffset != 0)
			{
				wsprintfA(szWithOffset, "%s + %d bytes", pszSymbol, siSymbol.dwOffset);
				pszSymbol = szWithOffset;
			}
		}
		else
			pszSymbol = "<no symbol>";
	}
	__except (EXCEPTION_EXECUTE_HANDLER)
	{
		pszSymbol = "<EX: no symbol>";
		siSymbol.dwOffset = dwAddress - mi.BaseOfImage;
	}

	lstrcpynA(siSymbol.szSymbol, pszSymbol, _countof(siSymbol.szSymbol));
	return siSymbol;
}

static DWORD __stdcall GetModuleBase( HANDLE hProcess, DWORD dwReturnAddress )
{
	IMAGEHLP_MODULE moduleInfo;

	if ( SymGetModuleInfo( hProcess, dwReturnAddress, &moduleInfo ) )
		return moduleInfo.BaseOfImage;

	MEMORY_BASIC_INFORMATION mbInfo;
	if ( VirtualQueryEx( hProcess, (LPCVOID)dwReturnAddress, &mbInfo, sizeof( mbInfo ) ) )
	{
		char szFile[ MAX_PATH ] = { 0 };
		DWORD cch = GetModuleFileNameA( (HINSTANCE)mbInfo.AllocationBase, szFile, MAX_PATH );

		// Ignore the return code since we can't do anything with it.
		if ( SymLoadModule( hProcess, NULL, cch ? szFile : NULL, NULL,
			(DWORD)mbInfo.AllocationBase, 0 ) )

			return (DWORD) mbInfo.AllocationBase;
	}

	return SymGetModuleBase( hProcess, dwReturnAddress );
}


static bool DumpStackVector(vector<tagSYMBOL_INFO> &CallStacks)
{
	HANDLE hProcess = ::GetCurrentProcess();
	if (!::SymInitialize(hProcess, NULL, FALSE))
		return false;

	// force undecorated names to get params
	DWORD dw = ::SymGetOptions();
	dw &= ~SYMOPT_UNDNAME;
	::SymSetOptions(dw);

	HANDLE hThread = ::GetCurrentThread();
	CONTEXT threadContext;

	threadContext.ContextFlags = CONTEXT_FULL;

	if (::GetThreadContext(hThread, &threadContext))
	{
		STACKFRAME stackFrame;
		memset(&stackFrame, 0, sizeof(stackFrame));
		stackFrame.AddrPC.Mode = AddrModeFlat;

		DWORD dwMachType(0);

		dwMachType					= IMAGE_FILE_MACHINE_I386;

		// program counter, stack pointer, and frame pointer
		stackFrame.AddrPC.Offset    = threadContext.Eip;
		stackFrame.AddrStack.Offset = threadContext.Esp;
		stackFrame.AddrStack.Mode   = AddrModeFlat;
		stackFrame.AddrFrame.Offset = threadContext.Ebp;
		stackFrame.AddrFrame.Mode   = AddrModeFlat;

		for (int nFrame = 0; nFrame < 1024; nFrame++)
		{
			if (!::StackWalk(dwMachType, hProcess, hThread,
				&stackFrame, &threadContext, NULL,
				SymFunctionTableAccess, GetModuleBase, NULL))
			{
				break;
			}
			CallStacks.push_back(ResolveSymbol(hProcess, stackFrame.AddrPC.Offset));
		}
	}
	return true;
}

CSimpleStringA DumpStack(DWORD nSkipLevel)
{
	CSimpleStringA strStack;
	vector<tagSYMBOL_INFO> CallStacks;
	if(!DumpStackVector(CallStacks))
	{
		strStack.Format("DumpStack Error: IMAGEHLP.DLL wasn't found. GetLastError() returned 0x%8.8X||",GetLastError());
		return strStack;
	}
	strStack="===begin===";
	// dump it out now
	int cAddresses = CallStacks.size();
	for (DWORD nAddress = nSkipLevel+1; nAddress < (DWORD)cAddresses; nAddress++)
	{
		tagSYMBOL_INFO info = CallStacks[nAddress];
		CSimpleStringA strLine;
		strLine.Format("%8.8X: %s%s||",info.dwAddress,info.szModule,info.szSymbol);
		strStack+=strLine;
	}
	strStack+="===end===";
	return strStack;
}

CSimpleStringA GetCallerFunctionName(DWORD nSkipLevel)
{
	vector<tagSYMBOL_INFO> CallStacks;
	if(!DumpStackVector(CallStacks))
		return "";
	tagSYMBOL_INFO info=CallStacks[nSkipLevel];
	CSimpleStringA strStack;
	strStack.Format("Address:%8.8X,Modual:\"%s\",Function:\"%s\"",info.dwAddress,info.szModule,info.szSymbol);
	return strStack;
}

typedef enum _THREADINFOCLASSEX
{
    ThreadBasicInformation,
    ThreadTimes,
    ThreadPriority,
    ThreadBasePriority,
    ThreadAffinityMask,
    ThreadImpersonationToken,
    ThreadDescriptorTableEntry,
    ThreadEnableAlignmentFaultFixup,
    ThreadEventPair_Reusable,
    ThreadQuerySetWin32StartAddress,
    ThreadZeroTlsCell,
    ThreadPerformanceCount,
    ThreadAmILastThread,
    ThreadIdealProcessor,
    ThreadPriorityBoost,
    ThreadSetTlsArrayAddress,
    ThreadIsIoPending_Reserved,
    ThreadHideFromDebugger,
    ThreadBreakOnTermination,
    MaxThreadInfoClass
}  THREADINFOCLASSEX;

#if defined(_MSC_VER) && (_MSC_VER < 1930)
typedef struct _CLIENT_ID
{
    HANDLE UniqueProcess;
    HANDLE UniqueThread;
} CLIENT_ID;
typedef CLIENT_ID* PCLIENT_ID;

#endif //_MSC_VER



typedef struct _THREAD_BASIC_INFORMATION { // Information Class 0
	LONG     ExitStatus;
	PVOID    TebBaseAddress;
	CLIENT_ID ClientId;
	LONG AffinityMask;
	LONG Priority;
	LONG BasePriority;
} THREAD_BASIC_INFORMATION, *PTHREAD_BASIC_INFORMATION;

extern "C" LONG (__stdcall *ZwQueryInformationThread) (
	IN HANDLE ThreadHandle,
	IN THREADINFOCLASSEX ThreadInformationClass,
	OUT PVOID ThreadInformation,
	IN ULONG ThreadInformationLength,
	OUT PULONG ReturnLength OPTIONAL
	) = NULL;


extern "C" LONG (__stdcall *pFnRtlNtStatusToDosError) (IN  ULONG status) = NULL;

BOOL ShowThreadInfo (DWORD tid)
{
	static BOOL isInit = FALSE;
	static std::map<DWORD, string> threadInModule;
	if (!isInit)
	{
		setlocale (LC_ALL, ".ACP");

		HINSTANCE hNTDLL = ::GetModuleHandle (TEXT ("ntdll"));

		(FARPROC&)ZwQueryInformationThread  =
			::GetProcAddress (hNTDLL, "ZwQueryInformationThread");

		(FARPROC&)pFnRtlNtStatusToDosError = ::GetProcAddress (hNTDLL, "RtlNtStatusToDosError");
		threadInModule.clear();
		isInit = TRUE;
	}

	std::map<DWORD, string>::iterator it = threadInModule.find(tid);
	if (threadInModule.end() != it)
		return TRUE;


	THREAD_BASIC_INFORMATION    tbi;
	PVOID                       startaddr;
	LONG                        status;
	HANDLE                      thread, process;

	thread = ::OpenThread (THREAD_ALL_ACCESS, FALSE, tid);
	if (thread == NULL)
		return FALSE;

	//thread 起始地址
	status = ZwQueryInformationThread (thread, ThreadQuerySetWin32StartAddress,	&startaddr,	sizeof (startaddr),	NULL);	

	if (status < 0)
	{
		CloseHandle (thread);
		SetLastError (pFnRtlNtStatusToDosError(status));
		return FALSE;
	};

	status = ZwQueryInformationThread (thread, ThreadBasicInformation, &tbi, sizeof (tbi), NULL);

	if (status < 0)
	{
		CloseHandle (thread);
		SetLastError (pFnRtlNtStatusToDosError(status));
		return FALSE;
	};

	process = ::OpenProcess (PROCESS_ALL_ACCESS, FALSE, (DWORD)tbi.ClientId.UniqueProcess);

	if (process == NULL)
	{
		DWORD error = ::GetLastError ();
		CloseHandle (thread);
		SetLastError (error);
		return FALSE;
	};

	

	TCHAR modname [0x100];
	::GetModuleFileNameEx (process, NULL, modname, 0x100);

	GetMappedFileName(process, startaddr, modname, 0x100);



	vector<tagSYMBOL_INFO> callStack;
	DumpStackVector(callStack);
	
	//MessageBox(NULL, NULL, NULL, 0);
	threadInModule[tid] = std::string(modname);
	CloseHandle (process);
	CloseHandle (thread);
	return TRUE;
};


class stackTrace
{
public:
	static void InitTrack();
	static CSimpleStringA StackTrack();
	static void UninitTrack();

	static HANDLE g_hHandle;
	static CONTEXT g_context;
};

HANDLE stackTrace::g_hHandle = NULL;
CONTEXT stackTrace::g_context = { CONTEXT_FULL };


CSimpleStringA TraceStack()
{
	static CRITICAL_SECTION cs_;
	static bool isInit = false;
	if (!isInit)
		InitializeCriticalSection(&cs_);
	EnterCriticalSection(&cs_);

	HANDLE hThread = GetCurrentThread();
	GetThreadContext(hThread, &stackTrace::g_context);
	__asm{call $ + 5}
	__asm{pop eax}
	__asm{mov stackTrace::g_context.Eip, eax}
	__asm{mov stackTrace::g_context.Ebp, ebp}
	__asm{mov stackTrace::g_context.Esp, esp}
	stackTrace::InitTrack();
	CSimpleStringA resultModule = stackTrace::StackTrack();
	stackTrace::UninitTrack();

	LeaveCriticalSection(&cs_);
	return resultModule;
}
void stackTrace::InitTrack()
{
	g_hHandle = GetCurrentProcess();

	SymInitialize(g_hHandle, NULL, TRUE);
}

CSimpleStringA stackTrace::StackTrack()
{
	HANDLE m_hThread = GetCurrentThread();
	STACKFRAME sf = { 0 };

	sf.AddrPC.Offset = g_context.Eip;
	sf.AddrPC.Mode = AddrModeFlat;

	sf.AddrFrame.Offset = g_context.Ebp;
	sf.AddrFrame.Mode = AddrModeFlat;

	sf.AddrStack.Offset = g_context.Esp;
	sf.AddrStack.Mode = AddrModeFlat;

	typedef struct tag_SYMBOL_INFO
	{
		IMAGEHLP_SYMBOL symInfo;
		TCHAR szBuffer[MAX_PATH];
	} SYMBOL_INFO, *LPSYMBOL_INFO;

	DWORD dwDisplament = 0;
	SYMBOL_INFO stack_info = { 0 };
	PIMAGEHLP_SYMBOL pSym = (PIMAGEHLP_SYMBOL)&stack_info;
	pSym->SizeOfStruct = sizeof(IMAGEHLP_SYMBOL);
	pSym->MaxNameLength = sizeof(SYMBOL_INFO) - offsetof(SYMBOL_INFO, symInfo.Name);
	IMAGEHLP_LINE ImageLine = { 0 };
	ImageLine.SizeOfStruct = sizeof(IMAGEHLP_LINE);

	while (StackWalk(IMAGE_FILE_MACHINE_I386, g_hHandle, m_hThread, &sf, &g_context, NULL, SymFunctionTableAccess, SymGetModuleBase, NULL))
	{
		IMAGEHLP_MODULE mi;

		memset(&mi,0,sizeof(mi));

		mi.SizeOfStruct = sizeof(IMAGEHLP_MODULE);
		SymGetModuleInfo(g_hHandle, sf.AddrPC.Offset, &mi);

		string moduleName = mi.ModuleName;
		if (-1 != moduleName.find("mod_"))
			return CSimpleStringA(mi.ModuleName);
	}
	return CSimpleStringA("");
}

void stackTrace::UninitTrack()
{
	HANDLE hProcess = GetCurrentProcess();
	SymCleanup(g_hHandle);
	g_hHandle = NULL;
}