【原创】另一种基于AVX2/SSE2的高效模式匹配算法在内存搜索中的应用-By.Haogl
本帖最后由 xxHxx 于 2024-9-24 16:51 编辑0x00. 前言 最近为对某软件的指定模块打特征码补丁，研究学习了各种搜索算法（Sunday、Shift_And、BNDM等），上一篇《对SSE2模式匹配算法SSE2PatternFind的一点改造优化》中的算法，主要是利用SSE2指令找到特征码中不是通配符的第一个字节，再基于找到的第一个字节用常规指令搜索后面的字节序列，不足之处是除了第一个字节利用了SSE2大位宽的单指令多数据处理优势，而后面的字节搜索用不上SSE2指令集的优势，基于这个问题，这两天思考后又想到了一个更好的特征码匹配方式，尝试了一下，发现效果很好，写下来和大家分享！ 上一篇文章《对SSE2模式匹配算法SSE2PatternFind的一点改造优化》链接： https://www.chinapyg.com/thread-155180-1-1.html
0x01. 新的特征码匹配算法特点
[*]（1）逻辑简单好理解；
[*]（2）充分利用AVX2、SSE2指令集的大位宽单指令多数据的优势，理论上指令集位宽翻倍、算法搜索速度翻倍；
[*]（3）支持前、中、后通配符‘??’，支持半字节‘?’，如："?? 5? 77 ?? 88 ?? ?A ??"
[*]（4）代码易于扩展到不同位宽的指令集。
0x02. 算法原理
（1）特征码字符串初始化处理 根据传入的std::string类型的特征码字符串，检查myPattern是否为空，并使用std::remove和erase去除特征码字符串中的所有空格字符；检查特征码中的每个字符是否是?或十六进制字符，如果有非法的字符（既不是?、也不是十六进制字符）则返回FALSE；检查特征码字符串长度是否为偶数，如果不是则返回FALSE；检查如果vecIdx大小为0，说明没有有效的特征码字节（如：所有的特征码字符都是通配符??），返回FALSE。 将传入的特征码字符串按每两个字符（一个字节）进行分割并依次处理： ① 判特征码字符是否含有?或??来处理传入的特征码字符串，即半字节中的?通过 0xFF << 4 或0xFF >> 4将左、右半字节的二进制位替换为1111，双问号??替换为0xFF（即二进制位全为1），其它保持原特征码字节数据，得到vecPtn特征码字节序列； ② 根据特征码字符是否含有?或??，生成相应的二进制掩码vecMsk，即半字节中的?通过 0xFF << 4 或0xFF >> 4将左、右半字节的二进制位替换为1111，双问号??替换为0xFF（即二进制位全为1），其它全为0； ③ 记录不是??（双问号）的特征码字节在原始特征码字节序列（传入的有??的特征码）中的索引下标（vecIdx）。 特征码字符串初始化代码实现如下：/****************************************************************************************
* @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;
}（2）特征码匹配算法 ① 首先，从要搜索的内存中取得第1组32字节(__m256i)的数据curMemByte，将取到的32字节数据curMemByte中的每一个字节与上面初始化得到的不是??的特征字节对应的二进制掩码vecMsk进行位或运算，以修正特征码字节序列中第1个特征字节vecPtn存在半字节时的情况（如果存在半字节，则对应的半字节?处的二进制位替换为1），修正后得到curByteCorr；将curByteCorr中的每一个字节分别和vecPtn（即m256VecPtn.at(0)中的一个字节）进行比较，即在curByteCorr中查找所有可能的vecPtn，得到查找结果curCmp；从curCmp中提取出找到的vecPtn，得到得到vecPtn.at(0)对应的curBit； ② 再次，从内存取出第2组32字节(__m256i)的数据curMemByte，步骤和前面第1组方式相同，这里需要说明：第2组32字节数据从内存地址0x03处开始取，因为我们要查找的第2个特征字节vecPtn.at(1)在原始特征码字符串中的索引下标是3，与vecPtn.at(0)之间有一个??通配符位置需要留出来；通过和前面第1组数据的完全相同方式得到vecPtn.at(1)对应的curBit； ③ 将两次得到的curBit 进行位与(&)运算，找出vecPtn.at(0)的curBit(即prevCmpBit)与vecPtn.at(1)的curBit相同索引位都为1的位置，这些位置就是内存中找的vecPtn前两个特征码字节。 这里的位与(&)运算结果curBit不为0说明找到相应的特征码字节vecPtn.at(i)，需要继续查找下一个vecPtn.at(i+1)，直到vecPtn的每个元素都遍历完时curBit还不为0，说明找到了整个特征码序列，此时curBit中一个为1的二进制位就是找到的一个特征码序列的位置标记（一个curBit中可能找到多个特征码序列）。 当这里的任何一次位与(&)运算结果curBit为0时，说明没有找到vecPtn的对应元素，本次查找结束，当前的内存查找指针+32Byte，进入内存的下一个32字节搜索。
重复上述步骤，直到整个要搜索的内存区域遍历完成。
④ 上述①~③步遍历完成后，给定内存区域末尾还会剩下少于32字节的内存区域没有被搜索，需要判断要搜索的特征码字节序列长度是否小于内存区域末尾剩余的字节数，如果小于，则需要单独对末尾剩余的这部分内存进行特征码字节序列的搜索。
特征码匹配算法的具体流程如下： 基于AVX2指令集的特征码字节序列遍历匹配的算法代码实现如下：
/****************************************************************************************
* @brief ▲▲▲BFPatternFind-暴力搜索给定内存范围内的特征码(用于搜索内存区域最末尾的少部分字节序列)-By.Haogl-20240906
* @param startAddr 需要搜索的起始内存地址
* @param searchSize 要搜索的内存块的大小,可上负下正（以字节为单位）
* @param vecPtn 已按特定格式转换后的特征码字节序列
* @param vecMsk 特征字符串中'??'、'?'的掩码('??'和'?'用1替代，非'?'为0)
* @param vecIdx 不是'??'的特征码字节在原始特征码字节序列(传入的有'??'的特征码)中的索引下标
* @return 返回搜索到的内存地址(ULONGLONG类型)
*/
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; // 特征码字节序列的首地址
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];
currPattern = currPtnCh + 0x1;// 把自己改得不是原来的自己，防止搜索到自己
// currEqual为0时表示两个字节相同(包含半字节'?'的判断)
currEqual = ((currAddress] | vecMsk.at(vecIdx)) ^ currPtnCh);
currPattern = 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
*/
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为负时(从给定地址的前面开始搜索)，需要修正为真正的起始地址
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)));
m256VecMsk.push_back(_mm256_set1_epi8(vecMsk.at(vecIdx)));
}
UCHAR bakVecPtnCh = vecPtn.at(vecIdx);
vecPtn.at(vecIdx) += 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);
if (std::abs(searchSize) < 32) { goto lessThan32Byte; }// 要搜索的内存区域小于32字节
for (size_t i = vecIdx; i <= maxEndSize; i += 32)
{
PUCHAR baseMemAddr = (PUCHAR)(realStartAddr + i - vecIdx);
size_t prevCmpBit = 0xFFFFFFFF;// prevCmpBit记录上一个特征字节的bit位比较结果
for (size_t j = 0; j < vecIdx.size(); j++)
{
curMemByte = _mm256_loadu_si256((__m256i*)(baseMemAddr + vecIdx));
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) = bakVecPtnCh;// 恢复原来的自己
lessThan32Byte:
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);
if (searchNum != 0 && retList.size() >= searchNum) { return TRUE; }
}
}
}
return TRUE;
}
0x03. 结果测试 对 VS Code 的主程序Code.exe的代码段".text"（大小131M）做如下搜索测试：// 特征码含前、中、后通配符‘??’及半字节‘?’
std::string myPattern = "?? 89 ?9 E8 ?? ?? ?? ?? 83 7B ?? ?? 0F 85 ?? ?? ?? ?? 48 8D 5C 24 ?? 4C 8? 73 ?? 0F 29 ??";
[*]算法基于SSE2指令集（位宽128bit）的搜索，每次搜索时间121~152毫秒，与上一篇《对SSE2模式匹配算法SSE2PatternFind的一点改造优化》中的算法搜索速度基本一致；
[*]算法基于AVX2指令集（位宽256bit）的搜索，每次搜索时间50~80毫秒，指令集位宽翻倍、搜索速度也翻倍。
0x04. 总结 算法原理简单高效，代码易于实现、易于扩展；只搜索特征码中不是通配符的特征字节，优化搜索字节数，搜索速度快；算法主要利用位操作对特征码进行比对，充分利用了AVX2、SSE2指令集的大位宽、单指令多数据的优势；采用掩码的方式实现通配符(含半字节)特征码的搜索支持。 另：如果采用AVX512指令集（位宽512bit）一次可比对的字节数达到64Byte，且AVX2中的__mm256_cmpeq_epi8和_mm256_movemask_epi8指令在AVX512中可简化一条指令_mm512_cmpeq_epi8_mask，理论上速度还会显著提升，机器支持AVX512指令集的可以修改代码后自行测试。
20240908修改：---①修改文章说法和代码逻辑有差别的问题；---②上传源码文件，源码中提供AVX2PatternFind256、SSE2PatternFind128、判断当前机器支持指令集情况自动选择对应搜索函数；---③代码增加进行内存搜索时不要搜到自己特征码字节序列的处理。
20240924修改：---修复搜索的内存尺寸小于32Byte时奔溃的问题。
源码文件：
图片看不清的可以下载下面的pdf：**** Hidden Message *****
感谢分享，学习了！ 太牛了。。。。
感谢分享，学习了！ 学习了，多谢分享 学习了，多谢分享 学习了，多谢分享!! 进来学习了 进来学习了 刚从吾爱下载完代码，就发现帖子被删了