loader: keep the caller's flags in the TLS stubs (#1571)

A patched TLS load calls TlsRegStub and SafeCall, which moved RSP past the
red zone with sub/add rsp,0x80. ADD and SUB write the arithmetic flags, so
the sub ran before SafeCall's pushfq saved them and the add ran after its
popfq restored them: guest code that compares, loads fs:[0] and then
branches on the flags could branch the wrong way. lea moves RSP without
touching the flags. The entry lea plus a two-byte nop keeps every patched
offset in both stubs; the exit lea is one byte longer and moves only ret.

TlsStubsPreserveFlags calls the stubs with two flag patterns and checks
the status flags, the output register and RAX; it fails with the old
bytes.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

src/loader/jit.h

@@ -58,16 +58,17 @@ struct TlsRegStub {
 	}
 	static uint64_t GetSize() { return 32; }
 
-	// sub rsp,0x80
+	// lea rsp,[rsp-0x80]
+	// nop
 	// push rax
 	// call safe_call
 	// mov <reg>,rax
 	// pop rax
-	// add rsp,0x80
+	// lea rsp,[rsp+0x80]
 	// ret
-	uint8_t code[32] = {0x48, 0x81, 0xEC, 0x80, 0x00, 0x00, 0x00, 0x50, 0xE8, 0x00, 0x00,
-	                    0x00, 0x00, 0x48, 0x89, 0xC0, 0x58, 0x48, 0x81, 0xC4, 0x80, 0x00,
-	                    0x00, 0x00, 0xC3, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90};
+	uint8_t code[32] = {0x48, 0x8D, 0x64, 0x24, 0x80, 0x66, 0x90, 0x50, 0xE8, 0x00, 0x00,
+	                    0x00, 0x00, 0x48, 0x89, 0xC0, 0x58, 0x48, 0x8D, 0xA4, 0x24, 0x80,
+	                    0x00, 0x00, 0x00, 0xC3, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90};
 };
 
 struct SafeCall {
@@ -78,7 +79,8 @@ struct SafeCall {
 	static uint64_t GetSize() { return 0x1000; }
 
 	uint8_t code[0x90] = {
-	    /*00*/ 0x48, 0x81, 0xec, 0x80, 0x00, 0x00, 0x00, // sub    rsp,0x80
+	    /*00*/ 0x48, 0x8d, 0x64, 0x24, 0x80,             // lea    rsp,[rsp-0x80]
+	    /*05*/ 0x66, 0x90,                               // nop
 	    /*07*/ 0x9c,                                     // pushfq
 	    /*08*/ 0x51,                                     // push   rcx
 	    /*09*/ 0x52,                                     // push   rdx
@@ -108,9 +110,9 @@ struct SafeCall {
 	    /*40*/ 0x5a,             // pop    rdx
 	    /*41*/ 0x59,             // pop    rcx
 	    /*42*/ 0x9d,             // popfq
-	    /*43*/ 0x48, 0x81, 0xc4, 0x80, 0x00, 0x00,
-	    0x00,        // add    rsp,0x80
-	    /*4a*/ 0xc3, // ret
+	    /*43*/ 0x48, 0x8d, 0xa4, 0x24, 0x80, 0x00, 0x00,
+	    0x00,        // lea    rsp,[rsp+0x80]
+	    /*4b*/ 0xc3, // ret
 	};
 };
 

tests/VirtualMemoryAllocationTests.cpp

@@ -8,6 +8,7 @@
 #include "kernel/pthread.h"
 #include "libs/errno.h"
 #include "loader/guestInstructionPatcher.h"
+#include "loader/jit.h"
 #include "loader/runtimeLinker.h"
 #include "loader/systemContent.h"
 #include "loader/x64InstructionEmulator.h"
@@ -4775,6 +4776,59 @@ void TestSmallFiberStacksAndMigration() {
 
 } // namespace
 
+#if defined(__linux__) || KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
+KYTY_MS_ABI uint8_t* TlsStubTestAddress() {
+	return reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(0x123456789abcull));
+}
+
+void TestTlsStubsPreserveFlags() {
+	const char*    test      = "TlsStubsPreserveFlags";
+	const uint64_t stub_size = Loader::Jit::SafeCall::GetSize();
+	constexpr uint64_t CODE_SIZE = 0x1000;
+	constexpr uint8_t  REG       = 1;
+	const auto         mapping   = Libs::LibKernel::Memory::AllocateProgramMemory(
+	    0x0000000906000000ull, stub_size + CODE_SIZE, Common::VirtualMemory::Mode::ExecuteReadWrite,
+	    "tls_stub_test");
+	Check(test, mapping != 0, "failed to allocate TLS stub test code");
+	auto* safe_call = new (reinterpret_cast<void*>(mapping)) Loader::Jit::SafeCall;
+	safe_call->SetFunc(TlsStubTestAddress);
+	const auto stub_address = mapping + Loader::Jit::TlsRegStub::GetOffset(REG);
+	auto*      stub         = new (reinterpret_cast<void*>(stub_address)) Loader::Jit::TlsRegStub;
+	stub->SetFunc(mapping);
+	stub->SetOutputReg(REG);
+
+	using Fixture  = uint64_t(KYTY_SYSV_ABI*)(uint64_t*);
+	bool preserved = true;
+	for (const auto [input, status]: {std::pair {uint8_t {0x7f}, uint64_t {0x890}},
+	                                  std::pair {uint8_t {0xff}, uint64_t {0x055}}}) {
+		Xbyak::CodeGenerator code(CODE_SIZE, reinterpret_cast<void*>(mapping + stub_size));
+		code.push(code.rbx);
+		code.mov(code.rbx, code.rdi);
+		code.mov(code.rax, 0x5555555555555555ull);
+		code.mov(code.al, input);
+		code.add(code.al, 1);
+		code.mov(code.r11, stub_address);
+		code.call(code.r11);
+		code.mov(code.qword[code.rbx], code.rcx);
+		code.mov(code.qword[code.rbx + 8], code.rax);
+		code.pushf();
+		code.pop(code.rax);
+		code.pop(code.rbx);
+		code.ret();
+		Check(test, Xbyak::GetError() == 0, "failed to generate TLS stub fixture");
+		Check(test, Common::VirtualMemory::FlushInstructionCache(mapping, stub_size + CODE_SIZE),
+		      "failed to flush TLS stub test code");
+		std::array<uint64_t, 2> registers {};
+		const auto flags = reinterpret_cast<Fixture>(mapping + stub_size)(registers.data());
+		const auto rax   = (0x5555555555555555ull & ~0xffull) | static_cast<uint8_t>(input + 1u);
+		preserved &= (flags & 0x8d5u) == status && registers[0] == 0x123456789abcull &&
+		             registers[1] == rax;
+	}
+	Libs::LibKernel::Memory::FreeGuestMemory(mapping, stub_size + CODE_SIZE);
+	Check(test, preserved, "TLS stubs changed the caller's flags or registers");
+}
+#endif
+
 int main(int argc, char** argv) {
 	InitSubsystems();
 	if (argc == 2 && std::strcmp(argv[1], "--virtual-ranges-only") == 0) {
@@ -4919,6 +4973,9 @@ int main(int argc, char** argv) {
 	RunTest(TestMemoryPoolCommitDecommitQueryFlags);
 	RunTest(TestProgramMemoryAllocationAndProtection);
 	RunTest(TestModuleRelocationUsesWritableHostMapping);
+#if defined(__linux__) || KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
+	RunTest(TestTlsStubsPreserveFlags);
+#endif
 
 	if (g_failed_tests != 0) {
 		std::printf("VirtualMemoryAllocationTests: %d case(s) failed\n", g_failed_tests);