admin/awesome_anti_virus_engine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve Rich Header detection with enhanced safety checks

Browse Source 
- Add robust boundary and size validation for Rich header parsing
- Prevent potential buffer overread with controlled search area
- Implement additional checks to handle maliciously crafted PE files
- Limit Rich header search to a maximum reasonable size
- Ensure safe memory access during signature detection
This commit is contained in:
Huoji's
2025-03-09 03:31:11 +08:00
parent 4d1ccb16aa
commit f00b0b2037
1 changed files with 25 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
ai_anti_malware/ml.cpp
View File

@@ -346,18 +346,34 @@ std::vector<double> MachineLearning::ExtractFeatures(const uint8_t* buffer,
dataDir = peconv::get_directory_entry(peBuffer, IMAGE_DIRECTORY_ENTRY_IAT); dataDir = peconv::get_directory_entry(peBuffer, IMAGE_DIRECTORY_ENTRY_IAT);
peInfo.hasEntryIat = dataDir && dataDir->VirtualAddress != 0; peInfo.hasEntryIat = dataDir && dataDir->VirtualAddress != 0;
// Rich头部检测 // Rich头部检测 - 安全实现
peInfo.hasRich = false; peInfo.hasRich = false;
PIMAGE_DOS_HEADER dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(peBuffer); PIMAGE_DOS_HEADER dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(peBuffer);
if (dosHeader && dosHeader->e_magic == IMAGE_DOS_SIGNATURE) { if (dosHeader && dosHeader->e_magic == IMAGE_DOS_SIGNATURE) {
const uint32_t* richPtr = reinterpret_cast<const uint32_t*>( // 确保e_lfanew值合理
peBuffer + sizeof(IMAGE_DOS_HEADER)); if (dosHeader->e_lfanew > sizeof(IMAGE_DOS_HEADER) &&
size_t maxLen = dosHeader->e_lfanew - sizeof(IMAGE_DOS_HEADER); dosHeader->e_lfanew < v_size) { // 确保在PE文件大小范围内
if (maxLen > 0 && richPtr != nullptr) {
for (size_t i = 0; i < maxLen / 4 - 1; i++) { size_t maxLen = dosHeader->e_lfanew - sizeof(IMAGE_DOS_HEADER);
if (richPtr[i] == 0x68636952) { // "Rich" // 确保搜索区域不会太大（预防恶意构造的文件）
peInfo.hasRich = true; const size_t MAX_RICH_SEARCH_SIZE = 1024; // 合理的Rich头最大区域
break; if (maxLen > MAX_RICH_SEARCH_SIZE) {
maxLen = MAX_RICH_SEARCH_SIZE;
}
// 确保不会越界
if (sizeof(IMAGE_DOS_HEADER) + maxLen <= v_size) {
const uint32_t* richPtr = reinterpret_cast<const uint32_t*>(
peBuffer + sizeof(IMAGE_DOS_HEADER));
// 确保剩余长度至少能容纳一个uint32_t
for (size_t i = 0;
i < maxLen / 4 - 1 && (i + 1) * sizeof(uint32_t) <= maxLen;
i++) {
if (richPtr[i] == 0x68636952) { // "Rich"
peInfo.hasRich = true;
break;
}
} }
} }
} }