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

[via llvm-commits]

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

Since we may copy code (see CopyInstructions) 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 anything that function may have rip-relative accesses to, i.e. within 2 GB of that function's whole module.

This fixes interception failures like the following scenario:

1. Intercept `CreateProcess` in kernel32.dll, allocating a trampoline region right after
2. Start intercepting `memcpy` in the main executable, which is loaded at a lower address than kernel32.dll, but still within 2 GB of the trampoline region so we keep using it.
3. Try to copy instructions from `memcpy` to the trampoline. Turns out one instruction references data that is more than 2GB away from the trampoline, so it can't be relocated.
4. The process exits due to a CHECK failure

(Full story at https://crbug.com/341936875#comment45 and following.)

>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] [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) {