相信大家一定经常碰到内存泄漏的问题,在诊断过程中常用的一个工具就是LeakDiag,该工具可以让用户制定要查看哪些堆中的内存分配情况。碰巧我在项目中也需要类似的功能,但是由于当时不知道有LeakDiag这个工具,所以自己开发了一个。这里把方法拿出来和大家分享一下,希望能对大家有所帮助。
其实此类工具都是将负责内存分配的API进行Hook,让相应的函数调用先到自己的函数中,在那里完成一些统计信息的记录,然后在把调用转到真正的内存分配函数,再把结果返回给原始调用者。
要理解这个功能就要了解Windows是如何实现对DLL中函数的调用。该实现通常需要一下几个部分的支持:
1. PE文件的支持
当程序中声明一个输出函数,无论是通过.def文件还是__declspec(dllexport),在进行链接的时候,编译器会产生两个文件:
l Import Library文件 (<output>.lib)
该文件保存每个输出函数的stub函数。如果应用程序希望做Loading-time动态链接,那么链接器会把对每个输出函数的调用链接到stub函数。每个stub函数都很简单,只有一条跳转语句,即跳转到函数的真实地址。该地址由当前DLL中导入表中的记录提供。其具体过程见下面的介绍。
l PE文件 (<output>.dll)
对于调用者所在的可执行文件或者DLL文件,编译器会插入一个导入函数表;对于提供函数实现的DLL文件,编译器会插入一个导出函数表,这两个表的结构如图1所示。
2. Image Loader的支持
当应用程序启动的时候,Image Loader会依次执行下列操作:
l 递归加载DLL导入表中指定的DLL
l 为导入表中的每个函数查找其所属DLL的导出表的记录的函数的真实地址,并把该地址拷贝到调用者的导入表中
所以我们的Hooking原理就是把导入表和导出表的相应函数地址替换成自己的函数即可。这样在调用时候,函数会先跳到我们的函数,然后再到真实函数。
下面是我程序中的关键部分代码,供大家参考。另外,虽然这个实现是用于诊断内存泄漏的,但是只要运用得当,可以用于任何可能的地方,大家可以充分发挥自己的想象力。
// Hook输出表中的函数
VOID CAPIHook::HookEATEntryInModule(PROC pfnHook)
{
HMODULE hCallee = GetModuleHandleA(m_pszCalleeModuleName.c_str());
ULONG ulSize;
PIMAGE_EXPORT_DIRECTORY pExportDir = NULL;
__try
{
pExportDir = (PIMAGE_EXPORT_DIRECTORY)ImageDirectoryEntryToData(
hCallee,
TRUE,
IMAGE_DIRECTORY_ENTRY_EXPORT,
&ulSize );
}
__except(EXCEPTION_EXECUTE_HANDLER) {}
if( pExportDir == NULL )
{
ATLASSERT(FALSE && "Invalid callee module accessed");
return;
}
PDWORD pdwNameRvas = (PDWORD)((PBYTE)hCallee + pExportDir->AddressOfNames);
PDWORD pdwNameOrdinals = (PDWORD)((PBYTE)hCallee + pExportDir->AddressOfNameOrdinals);
PDWORD pdwNameFunctions = (PDWORD)((PBYTE)hCallee + pExportDir->AddressOfFunctions);
for( DWORD i = 0; i < pExportDir->NumberOfNames; ++i )
{
PSTR pszFunctionName = (PSTR)((PBYTE)hCallee + pdwNameRvas[i]);
if( stricmp(pszFunctionName,m_pszFunctionName.c_str()) != 0 )
{
continue;
}
PROC *ppfn = (PROC*)&pdwNameFunctions[pdwNameOrdinals[i]];
PROC pfnRva = (PROC)((PBYTE)pfnHook - (PBYTE)hCallee);
if( !WriteProcessMemory(GetCurrentProcess(),
ppfn,
&pfnRva,
sizeof(pfnRva),
NULL ) && GetLastError() == ERROR_NOACCESS )
{
DWORD dwOldProtect;
if( VirtualProtect( ppfn,
sizeof(pfnRva),
PAGE_WRITECOPY,
&dwOldProtect) )
{
WriteProcessMemory(GetCurrentProcess(),
ppfn,
&pfnRva,
sizeof(pfnRva),
NULL );
VirtualProtect(ppfn,
sizeof(pfnRva),
dwOldProtect,
&dwOldProtect);
}
}
break;
}
}
// Hook输入表中的函数
VOID CAPIHook::HookIATEntryInModule(HMODULE hCaller,PROC pfnHook,PROC pfnOrigin)
{
ULONG ulSize;
PIMAGE_IMPORT_DESCRIPTOR pImportDesc = NULL;
__try
{
pImportDesc = (PIMAGE_IMPORT_DESCRIPTOR)ImageDirectoryEntryToData(
hCaller,
TRUE,
IMAGE_DIRECTORY_ENTRY_IMPORT,
&ulSize );
}
__except(EXCEPTION_EXECUTE_HANDLER) {}
if( pImportDesc == NULL )
{
// this might not be an error because some dlls like ntdll.dll,
// which do not have import section will also return NULL
return;
}
for( ; pImportDesc->Name; pImportDesc++ )
{
PSTR pszModuleName = (PSTR)((PBYTE)hCaller + pImportDesc->Name);
if( stricmp(pszModuleName,m_pszCalleeModuleName.c_str()) == 0 )
{
for( PIMAGE_THUNK_DATA pThunk = (PIMAGE_THUNK_DATA)((PBYTE)hCaller + pImportDesc->FirstThunk);
pThunk->u1.Function;
pThunk++ )
{
PROC *ppfn = (PROC*)&pThunk->u1.Function;
if( *ppfn == pfnOrigin )
{
if( !WriteProcessMemory(GetCurrentProcess(),
ppfn,
&pfnHook,
sizeof(pfnHook),
NULL ) && GetLastError() == ERROR_NOACCESS )
{
DWORD dwOldProtect;
if( VirtualProtect(ppfn, sizeof(pfnHook),PAGE_WRITECOPY,&dwOldProtect) )
{
WriteProcessMemory(GetCurrentProcess(),
ppfn,
&pfnHook,
sizeof(pfnHook),
NULL);
VirtualProtect(ppfn,
sizeof(pfnHook),
dwOldProtect,
&dwOldProtect);
}
}
// return if we found what we want
return;
}
}
// since one image might have multiple import entries, so go on searching here
}
}
}