diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp --- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp +++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp @@ -258,6 +258,11 @@ cl::desc("Emit homogeneous prologue and epilogue for the size " "optimization (default = off)")); +cl::opt CheckAuthenticatedLRByLoad( + "aarch64-check-authenticated-lr-by-load", cl::Hidden, cl::init(true), + cl::desc("When performing a tail call with authenticated LR, " + "use a load instruction to check the LR")); + STATISTIC(NumRedZoneFunctions, "Number of functions using red zone"); /// Returns how much of the incoming argument stack area (in bytes) we should @@ -1946,6 +1951,106 @@ } } +// Checks authenticated LR just before the TCRETURN* instruction. +// This function may split MBB - in that case, the returned basic block +// is the new block containing the return instruction. +static MachineBasicBlock &checkAuthenticatedLR(MachineFunction &MF, + MachineBasicBlock &MBB) { + const auto &MFnI = *MF.getInfo(); + if (!MFnI.shouldSignReturnAddress(MF)) + return MBB; + + auto TI = MBB.getFirstTerminator(); + if (TI == MBB.end()) + return MBB; + + // Only explicitly check LR if we are performing tail call. + Register TmpReg = AArch64::X16; + switch (TI->getOpcode()) { + case AArch64::TCRETURNdi: + break; + case AArch64::TCRETURNri: + case AArch64::TCRETURNriBTI: + case AArch64::TCRETURNriALL: + if (TI->getOperand(0).getReg() == AArch64::X16) + TmpReg = AArch64::X17; + break; + default: + return MBB; + } + + // The following code may create a signing oracle: + // + // + // TCRETURN ; the callee may sign and spill the LR in its prologue + // + // Turn it into: + // + // + // mov tmp, lr + // xpaclri ; encoded as "hint #7" + // ; Note: at this point, the LR register contains the return address as if + // ; the authentication succeeded and the temporary register contains the + // ; *real* result of authentication. + // cmp tmp, lr + // b.ne break_block + // ret_block: + // TCRETURN + // break_block: + // brk 0xc471 + // + // or just + // + // + // ldr tmp, [lr] + // TCRETURN + + const BasicBlock *BB = MBB.getBasicBlock(); + const AArch64Subtarget &Subtarget = MF.getSubtarget(); + const TargetInstrInfo *TII = Subtarget.getInstrInfo(); + DebugLoc DL = TI->getDebugLoc(); + + if (CheckAuthenticatedLRByLoad) { + BuildMI(MBB, TI, DL, TII->get(AArch64::LDRXui), TmpReg) + .addReg(AArch64::LR) + .addImm(0) + .setMIFlags(MachineInstr::FrameDestroy); + return MBB; + } + + // Auth instruction was previously added to this basic block. + assert(TI != MBB.begin() && "Non-terminator instructions expected"); + auto &LastNonTerminator = *std::prev(TI); + MachineBasicBlock *RetBlock = MBB.splitAt(LastNonTerminator); + + MachineBasicBlock *BreakBlock = MF.CreateMachineBasicBlock(BB); + MF.push_back(BreakBlock); + MBB.splitSuccessor(RetBlock, BreakBlock); + + MachineBasicBlock *AuthBlock = &MBB; + BuildMI(AuthBlock, DL, TII->get(TargetOpcode::COPY), TmpReg) + .addReg(AArch64::LR) + .setMIFlags(MachineInstr::FrameDestroy); + BuildMI(AuthBlock, DL, TII->get(AArch64::XPACLRI)) + .setMIFlags(MachineInstr::FrameDestroy); + BuildMI(AuthBlock, DL, TII->get(AArch64::SUBSXrs), AArch64::XZR) + .addReg(TmpReg) + .addReg(AArch64::LR) + .addImm(0) + .setMIFlags(MachineInstr::FrameDestroy); + BuildMI(AuthBlock, DL, TII->get(AArch64::Bcc)) + .addImm(AArch64CC::NE) + .addMBB(BreakBlock) + .setMIFlags(MachineInstr::FrameDestroy); + assert(AuthBlock->getFallThrough() == RetBlock); + + BuildMI(BreakBlock, DL, TII->get(AArch64::BRK)) + .addImm(0xc471) + .setMIFlags(MachineInstr::FrameDestroy); + + return *RetBlock; +} + static bool isFuncletReturnInstr(const MachineInstr &MI) { switch (MI.getOpcode()) { default: @@ -1976,13 +2081,17 @@ } auto FinishingTouches = make_scope_exit([&]() { + MachineBasicBlock *RetMBB = &MBB; InsertReturnAddressAuth(MF, MBB, NeedsWinCFI, &HasWinCFI); if (needsShadowCallStackPrologueEpilogue(MF)) emitShadowCallStackEpilogue(*TII, MF, MBB, MBB.getFirstTerminator(), DL); + else + RetMBB = &checkAuthenticatedLR(MF, MBB); + // checkAuthenticatedLR() may have split MBB at this point. if (EmitCFI) - emitCalleeSavedGPRRestores(MBB, MBB.getFirstTerminator()); + emitCalleeSavedGPRRestores(*RetMBB, RetMBB->getFirstTerminator()); if (HasWinCFI) - BuildMI(MBB, MBB.getFirstTerminator(), DL, + BuildMI(*RetMBB, RetMBB->getFirstTerminator(), DL, TII->get(AArch64::SEH_EpilogEnd)) .setMIFlag(MachineInstr::FrameDestroy); }); diff --git a/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll b/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll @@ -0,0 +1,134 @@ +; RUN: llc -mtriple=aarch64 -asm-verbose=0 < %s | FileCheck -DAUT="hint #29" --check-prefixes=COMMON,LDR %s +; RUN: llc -mtriple=aarch64 -asm-verbose=0 -mattr=v8.3a < %s | FileCheck -DAUT="autiasp" --check-prefixes=COMMON,LDR %s +; RUN: llc -mtriple=aarch64 -asm-verbose=0 -aarch64-check-authenticated-lr-by-load=0 < %s | FileCheck -DAUT="hint #29" -DXPAC="hint #7" --check-prefixes=COMMON,XPAC %s +; RUN: llc -mtriple=aarch64 -asm-verbose=0 -aarch64-check-authenticated-lr-by-load=0 -mattr=v8.3a < %s | FileCheck -DAUT="autiasp" -DXPAC="xpaclri" --check-prefixes=COMMON,XPAC %s + +define i32 @tailcall_direct() "sign-return-address"="non-leaf" { +; COMMON-LABEL: tailcall_direct: +; COMMON: str x30, [sp, #-16]! +; COMMON: ldr x30, [sp], #16 +; +; LDR-NEXT: [[AUT]] +; LDR-NEXT: .cfi_negate_ra_state +; LDR-NEXT: ldr x16, [x30] +; LDR-NEXT: b callee +; +; XPAC-NEXT: [[AUT]] +; XPAC-NEXT: .cfi_negate_ra_state +; XPAC-NEXT: mov x16, x30 +; XPAC-NEXT: [[XPAC]] +; XPAC-NEXT: cmp x16, x30 +; XPAC-NEXT: b.ne .[[FAIL:LBB[_0-9]+]] +; XPAC-NEXT: b callee +; XPAC-NEXT: .[[FAIL]]: +; XPAC-NEXT: brk #0xc471 + tail call void asm sideeffect "", "~{lr}"() + %call = tail call i32 @callee() + ret i32 %call +} + +define i32 @tailcall_indirect(ptr %fptr) "sign-return-address"="non-leaf" { +; COMMON-LABEL: tailcall_indirect: +; COMMON: str x30, [sp, #-16]! +; COMMON: ldr x30, [sp], #16 +; +; LDR-NEXT: [[AUT]] +; LDR-NEXT: .cfi_negate_ra_state +; LDR-NEXT: ldr x16, [x30] +; LDR-NEXT: br x0 +; +; XPAC-NEXT: [[AUT]] +; XPAC-NEXT: .cfi_negate_ra_state +; XPAC-NEXT: mov x16, x30 +; XPAC-NEXT: [[XPAC]] +; XPAC-NEXT: cmp x16, x30 +; XPAC-NEXT: b.ne .[[FAIL:LBB[_0-9]+]] +; XPAC-NEXT: br x0 +; XPAC-NEXT: .[[FAIL]]: +; XPAC-NEXT: brk #0xc471 + tail call void asm sideeffect "", "~{lr}"() + %call = tail call i32 %fptr() + ret i32 %call +} + +define i32 @tailcall_direct_noframe() "sign-return-address"="non-leaf" { +; COMMON-LABEL: tailcall_direct_noframe: +; COMMON-NEXT: .cfi_startproc +; COMMON-NEXT: b callee + %call = tail call i32 @callee() + ret i32 %call +} + +define i32 @tailcall_indirect_noframe(ptr %fptr) "sign-return-address"="non-leaf" { +; COMMON-LABEL: tailcall_indirect_noframe: +; COMMON-NEXT: .cfi_startproc +; COMMON-NEXT: br x0 + %call = tail call i32 %fptr() + ret i32 %call +} + +define i32 @tailcall_direct_noframe_sign_all() "sign-return-address"="all" { +; COMMON-LABEL: tailcall_direct_noframe_sign_all: +; COMMON-NOT: str{{.*}}x30 +; COMMON-NOT: ldr{{.*}}x30 +; +; LDR: [[AUT]] +; LDR-NEXT: .cfi_negate_ra_state +; LDR-NEXT: ldr x16, [x30] +; LDR-NEXT: b callee +; +; XPAC: [[AUT]] +; XPAC-NEXT: .cfi_negate_ra_state +; XPAC-NEXT: mov x16, x30 +; XPAC-NEXT: [[XPAC]] +; XPAC-NEXT: cmp x16, x30 +; XPAC-NEXT: b.ne .[[FAIL:LBB[_0-9]+]] +; XPAC-NEXT: b callee +; XPAC-NEXT: .[[FAIL]]: +; XPAC-NEXT: brk #0xc471 + %call = tail call i32 @callee() + ret i32 %call +} + +define i32 @tailcall_indirect_noframe_sign_all(ptr %fptr) "sign-return-address"="all" { +; COMMON-LABEL: tailcall_indirect_noframe_sign_all: +; COMMON-NOT: str{{.*}}x30 +; COMMON-NOT: ldr{{.*}}x30 +; +; LDR: [[AUT]] +; LDR-NEXT: .cfi_negate_ra_state +; LDR-NEXT: ldr x16, [x30] +; LDR-NEXT: br x0 +; +; XPAC: [[AUT]] +; XPAC-NEXT: .cfi_negate_ra_state +; XPAC-NEXT: mov x16, x30 +; XPAC-NEXT: [[XPAC]] +; XPAC-NEXT: cmp x16, x30 +; XPAC-NEXT: b.ne .[[FAIL:LBB[_0-9]+]] +; XPAC-NEXT: br x0 +; XPAC-NEXT: .[[FAIL]]: +; XPAC-NEXT: brk #0xc471 + %call = tail call i32 %fptr() + ret i32 %call +} + +; Do not emit any LR checks when Shadow Call Stack is enabled +define i32 @tailcall_scs(ptr %fptr) "sign-return-address"="all" shadowcallstack "target-features"="+reserve-x18" { +; COMMON-LABEL: tailcall_scs: +; COMMON: str x30, [sp, #-16]! +; COMMON: ldr x30, [sp], #16 +; +; COMMON-NOT: ldr {{.*}}, [x30] +; COMMON-NOT: xpac +; COMMON-NOT: hint #7 +; COMMON-NOT: brk +; +; Match the end of function: +; COMMON: .size tailcall_scs, + tail call void asm sideeffect "", "~{lr}"() + %call = tail call i32 %fptr() + ret i32 %call +} + +declare i32 @callee() diff --git a/llvm/test/CodeGen/AArch64/sign-return-address.ll b/llvm/test/CodeGen/AArch64/sign-return-address.ll --- a/llvm/test/CodeGen/AArch64/sign-return-address.ll +++ b/llvm/test/CodeGen/AArch64/sign-return-address.ll @@ -176,6 +176,7 @@ ; COMPAT-NEXT: ldr x30, [sp], #16 // 8-byte Folded Reload ; COMPAT-NEXT: hint #29 ; COMPAT-NEXT: .cfi_negate_ra_state +; COMPAT-NEXT: ldr x16, [x30] ; COMPAT-NEXT: b bar ; ; V83A-LABEL: spill_lr_and_tail_call: @@ -191,6 +192,7 @@ ; V83A-NEXT: ldr x30, [sp], #16 // 8-byte Folded Reload ; V83A-NEXT: autiasp ; V83A-NEXT: .cfi_negate_ra_state +; V83A-NEXT: ldr x16, [x30] ; V83A-NEXT: b bar call void asm sideeffect "mov x30, $0", "r,~{lr}"(i64 %x) #1 tail call fastcc i64 @bar(i64 %x)