#include "StdAfx.h"
#include "DumpCollect.h"

namespace DumpCollect
{
	TCHAR _szDrive[_MAX_DRIVE] = { 0 };
	TCHAR _szDir[_MAX_DIR] = { 0 };
	TCHAR _szFna[_MAX_DIR] = { 0 };
	TCHAR _szExt[_MAX_DIR] = { 0 };
	TCHAR _szModulePath[MAX_PATH] = {0};
	TCHAR _szModuleFileName[MAX_PATH] = {0};

	void GetName()
	{
		static BOOL bGetName = FALSE;
		if ( bGetName )
			return;

		::GetModuleFileName(NULL, _szModulePath, sizeof(_szModulePath) / sizeof(TCHAR));
		_stprintf_s(_szModuleFileName, _T("%s"), _szModulePath);

		_tsplitpath_s(_szModulePath, _szDrive, _szDir, _szFna, _szExt);
		_tcscpy_s(_szModulePath, _szDrive);
		_tcscat_s(_szModulePath, _szDir);
	}

 	void CreateDumpFile( IN LPCTSTR lpstrDumpFilePath, IN EXCEPTION_POINTERS *pException )
	{
		// 创建Dump文件;
		HANDLE hDumpFile = CreateFile( lpstrDumpFilePath, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_WRITE, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL );

		if ( (hDumpFile != NULL) && (hDumpFile != INVALID_HANDLE_VALUE) )
		{
			// Dump信息;
			MINIDUMP_EXCEPTION_INFORMATION	dumpInfo;
			ZeroMemory(&dumpInfo, sizeof(MINIDUMP_EXCEPTION_INFORMATION));

			dumpInfo.ThreadId			= GetCurrentThreadId();
			dumpInfo.ExceptionPointers	= pException;
			dumpInfo.ClientPointers		= TRUE;

			// 写入Dump文件内容;
			MINIDUMP_CALLBACK_INFORMATION	mci;
			ZeroMemory(&mci, sizeof(MINIDUMP_CALLBACK_INFORMATION));

			mci.CallbackRoutine		= (MINIDUMP_CALLBACK_ROUTINE)MiniDumpCallback;
			mci.CallbackParam		= 0;

			::MiniDumpWriteDump( GetCurrentProcess(), GetCurrentProcessId(), hDumpFile, MiniDumpNormal, (pException != NULL) ? &dumpInfo : NULL, NULL, &mci );

			CloseHandle(hDumpFile);
		}
	}

	BOOL CALLBACK MiniDumpCallback( IN PVOID pParam, IN const PMINIDUMP_CALLBACK_INPUT pInput, IN PMINIDUMP_CALLBACK_OUTPUT pOutput )
	{
		if ( pInput == NULL || pOutput == NULL )
		{
			return FALSE;
		}

		switch(pInput->CallbackType)
		{
		case ModuleCallback:
			if ( pOutput->ModuleWriteFlags & ModuleWriteDataSeg )
			{
				if ( !IsDataSectionNeeded(pInput->Module.FullPath) )
				{
					pOutput->ModuleWriteFlags &= (~ModuleWriteDataSeg);
				}
			}
		case IncludeModuleCallback:
		case IncludeThreadCallback:
		case ThreadCallback:
		case ThreadExCallback:
			return TRUE;
		default:
			break;
		}

		return FALSE;
	}

	BOOL IsDataSectionNeeded( IN const WCHAR* pModuleName )
	{
		if ( pModuleName == NULL || pModuleName[0] == '\0' )
		{
			return FALSE;
		}

		WCHAR szFileName[_MAX_FNAME] = L"";
		_wsplitpath(pModuleName, NULL, NULL, szFileName, NULL);

		if ( _wcsicmp( szFileName, L"ntdll" ) == 0 )
		{
			return TRUE;
		}

		return FALSE;
	}

	LONG CALLBACK MyUnhandledExceptionFilter( IN PEXCEPTION_POINTERS pExceptionInfo )
	{
		TCHAR szDumpFile[MAX_PATH] = {0};
		_stprintf_s(szDumpFile, MAX_PATH, _T("%s%s[%s].dmp"), _szModulePath, _szFna, CTime::GetCurrentTime().Format(_T("%Y%m%d%H%M%S")));
		CreateDumpFile(szDumpFile, pExceptionInfo);
		return EXCEPTION_EXECUTE_HANDLER;
	}

	// 此函数一旦成功调用，之后对SetUnhandledExceptionFilter的调用将无效;
	void DisableSetUnhandledExceptionFilter()
	{
#if 1
		void *addr = (void*)GetProcAddress(LoadLibrary(_T("kernel32.dll")), "SetUnhandledExceptionFilter");

		if ( addr )
		{
			BYTE byCode[16] = {0};
			INT nSize = 0;
			byCode[nSize++] = 0x33;
			byCode[nSize++] = 0xC0;
			byCode[nSize++] = 0xC2;
			byCode[nSize++] = 0x04;
			byCode[nSize++] = 0x00;

			DWORD dwOldFlag, dwTempFlag;
			VirtualProtect(addr, nSize, PAGE_READWRITE, &dwOldFlag);
			WriteProcessMemory(GetCurrentProcess(), addr, byCode, nSize, NULL);
			VirtualProtect(addr, nSize, dwOldFlag, &dwTempFlag);
		}
#else
		HMODULE hkernel32 = LoadLibrary(_T("kernel32.dll"));

		if ( hkernel32 == NULL)
		{
			return ;
		}

		void *pOrgEntry = GetProcAddress(hkernel32, "SetUnhandledExceptionFilter");
		if ( pOrgEntry == NULL )
			return;

		BYTE newJump[100];
		DWORD dwOrgEntryAddr = (DWORD)pOrgEntry;

		dwOrgEntryAddr += 5;

		void *pNewFunc = &MyUnhandledExceptionFilter;
		DWORD dwNewEntryAddr = (DWORD)pNewFunc;
		DWORD dwRelativeAddr = dwNewEntryAddr - dwOrgEntryAddr;

		newJump[0] = 0xE9;	// Jmp absolute
		memcpy(&newJump[1], &dwRelativeAddr, sizeof(pNewFunc));
		SIZE_T bytesWritten;
		BOOL bRet = WriteProcessMemory(GetCurrentProcess(), pOrgEntry, newJump, sizeof(pNewFunc) + 1, &bytesWritten);
#endif
	}

	void InitMinDump()
	{
		GetName();
		// 注册异常处理函数;
		SetUnhandledExceptionFilter((LPTOP_LEVEL_EXCEPTION_FILTER)MyUnhandledExceptionFilter);

		// 使VS2005 自动设置SetUnhandledExceptionFilter无效;
		DisableSetUnhandledExceptionFilter();
	}
}