From 8e0fdfc325359d99c53e258eee6c501c1c3cf228 Mon Sep 17 00:00:00 2001
From: Billy Laws <blaws05@gmail.com>
Date: Thu, 11 Jul 2024 13:17:11 +0100
Subject: [PATCH] ARM64EC: Add a helper to lookup the redirected address of an
 export

FEX is unable to deal with reentrant compilation of any x64 hotpatches
so they need to be ignored by bypassing FFSs and calling directly into
the native target.
---
 Source/Windows/ARM64EC/Module.cpp | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/Source/Windows/ARM64EC/Module.cpp b/Source/Windows/ARM64EC/Module.cpp
index ee578f4ed..7e5c22673 100644
--- a/Source/Windows/ARM64EC/Module.cpp
+++ b/Source/Windows/ARM64EC/Module.cpp
@@ -41,6 +41,7 @@ $end_info$
 #include <ntstatus.h>
 #include <windef.h>
 #include <winternl.h>
+#include <winnt.h>
 #include <wine/debug.h>
 
 class ECSyscallHandler;
@@ -109,6 +110,30 @@ bool IsDispatcherAddress(uint64_t Address) {
   const auto& Config = SignalDelegator->GetConfig();
   return Address >= Config.DispatcherBegin && Address < Config.DispatcherEnd;
 }
+
+// GetProcAddress on ARM64EC returns a pointer to an x64 fast forward sequence to allow for redirecting to the JIT if functions are hotpatched.
+// This looks up the procedure address of the native code even if the fast forward sequence has been patched.
+uintptr_t GetRedirectedProcAddress(HMODULE Module, const char* ProcName) {
+  const uintptr_t Proc = reinterpret_cast<uintptr_t>(GetProcAddress(Module, ProcName));
+  if (!Proc) {
+    return 0;
+  }
+
+  ULONG Size;
+  const auto* LoadConfig =
+    reinterpret_cast<_IMAGE_LOAD_CONFIG_DIRECTORY64*>(RtlImageDirectoryEntryToData(Module, true, IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG, &Size));
+  const auto* CHPEMetadata = reinterpret_cast<IMAGE_ARM64EC_METADATA*>(LoadConfig->CHPEMetadataPointer);
+  const uintptr_t ModuleBase = reinterpret_cast<uintptr_t>(Module);
+  const uintptr_t ProcRVA = Proc - ModuleBase;
+  const auto* RedirectionTableBegin = reinterpret_cast<IMAGE_ARM64EC_REDIRECTION_ENTRY*>(ModuleBase + CHPEMetadata->RedirectionMetadata);
+  const auto* RedirectionTableEnd = RedirectionTableBegin + CHPEMetadata->RedirectionMetadataCount;
+  const auto* It =
+    std::lower_bound(RedirectionTableBegin, RedirectionTableEnd, ProcRVA, [](const auto& Entry, uintptr_t RVA) { return Entry.Source < RVA; });
+  if (It->Source != ProcRVA) {
+    return 0;
+  }
+  return ModuleBase + It->Destination;
+}
 } // namespace
 
 namespace Exception {