cpu_patches: Patch stack canary accesses (#3857)

* Patch stack checks done using fs:[0x28]

Additionally adds support for multiple patches per instruction, since this makes two separate patches we need to conditionally perform for mov instructions.

* Missing include

* Disable patches for Apple

Mac can use their native FS segment directly, so these patches aren't needed

* Oops
This commit is contained in:
Stephen Miller
2025-12-07 00:47:39 -06:00
committed by GitHub
parent d3ad728ac0
commit 391d30cbb1

View File

@@ -5,6 +5,7 @@
#include <memory> #include <memory>
#include <mutex> #include <mutex>
#include <set> #include <set>
#include <vector>
#include <Zydis/Zydis.h> #include <Zydis/Zydis.h>
#include <xbyak/xbyak.h> #include <xbyak/xbyak.h>
#include <xbyak/xbyak_util.h> #include <xbyak/xbyak_util.h>
@@ -122,6 +123,30 @@ static void GenerateTcbAccess(void* /* address */, const ZydisDecodedOperand* op
#endif #endif
} }
static bool FilterStackCheck(const ZydisDecodedOperand* operands) {
const auto& dst_op = operands[0];
const auto& src_op = operands[1];
// Some compilers emit stack checks by starting a function with
// 'mov (64-bit register), fs:[0x28]', then checking with `xor (64-bit register), fs:[0x28]`
return src_op.type == ZYDIS_OPERAND_TYPE_MEMORY && src_op.mem.segment == ZYDIS_REGISTER_FS &&
src_op.mem.base == ZYDIS_REGISTER_NONE && src_op.mem.index == ZYDIS_REGISTER_NONE &&
src_op.mem.disp.value == 0x28 && dst_op.reg.value >= ZYDIS_REGISTER_RAX &&
dst_op.reg.value <= ZYDIS_REGISTER_R15;
}
static void GenerateStackCheck(void* /* address */, const ZydisDecodedOperand* operands,
Xbyak::CodeGenerator& c) {
const auto dst = ZydisToXbyakRegisterOperand(operands[0]);
c.xor_(dst, 0);
}
static void GenerateStackCanary(void* /* address */, const ZydisDecodedOperand* operands,
Xbyak::CodeGenerator& c) {
const auto dst = ZydisToXbyakRegisterOperand(operands[0]);
c.mov(dst, 0);
}
static bool FilterNoSSE4a(const ZydisDecodedOperand*) { static bool FilterNoSSE4a(const ZydisDecodedOperand*) {
Cpu cpu; Cpu cpu;
return !cpu.has(Cpu::tSSE4a); return !cpu.has(Cpu::tSSE4a);
@@ -440,18 +465,26 @@ struct PatchInfo {
bool trampoline; bool trampoline;
}; };
static const std::unordered_map<ZydisMnemonic, PatchInfo> Patches = { static const std::unordered_map<ZydisMnemonic, std::vector<PatchInfo>> Patches = {
// SSE4a // SSE4a
{ZYDIS_MNEMONIC_EXTRQ, {FilterNoSSE4a, GenerateEXTRQ, true}}, {ZYDIS_MNEMONIC_EXTRQ, {{FilterNoSSE4a, GenerateEXTRQ, true}}},
{ZYDIS_MNEMONIC_INSERTQ, {FilterNoSSE4a, GenerateINSERTQ, true}}, {ZYDIS_MNEMONIC_INSERTQ, {{FilterNoSSE4a, GenerateINSERTQ, true}}},
{ZYDIS_MNEMONIC_MOVNTSS, {FilterNoSSE4a, ReplaceMOVNTSS, false}}, {ZYDIS_MNEMONIC_MOVNTSS, {{FilterNoSSE4a, ReplaceMOVNTSS, false}}},
{ZYDIS_MNEMONIC_MOVNTSD, {FilterNoSSE4a, ReplaceMOVNTSD, false}}, {ZYDIS_MNEMONIC_MOVNTSD, {{FilterNoSSE4a, ReplaceMOVNTSD, false}}},
#if !defined(__APPLE__)
// FS segment patches
// These first two patches are for accesses to the stack canary, fs:[0x28]
{ZYDIS_MNEMONIC_XOR, {{FilterStackCheck, GenerateStackCheck, false}}},
{ZYDIS_MNEMONIC_MOV,
{{FilterStackCheck, GenerateStackCanary, false},
#if defined(_WIN32) #if defined(_WIN32)
// Windows needs a trampoline. // Windows needs a trampoline for Tcb accesses.
{ZYDIS_MNEMONIC_MOV, {FilterTcbAccess, GenerateTcbAccess, true}}, {FilterTcbAccess, GenerateTcbAccess, true}
#elif !defined(__APPLE__) #else
{ZYDIS_MNEMONIC_MOV, {FilterTcbAccess, GenerateTcbAccess, false}}, {FilterTcbAccess, GenerateTcbAccess, false}
#endif
}},
#endif #endif
}; };
@@ -503,51 +536,53 @@ static std::pair<bool, u64> TryPatch(u8* code, PatchModule* module) {
} }
if (Patches.contains(instruction.mnemonic)) { if (Patches.contains(instruction.mnemonic)) {
const auto& patch_info = Patches.at(instruction.mnemonic); const auto& patches = Patches.at(instruction.mnemonic);
bool needs_trampoline = patch_info.trampoline; for (const auto& patch_info : patches) {
if (patch_info.filter(operands)) { bool needs_trampoline = patch_info.trampoline;
auto& patch_gen = module->patch_gen; if (patch_info.filter(operands)) {
auto& patch_gen = module->patch_gen;
if (needs_trampoline && instruction.length < 5) { if (needs_trampoline && instruction.length < 5) {
// Trampoline is needed but instruction is too short to patch. // Trampoline is needed but instruction is too short to patch.
// Return false and length to signal to AOT compilation that this instruction // Return false and length to signal to AOT compilation that this instruction
// should be skipped and handled at runtime. // should be skipped and handled at runtime.
return std::make_pair(false, instruction.length); return std::make_pair(false, instruction.length);
} }
// Reset state and move to current code position. // Reset state and move to current code position.
patch_gen.reset(); patch_gen.reset();
patch_gen.setSize(code - patch_gen.getCode()); patch_gen.setSize(code - patch_gen.getCode());
if (needs_trampoline) { if (needs_trampoline) {
auto& trampoline_gen = module->trampoline_gen; auto& trampoline_gen = module->trampoline_gen;
const auto trampoline_ptr = trampoline_gen.getCurr(); const auto trampoline_ptr = trampoline_gen.getCurr();
patch_info.generator(code, operands, trampoline_gen); patch_info.generator(code, operands, trampoline_gen);
// Return to the following instruction at the end of the trampoline. // Return to the following instruction at the end of the trampoline.
trampoline_gen.jmp(code + instruction.length); trampoline_gen.jmp(code + instruction.length);
// Replace instruction with near jump to the trampoline. // Replace instruction with near jump to the trampoline.
patch_gen.jmp(trampoline_ptr, Xbyak::CodeGenerator::LabelType::T_NEAR); patch_gen.jmp(trampoline_ptr, Xbyak::CodeGenerator::LabelType::T_NEAR);
} else { } else {
patch_info.generator(code, operands, patch_gen); patch_info.generator(code, operands, patch_gen);
} }
const auto patch_size = patch_gen.getCurr() - code; const auto patch_size = patch_gen.getCurr() - code;
if (patch_size > 0) { if (patch_size > 0) {
ASSERT_MSG(instruction.length >= patch_size, ASSERT_MSG(instruction.length >= patch_size,
"Instruction {} with length {} is too short to replace at: {}", "Instruction {} with length {} is too short to replace at: {}",
ZydisMnemonicGetString(instruction.mnemonic), instruction.length, ZydisMnemonicGetString(instruction.mnemonic), instruction.length,
fmt::ptr(code)); fmt::ptr(code));
// Fill remaining space with nops. // Fill remaining space with nops.
patch_gen.nop(instruction.length - patch_size); patch_gen.nop(instruction.length - patch_size);
module->patched.insert(code); module->patched.insert(code);
LOG_DEBUG(Core, "Patched instruction '{}' at: {}", LOG_DEBUG(Core, "Patched instruction '{}' at: {}",
ZydisMnemonicGetString(instruction.mnemonic), fmt::ptr(code)); ZydisMnemonicGetString(instruction.mnemonic), fmt::ptr(code));
return std::make_pair(true, instruction.length); return std::make_pair(true, instruction.length);
}
} }
} }
} }