11#include " stdafx.h"
2- #include < tchar.h>
3- #include < algorithm>
4-
5- #if _MSC_VER
6- #pragma warning(disable:4055)
7- #pragma warning(error: 4244)
8- #pragma warning(error: 4267)
9- #pragma warning(disable:4996)
10- #define inline __inline
11- #endif
122
133#ifdef _WIN64
144#define HOST_MACHINE IMAGE_FILE_MACHINE_AMD64
@@ -24,7 +14,7 @@ PMEMORYMODULE WINAPI MapMemoryModuleHandle(HMEMORYMODULE hModule) {
2414 if (!dos)return nullptr ;
2515 PIMAGE_NT_HEADERS nt = (PIMAGE_NT_HEADERS)((LPBYTE)hModule + dos->e_lfanew );
2616 if (!nt)return nullptr ;
27- PMEMORYMODULE pModule = (PMEMORYMODULE)((LPBYTE)hModule + nt->OptionalHeader .SizeOfHeaders );
17+ PMEMORYMODULE pModule = (PMEMORYMODULE)((LPBYTE)hModule + nt->OptionalHeader .SizeOfHeaders - sizeof (MEMORYMODULE) );
2818 if (!_ProbeForRead (pModule, sizeof (MEMORYMODULE)))return nullptr ;
2919 if (pModule->Signature != MEMORY_MODULE_SIGNATURE || (size_t )pModule->codeBase != nt->OptionalHeader .ImageBase )return nullptr ;
3020 return pModule;
@@ -290,10 +280,25 @@ NTSTATUS MemoryLoadLibrary(
290280 );
291281 new_header->OptionalHeader .ImageBase = (size_t )base;
292282
283+ // https://learn.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-image_optional_header32
284+ int sizeOfHeaders = dos_header->e_lfanew + // e_lfanew member of IMAGE_DOS_HEADER
285+ 4 + // 4 byte signature
286+ sizeof (IMAGE_FILE_HEADER) + // size of IMAGE_FILE_HEADER
287+ sizeof (IMAGE_OPTIONAL_HEADER) + // size of optional header
288+ sizeof (IMAGE_SECTION_HEADER) * old_header->FileHeader .NumberOfSections ; // size of all section headers
289+
290+ //
291+ // Make sure there have enough free space to embed our structure.
292+ //
293+ if (sizeOfHeaders + sizeof (MEMORYMODULE) > old_header->OptionalHeader .SizeOfHeaders ) {
294+ status = STATUS_NOT_SUPPORTED;
295+ return status;
296+ }
297+
293298 //
294299 // Setup MemoryModule structure.
295300 //
296- PMEMORYMODULE hMemoryModule = (PMEMORYMODULE)(base + old_header->OptionalHeader .SizeOfHeaders );
301+ PMEMORYMODULE hMemoryModule = (PMEMORYMODULE)(base + old_header->OptionalHeader .SizeOfHeaders - sizeof (MEMORYMODULE) );
297302 RtlZeroMemory (hMemoryModule, sizeof (MEMORYMODULE));
298303 hMemoryModule->codeBase = base;
299304 hMemoryModule->dwImageFileSize = size;
0 commit comments