From 9e295953a8b056036525c4274fb1ddf91dbb85c7 Mon Sep 17 00:00:00 2001
From: Joachim Bauch <mail@joachim-bauch.de>
Date: Wed, 14 Feb 2018 20:31:28 +0100
Subject: [PATCH] Prevent memory of library from spanning over 4GB bounaries.

Based on code from #67 / #79, fixes #63.
---
 MemoryModule.c | 65 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/MemoryModule.c b/MemoryModule.c
index 59a5844..05646e1 100644
--- a/MemoryModule.c
+++ b/MemoryModule.c
@@ -70,6 +70,13 @@ struct ExportNameEntry {
 typedef BOOL (WINAPI *DllEntryProc)(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpReserved);
 typedef int (WINAPI *ExeEntryProc)(void);
 
+#ifdef _WIN64
+typedef struct POINTER_LIST {
+    struct POINTER_LIST *next;
+    void *address;
+} POINTER_LIST;
+#endif
+
 typedef struct {
     PIMAGE_NT_HEADERS headers;
     unsigned char *codeBase;
@@ -87,6 +94,9 @@ typedef struct {
     void *userdata;
     ExeEntryProc exeEntry;
     DWORD pageSize;
+#ifdef _WIN64
+    POINTER_LIST *blockedMemory;
+#endif
 } MEMORYMODULE, *PMEMORYMODULE;
 
 typedef struct {
@@ -137,6 +147,21 @@ OutputLastError(const char *msg)
 #endif
 }
 
+#ifdef _WIN64
+static void
+FreePointerList(POINTER_LIST *head, CustomFreeFunc freeMemory, void *userdata)
+{
+    POINTER_LIST *node = head;
+    while (node) {
+        POINTER_LIST *next;
+        freeMemory(node->address, 0, MEM_RELEASE, userdata);
+        next = node->next;
+        free(node);
+        node = next;
+    }
+}
+#endif
+
 static BOOL
 CheckSize(size_t size, size_t expected) {
     if (size < expected) {
@@ -535,6 +560,9 @@ HMEMORYMODULE MemoryLoadLibraryEx(const void *data, size_t size,
     size_t optionalSectionSize;
     size_t lastSectionEnd = 0;
     size_t alignedImageSize;
+#ifdef _WIN64
+    POINTER_LIST *blockedMemory = NULL;
+#endif
 
     if (!CheckSize(size, sizeof(IMAGE_DOS_HEADER))) {
         return NULL;
@@ -610,9 +638,40 @@ HMEMORYMODULE MemoryLoadLibraryEx(const void *data, size_t size,
         }
     }
 
+#ifdef _WIN64
+    // Memory block may not span 4 GB boundaries.
+    while ((((uintptr_t) code) >> 32) < (((uintptr_t) (code + alignedImageSize)) >> 32)) {
+        POINTER_LIST *node = (POINTER_LIST*) malloc(sizeof(POINTER_LIST));
+        if (!node) {
+            freeMemory(code, 0, MEM_RELEASE, userdata);
+            FreePointerList(blockedMemory, freeMemory, userdata);
+            SetLastError(ERROR_OUTOFMEMORY);
+            return NULL;
+        }
+
+        node->next = blockedMemory;
+        node->address = code;
+        blockedMemory = node;
+
+        code = (unsigned char *)allocMemory(NULL,
+            alignedImageSize,
+            MEM_RESERVE | MEM_COMMIT,
+            PAGE_READWRITE,
+            userdata);
+        if (code == NULL) {
+            FreePointerList(blockedMemory, freeMemory, userdata);
+            SetLastError(ERROR_OUTOFMEMORY);
+            return NULL;
+        }
+    }
+#endif
+
     result = (PMEMORYMODULE)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(MEMORYMODULE));
     if (result == NULL) {
         freeMemory(code, 0, MEM_RELEASE, userdata);
+#ifdef _WIN64
+        FreePointerList(blockedMemory, freeMemory, userdata);
+#endif
         SetLastError(ERROR_OUTOFMEMORY);
         return NULL;
     }
@@ -626,6 +685,9 @@ HMEMORYMODULE MemoryLoadLibraryEx(const void *data, size_t size,
     result->freeLibrary = freeLibrary;
     result->userdata = userdata;
     result->pageSize = sysInfo.dwPageSize;
+#ifdef _WIN64
+    result->blockedMemory = blockedMemory;
+#endif
 
     if (!CheckSize(size, old_header->OptionalHeader.SizeOfHeaders)) {
         goto error;
@@ -823,6 +885,9 @@ void MemoryFreeLibrary(HMEMORYMODULE mod)
         module->free(module->codeBase, 0, MEM_RELEASE, module->userdata);
     }
 
+#ifdef _WIN64
+    FreePointerList(module->blockedMemory, module->free, module->userdata);
+#endif
     HeapFree(GetProcessHeap(), 0, module);
 }