From 2fed2d5bae66934ffc2a0f0baeff8c2e6824a33f Mon Sep 17 00:00:00 2001 From: Huoji's <1296564236@qq.com> Date: Sun, 9 Mar 2025 03:25:29 +0800 Subject: [PATCH] Enhance PE file Rich header detection with null and boundary checks - Added null pointer and boundary checks in ExtractFeatures method - Prevent potential buffer overread when searching for Rich header signature - Improve robustness of feature extraction for PE file analysis --- ai_anti_malware/ml.cpp | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/ai_anti_malware/ml.cpp b/ai_anti_malware/ml.cpp index aa2f004..861af55 100644 --- a/ai_anti_malware/ml.cpp +++ b/ai_anti_malware/ml.cpp @@ -289,10 +289,12 @@ std::vector MachineLearning::ExtractFeatures(const uint8_t* buffer, const uint32_t* richPtr = reinterpret_cast( peBuffer + sizeof(IMAGE_DOS_HEADER)); size_t maxLen = dosHeader->e_lfanew - sizeof(IMAGE_DOS_HEADER); - for (size_t i = 0; i < maxLen / 4 - 1; i++) { - if (richPtr[i] == 0x68636952) { // "Rich" - peInfo.hasRich = true; - break; + if (maxLen > 0 && richPtr != nullptr) { + for (size_t i = 0; i < maxLen / 4 - 1; i++) { + if (richPtr[i] == 0x68636952) { // "Rich" + peInfo.hasRich = true; + break; + } } } }