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

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;
}