From 4877bb3f19fdb606c7231975d3c343ea20132981 Mon Sep 17 00:00:00 2001
From: Billy Laws <blaws05@gmail.com>
Date: Tue, 23 Jul 2024 18:08:59 +0000
Subject: [PATCH] ARM64EC: Support directly issuing the NtContinue syscall

This is required for handling SMC with the ResetToConsistentState
arguments as used in Windows, as using the NTDLL exported NtContinue
would wipe out any reserved registers in the ARM64EC ABI.

For Windows the syscall numbers are somewhat stable, and the SVC
instruction can be called directly. Since wine doesn't handle that on
ARM64, hardcode the system call number and manually call into wine
dispatcher. Once wine gains proper syscall thunks, those can be
parsed to get the number and the hardcoding dropped.
---
 Source/Windows/ARM64EC/Module.S   | 17 +++++++++++++++++
 Source/Windows/ARM64EC/Module.cpp | 18 +++++++++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/Source/Windows/ARM64EC/Module.S b/Source/Windows/ARM64EC/Module.S
index 9d8e169ee..6510dd447 100644
--- a/Source/Windows/ARM64EC/Module.S
+++ b/Source/Windows/ARM64EC/Module.S
@@ -81,3 +81,20 @@ ret_sp_misaligned:
   adrp lr, X64ReturnInstr
   ldr lr, [lr, #:lo12:X64ReturnInstr]
   br x17
+
+  // Calls NtContinue directly to allow continuing from a full native context, as the NTDLL NtContinue export takes in
+  // an x64 context with EC and the conversion to that loses the ARM64EC ABI-disallowed registers that FEX uses.
+.global "#NtContinueNative"
+"#NtContinueNative":
+  adrp x16, WineSyscallDispatcher
+  ldr x16, [x16, #:lo12:WineSyscallDispatcher]
+  cbz x16, direct_syscall
+wine_syscall:
+  mov x9, x30
+  adrp x8, WineNtContinueSyscallId
+  ldr x8, [x8, #:lo12:WineNtContinueSyscallId]
+  blr x16
+  ret
+direct_syscall:
+  svc #0x43
+  ret
diff --git a/Source/Windows/ARM64EC/Module.cpp b/Source/Windows/ARM64EC/Module.cpp
index 8b6a0195e..8584afcf9 100644
--- a/Source/Windows/ARM64EC/Module.cpp
+++ b/Source/Windows/ARM64EC/Module.cpp
@@ -46,9 +46,17 @@ $end_info$
 #include <wine/debug.h>
 
 class ECSyscallHandler;
+
+extern "C" {
 void* X64ReturnInstr; // See Module.S
 extern void* ExitFunctionEC;
 
+// Wine doesn't support issuing direct system calls with SVC, and unlike Windows it doesn't have a 'stable' syscall number for NtContinue
+void* WineSyscallDispatcher;
+// TODO: this really shouldn't be hardcoded, once wine gains proper syscall thunks this can be dropped.
+uint64_t WineNtContinueSyscallId = 0x1a;
+}
+
 struct ThreadCPUArea {
   static constexpr size_t TEBCPUAreaOffset = 0x1788;
   CHPE_V2_CPU_AREA_INFO* Area;
@@ -85,6 +93,9 @@ struct ThreadCPUArea {
   }
 };
 
+
+extern "C" NTSTATUS NtContinueNative(ARM64_NT_CONTEXT* NativeContext, BOOLEAN Alert);
+
 namespace {
 fextl::unique_ptr<FEXCore::Context::Context> CTX;
 fextl::unique_ptr<FEX::DummyHandlers::DummySignalDelegator> SignalDelegator;
@@ -437,7 +448,12 @@ void ProcessInit() {
   X64ReturnInstr = ::VirtualAlloc(nullptr, FEXCore::Utils::FEX_PAGE_SIZE, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
   *reinterpret_cast<uint8_t*>(X64ReturnInstr) = 0xc3;
 
-  Exception::KiUserExceptionDispatcher = GetRedirectedProcAddress(GetModuleHandle("ntdll.dll"), "KiUserExceptionDispatcher");
+  const auto NtDll = GetModuleHandle("ntdll.dll");
+  Exception::KiUserExceptionDispatcher = GetRedirectedProcAddress(NtDll, "KiUserExceptionDispatcher");
+  const auto WineSyscallDispatcherPtr = reinterpret_cast<void**>(GetProcAddress(NtDll, "__wine_syscall_dispatcher"));
+  if (WineSyscallDispatcherPtr) {
+    WineSyscallDispatcher = *WineSyscallDispatcherPtr;
+  }
 }
 
 void ProcessTerm() {}