Wednesday, March 30, 2011

LeakMon :- Part 2 Under the hood

This post explains how LeakMon application tracks various kind of resource allocation and report them.

Download Source :- LeakMon_source.zip

Basic working principle

Consider the case of tracking memory allocation. In this case, Leakmon will hook all the possible memory allocation and de-allocation functions. So when ever a function allocates some memory it will keep the pointer to the newly allocated memory as a key in a map and the call stack as value. If this memory is de-allocated, this entry in the map is removed. So at any point all the entries in the map are those memory pointers which haven't de-allocted. When you ask the Leakmon to Dump the leak, all it does is dumping the entries in the map.

This application has the following distinguishing features from the other resource leak tracking applications in codeproject and around.

1. You dont have to modify any single line of code to start tracking
2. It will track the allocation from all the binaries in the process. Means the tracking is not restricted to only the EXE binary of the process. It will track allocation from all the dependent DLL's as well

OK, Let's start with the working of LeakMon from the first binary, Injector.exe

Injector application



The Injector.exe allows you to select the process of which you want to track leak. It will list all the process in the system with the help of CreateToolhelp32Snapshot, Process32First and Process32Next functions.

So, when if you select a process and press the "Inject" button, what it does is, It will inject the "HookDll.dll" to the traget process' address space. This is done with the help of CreateRemoteThread API. Injecting a DLL using CreateRemoteThread API is explained detailed in the article Three Ways to Inject Your Code into Another Process. Anyway in Injector.exe the following code part does the injection.
void CInjectorDlg::OnInject()
{
......
CString csPid = m_List.GetItemText( nSelected, 0 );
DWORD dwPID = _ttoi( csPid );
HANDLE hProcess = OpenProcess( PROCESS_CREATE_THREAD|PROCESS_QUERY_INFORMATION|
PROCESS_VM_OPERATION|PROCESS_VM_WRITE|PROCESS_VM_READ,
FALSE, dwPID );
if( !hProcess)
{
      AfxMessageBox( _T("Failed to open the process" ));
      return;
}
HINSTANCE hLib = LoadLibrary( "Kernel32.dll" );
PROC pLoadLib = (PROC)GetProcAddress( hLib, "LoadLibraryA" );
void* pLibRemote = ::VirtualAllocEx( hProcess, NULL, csPath.GetLength(),
MEM_COMMIT, PAGE_READWRITE );
::WriteProcessMemory( hProcess, pLibRemote, (void*)csPath.operator LPCTSTR(),
csPath.GetLength(), NULL );
if( !CreateRemoteThread( hProcess, 0, 0, (LPTHREAD_START_ROUTINE)pLoadLib, pLibRemote, 0, 0 ))
{
          AfxMessageBox( "Create Remote thread Failed" );
          :VirtualFreeEx( hProcess, pLibRemote,csPath.GetLength(), MEM_RELEASE );
}
}
The HookDll

HookDll.dll is the core part of the LeakMon. It performs almost all of the resource tracking stuffs. HookDll starts its working once it got injected into the target process. The BOOL CHookDllApp::InitInstance function gets called when this dll is loaded in the target process. This function however doesn't do much other than just creating a worker thread, DumpController.
BOOL CHookDllApp::InitInstance()
{
       HANDLE hThread = ::CreateThread( 0,0,DumpController, 0,0, 0 );
      CloseHandle( hThread );
      return CWinApp::InitInstance();
}
Now when the DumpController starts, it will first show the configuration dialog, where you can,
1. Select what kind of resource allocation you want to track
2. Select the path of the PDB's of the application
3. Stack depth


DebugHelp Functions
HookDll creates the call the stack, with the help of DebugHelp Functions. For the DebugHelp function to return the call stack correctly, it is necessary that you specify the correct PDB path of the application. Also it always good to use the latest version of the dbghelp.dll that comes with latest version of Debugging tool for windows. Another good feature that LeakMon support is integration with symbol server. However, adding the Symbol server functionality is very easy. All you have to do is to ensure that, the symsrv.dll is present in the same directory from where dbghelp.dll is loaded.
So now once you click the OK button, in the config dialog, the HookDll will start initializing symbol handler. To initialize symbol handler we have to call SymInitialize. When you call SymInitialize with the value of fInvadeProcess as TRUE, it will load symbol table for all the dlls in the process. But wait, we do better, we can show the progress of loading the symbols of each dll. Showing the progress of symbol loading is important especially when you are downloading a symbol from symbol server, which may take significant time.You might have seen this feature in Visual studio. when you start debugging an application in VS, it will show the symbol loading progress of each dll in the status bar.
To show the progress, we can register one call back with the function SymRegisterCallback64. Our callback function will get called before loading each symbol.
void ConfigDlg::OnOk()
{
..................
       SymInitialize(GetCurrentProcess(), (LPTSTR)csWholePath.operator LPCTSTR() , FALSE );
       SymRegisterCallback64( GetCurrentProcess(),SymRegisterCallbackProc64,(ULONG64 )this );
..................
       if( pSymRefreshModuleList )
       {
              pSymRefreshModuleList( GetCurrentProcess());
       }
       ..................
}


BOOL CALLBACK ConfigDlg::SymRegisterCallbackProc64(HANDLE hProcess,
                                                                                                ULONG ActionCode,
                                                                                                ULONG64 CallbackData,
                                                                                                ULONG64 UserContext)
{
       if( CBA_DEFERRED_SYMBOL_LOAD_START == ActionCode )
       {
              PIMAGEHLP_DEFERRED_SYMBOL_LOAD64 pSybolLoadInfo =(PIMAGEHLP_DEFERRED_SYMBOL_LOAD64)CallbackData;
              ConfigDlg* pDlg = (ConfigDlg*)UserContext;
              CString csLoadtext = _T("Loading symbol for file: ");
              csLoadtext += pSybolLoadInfo->FileName;
              pDlg->m_ProgressDlg.SetDlgItemText( IDC_LOAD_INFO, csLoadtext );
       }
       // return false to indicate that we are not doing any symbol loading by ourself
       return FALSE;
}
API Hooking

To intercept the allocation and de-allocation calls(functions) made by dll's, LeakMon relies on User mode API hooking though IAT patching. I am not going to explain the IAT patching because there are already several articles in codeproject and other sites having a detailed explanation. Here is one good article to start with API hooking revealed.
How ever one thing that is different in LeakMon with usual custom is that, it doesn't walk through the dependency chain of dll. I mean, if A.exe has statically linked to B.dll and again if B.dll is linked to C.dll, what usually does for API hooking is, it will hook the API imported by A.exe, then find out the dependency of A.exe, which is B.dll, so hook its imported functions. Then it find out the depended dll of B.dll which is C.dll, hook it functions and so on. The problem with this approach is that only statically linked dlls will be hooked. If you are making some calls to a COM dll or if I you are creating an Active-X control in your application, those dll will never get hooked. So instead of walking through the dependency chain, Leakmon, enumerates all the dll's in that process and hooks it. In fact this approach is a little more easier too. It make use of the Module32First and Module32Next set of functions for this purpose.

As you can see, this application can track three kind of resource allocation and de-allocation namely Memory, GDI objects and HANDLES. Windows have a wider set of function for allocation and de-allocation for each of of resources. The following table shows the list of API's LeakMon hooks for tracking allocations/de-allocation of resources.

Memory allocation and de-allocation functions
HeapAllocHeapFreeHeapReAllocVirtualAllocEx
VirtualFreeExGlobalAllocGlobalReAllocGlobalFree
LocalAllocLocalReAllocLocalFree

GDI object creation and deletion functions
Bitmaps
LoadBitmapALoadBitmapWLoadImageA
LoadImageWCreateBitmapCreateBitmapIndirect
CreateCompatibleBitmapCreateDIBitmapCreateDIBSection
CreateDiscardableBitmapCopyImageGetIconInfo
GetIconInfoExAGetIconInfoExW
Icons
CopyIconCreateIconCreateIconFromResource
CreateIconFromResourceExCreateIconIndirectDestroyIcon
DuplicateIconExtractAssociatedIconAExtractAssociatedIconW
ExtractAssociatedIconExAExtractAssociatedIconExWExtractIconA
ExtractIconWExtractIconExAExtractIconExW
LoadIconALoadIconWPrivateExtractIconsA
PrivateExtractIconsW
Cursor
CreateCursorDestroyCursorLoadCursorA
LoadCursorWLoadCursorFromFileALoadCursorFromFileW
Brush
CreateBrushIndirectCreateSolidBrushCreatePatternBrush
CreateDIBPatternBrushCreateDIBPatternBrushPtCreateHatchBrush
Device context
CreateCompatibleDCCreateDCACreateDCW
CreateICACreateICWGetDC
GetDCExGetWindowDCReleaseDC
DeleteDC
Font
CreateFontACreateFontWCreateFontIndirectA
CreateFontIndirectW
Metafile
CreateMetaFileACreateMetaFileWCreateEnhMetaFileA
CreateEnhMetaFileWGetEnhMetaFileAGetEnhMetaFileW
GetMetaFileAGetMetaFileWDeleteMetaFile
DeleteEnhMetaFileCopyEnhMetaFileACopyEnhMetaFileW
CloseEnhMetaFileCloseMetaFile
Pen
CreatePenCreatePenIndirectExtCreatePen
Region
PathToRegionCreateEllipticRgnCreateEllipticRgnIndirect
CreatePolygonRgnCreatePolyPolygonRgnCreateRectRgn
CreateRectRgnIndirectCreateRoundRectRgnExtCreateRegion
Palette
CreateHalftonePaletteCreatePalette
Common Function
DeleteObject

Handle creation and deletion functions
Synchronization objects
CreateEventACreateEventWCreateEventExA
CreateEventExWOpenEventAOpenEventW
CreateMutexACreateMutexWCreateMutexExA
CreateMutexExWOpenMutexAOpenMutexW
CreateSemaphoreACreateSemaphoreWCreateSemaphoreExA
CreateSemaphoreExWOpenSemaphoreAOpenSemaphoreW
CreateWaitableTimerACreateWaitableTimerWCreateWaitableTimerExA
CreateWaitableTimerExWOpenWaitableTimerAOpenWaitableTimerW
File function
CreateFileACreateFileWCreateFileTransactedA
CreateFileTransactedWFindFirstFileAFindFirstFileW
FindFirstFileExAFindFirstFileExWFindFirstFileNameTransactedW
FindFirstFileNameWFindFirstFileTransactedAFindFirstFileTransactedW
FindFirstStreamTransactedWFindFirstStreamWFindClose
OpenFileByIdReOpenFileCreateIoCompletionPort
Authorization function
CreateRestrictedTokenDuplicateTokenDuplicateTokenEx
OpenProcessTokenOpenThreadToken
Directory management
FindFirstChangeNotificationAFindFirstChangeNotificationWFindCloseChangeNotification
File mapping
CreateMemoryResourceNotificationCreateFileMappingACreateFileMappingW
CreateFileMappingNumaACreateFileMappingNumaWOpenFileMappingA
OpenFileMappingW
Memory
HeapCreateHeapDestroyGlobalAlloc
GlobalReAllocGlobalFreeLocalAlloc
LocalReAllocLocalFree
Process and thread
CreateProcessACreateProcessWCreateProcessAsUserA
CreateProcessAsUserWCreateProcessWithLogonWCreateProcessWithTokenW
OpenProcessCreateThreadCreateRemoteThread
OpenThreadCreateJobObjectACreateJobObjectW
Mail slot
CreateMailslotACreateMailslotW
pipe
CreatePipeCreateNamedPipeACreateNamedPipeW
Registry
RegCreateKeyExARegCreateKeyExWRegCreateKeyTransactedA
RegCreateKeyTransactedWRegOpenCurrentUserRegOpenKeyA
RegOpenKeyWRegOpenKeyExARegOpenKeyExW
RegOpenKeyTransactedARegOpenKeyTransactedWRegOpenUserClassesRoot
RegCreateKeyARegCreateKeyWRegCloseKey
Common functions
DuplicateHandleCloseHandle

Note:- I have tried my best to include all the functions I know which allocates or de-allocate resources. Still, if know some function that i have missed, please let me know.

Now lets see how the hooking works. For each of the function in the above table LeakMon have a dummy function ( Which is actually 199 dummy functions ). So after hooking, the application will be calling my dummy function instead of the original function. In the dummy function what it does is,

1) Call the original function
2) Create the call stack and store it
3) Return the return value of retured by original function

Let's check the example of HeapAlloc function ( Remember that all the memory allocation functions such as new, malloc etc will finally ends up in calling HeapAlloc function to allocate memory. Similarly free, delete functions calls HeapFree ).
LPVOID WINAPI MyHeapAlloc( HANDLE hHeap,
                                                         DWORD dwFlags,
                                                         SIZE_T dwBytes )
{
           LPVOID lMem = pOrgHeapAlloc( hHeap, dwFlags, dwBytes );
           CreateCallStack( lMem, dwBytes );
           return lMem;
}
The CreateCallStack function creates the current function call stack and save it in a map. The key of the map is the memory address returned by original HeapAlloc and value is a structure which holds the call stack and size of memory allocated.

Now when the application calls HeapFree, the dummy function just removes the entry from the map and then call the original HeapFree function.
BOOL WINAPI MyHeapFree( HANDLE hHeap, DWORD dwFlags, LPVOID lpMem )
{
       RemovCallStack( lpMem );
       return pOrgHeapFree( hHeap, dwFlags, lpMem );
}
Creating the Call Stack

This is the only part which actually has different processing for 64 bit and 32 bit. In 32 bit applications the call stack is created using the StackWalk64 function. You can either check the code ( StackDump function ) or read this nice article Walking the callstack, to learn more about the StackWalk64 function.

The stack walking in x64 machine was bit difficult for me. I actually had no idea how to do it. Thanks to Ken Johnson for posting an example of stack walking in x64 machine which practically made me to port this application to 64 bit also.

Dumping the leaks

After the hook dll have setup all the hooks, it will wait for couple of events. One among those event is Dump event. When user clicks on the Dump button in the Injector application, it will set the dump event to signalled state. When this event is signalled the thread in the HookDll will start dumping each entry in the map to a user specified file. The format of the dumpfile is mentioned in the previous post.

That's all the main thing about the Injector and HookDll. There is one more application DumpViewer, which actually has some nice features done using Visual Studio automation. I will explain those things in another post.

Open Source

I got couple of request to make this project open source, so I have opened a project in code.google.com and is available at http://code.google.com/p/leakmon/ . So if you have some idea or got some improvement points you are welcome to join me there. Suggestions are also welcome.