一种C语言"打桩"的源码实现

实际工作中，在我们写好源代码后，通常需要对代码进行UT、FT测试，这个时候我们经常需要“打桩”，考虑以下情形：

1、本模块A的正常业务过程需要调用模块B的函数b1，但函数b1有可能还未实现（或者系统还未集成模块A无法调用b1），这个时侯为了顺利的进行UT，我们就可以对函数b1进行打桩。

2、模块A正常业务过程会向模块C发送消息，而我们想查看消息的内容是否正确，这个时侯就可以对发送消息的函数打桩，改变其行为，打桩后测试过程中模块A不会向C发送消息，而会将消息码流打印到屏幕（或写到文件，这个要看桩函数的实现）。

3、模拟UT、FT测试过程中无法实现的场景，我们的代码肯定都是针对实际运行环境，如果我们代码中有关于数据库的操作、文件的操作，我们不会真的去操作数据库和写文件，而是使用桩将这个场景屏蔽或替代成其它过程，以满足我们的测试需要。

“废话”说了一大通，下面能一段简单的代码了解一下打桩和桩函数：

编译环境：window10 + VS2015

先来看一下”打桩”的效果：

// stub_test.c : 定义控制台应用程序的入口点。
//
#include "stub.h"
#include 
void add(int i)
{
    printf("add(%d)\n",i);
}

void add_stub(int i)
{
    printf("add_stub(%d)\n",i);
}

int main()
{
    INSTALL_STUB(add,add_stub);
    add(12);
    REMOVE_STUB(add_stub);
    add(11);
    return 0;
}　

编译运行结果：

add_stub(12)
add(11)

本文就是探究一下这个打桩(INSTALL_STUB)和移除桩(REMOVE_STUB)过程的源码实现，源码中的一些关键点都做了详细注释。

stub.h源码如下：

#pragma once
#ifdef __cplusplus
extern "C" {
#endif
    int uninstall_stub(void* stub_f);
    int install_stub(void *orig_f, void *stub_f, char *desc);

#define INSTALL_STUB(o,s) install_stub((void*)o,(void*)s,(char*)#o"->"#s)
#define REMOVE_STUB(s) uninstall_stub((void*)s)

#ifdef __cplusplus
}
#endif

stub_list.h源码如下：

#pragma once
#define __inline__

typedef struct list_head
{
    struct list_head *next, *prev;

}list_head_t;

#define LIST_HEAD_INIT(name) { &(name),&(name) }
#define LIST_HEAD(name) \
struct list_head name = LIST_HEAD_INIT(name)

#define INIT_LIST_HEAD(ptr) do { \
(ptr)->next = (ptr);(ptr)->prev = (ptr);\
} while (0)

static __inline__ void __list_add(struct list_head *new,
    struct list_head *prev,
    struct list_head *next)
{
    next->prev = new;
    new->next = next;
    new->prev = prev;
    prev->next = new;
}

static __inline__ void list_add(struct list_head *new, struct list_head *head)
{
    __list_add(new, head, head->next);
}

static __inline__ void __list_del(struct list_head *prev, struct list_head *next)
{
    next->prev = prev;
    prev->next = next;
}

static __inline__ void list_del(struct list_head *entry)
{
    __list_del(entry->prev, entry->next);
}

#define list_entry(ptr,type,member) \
   ((type*)((char*)(ptr)-(unsigned long)(&((type*)0)->member)))
#define list_for_each(pos,head) \
for(pos=(head)->next;pos!=(head);pos=pos->next)

static __inline__ int list_count(struct list_head *head)
{
    struct list_head *pos;
    int count = 0;
    list_for_each(pos, head)
    {
        count++;
    }
    return count;
}

stub.c源码如下:

#define _CRT_SECURE_NO_WARNINGS
#include 

#include "stub_list.h"
LIST_HEAD(head);
struct stub
{
    struct list_head node;
    char             desc[256];
    void             *orig_f;
    void             *stub_f;
    unsigned int     stubpath;
    unsigned int     old_flg;
    unsigned char    assm[5];
};

HANDLE mutex;
int init_lock_flag = 0;
void initlock()
{
    mutex = CreateSemaphore(NULL, 1, 1, NULL);
}
void lock()
{
    (init_lock_flag == 0) && (initlock(), init_lock_flag = 1);
    while (WaitForSingleObject(mutex, INFINITE) != WAIT_OBJECT_0)
        ;
}

void unlock()
{
    ReleaseSemaphore(mutex, 1, NULL);
}

#define _PAGESIZE 4096

static int set_mprotect(struct stub* pstub)
{
    void *p;
    /****************************************************************************
    *函数VirtualProtectExVirtualProtectEx用来设置内存区域的保护属性,可以将原始函数
    *所在的内存设置为可读、可写、可执行，这样就可以改变原始函数所在内存的代码指令。
    *由于 VirtualProtectExVirtualProtectEx只能设置从内存页大小(4096)的整数倍开始的
    *地址，因此利用(long)pstub->orig_f& ~(_PAGESIZE - 1) 计算出一个比原始函数地址小
    *且为内存页大小(4096)的整数倍的地址，把它当作VirtualProtectExVirtualProtectEx的
    *作用的起始地址，内存大小为两个内存页_PAGESIZE << 1
    *****************************************************************************/

    p = (void*)((long)pstub->orig_f& ~(_PAGESIZE - 1));
    return TRUE - VirtualProtectEx((HANDLE)-1,
        p, _PAGESIZE << 1,
        PAGE_EXECUTE_READWRITE,/* 设置内存属性为可读、可写、可执行 */
        &pstub->old_flg);/* 将该内存的之前的属性保存下来，以便移除桩函数时，恢复原始函数 */
}

static int set_asm_jmp(struct stub* pstub)
{
    unsigned int offset;
    /* 保存从原始函数地址开始的5个字节，因为之后我们会改写这块区域 */
    memcpy(pstub->assm, pstub->orig_f, sizeof(pstub->assm));
    *((char*)pstub->orig_f) = 0xE9;/* 这个是相对跳转指令jmp */
    /**************************************************************
     *计算出桩函数与原始函数之间的相对地址，注意要减去jmp指令的
     *5个字节(0xE9加上一个4字节的相对地址)，然后用这条jmp指令，改写
     *原始函数地址开始的5个字节，这样调用原始函数，就会自动跳到桩函数
     **************************************************************/
    offset = (unsigned int)((long)pstub->stub_f - ((long)pstub->orig_f + 5));
    *((unsigned int*)((char*)pstub->orig_f + 1)) = offset;
    return 0;
}

static void restore_asm(struct stub* pstub)
{
    /* 恢复原始函数地址开始的5个字节 */
    memcpy(pstub->orig_f, pstub->assm, sizeof(pstub->assm));
}


int install_stub(void *orig_f, void *stub_f, char *desc)
{
    struct stub *pstub;
    pstub = (struct stub*)malloc(sizeof(struct stub));
    pstub->orig_f = orig_f;
    pstub->stub_f = stub_f;
    do
    {
        /* 设置该内存段属性 */
        if (set_mprotect(pstub))
        {
            break;
        }
        /* 用jmp指令去覆盖orig_f开始的5个字节 */
        if (set_asm_jmp(pstub))
        {
            break;
        }

        if (desc)
        {
            strncpy(pstub->desc, desc, sizeof(pstub->desc));
        }
        lock(); /* 如果有多个线程同时操作链表，要使用锁进行同步 */
        /* 如果对多个函数打桩，就需要保存多个函数的相关信息，这里使用链表储存 */
        list_add(&pstub->node, &head);
        unlock(); /* 操作完成，释放锁 */
        return 0;
    } while (0);

    free(pstub);
    return -1;
}

int uninstall_stub(void* stub_f)
{
    struct stub *pstub;
    struct list_head *pos;
    /* 移除桩函数就是将原始函数地址开始的5个字节恢复，然后将该
       函数的信息从链表中移除同时释放之前动态申请的内存 */
    list_for_each(pos, &head)
    {
        pstub = list_entry(pos, struct stub, node);
        if (pstub->stub_f == stub_f)
        {
            restore_asm(pstub);
            lock();
            list_del(&pstub->node);
            unlock();
            free(pstub);
            return 0;
        }
    }
    return -1;
}