基于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     = ;
     P_AC_NODE       pRoot       = NULL;
     unsigned int    uiPattIdx   = ;
     unsigned int    uiUCharIdx  = ;
     uint16_t        u16Idx      = ;

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

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

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

     for (uiPattIdx = ; uiPattIdx < In_vctPattern.size(); uiPattIdx ++)
     {
         P_AC_NODE   pCurNode    = pRoot;
         for (uiUCharIdx = ; 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(, sizeof(AC_NODE));
                 if (pNode == NULL)
                 {
                     iRetVal = -;
                     goto fun_ret;
                 }

                 pNode->u64Depth = uiUCharIdx + ;
                 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() - )
                 pCurNode->bIsMathed = true;
         }
     }

 fun_ret:
     return iRetVal;
 }

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

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

     quNodes.push(Out_vctACNodes[]);
     while (!quNodes.empty())
     {
         std::map<unsigned char, struct ACNode *>::iterator  itGoto;
         P_AC_NODE   pNode = quNodes.front();
         quNodes.pop();
         if (pNode->u64Depth <= )
             pNode->pFail = Out_vctACNodes[];
         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[])
                 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 == )
         goto fun_ret;

     printf("<<<<<<<<<<FUCKOFF:%x\n", 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     = ;
     P_AC_NODE   pCurrent    = NULL;
     uint64_t    u64Idx      = ;

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

     pCurrent = In_pRoot;
     for (u64Idx = ; 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   = ;
     for (uiIdx = ; 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   = ;
     time_t          tACBegin    = {};
     double          dMemSec     = 0.0;

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

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