From a432f11a52dd5ec21a3438bdaa8f623e32a3234c Mon Sep 17 00:00:00 2001 From: Lang Hames Date: Thu, 21 Nov 2024 11:57:15 +1100 Subject: [PATCH] [JITLink][arm64] Support arm64e JIT'd code (initially enabled for MachO only). Adds two new JITLink passes to create and populate a pointer-signing function that can be called via an allocation-action attached to the LinkGraph: * createEmptyPointerSigningFunction creates a pointer signing function in a custome section, reserving sufficient space for the signing code. It should be run as a post-prune pass (to ensure that memory is reserved prior to allocation). * lowerPointer64AuthEdgesToSigningFunction pass populates the signing function by walking the graph, decoding the ptrauth info (encoded in the edge addend) and writing an instruction sequence to sign all ptrauth fixup locations. rdar://61956998 --- .../include/llvm/ExecutionEngine/JITLink/aarch64.h | 53 ++++ llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp | 36 ++- llvm/lib/ExecutionEngine/JITLink/aarch64.cpp | 276 +++++++++++++++++++++ .../JITLink/AArch64/MachO_ptrauth-globals.s | 158 ++++++++++++ 4 files changed, 520 insertions(+), 3 deletions(-) create mode 100644 llvm/test/ExecutionEngine/JITLink/AArch64/MachO_ptrauth-globals.s diff --git a/llvm/include/llvm/ExecutionEngine/JITLink/aarch64.h b/llvm/include/llvm/ExecutionEngine/JITLink/aarch64.h index e8c3e3414dce..db440c378d24 100644 --- a/llvm/include/llvm/ExecutionEngine/JITLink/aarch64.h +++ b/llvm/include/llvm/ExecutionEngine/JITLink/aarch64.h @@ -31,6 +31,36 @@ enum EdgeKind_aarch64 : Edge::Kind { /// Pointer64 = Edge::FirstRelocation, + /// An arm64e authenticated pointer relocation. The addend contains a 64-bit + /// struct containing the authentication parameters: + /// + /// Addend encoding: + /// int32_t addend; + /// uint16_t diversityData; + /// uint16_t hasAddressDiversity : 1; + /// uint16_t key : 2; + /// uint16_t zeroes : 12; + /// uint16_t authenticated : 1; + /// + /// Note: This means that the addend cannot be interpreted as a plain offset + /// prior to lowering. + /// + /// Authenticated pointer edges cannot be fixed up directly by JITLink as the + /// signing keys are held in the executing process. They can be removed from + /// the graph by a combination of the createEmptyPointerSigningFunction pass + /// (post-prune) and the lowerPointer64AuthEdgesToSigningFunction pass + /// (pre-fixup). Together these passes construct a signing function that will + /// be run in the executing process to write the signed pointers to the fixup + /// locations. + /// + /// Fixup expression: + /// NONE + /// + /// Errors: + /// - Failure to handle edges of this kind prior to the fixup phase will + /// result in an unsupported error during the fixup phase. + Pointer64Authenticated, + /// A plain 32-bit pointer value relocation. /// /// Fixup expression: @@ -832,6 +862,29 @@ public: Section *StubsSection = nullptr; }; +/// Returns the name of the pointer signing function section. +const char *getPointerSigningFunctionSectionName(); + +/// Creates a pointer signing function section, block, and symbol to reserve +/// space for a signing function for this LinkGraph. Clients should insert this +/// pass in the post-prune phase, and add the paired +/// lowerPointer64AuthEdgesToSigningFunction pass to the pre-fixup phase. +/// +/// No new Pointer64Auth edges can be inserted into the graph between when this +/// pass is run and when the pass below runs (since there will not be sufficient +/// space reserved in the signing function to write the signing code for them). +Error createEmptyPointerSigningFunction(LinkGraph &G); + +/// Given a LinkGraph containing Pointer64Authenticated edges, transform those +/// edges to Pointer64 and add signing code to the pointer signing function +/// (which must already have been created by the +/// createEmptyPointerSigningFunction pass above). +/// +/// This function will add a $__ptrauth_sign section with finalization-lifetime +/// containing an anonymous function that will sign all pointers in the graph. +/// An allocation action will be added to run this function during finalization. +Error lowerPointer64AuthEdgesToSigningFunction(LinkGraph &G); + } // namespace aarch64 } // namespace jitlink } // namespace llvm diff --git a/llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp b/llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp index 125c6373f82d..5607963e3774 100644 --- a/llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp +++ b/llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp @@ -28,8 +28,8 @@ class MachOLinkGraphBuilder_arm64 : public MachOLinkGraphBuilder { public: MachOLinkGraphBuilder_arm64(const object::MachOObjectFile &Obj, SubtargetFeatures Features) - : MachOLinkGraphBuilder(Obj, Triple("arm64-apple-darwin"), - std::move(Features), aarch64::getEdgeKindName), + : MachOLinkGraphBuilder(Obj, getObjectTriple(Obj), std::move(Features), + aarch64::getEdgeKindName), NumSymbols(Obj.getSymtabLoadCommand().nsyms) {} private: @@ -38,6 +38,7 @@ private: MachOPointer32, MachOPointer64, MachOPointer64Anon, + MachOPointer64Authenticated, MachOPage21, MachOPageOffset12, MachOGOTPage21, @@ -53,6 +54,18 @@ private: MachONegDelta64, }; + static Triple getObjectTriple(const object::MachOObjectFile &Obj) { + // Get the CPU sub-type from the header. + // jitLink_MachO should already have validated that the buffer is big enough + // to cover a mach_header64 so this is safe. + uint32_t CPUSubType = + *(const support::ulittle32_t *)(Obj.getData().data() + 8); + CPUSubType &= ~MachO::CPU_SUBTYPE_MASK; + if (CPUSubType == MachO::CPU_SUBTYPE_ARM64E) + return Triple("arm64e-apple-darwin"); + return Triple("arm64-apple-darwin"); + } + static Expected getRelocationKind(const MachO::relocation_info &RI) { switch (RI.r_type) { @@ -103,6 +116,10 @@ private: if (!RI.r_pcrel && !RI.r_extern && RI.r_length == 2) return MachOPairedAddend; break; + case MachO::ARM64_RELOC_AUTHENTICATED_POINTER: + if (!RI.r_pcrel && RI.r_extern && RI.r_length == 3) + return MachOPointer64Authenticated; + break; case MachO::ARM64_RELOC_TLVP_LOAD_PAGE21: if (RI.r_pcrel && RI.r_extern && RI.r_length == 2) return MachOTLVPage21; @@ -366,12 +383,15 @@ private: Kind = aarch64::Pointer32; break; case MachOPointer64: + case MachOPointer64Authenticated: if (auto TargetSymbolOrErr = findSymbolByIndex(RI.r_symbolnum)) TargetSymbol = TargetSymbolOrErr->GraphSymbol; else return TargetSymbolOrErr.takeError(); Addend = *(const ulittle64_t *)FixupContent; - Kind = aarch64::Pointer64; + Kind = *MachORelocKind == MachOPointer64 + ? aarch64::Pointer64 + : aarch64::Pointer64Authenticated; break; case MachOPointer64Anon: { orc::ExecutorAddr TargetAddress(*(const ulittle64_t *)FixupContent); @@ -493,6 +513,8 @@ private: return "MachOPointer64"; case MachOPointer64Anon: return "MachOPointer64Anon"; + case MachOPointer64Authenticated: + return "MachOPointer64Authenticated"; case MachOPage21: return "MachOPage21"; case MachOPageOffset12: @@ -601,6 +623,14 @@ void link_MachO_arm64(std::unique_ptr G, // Add an in-place GOT/Stubs pass. Config.PostPrunePasses.push_back(buildTables_MachO_arm64); + + // If this is an arm64e graph then add pointer signing passes. + if (G->getTargetTriple().isArm64e()) { + Config.PostPrunePasses.push_back( + aarch64::createEmptyPointerSigningFunction); + Config.PreFixupPasses.push_back( + aarch64::lowerPointer64AuthEdgesToSigningFunction); + } } if (auto Err = Ctx->modifyPassConfig(*G, Config)) diff --git a/llvm/lib/ExecutionEngine/JITLink/aarch64.cpp b/llvm/lib/ExecutionEngine/JITLink/aarch64.cpp index 4d3c19574a23..a79dbd5e4494 100644 --- a/llvm/lib/ExecutionEngine/JITLink/aarch64.cpp +++ b/llvm/lib/ExecutionEngine/JITLink/aarch64.cpp @@ -12,6 +12,8 @@ #include "llvm/ExecutionEngine/JITLink/aarch64.h" +#include "llvm/Support/BinaryStreamWriter.h" + #define DEBUG_TYPE "jitlink" namespace llvm { @@ -80,6 +82,280 @@ const char *getEdgeKindName(Edge::Kind R) { } } +// Write a 64-bit GPR -> GPR move. +template +static Error writeMovRegRegSeq(AppendFtor &Append, uint64_t DstReg, + uint64_t SrcReg) { + assert(DstReg < 32 && "Dst reg out of range"); + assert(SrcReg < 32 && "Src reg out of range"); + + if (DstReg == SrcReg) + return Error::success(); + + constexpr uint32_t MOVGPR64Template = 0xaa0003e0; + constexpr uint32_t DstRegIndex = 0; + constexpr uint32_t SrcRegIndex = 16; + uint32_t Instr = MOVGPR64Template; + Instr |= DstReg << DstRegIndex; + Instr |= SrcReg << SrcRegIndex; + return Append(Instr); +} + +// Generate a sequence of imm writes to assign the given value. +template +static Error writeMovRegImm64Seq(AppendFtor &Append, uint64_t Reg, + uint64_t Imm) { + assert(Reg < 32 && "Invalid register number"); + + constexpr uint32_t MovRegImm64Template = 0xd2800000; + constexpr unsigned PreserveBitIndex = 29; + constexpr unsigned ShiftBitsIndex = 21; + constexpr unsigned ImmBitsIndex = 5; + + bool PreserveRegValue = false; + for (unsigned I = 0; I != 4; ++I) { + uint32_t ImmBits = Imm & 0xffff; + Imm >>= 16; + + // Skip any all-zero immediates after the first one. + if (PreserveRegValue && !ImmBits) + continue; + + uint32_t Instr = MovRegImm64Template; + Instr |= PreserveRegValue << PreserveBitIndex; + Instr |= (I << ShiftBitsIndex); + Instr |= ImmBits << ImmBitsIndex; + Instr |= Reg; + if (auto Err = Append(Instr)) + return Err; + PreserveRegValue = true; + } + + return Error::success(); +} + +template +static Error +writePACSignSeq(AppendFtor &Append, unsigned DstReg, orc::ExecutorAddr RawAddr, + unsigned RawAddrReg, unsigned DiscriminatorReg, unsigned Key, + uint64_t EncodedDiscriminator, bool AddressDiversify) { + assert(DstReg < 32 && "DstReg out of range"); + assert(RawAddrReg < 32 && "AddrReg out of range"); + assert(DiscriminatorReg < 32 && "DiscriminatorReg out of range"); + assert(EncodedDiscriminator < 0x10000 && "EncodedDiscriminator out of range"); + + if (AddressDiversify) { + // Move the address into the discriminator register. + if (auto Err = writeMovRegRegSeq(Append, DiscriminatorReg, RawAddrReg)) + return Err; + // Blend encoded discriminator if there is one. + if (EncodedDiscriminator) { + constexpr uint32_t MOVKTemplate = 0xf2e00000; + constexpr unsigned ImmIndex = 5; + uint32_t BlendInstr = MOVKTemplate; + BlendInstr |= EncodedDiscriminator << ImmIndex; + BlendInstr |= DiscriminatorReg; + if (auto Err = Append(BlendInstr)) + return Err; + } + } else if (EncodedDiscriminator) { + // Move the encoded discriminator into the discriminator register. + if (auto Err = + writeMovRegImm64Seq(Append, DiscriminatorReg, EncodedDiscriminator)) + return Err; + } else + DiscriminatorReg = 31; // WZR + + constexpr uint32_t PACTemplate = 0xdac10000; + constexpr unsigned ZBitIndex = 13; + constexpr unsigned KeyIndex = 10; + constexpr unsigned DiscriminatorRegIndex = 5; + + uint32_t Instr = PACTemplate; + Instr |= (DiscriminatorReg == 31) << ZBitIndex; + Instr |= Key << KeyIndex; + Instr |= DiscriminatorReg << DiscriminatorRegIndex; + Instr |= DstReg; + + return Append(Instr); +} + +template +static Error writeStoreRegSeq(AppendFtor &Append, unsigned DstLocReg, + unsigned SrcReg) { + assert(DstLocReg < 32 && "DstLocReg out of range"); + assert(SrcReg < 32 && "SrcReg out of range"); + + constexpr uint32_t STRTemplate = 0xf9000000; + constexpr unsigned DstLocRegIndex = 5; + constexpr unsigned SrcRegIndex = 0; + + uint32_t Instr = STRTemplate; + Instr |= DstLocReg << DstLocRegIndex; + Instr |= SrcReg << SrcRegIndex; + + return Append(Instr); +} + +const char *getPointerSigningFunctionSectionName() { return "$__ptrauth_sign"; } + +/// Creates a pointer signing function section, block, and symbol to reserve +/// space for a signing function for this LinkGraph. Clients should insert this +/// pass in the post-prune phase, and add the paired +/// lowerPointer64AuthEdgesToSigningFunction pass to the pre-fixup phase. +Error createEmptyPointerSigningFunction(LinkGraph &G) { + LLVM_DEBUG({ + dbgs() << "Creating empty pointer signing function for " << G.getName() + << "\n"; + }); + + // FIXME: We could put a tighter bound on this if we inspected the ptrauth + // info encoded in the addend -- the only actually unknown quantity is the + // fixup location, and we can probably put constraints even on that. + size_t NumPtrAuthFixupLocations = 0; + for (auto *B : G.blocks()) + for (auto &E : B->edges()) + NumPtrAuthFixupLocations += + E.getKind() == aarch64::Pointer64Authenticated; + + constexpr size_t MaxPtrSignSeqLength = + 4 + // To materialize the value to sign. + 4 + // To materialize the fixup location. + 3 + // To copy, blend discriminator, and sign + 1; // To store the result. + + // The maximum number of signing instructions required is the maximum per + // location, times the number of locations, plus three instructions to + // materialize the return value and return. + size_t NumSigningInstrs = NumPtrAuthFixupLocations * MaxPtrSignSeqLength + 3; + + // Create signing function section. + auto &SigningSection = + G.createSection(getPointerSigningFunctionSectionName(), + orc::MemProt::Read | orc::MemProt::Exec); + SigningSection.setMemLifetime(orc::MemLifetime::Finalize); + + size_t SigningFunctionSize = NumSigningInstrs * 4; + auto &SigningFunctionBlock = G.createMutableContentBlock( + SigningSection, G.allocateBuffer(SigningFunctionSize), + orc::ExecutorAddr(), 4, 0); + G.addAnonymousSymbol(SigningFunctionBlock, 0, SigningFunctionBlock.getSize(), + true, true); + + LLVM_DEBUG({ + dbgs() << " " << NumPtrAuthFixupLocations << " location(s) to sign, up to " + << NumSigningInstrs << " instructions required (" + << formatv("{0:x}", SigningFunctionBlock.getSize()) << " bytes)\n"; + }); + + return Error::success(); +} + +/// Given a LinkGraph containing Pointer64Auth edges, transform those edges to +/// Pointer64 and add code to sign the pointers in the executor. +/// +/// This function will add a $__ptrauth_sign section with finalization-lifetime +/// containing an anonymous function that will sign all pointers in the graph. +/// An allocation action will be added to run this function during finalization. +Error lowerPointer64AuthEdgesToSigningFunction(LinkGraph &G) { + LLVM_DEBUG({ + dbgs() << "Writing pointer signing function for " << G.getName() << "\n"; + }); + + constexpr unsigned Reg1 = 8; // Holds pointer value to sign. + constexpr unsigned Reg2 = 9; // Holds fixup address. + constexpr unsigned Reg3 = 10; // Temporary for discriminator value if needed. + + // Find the signing function. + auto *SigningSection = + G.findSectionByName(getPointerSigningFunctionSectionName()); + assert(SigningSection && "Siging section missing"); + assert(SigningSection->blocks_size() == 1 && + "Unexpected number of blocks in signing section"); + assert(SigningSection->symbols_size() == 1 && + "Unexpected number of symbols in signing section"); + + auto &SigningFunctionSym = **SigningSection->symbols().begin(); + auto &SigningFunctionBlock = SigningFunctionSym.getBlock(); + auto SigningFunctionBuf = SigningFunctionBlock.getAlreadyMutableContent(); + + // Write the instructions to the block content. + BinaryStreamWriter InstrWriter( + {reinterpret_cast(SigningFunctionBuf.data()), + SigningFunctionBuf.size()}, + G.getEndianness()); + + auto AppendInstr = [&](uint32_t Instr) { + return InstrWriter.writeInteger(Instr); + }; + + for (auto *B : G.blocks()) { + for (auto EI = B->edges().begin(); EI != B->edges().end();) { + auto &E = *EI; + if (E.getKind() == aarch64::Pointer64Authenticated) { + uint64_t EncodedInfo = E.getAddend(); + int32_t RealAddend = (uint32_t)(EncodedInfo & 0xffffffff); + uint32_t InitialDiscriminator = (EncodedInfo >> 32) & 0xffff; + bool AddressDiversify = (EncodedInfo >> 48) & 0x1; + uint32_t Key = (EncodedInfo >> 49) & 0x3; + uint32_t HighBits = EncodedInfo >> 51; + auto ValueToSign = E.getTarget().getAddress() + RealAddend; + + if (HighBits != 0x1000) + return make_error( + "Pointer64Auth edge at " + + formatv("{0:x}", B->getFixupAddress(E).getValue()) + + " has invalid encoded addend " + formatv("{0:x}", EncodedInfo)); + +#ifndef NDEBUG + const char *const KeyNames[] = {"IA", "IB", "DA", "DB"}; +#endif // NDEBUG + LLVM_DEBUG({ + dbgs() << " " << B->getFixupAddress(E) << " <- " << ValueToSign + << " : key = " << KeyNames[Key] << ", discriminator = " + << formatv("{0:x4}", InitialDiscriminator) + << ", address diversified = " + << (AddressDiversify ? "yes" : "no") << "\n"; + }); + + // Materialize pointer value. + cantFail( + writeMovRegImm64Seq(AppendInstr, Reg1, ValueToSign.getValue())); + + // Materialize fixup pointer. + cantFail(writeMovRegImm64Seq(AppendInstr, Reg2, + B->getFixupAddress(E).getValue())); + + // Write signing instruction(s). + cantFail(writePACSignSeq(AppendInstr, Reg1, ValueToSign, Reg2, Reg3, + Key, InitialDiscriminator, AddressDiversify)); + + // Store signed pointer. + cantFail(writeStoreRegSeq(AppendInstr, Reg2, Reg1)); + + // Remove this edge. + EI = B->removeEdge(EI); + } else + ++EI; + } + } + + // Write epilogue. x0 = 0, x1 = 1 is an SPS serialized Error::success value. + constexpr uint32_t RETInstr = 0xd65f03c0; + cantFail(writeMovRegImm64Seq(AppendInstr, 0, 0)); // mov x0, #0 + cantFail(writeMovRegImm64Seq(AppendInstr, 1, 1)); // mov x1, #1 + cantFail(AppendInstr(RETInstr)); // ret + + // Add an allocation action to call the signing function. + using namespace orc::shared; + G.allocActions().push_back( + {cantFail(WrapperFunctionCall::Create>( + SigningFunctionSym.getAddress())), + {}}); + + return Error::success(); +} + } // namespace aarch64 } // namespace jitlink } // namespace llvm diff --git a/llvm/test/ExecutionEngine/JITLink/AArch64/MachO_ptrauth-globals.s b/llvm/test/ExecutionEngine/JITLink/AArch64/MachO_ptrauth-globals.s new file mode 100644 index 000000000000..1a4939f3a25c --- /dev/null +++ b/llvm/test/ExecutionEngine/JITLink/AArch64/MachO_ptrauth-globals.s @@ -0,0 +1,158 @@ +# RUN: llvm-mc -triple=arm64e-apple-macosx -filetype=obj -o %t.o %s +# RUN: llvm-jitlink %t.o +# +# REQUIRES: native && system-darwin +# +# Check that arm64e ptrauth relocations are handled correctly. +# +# This test contains eight global pointers with different signing schemes +# (IA vs DA key, with and without address diversity, and with 0 or 0xa5a5 as +# the additional diversity value). If all pointers pass authentication at +# runtime then the test returns zero. +# +# This test requires execution since the signed pointers are written by a +# signing function attached to the graph. +# +# TODO: Write an out-of-process version. This will probably need to be added to +# the ORC runtime. + + .section __TEXT,__text,regular,pure_instructions + .build_version macos, 13, 0 sdk_version 13, 3 + .globl _main + .p2align 2 +_main: + adrp x8, _p1@PAGE + ldr x16, [x8, _p1@PAGEOFF] + autiza x16 + + adrp x9, _p2@PAGE + add x9, x9, _p2@PAGEOFF + ldr x16, [x9] + autia x16, x9 + + adrp x10, _p3@PAGE + ldr x16, [x10, _p3@PAGEOFF] + mov x17, #23130 + autia x16, x17 + + adrp x9, _p4@PAGE + add x9, x9, _p4@PAGEOFF + ldr x16, [x9] + mov x17, x9 + movk x17, #23130, lsl #48 + autia x16, x17 + + adrp x10, _p5@PAGE + ldr x10, [x10, _p5@PAGEOFF] + ldraa x10, [x10] + + adrp x9, _p6@PAGE + add x9, x9, _p6@PAGEOFF + ldr x16, [x9] + autda x16, x9 + + adrp x10, _p7@PAGE + ldr x16, [x10, _p7@PAGEOFF] + mov x17, #23130 + autda x16, x17 + + adrp x9, _p8@PAGE + add x9, x9, _p8@PAGEOFF + ldr x16, [x9] + mov x17, x9 + movk x17, #23130, lsl #48 + autda x16, x17 + + mov w0, #0 + ret + + .private_extern _a + .section __DATA,__data + .globl _a + .p2align 3 +_a: + .quad 1 + + .private_extern _b + .globl _b + .p2align 3 +_b: + .quad 2 + + .private_extern _c + .globl _c + .p2align 3 +_c: + .quad 3 + + .private_extern _d + .globl _d + .p2align 3 +_d: + .quad 4 + + .private_extern _e + .globl _e + .p2align 3 +_e: + .quad 5 + + .private_extern _f + .globl _f + .p2align 3 +_f: + .quad 6 + + .private_extern _g + .globl _g + .p2align 3 +_g: + .quad 7 + + .private_extern _h + .globl _h + .p2align 3 +_h: + .quad 8 + + .globl _p1 + .p2align 3 +_p1: + .quad _a@AUTH(ia,0) + + .globl _p2 + .p2align 3 +_p2: + .quad _b@AUTH(ia,0,addr) + + .globl _p3 + .p2align 3 +_p3: + .quad _c@AUTH(ia,23130) + + .globl _p4 + .p2align 3 +_p4: + .quad _d@AUTH(ia,23130,addr) + + .globl _p5 + .p2align 3 +_p5: + .quad _e@AUTH(da,0) + + .globl _p6 + .p2align 3 +_p6: + .quad _f@AUTH(da,0,addr) + + .globl _p7 + .p2align 3 +_p7: + .quad _g@AUTH(da,23130) + + .globl _p8y + .p2align 3 +_p8: + .quad _h@AUTH(da,23130,addr) + +.subsections_via_symbols -- 2.11.4.GIT