From df3e51fc8c71a996bb048a0d88a341808d202d2d Mon Sep 17 00:00:00 2001
From: Tony Wasserka <neobrain@users.noreply.github.com>
Date: Mon, 5 Feb 2024 17:13:32 +0100
Subject: [PATCH] Library Forwarding: Allocate packed arguments on the guest
 stack if needed

This is required for host-side calls to guest functions on 32-bit guests.
Since the host stack is allocated before FEX blocks memory inaccessible to
the guest, the guest would otherwise fail to read the packed argument data.
---
 .../Source/Interface/HLE/Thunks/Thunks.cpp    | 40 ++++++++++-------
 ThunkLibs/include/common/Host.h               | 43 +++++++++++++++++++
 2 files changed, 67 insertions(+), 16 deletions(-)

diff --git a/FEXCore/Source/Interface/HLE/Thunks/Thunks.cpp b/FEXCore/Source/Interface/HLE/Thunks/Thunks.cpp
index 9fcdb2a6a3..48faeee99a 100644
--- a/FEXCore/Source/Interface/HLE/Thunks/Thunks.cpp
+++ b/FEXCore/Source/Interface/HLE/Thunks/Thunks.cpp
@@ -181,23 +181,11 @@ namespace FEXCore {
             Thread->CurrentFrame->State.gregs[FEXCore::X86State::REG_RDI] = (uintptr_t)arg0;
             Thread->CurrentFrame->State.gregs[FEXCore::X86State::REG_RSI] = (uintptr_t)arg1;
           } else {
-            // The argument pointer comes from the host stack, so it's not
-            // accessible by the guest. Allocate a thread-local chunk of memory
-            // to relocate the argument data.
-            // TODO: Directly store the arguments in a guest-accessible location instead
-            // TODO: FEXCore::Allocator::malloc() still returns pointers inaccessible from 32-bit guests here
-            thread_local void* local_args =
-              mmap( 0, 128, PROT_READ | PROT_WRITE,
-                    MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
-
-            // We don't know how much argument data is on the stack, so we
-            // unconditionally copy a fixed amount and leave everything else
-            // uninitialized.
-            // TODO: This breaks functions with large argument counts.
-            memcpy(local_args, arg1, 128);
-
+            if ((reinterpret_cast<uintptr_t>(arg1) >> 32) != 0) {
+              ERROR_AND_DIE_FMT("Tried to call guest function with arguments packed to a 64-bit address");
+            }
             Thread->CurrentFrame->State.gregs[FEXCore::X86State::REG_RCX] = (uintptr_t)arg0;
-            Thread->CurrentFrame->State.gregs[FEXCore::X86State::REG_RDX] = (uintptr_t)local_args;
+            Thread->CurrentFrame->State.gregs[FEXCore::X86State::REG_RDX] = (uintptr_t)arg1;
           }
 
           Thread->CTX->HandleCallback(Thread, (uintptr_t)callback);
@@ -499,6 +487,26 @@ namespace FEXCore {
       }
     }
 
+    FEX_DEFAULT_VISIBILITY void* GetGuestStack() {
+      if (!Thread) {
+        ERROR_AND_DIE_FMT("Thunked library attempted to query guest stack pointer asynchronously");
+      }
+
+      return (void*)(uintptr_t)((Thread->CurrentFrame->State.gregs[FEXCore::X86State::REG_RSP]));
+    }
+
+    FEX_DEFAULT_VISIBILITY void MoveGuestStack(uintptr_t NewAddress) {
+      if (!Thread) {
+        ERROR_AND_DIE_FMT("Thunked library attempted to query guest stack pointer asynchronously");
+      }
+
+      if (NewAddress >> 32) {
+         ERROR_AND_DIE_FMT("Tried to set stack pointer for 32-bit guest to a 64-bit address");
+      }
+
+      Thread->CurrentFrame->State.gregs[FEXCore::X86State::REG_RSP] = NewAddress;
+    }
+
 #else
     fextl::unique_ptr<ThunkHandler> ThunkHandler::Create() {
       ERROR_AND_DIE_FMT("Unsupported");
diff --git a/ThunkLibs/include/common/Host.h b/ThunkLibs/include/common/Host.h
index 8c3324dff0..0c48f9e567 100644
--- a/ThunkLibs/include/common/Host.h
+++ b/ThunkLibs/include/common/Host.h
@@ -30,6 +30,12 @@ namespace FEXCore {
   __attribute__((weak))
   HostToGuestTrampolinePtr*
   FinalizeHostTrampolineForGuestFunction(HostToGuestTrampolinePtr*, void* HostPacker);
+
+  __attribute__((weak))
+  void* GetGuestStack();
+
+  __attribute__((weak))
+  void MoveGuestStack(uintptr_t NewAddress);
 }
 
 template<typename Fn>
@@ -447,16 +453,53 @@ auto Projection(guest_layout<T>& data) {
   }
 }
 
+#ifdef IS_32BIT_THUNK
+/**
+ * Helper class to manage guest stack memory from a host function.
+ *
+ * The current guest stack position is saved upon construction and bumped
+ * for each object construction. Upon destruction, the old guest stack is
+ * restored.
+ */
+class GuestStackBumpAllocator final {
+  uintptr_t Top = reinterpret_cast<uintptr_t>(FEXCore::GetGuestStack());
+  uintptr_t Next = Top;
+
+public:
+  ~GuestStackBumpAllocator() {
+    FEXCore::MoveGuestStack(Top);
+  }
+
+  template<typename T, typename... Args>
+  T* New(Args&&... args) {
+    Next -= sizeof(T);
+    Next &= ~uintptr_t { alignof(T) - 1 };
+    FEXCore::MoveGuestStack(Next);
+    return new (reinterpret_cast<void*>(Next)) T {
+      std::forward<Args>(args)...
+    };
+  }
+};
+#endif
+
 template<typename Result, typename... Args>
 struct CallbackUnpack<Result(Args...)> {
   static Result CallGuestPtr(Args... args) {
     GuestcallInfo *guestcall;
     LOAD_INTERNAL_GUESTPTR_VIA_CUSTOM_ABI(guestcall);
 
+#ifndef IS_32BIT_THUNK
     PackedArguments<Result, guest_layout<Args>...> packed_args = {
       to_guest(to_host_layout(args))...
     };
+#else
+    GuestStackBumpAllocator GuestStack;
+    auto& packed_args = *GuestStack.New<PackedArguments<Result, guest_layout<Args>...>>(
+      to_guest(to_host_layout(args))...
+    );
+#endif
     guestcall->CallCallback(guestcall->GuestUnpacker, guestcall->GuestTarget, &packed_args);
+
     if constexpr (!std::is_void_v<Result>) {
       return packed_args.rv;
     }