memory leak buster - primitive: blog

VisualC++ のデバッグ版 CRT (C 標準ライブラリ) にはメモリリーク検出機能が備わっています。(参考)
しかし、これが出してくれる情報は、どのアドレスで何 byte リークしてますよ、というだけでデバッグの助けにはあまりなりません。
これがものすごく不満だったので、もっと親切なリークチェッカを作ってみました。

https://github.com/i-saint/scribble/blob/master/MemoryLeakBuster2.cpp

プロジェクトに含めるだけで有効になり、プログラム終了時にリーク箇所とそれを確保した時の callstack をデバッグ出力に表示します。

実装には、HeapAlloc に hook を仕掛けるテクニックを用いました。
CRT の malloc() 一族は、WinAPI の HeapAlloc() で実装されています。operator new は malloc() で実装されており、CRT のメモリ確保関数は全て HeapAlloc() に行き着きます。なので、HeapAlloc() を乗っ取れば CRT の全てのメモリ確保を盗み見ることができそうです。
hook を仕掛ける方法は色々ありますが、今回は import 関数テーブルを塗り替える方法を用いました。
外部 dll の関数は実行時までアドレスが決まらないため、直接呼ぶことができません。このため、exe や dll には import 関数の名前、関数へのポインタを保持する領域があります。その情報を元に、exe 起動時や dll ロード時に Windows が関数ポインタを適切に設定します。実行コードはその関数ポインタを経由することで、アドレスの不確定性を意識することなく dll の関数を呼ぶことができるようになっています。
(dll の関数を呼ぶ箇所をデバッガで見ると、__imp__malloc のような、頭に __imp_ が付くポインタ変数を経由して呼んでいることがわかります。これが import 関数ポインタです)
なので、その import 関数ポインタを書き換えてやれば、dll の関数呼び出しをいくらでも好きなようにリダイレクトさせることができます。HeapAlloc と HeapFree をカスタム関数にリダイレクトさせて情報を掠め取ってリークチェッカを実現したのが今回のプログラムです。

今回の肝である import 関数テーブルを巡回したり書き換えたりするのは、以下のようなコードで実現できます。(この記事にとても助けられました)

#include <windows.h>

// write protect がかかったメモリ領域を強引に書き換える
template<class T>
inline void ForceWrite(T &dst, const T &src)
{
    DWORD old_flag;
    VirtualProtect(&dst, sizeof(T), PAGE_EXECUTE_READWRITE, &old_flag);
    dst = src;
    VirtualProtect(&dst, sizeof(T), old_flag, &old_flag);
}

void HookHeapAlloc(HMODULE module)
{
    if(module==0) { return; }

    size_t ImageBase = (size_t)module;
    PIMAGE_DOS_HEADER pDosHeader = (PIMAGE_DOS_HEADER)ImageBase;
    PIMAGE_NT_HEADERS pNTHeader = (PIMAGE_NT_HEADERS)(ImageBase + pDosHeader->e_lfanew);

    size_t RVAImports = pNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
    IMAGE_IMPORT_DESCRIPTOR *pImportDesc = (IMAGE_IMPORT_DESCRIPTOR*)(ImageBase + RVAImports);
    while(pImportDesc->Name != 0) {
        // HeapAlloc は kernel32.dll の中にあるので、kernel32.dll から import している関数群を探す
        if(stricmp((const char*)(ImageBase+pImportDesc->Name), "kernel32.dll")==0) {
            const char *dllname = (const char*)(ImageBase+pImportDesc->Name);
            IMAGE_IMPORT_BY_NAME **func_names = (IMAGE_IMPORT_BY_NAME**)(ImageBase+pImportDesc->Characteristics);
            void **import_table = (void**)(ImageBase+pImportDesc->FirstThunk);
            // HeapAlloc / HeapFree を探してハイジャック
            for(size_t i=0; ; ++i) {
                if((size_t)func_names[i] == 0) { break;}
                const char *funcname = (const char*)(ImageBase+(size_t)func_names[i]->Name);
                if(strcmp(funcname, "HeapAlloc")==0) {
                    //ForceWrite<void*>(import_table[i], HeapAlloc_Hooked);
                }
                else if(strcmp(funcname, "HeapFree")==0) {
                    //ForceWrite<void*>(import_table[i], HeapFree_Hooked);
                }
            }
        }
        ++pImportDesc;
    }
}

初期化とリーク情報出力は、global オブジェクトのコンストラクタ/デストラクタで実装しています。これにより .cpp をプロジェクトに含めるだけで有効になるのを実現していますが、タイミングの問題が残っています。
dll アンロード時にそれまで確保したメモリをまとめて開放するお行儀の悪いライブラリがあります (CRT がそう！)。 dll のアンロードが行われるのは global オブジェクトの破棄の後なので、誤認リークが発生してしまうわけです。
これを解決したい場合、リークチェッカ自身を dll として実装し、起動後できるだけ早く load、終了時できるだけ遅く unload する必要がありそうですが、dll のロード順をコントロールする手段がよくわからないし、dll 化すると使用が面倒になるしで今回はここまでにしました。

おまけ。HeapAlloc() 乗っとり方式を思いつく前に作った operator new オーバーロードによる実装。こっちの方が色々非力ですがよりポータブル…かもしれません。
https://github.com/i-saint/scribble/blob/master/MemoryLeakBuster.cpp

[間連: memory leak buster (2)]