[compiler-rt] [win/asan] AllocateMemoryForTrampoline within 2 GB of the module's base address (PR #108822)

[via llvm-commits]

https://github.com/zmodem updated https://github.com/llvm/llvm-project/pull/108822

>From bb222ba6859e493e7d61e1961b42b94179733e10 Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans at chromium.org>
Date: Mon, 16 Sep 2024 13:51:38 +0200
Subject: [PATCH 1/2] [win/asan] AllocateMemoryForTrampoline within 2 GB of the
 module's base address

---
 .../lib/interception/interception_win.cpp     | 25 ++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/compiler-rt/lib/interception/interception_win.cpp b/compiler-rt/lib/interception/interception_win.cpp
index a638e66eccee58..b7979b271f468f 100644
--- a/compiler-rt/lib/interception/interception_win.cpp
+++ b/compiler-rt/lib/interception/interception_win.cpp
@@ -130,6 +130,7 @@
 #include "sanitizer_common/sanitizer_platform.h"
 #define WIN32_LEAN_AND_MEAN
 #include <windows.h>
+#include <psapi.h>
 
 namespace __interception {
 
@@ -385,7 +386,29 @@ void TestOnlyReleaseTrampolineRegions() {
   }
 }
 
-static uptr AllocateMemoryForTrampoline(uptr image_address, size_t size) {
+static uptr AllocateMemoryForTrampoline(uptr func_address, size_t size) {
+  uptr image_address = func_address;
+
+#if SANITIZER_WINDOWS64
+  // Since we may copy code to the trampoline which could reference data
+  // inside the original module, we really want the trampoline to be within
+  // 2 GB of not just the original function, but within 2 GB of that function's
+  // whole module. Since the allocated trampoline's address is always greater
+  // than image_address, we achieve this by setting image_address to the base
+  // address of the module, if we can find it (which is not the case if
+  // func_address is in mmap'ed memory for example).
+  HMODULE module;
+  if (::GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
+                           GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+                           (LPCWSTR)func_address, &module)) {
+    MODULEINFO module_info;
+    if (::GetModuleInformation(::GetCurrentProcess(), module,
+                                &module_info, sizeof(module_info))) {
+      image_address = (uptr)module_info.lpBaseOfDll;
+    }
+  }
+#endif
+
   // Find a region within 2G with enough space to allocate |size| bytes.
   TrampolineMemoryRegion *region = nullptr;
   for (size_t bucket = 0; bucket < kMaxTrampolineRegion; ++bucket) {

>From 32b02e0cfce2e6432b237e92e378f9cbc8a6ff7b Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans at chromium.org>
Date: Tue, 17 Sep 2024 11:10:32 +0200
Subject: [PATCH 2/2] rework the comment

---
 compiler-rt/lib/interception/interception_win.cpp | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/compiler-rt/lib/interception/interception_win.cpp b/compiler-rt/lib/interception/interception_win.cpp
index b7979b271f468f..a0ff124a89c9ed 100644
--- a/compiler-rt/lib/interception/interception_win.cpp
+++ b/compiler-rt/lib/interception/interception_win.cpp
@@ -390,13 +390,14 @@ static uptr AllocateMemoryForTrampoline(uptr func_address, size_t size) {
   uptr image_address = func_address;
 
 #if SANITIZER_WINDOWS64
-  // Since we may copy code to the trampoline which could reference data
-  // inside the original module, we really want the trampoline to be within
-  // 2 GB of not just the original function, but within 2 GB of that function's
-  // whole module. Since the allocated trampoline's address is always greater
-  // than image_address, we achieve this by setting image_address to the base
-  // address of the module, if we can find it (which is not the case if
-  // func_address is in mmap'ed memory for example).
+  // Allocate memory after the module (DLL or EXE file), but within 2GB
+  // of the start of the module so that any address within the module can be
+  // referenced with PC-relative operands.
+  // This allows us to not just jump to the trampoline with a PC-relative
+  // offset, but to relocate any instructions that we copy to the trampoline
+  // which have references to the original module. If we can't find the base
+  // address of the module (e.g. if func_address is in mmap'ed memory), just
+  // use func_address as is.
   HMODULE module;
   if (::GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
                            GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,