动态获取桩函数的执行次数

    UT中经常会用到打桩技术，保证测试用例顺利执行。但在有些情况下需要知道这个桩函数有没有被执行，或者执行了几次，这时就需要通过某种方式来获取函数的执行次数。一般的做法就是定义一个全局变量，在测试之前将变量赋值为0，然后在桩函数中对该变量进行加1，代码执行完后该变量中存放的就是桩函数的执行次数。这种方法简单也易于实现，但如果有大量的这种需求，就要定义大量的变量，测试代码变得非常臃肿，后期维护也不是很方便。因此自己就动手，做了个能够在代码执行时自动对桩函数执行次数计数的工具类，方便以后的工作。方法就是，每次在对函数打桩时先申请一小块内存，修改原函数前5个字节为jmp到这块内存的指令，然后在这块内存中存放对特定变量加1的指令，再跳到桩函数处执行。

 #ifndef STUB_HEADER_H #define STUB_HEADER_H #include <windows.h> #include <assert.h> #define JMP_CODE_LENGTH 5 //拷贝替换的长度，为5个字节：JMP xxx //定义用于计数的跳转指令集 BYTE Trampoline[] = { 0x50, //push eax; 0x53, //push ebx; 0xBB,0x0,0x0,0x0,0x0, //mov ebx,xxx 压入对应的count指针,字节序 0x83,0x3,0x1, //add dword ptr [ebx],1; 0x5B, //pop ebx; 0x58, //pop eax; 0xE9,0x0,0x0,0x0,0x0 //jmp xxx; }; //打桩类 class Stub { public: Stub(void* lpOrigFuncAddr,void* lpStubFuncAddr) { memset(m_abOrigCode,0,JMP_CODE_LENGTH); m_lpOrigFuncAddr = lpOrigFuncAddr; m_lpStubFuncAddr = lpStubFuncAddr; m_bInstallStub = FALSE; m_dwExecCount = 0; InstallStub(); } Stub(LPCTSTR lpModuleName, LPCSTR lpProcName, void* lpStubFuncAddr) { assert(lpModuleName != NULL); assert(lpProcName != NULL); m_bInstallStub = FALSE; HMODULE hModule = LoadLibrary(lpModuleName); PROC hProc = GetProcAddress(hModule, lpProcName); if(NULL == hProc) return; m_lpOrigFuncAddr = (void*)hProc; m_lpStubFuncAddr = lpStubFuncAddr; memset(m_abOrigCode,0,JMP_CODE_LENGTH); m_dwExecCount = 0; InstallStub(); } ~Stub() { UnInstallStub(); } BOOL IsStub() { return m_bInstallStub; } DWORD GetStubExecuteCount() { return m_dwExecCount; } private: void InstallStub() { BYTE abJmpCode[JMP_CODE_LENGTH]; if (m_bInstallStub) return ; memcpy(m_abOrigCode,m_lpOrigFuncAddr,JMP_CODE_LENGTH); DWORD dwProcessId = GetCurrentProcessId(); HANDLE hProcess = OpenProcess(PROCESS_VM_READ|PROCESS_VM_WRITE|PROCESS_VM_OPERATION , FALSE, dwProcessId); if (NULL == hProcess) return ; LPVOID lpvTrampoline = VirtualAlloc(NULL, sizeof(Trampoline), MEM_COMMIT, PAGE_READWRITE); if (lpvTrampoline == NULL ) return ; *(DWORD*)(Trampoline+3) = (DWORD)(&m_dwExecCount); //设置为跳转到桩函数的jmp指令 *(DWORD*)(Trampoline+13) = (DWORD)m_lpStubFuncAddr-(DWORD)lpvTrampoline-sizeof(Trampoline); WriteProcessMemory(hProcess, lpvTrampoline, Trampoline, sizeof(Trampoline), NULL); VirtualProtectEx(hProcess, lpvTrampoline, sizeof(Trampoline), PAGE_EXECUTE, NULL); //重新设置原函数的前五个字节，设置为跳转到lpvTrampoline地址 abJmpCode[0] = 0xe9; *(DWORD*)(abJmpCode+1) = (DWORD)lpvTrampoline - (DWORD)m_lpOrigFuncAddr-JMP_CODE_LENGTH; DWORD dwOldFlag; VirtualProtectEx(hProcess, m_lpOrigFuncAddr, JMP_CODE_LENGTH, PAGE_READWRITE, &dwOldFlag); WriteProcessMemory(hProcess, m_lpOrigFuncAddr, abJmpCode, JMP_CODE_LENGTH, NULL); VirtualProtectEx(hProcess, m_lpOrigFuncAddr, JMP_CODE_LENGTH, dwOldFlag, NULL); m_bInstallStub = TRUE; } void UnInstallStub() { if(FALSE == m_bInstallStub) return ; DWORD dwProcessId = GetCurrentProcessId(); HANDLE hProcess = OpenProcess (PROCESS_VM_READ|PROCESS_VM_WRITE|PROCESS_VM_OPERATION , FALSE, dwProcessId); if (NULL == hProcess) return ; DWORD dwOldFlag; VirtualProtectEx(hProcess, m_lpOrigFuncAddr, JMP_CODE_LENGTH, PAGE_READWRITE, &dwOldFlag); WriteProcessMemory(hProcess, m_lpOrigFuncAddr, m_abOrigCode, JMP_CODE_LENGTH, NULL); VirtualProtectEx(hProcess, m_lpOrigFuncAddr, JMP_CODE_LENGTH, dwOldFlag, NULL); m_bInstallStub = FALSE; } private: DWORD m_dwExecCount; BYTE m_abOrigCode[JMP_CODE_LENGTH]; LPVOID m_lpOrigFuncAddr; LPVOID m_lpStubFuncAddr; BOOL m_bInstallStub; }; #endif

通过如下方式使用

 // cstub.cpp : 定义控制台应用程序的入口点。 // #include "stdafx.h" #include "stub.h" #pragma comment(lib,"ws2_32") int add(int a,int b) { return a+b; } int stub_add(int a,int b) { return 0; } int func(int a,int b) { add(a,b); add(a,b); add(a,b); return 0; } int WINAPI recvStub(__in SOCKET s,__out char* buf,__in int len, __in int flags) { return 0; } void test() { recv(0,0,0,0); } int _tmain(int argc, _TCHAR* argv[]) { Stub stub1(_T("Ws2_32.dll"),_T("recv"),recvStub); test(); printf("recvStub execute count %d/n",stub1.GetStubExecuteCount()); Stub stub2(add,stub_add); func(1,2); printf("stub_add execute count %d/n",stub2.GetStubExecuteCount()); return 0; }