From f00b0b203736cdbfa1bf2df3a571cf5676e66efa Mon Sep 17 00:00:00 2001
From: Huoji's <1296564236@qq.com>
Date: Sun, 9 Mar 2025 03:31:11 +0800
Subject: [PATCH] Improve Rich Header detection with enhanced safety checks

- 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
---
 ai_anti_malware/ml.cpp | 34 +++++++++++++++++++++++++---------
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/ai_anti_malware/ml.cpp b/ai_anti_malware/ml.cpp
index 16c8320..de32a9c 100644
--- a/ai_anti_malware/ml.cpp
+++ b/ai_anti_malware/ml.cpp
@@ -346,18 +346,34 @@ std::vector<double> MachineLearning::ExtractFeatures(const uint8_t* buffer,
     dataDir = peconv::get_directory_entry(peBuffer, IMAGE_DIRECTORY_ENTRY_IAT);
     peInfo.hasEntryIat = dataDir && dataDir->VirtualAddress != 0;
 
-    // Rich头部检测
+    // Rich头部检测 - 安全实现
     peInfo.hasRich = false;
     PIMAGE_DOS_HEADER dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(peBuffer);
     if (dosHeader && dosHeader->e_magic == IMAGE_DOS_SIGNATURE) {
-        const uint32_t* richPtr = reinterpret_cast<const uint32_t*>(
-            peBuffer + sizeof(IMAGE_DOS_HEADER));
-        size_t maxLen = dosHeader->e_lfanew - sizeof(IMAGE_DOS_HEADER);
-        if (maxLen > 0 && richPtr != nullptr) {
-            for (size_t i = 0; i < maxLen / 4 - 1; i++) {
-                if (richPtr[i] == 0x68636952) {  // "Rich"
-                    peInfo.hasRich = true;
-                    break;
+        // 确保e_lfanew值合理
+        if (dosHeader->e_lfanew > sizeof(IMAGE_DOS_HEADER) &&
+            dosHeader->e_lfanew < v_size) {  // 确保在PE文件大小范围内
+
+            size_t maxLen = dosHeader->e_lfanew - sizeof(IMAGE_DOS_HEADER);
+            // 确保搜索区域不会太大（预防恶意构造的文件）
+            const size_t MAX_RICH_SEARCH_SIZE = 1024;  // 合理的Rich头最大区域
+            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;
+                    }
                 }
             }
         }