From efe15ce3368d07c2bcc6aa03943866b2fe69ea01 Mon Sep 17 00:00:00 2001
From: Billy Laws <blaws05@gmail.com>
Date: Tue, 23 Jul 2024 17:38:10 +0000
Subject: [PATCH] ARM64EC: Handle direct syscall instructions

Most syscalls on Windows are done by calling into their NTDLL thunks,
however some DRMs parse out their numbers from NTDLL and directly call
them. Support this by redirecting to their entry thunks in the FEX
syscall handler.
---
 Source/Windows/ARM64EC/Module.cpp | 32 +++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/Source/Windows/ARM64EC/Module.cpp b/Source/Windows/ARM64EC/Module.cpp
index a0ff082c5..193e577e9 100644
--- a/Source/Windows/ARM64EC/Module.cpp
+++ b/Source/Windows/ARM64EC/Module.cpp
@@ -113,6 +113,8 @@ std::recursive_mutex ThreadCreationMutex;
 // Map of TIDs to their FEX thread state, `ThreadCreationMutex` must be locked when accessing
 std::unordered_map<DWORD, FEXCore::Core::InternalThreadState*> Threads;
 
+// Map from system call numbers to the relative addresses of their native implementations in ntdll
+std::vector<uint32_t> NtDllSyscallLUT;
 
 std::pair<NTSTATUS, ThreadCPUArea> GetThreadCPUArea(HANDLE Thread) {
   THREAD_BASIC_INFORMATION Info;
@@ -149,6 +151,34 @@ void FillNtDllLUTs() {
   for (auto It = RedirectionTableBegin; It != RedirectionTableEnd; It++) {
     NtDllRedirectionLUT[It->Source] = It->Destination;
   }
+
+  const auto* Exports = reinterpret_cast<IMAGE_EXPORT_DIRECTORY*>(RtlImageDirectoryEntryToData(NtDll, true, IMAGE_DIRECTORY_ENTRY_EXPORT, &Size));
+  const auto* FunctionTableBegin = reinterpret_cast<uint32_t*>(NtDllBase + Exports->AddressOfFunctions);
+  const auto* FunctionTableEnd = FunctionTableBegin + Exports->NumberOfFunctions;
+
+  NtDllSyscallLUT.reserve(0x200);
+  for (auto It = FunctionTableBegin; It != FunctionTableEnd; It++) {
+    const uint8_t* FunctionAddr = reinterpret_cast<uint8_t*>(NtDllBase + *It);
+    // Windows syscall thunks are as follows:
+    // 00: mov r10, rcx
+    // 03: mov eax, <NUM>
+    // <cont into MatchSeq>
+    static constexpr std::array<uint8_t, 16> MatchSeq {{
+      0xf6, 0x04, 0x25, 0x08, 0x03, 0xfe, 0x7f, 0x01, // 08: test byte ptr ds:7FFE0308h, 1
+      0x75, 0x03,                                     // 10: jnz short lbl
+      0x0f, 0x05,                                     // 12: syscall
+      0xc3,                                           // 14: retn
+      0xcd, 0x2e,                                     // 15: lbl: int 2Eh
+      0xc3                                            // 17: retn
+    }};
+
+    const uint8_t* MatchAddr = FunctionAddr + 8;
+    if (!memcmp(MatchSeq.data(), MatchAddr, MatchSeq.size())) {
+      const uint32_t SyscallNum = *reinterpret_cast<const uint32_t*>(FunctionAddr + 4);
+      NtDllSyscallLUT.resize(std::max<size_t>(NtDllSyscallLUT.size(), SyscallNum));
+      NtDllSyscallLUT[SyscallNum] = NtDllRedirectionLUT[*It];
+    }
+  }
 }
 
 } // namespace
@@ -408,6 +438,8 @@ public:
   }
 
   uint64_t HandleSyscall(FEXCore::Core::CpuStateFrame* Frame, FEXCore::HLE::SyscallArguments* Args) override {
+    Frame->State.rip = NtDllBase + NtDllSyscallLUT[Frame->State.gregs[FEXCore::X86State::REG_RAX]];
+    Frame->State.gregs[FEXCore::X86State::REG_RCX] = Frame->State.gregs[FEXCore::X86State::REG_R10];
     return 0;
   }