Driver: fortmm: Simplify IsPEHeaderValid()

This commit is contained in:
Nodir Temirkhodjaev 2023-05-09 17:10:33 +03:00
parent d086a9b8b0
commit 7e29bd622b

View File

@ -293,40 +293,57 @@ static BOOL CheckPEHeaderSections(const PIMAGE_NT_HEADERS pNtHeaders)
return TRUE; return TRUE;
} }
inline static BOOL CheckPEHeaderOptionalValid(const PIMAGE_NT_HEADERS pNtHeaders, DWORD dwSize)
{
/* Check NT header for valid signature */
if (pNtHeaders->Signature != IMAGE_NT_SIGNATURE)
return FALSE;
/* Check size of optional headerss */
if (dwSize < pNtHeaders->OptionalHeader.SizeOfHeaders)
return FALSE;
/* Check for the correct architecture */
if (pNtHeaders->FileHeader.Machine !=
#if defined(_M_ARM64)
IMAGE_FILE_MACHINE_ARM64
#elif defined(_WIN64)
IMAGE_FILE_MACHINE_AMD64
#else
IMAGE_FILE_MACHINE_I386
#endif
)
return FALSE;
/* Check to see if the image is really an executable file */
if ((pNtHeaders->FileHeader.Characteristics & (IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_DLL))
== 0)
return FALSE;
/* Check sections */
return CheckPEHeaderSections(pNtHeaders);
}
static BOOL IsPEHeaderValid(PVOID lpData, DWORD dwSize) static BOOL IsPEHeaderValid(PVOID lpData, DWORD dwSize)
{ {
const PIMAGE_DOS_HEADER pDosHeader = (PIMAGE_DOS_HEADER) lpData; const PIMAGE_DOS_HEADER pDosHeader = (PIMAGE_DOS_HEADER) lpData;
if (pDosHeader->e_magic != IMAGE_DOS_SIGNATURE /* Check DOS header for valid signature */ /* Check DOS header for valid signature */
/* Make sure size is at least size of headers */ if (pDosHeader->e_magic != IMAGE_DOS_SIGNATURE)
|| dwSize < (sizeof(IMAGE_DOS_HEADER) + sizeof(IMAGE_OPTIONAL_HEADER)) return FALSE;
|| dwSize < (pDosHeader->e_lfanew + sizeof(IMAGE_NT_HEADERS)))
/* Make sure size is at least size of headers */
if (dwSize < (sizeof(IMAGE_DOS_HEADER) + sizeof(IMAGE_OPTIONAL_HEADER)))
return FALSE;
if (dwSize < (pDosHeader->e_lfanew + sizeof(IMAGE_NT_HEADERS)))
return FALSE; return FALSE;
/* Check for optional headers */ /* Check for optional headers */
const PIMAGE_NT_HEADERS pNtHeaders = const PIMAGE_NT_HEADERS pNtHeaders =
(PIMAGE_NT_HEADERS) & ((PUCHAR) lpData)[pDosHeader->e_lfanew]; (PIMAGE_NT_HEADERS) & ((PUCHAR) lpData)[pDosHeader->e_lfanew];
if (pNtHeaders->Signature != IMAGE_NT_SIGNATURE /* Check NT header for valid signature */ return CheckPEHeaderOptionalValid(pNtHeaders, dwSize);
/* Check size of optional headerss */
|| dwSize < pNtHeaders->OptionalHeader.SizeOfHeaders
/* Check for the correct architecture */
|| pNtHeaders->FileHeader.Machine !=
#if defined(_M_ARM64)
IMAGE_FILE_MACHINE_ARM64
#elif defined(_WIN64)
IMAGE_FILE_MACHINE_AMD64
#else
IMAGE_FILE_MACHINE_I386
#endif
/* Check to see if the image is really an executable file */
|| (pNtHeaders->FileHeader.Characteristics
& (IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_DLL))
== 0)
return FALSE;
/* Check sections */
return CheckPEHeaderSections(pNtHeaders);
} }
static NTSTATUS InitializeModuleImage(PLOADEDMODULE pModule, const PIMAGE_NT_HEADERS lpNtHeaders, static NTSTATUS InitializeModuleImage(PLOADEDMODULE pModule, const PIMAGE_NT_HEADERS lpNtHeaders,