另一种基于AVX2/SSE2的高效模式匹配算法在内存搜索中的应用

/****************************************************************************************
 * @brief 特征码初始化，将传入的特征码字符串myPattern格式转化为目标特征码字节序列
 * @param myPattern 特征码文本字符串
 * @param vecPtn 存储转换后的特征码字节序列的向量
 * @param vecMsk 存储特征字符串中'??'、'?'的掩码('??'和'?'用1替代，非'?'为0)
 * @param vecIdx 存储不是'??'的特征码字节在原始特征码字节序列(传入的有'??'的特征码)中的索引下标
 * @return 转换成功返回TRUE，失败返回FALSE
 */
inline BOOL InitPattern(const std::string& myPattern, std::vector<UCHAR>& vecPtn, std::vector<UCHAR>& vecMsk, std::vector<ULONG>& vecIdx)
{
    std::string patternText = myPattern;
    if (patternText.empty()) { return FALSE; }
    // 去除所有空格
    patternText.erase(std::remove(patternText.begin(), patternText.end(), ' '), patternText.end());  

    for (char ch : patternText) {
        if (ch != '?' && !((ch >= '0' && ch <= '9') || (ch >= 'A' && ch <= 'F') || (ch >= 'a' && ch <= 'f'))) { return FALSE; }
    }

    if (patternText.length() % 2 != 0) { return FALSE; } // 特征码长度不能为单数
    ULONG len = patternText.length() / 2; // 原始特征码长度(字节)

    for (ULONG i = 0; i < len; i++)
    {
        std::string tmpS = patternText.substr(i * 2, 2); // 从patternText的第 j*2 个字符处开始，取2个字符。
        if ("??" != tmpS) // 不是"??"的特征字符
        {
            if ('?' == tmpS.at(0)) // 左半字节为'?'
            {
                tmpS.at(0) = 'F';
                vecMsk.push_back(UCHAR(0xFF) << 4);
            }
            else if ('?' == tmpS.at(1)) // 右半字节为'?'
            {
                tmpS.at(1) = 'F';
                vecMsk.push_back(UCHAR(0xFF) >> 4);
            }
            else
            {
                vecMsk.push_back(UCHAR(0x00)); // 需要做比较的特征字节掩码全为0
            }
            vecIdx.push_back(i);
        }
        if ("??" == tmpS) // 为"??"的特征字符
        {
            tmpS.at(0) = 'F';
            tmpS.at(1) = 'F';
            vecMsk.push_back(UCHAR(0xFF));
        }
        vecPtn.push_back(strtoul(tmpS.c_str(), nullptr, 16));
    }
    if (0 == vecIdx.size()) { return FALSE; }
    return TRUE;
}

/****************************************************************************************
 * @brief ▲▲▲BFPatternFind-暴力搜索给定内存范围内的特征码(用于搜索内存区域最末尾的少部分字节序列)-By.Haogl-20240906
 * @param startAddr 需要搜索的起始内存地址
 * @param searchSize 要搜索的内存块的大小,可上负下正（以字节为单位）
 * @param vecPtn 已按特定格式转换后的特征码字节序列
 * @param vecMsk 特征字符串中'??'、'?'的掩码('??'和'?'用1替代，非'?'为0)
 * @param vecIdx 不是'??'的特征码字节在原始特征码字节序列(传入的有'??'的特征码)中的索引下标
 * @return 返回搜索到的内存地址(ULONGLONG类型)
 */
DLL_API inline ULONGLONG BFPatternFind(const ULONGLONG startAddr, const ULONGLONG searchSize,
    const std::vector<UCHAR>& vecPtn, const std::vector<UCHAR>& vecMsk, const std::vector<ULONG>& vecIdx)
{
    if (searchSize < vecPtn.size()) { return 0; }
    PUCHAR maxAddress = (PUCHAR)(startAddr + searchSize);
    PUCHAR currPattern = (PUCHAR)&vecPtn[0]; // 特征码字节序列的首地址
    UCHAR currEqual;
    register UCHAR currPtnCh; 
    PUCHAR currAddress = (PUCHAR)startAddr;

    for (size_t iCh = 0; iCh < vecIdx.size() && (size_t)currAddress <= (size_t)maxAddress; iCh++)
    {
        currPtnCh = currPattern[vecIdx[iCh]];

        currPattern[vecIdx.at(iCh)] = currPtnCh + 0x1;  // 把自己改得不是原来的自己，防止搜索到自己
        // currEqual为0时表示两个字节相同(包含半字节'?'的判断)
        currEqual = ((currAddress[vecIdx[iCh]] | vecMsk.at(vecIdx[iCh])) ^ currPtnCh);
        currPattern[vecIdx.at(iCh)] = currPtnCh; // 过了判断后恢复原来的自己

        if (currEqual) { return 0; } // currEqual不为0时两个字节不相同

        if (iCh + 1 == vecIdx.size()) // 相同数据的个数等于特征码字节数长度，说明找到特征码字节序列
        {
            return (ULONGLONG)currAddress;
        }
    }
    return 0;
}


/****************************************************************************************
 * @brief ▲▲▲AVX2PatternFind256-使用AVX2加速搜索内存特征码(支持模糊匹配)-By.Haogl-20240906
 * @param retList 用于存储搜索到的特征码对应的内存地址（可以存储多个搜索到的内存地址）
 * @param searchStartAddr 需要搜索的起始内存地址
 * @param searchSize 要搜索的内存块的大小,可上负下正（以字节为单位）
 * @param myPattern 搜索特征码，支持通配符??、?，如："?? 5? 77 ?? 88 ?? ?A ??"
 * @param offsetSize 特征码地址离目标地址的偏移距离，上负下正，默认0不偏移
 * @param searchNum 搜索个数，0为搜索整个模块，默认为0
 * @return 成功返回TRUE，失败返回FALSE
*/
DLL_API BOOL AVX2PatternFind256(std::vector<ULONGLONG>& retList, const ULONGLONG searchStartAddr,
    const LONGLONG searchSize, const std::string& myPattern, const LONGLONG offsetSize, const ULONGLONG searchNum)
{
    if (0 == searchStartAddr || 0 == searchSize) { return FALSE; }
    
    ULONGLONG realStartAddr = searchStartAddr; // searchSize为负时(从给定的地址往上搜索)，searchStartAddr需要修正为真正的起始地址
    if ((searchSize < 0) && (searchStartAddr > std::abs(searchSize))) //searchSize可上负下正（以字节为单位）
    {
        realStartAddr = searchStartAddr - std::abs(searchSize);
    }

    std::vector<UCHAR> vecPtn;  // 存储转换后的特征码字节序列
    vecPtn.reserve(16);  // 设置容器的初始预留空间
    std::vector<UCHAR> vecMsk;  // 存储特征字符串中'??'、'?'的掩码('??'和'?'用1替代，非'?'为0)
    vecMsk.reserve(16);
    std::vector<ULONG> vecIdx;  // 存储不是'??'的特征码字节在原始特征码字节序列(传入的有'??'的特征码)中的索引下标
    vecIdx.reserve(8);
    if (!InitPattern(myPattern, vecPtn, vecMsk, vecIdx)) { return FALSE; }

    std::vector<__m256i> m256VecPtn;
    m256VecPtn.reserve(16);
    std::vector<__m256i> m256VecMsk;
    m256VecMsk.reserve(16);
    for (size_t k = 0; k < vecIdx.size(); k++)
    {
        m256VecPtn.push_back(_mm256_set1_epi8(vecPtn.at(vecIdx[k])));
        m256VecMsk.push_back(_mm256_set1_epi8(vecMsk.at(vecIdx[k])));
    }

    UCHAR bakVecPtnCh = vecPtn.at(vecIdx[0]);
    vecPtn.at(vecIdx[0]) += 1;  // 下面搜索前把自己改得不是自己，防止搜索到自己

    retList.clear();
    retList.reserve(16);
    __m256i curMemByte, curCmp, curByteCorr;
    register size_t curBit = 0;
    PUCHAR currMemAddr;
    size_t maxEndSize = min(std::abs(searchSize) - vecPtn.size(), std::abs(searchSize) - 32);
    for (size_t i = vecIdx[0]; i <= maxEndSize; i += 32)
    {
        PUCHAR baseMemAddr = (PUCHAR)(realStartAddr + i - vecIdx[0]);
        size_t prevCmpBit = 0xFFFFFFFF;  // prevCmpBit记录上一个特征字节的bit位比较结果
        for (size_t j = 0; j < vecIdx.size(); j++)
        {
            curMemByte = _mm256_loadu_si256((__m256i*)(baseMemAddr + vecIdx[j]));
            curByteCorr = _mm256_or_si256(curMemByte, m256VecMsk.at(j));	// 当前内存字节的'?'位置的二进制位修正为1
            curCmp = _mm256_cmpeq_epi8(m256VecPtn.at(j), curByteCorr); // 进行字节比较，并将比较结果存储在curCmp中。
            curBit = _mm256_movemask_epi8(curCmp); // 将curCmp中的每个字节的最高位提取出来，并将这些位组合成一个32位的整数curBit。
            curBit = curBit & prevCmpBit;

            if (0 == curBit) { break; }  // 在给定的内存区域中没有找到特征码
            prevCmpBit = curBit;

            if (j + 1 == vecIdx.size()) // 相同数据的个数等于特征码字节数长度，说明找到特征码字节序列
            {
                ULONG bitIdx = 0, n = 0;
                while (_BitScanForward(&bitIdx, curBit))  // 遍历不为0的二进制位(找到的特征码位置)
                {
                    currMemAddr = baseMemAddr + n + bitIdx;
                    retList.push_back((size_t)(currMemAddr + offsetSize));
                    if (searchNum != 0 && retList.size() >= searchNum) { return TRUE; }
                    ++bitIdx;
                    curBit = curBit >> bitIdx;
                    n += bitIdx;
                }
            }
        }
    }

    vecPtn.at(vecIdx[0]) = bakVecPtnCh;  // 恢复原来的自己
    if (vecPtn.size() < 32)  // 给定内存区域末尾搜剩的少于32字节的区域进行搜索
    {
        ULONGLONG tmpStarAddr = realStartAddr + maxEndSize + 1;
        ULONGLONG tmpSearchSize = std::abs(searchSize) - maxEndSize - 1;

        for (int i = 0; i <= tmpSearchSize - vecPtn.size(); i += vecPtn.size())
        {
            ULONGLONG tailPtnAddr = BFPatternFind(tmpStarAddr + i, tmpSearchSize - i, vecPtn, vecMsk, vecIdx);
            if (tailPtnAddr)
            {
                retList.push_back(tailPtnAddr);
            }
        }
    }
    return TRUE;
}

// 特征码含前、中、后通配符‘??’及半字节‘?’
std::string myPattern = "?? 89 ?9 E8 ?? ?? ?? ?? 83 7B ?? ?? 0F 85 ?? ?? ?? ?? 48 8D 5C 24 ?? 4C 8? 73 ?? 0F 29 ??";