基于AC有限状态机的多模匹配算法

参考链接：http://www.cnblogs.com/zzqcn/p/3525636.html

感谢原文作者。

花了两天半时间实现并测试了算法。

按照上文的思路实现了一遍，可能是原文中有些地方描述的不是特别清楚，导致一开始测试的时候发现了各种匹配遗漏的情况，后经过自己各种努力终于解决了各种遗漏。

同时在实现过程中也遇到了各种小问题，最后都解决了，总结起来主要有四个大坑，自己实现的时候需要注意，四个坑都在代码的注释里面了。

这里的实现虽然不会有遗漏的情况，但会有同一模式串在相同的偏移多次被命中的情况，但无伤大雅，至少没有遗漏不是吗。实际应用中只需对结果做去重就好了。

测试结论：对一个101.3MB的PE，从中随机抽取长度在[16-116)Bytes的模式串16个，分别用memcmp方式和AC自动机方式进行匹配，memcmp方式耗时33秒，AC方式耗时12秒，可见优势还是比较明显的。

代码中如有哪里不对，欢迎一起讨论。

#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <stdint.h>
#include <vector>
#include <map>
#include <queue>
#include <ctime>

typedef struct ACNode
{
    uint64_t        u64Depth;
    struct ACNode   *pFail;
    std::map<unsigned char, struct ACNode *>    *pmpGotoTab;
    struct ACParrent
    {
        struct ACNode   *pParent;
        unsigned char   ucCondition;
    } Parent;
    bool            bIsMathed;
} AC_NODE, *P_AC_NODE;

typedef void (__stdcall *P_AC_FOUND_CALLBACK)(const unsigned char *In_pucBuf, uint64_t In_u64EndPos, uint64_t In_u64Len);

int InitACGoto(const std::vector<const std::vector<unsigned char> *> &In_vctPattern,
    std::vector<P_AC_NODE> &Out_vctACNodes)
{
    int             iRetVal     = 0;
    P_AC_NODE       pRoot       = NULL;
    unsigned int    uiPattIdx   = 0;
    unsigned int    uiUCharIdx  = 0;
    uint16_t        u16Idx      = 0;

    if (In_vctPattern.empty())
    {
        iRetVal = -1;
        goto fun_ret;
    }

    pRoot = (P_AC_NODE)calloc(1, sizeof(AC_NODE));
    if (pRoot == NULL)
    {
        iRetVal = -2;
        goto fun_ret;
    }

    pRoot->pmpGotoTab = new std::map<unsigned char, struct ACNode *>();
    for (u16Idx = 0; u16Idx <= 0xff; u16Idx ++)
        pRoot->pmpGotoTab->insert(std::pair<unsigned char, struct ACNode *>((unsigned char)u16Idx, pRoot));
    Out_vctACNodes.push_back(pRoot);

    for (uiPattIdx = 0; uiPattIdx < In_vctPattern.size(); uiPattIdx ++)
    {
        P_AC_NODE   pCurNode    = pRoot;
        for (uiUCharIdx = 0; uiUCharIdx < In_vctPattern[uiPattIdx]->size(); uiUCharIdx ++)
        {
            unsigned char   ucCurUChar  = In_vctPattern[uiPattIdx]->at(uiUCharIdx);
            if (pCurNode->pmpGotoTab->find(ucCurUChar) == pCurNode->pmpGotoTab->end()
                || (pCurNode->pmpGotoTab->find(ucCurUChar) != pCurNode->pmpGotoTab->end()
                && pCurNode->pmpGotoTab->at(ucCurUChar) == pRoot))
            {
                P_AC_NODE   pNode = (P_AC_NODE)calloc(1, sizeof(AC_NODE));
                if (pNode == NULL)
                {
                    iRetVal = -3;
                    goto fun_ret;
                }

                pNode->u64Depth = uiUCharIdx + 1;
                pNode->Parent.pParent = pCurNode;
                pNode->Parent.ucCondition = ucCurUChar;
                pNode->pmpGotoTab = new std::map<unsigned char, struct ACNode *>();

                if (pCurNode->pmpGotoTab->find(ucCurUChar) != pCurNode->pmpGotoTab->end())
                    pCurNode->pmpGotoTab->erase(ucCurUChar);
                pCurNode->pmpGotoTab->insert(std::pair<unsigned char, struct ACNode *>(ucCurUChar, pNode));
                pCurNode = pNode;
                Out_vctACNodes.push_back(pNode);
            }
            else
                pCurNode = pCurNode->pmpGotoTab->at(ucCurUChar);

            if (uiUCharIdx == In_vctPattern[uiPattIdx]->size() - 1)
                pCurNode->bIsMathed = true;
        }
    }

fun_ret:
    return iRetVal;
}

int ACFail(std::vector<P_AC_NODE> &Out_vctACNodes)
{
    int                     iRetVal = 0;
    std::queue<P_AC_NODE>   quNodes;

    if (Out_vctACNodes.empty())
    {
        iRetVal = -1;
        goto fun_ret;
    }

    quNodes.push(Out_vctACNodes[0]);
    while (!quNodes.empty())
    {
        std::map<unsigned char, struct ACNode *>::iterator  itGoto;
        P_AC_NODE   pNode = quNodes.front();
        quNodes.pop();
        if (pNode->u64Depth <= 1)
            pNode->pFail = Out_vctACNodes[0];
        else
        {
            P_AC_NODE   pParentFail = pNode->Parent.pParent->pFail;
            while (pParentFail->pmpGotoTab->find(pNode->Parent.ucCondition) == pParentFail->pmpGotoTab->end())
                pParentFail = pParentFail->pFail;
            pNode->pFail = pParentFail->pmpGotoTab->at(pNode->Parent.ucCondition);
        }
        for (itGoto = pNode->pmpGotoTab->begin(); itGoto != pNode->pmpGotoTab->end(); itGoto ++)
        {
            if (itGoto->second != Out_vctACNodes[0])
                quNodes.push(itGoto->second);
        }
    }

fun_ret:
    return iRetVal;
}

void __stdcall ACFoundCallBack(const unsigned char *In_pucBuf, uint64_t In_u64EndPos, uint64_t In_u64Len)
{
    if (In_pucBuf == NULL || In_u64Len == 0)
        goto fun_ret;

    printf("<<<<<<<<<<FUCKOFF:%x
", In_u64EndPos - In_u64Len);

fun_ret:
    return;
}

int ACSearch(const P_AC_NODE In_pRoot, const unsigned char *In_pucBuf, uint64_t In_u64BufLen, P_AC_FOUND_CALLBACK In_pfCallBack)
{
    int         iRetVal     = 0;
    P_AC_NODE   pCurrent    = NULL;
    uint64_t    u64Idx      = 0;

    if (In_pRoot == NULL || In_pucBuf == NULL || In_u64BufLen == 0 || In_pfCallBack == NULL)
    {
        iRetVal = -1;
        goto fun_ret;
    }

    pCurrent = In_pRoot;
    for (u64Idx = 0; u64Idx < In_u64BufLen;)
    {
        P_AC_NODE   pFail   = NULL;
        if (pCurrent->pmpGotoTab->find(In_pucBuf[u64Idx]) != pCurrent->pmpGotoTab->end())
        {
            pCurrent = pCurrent->pmpGotoTab->at(In_pucBuf[u64Idx]);
            //坑1，出现匹配失败时不要前进，只在匹配成功时前进
            u64Idx ++;
        }
        else
            pCurrent = pCurrent->pFail;

        //坑3，每个节点都需要沿着失配指针一直向上找所有匹配到的结果，而不是
        //只在匹配成功时才这么做，否则会出现匹配遗漏（形如“abcd”和“bc”这样的特征串并存的情况）
        pFail = pCurrent->pFail;
        //坑4，一定要走到根，否则会出现匹配遗漏
        while (pFail != In_pRoot)
        {
            if (pFail->bIsMathed)
                In_pfCallBack(In_pucBuf, u64Idx, pFail->u64Depth);
            pFail = pFail->pFail;
        }
        //坑2，不管是否匹配成功，都要判断当前节点状态，因为出现失配后的
        //转移也有可能转到一个成功匹配的节点上
        if (pCurrent->bIsMathed)
            In_pfCallBack(In_pucBuf, u64Idx, pCurrent->u64Depth);
    }

fun_ret:
    return iRetVal;
}

void ReleaseACNodes(std::vector<P_AC_NODE> &Out_vctACNodes)
{
    unsigned int    uiIdx   = 0;
    for (uiIdx = 0; uiIdx < Out_vctACNodes.size(); uiIdx ++)
    {
        delete Out_vctACNodes[uiIdx]->pmpGotoTab;
        free(Out_vctACNodes[uiIdx]);
    }
    Out_vctACNodes.clear();
}

void main(int argc, char **argv)
{
    std::vector<P_AC_NODE>  vctNodes;
    std::vector<const std::vector<unsigned char> *> vctPatterns;
    unsigned char   *pucBuf = NULL;
    FILE            *pf     = NULL;
    long            lFileSize   = 0;
    time_t          tACBegin    = {0};
    double          dMemSec     = 0.0;

    pf = fopen(argv[1], "rb");
    fseek(pf, 0, SEEK_END);
    lFileSize = ftell(pf);
    fseek(pf, 0, SEEK_SET);
    pucBuf = (unsigned char *)calloc(lFileSize, 1);
    fread(pucBuf, 1, lFileSize, pf);
    fclose(pf);
    for (int i = 0; i < 1600; i ++)
    {
        std::vector<unsigned char>  *pvctPattern = new std::vector<unsigned char>();
        int iBegin  = rand() % (lFileSize - 128);
        int iLen    = rand() % 100 + 16;
        for (int j = 0; j < iLen; j ++)
            pvctPattern->push_back(pucBuf[j + iBegin]);
        vctPatterns.push_back(pvctPattern);
        printf("%x:%u
", iBegin, iLen);
        for (long j = 0; j < lFileSize - iLen; j ++)
        {
            time_t  tMemBegin   = time(NULL);
            if (memcmp(pucBuf + iBegin, pucBuf + j, iLen) == 0)
                printf(">>>>>>>>>>Off:%x
", j);
            dMemSec += difftime(time(NULL), tMemBegin);
        }
    }

    InitACGoto(vctPatterns, vctNodes);
    ACFail(vctNodes);
    tACBegin = time(NULL);
    ACSearch(vctNodes[0], pucBuf, lFileSize, ACFoundCallBack);
    printf("MemTime::%f
ACTime::%f
", dMemSec, difftime(time(NULL), tACBegin));
    ReleaseACNodes(vctNodes);
    return;
}