Index: lib/Target/AArch64/AArch64FrameLowering.cpp =================================================================== --- lib/Target/AArch64/AArch64FrameLowering.cpp +++ lib/Target/AArch64/AArch64FrameLowering.cpp @@ -98,6 +98,7 @@ #include "AArch64Subtarget.h" #include "AArch64TargetMachine.h" #include "MCTargetDesc/AArch64AddressingModes.h" +#include "llvm/ADT/ScopeExit.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/Statistic.h" #include "llvm/CodeGen/LivePhysRegs.h" @@ -279,6 +280,31 @@ return MBB.erase(I); } +static bool ShouldSignReturnAddress(MachineFunction &MF) { + // The function should be signed in the following situations: + // - sign-return-address=all + // - sign-return-address=non-leaf and the functions spills the LR + + const Function &F = MF.getFunction(); + if (!F.hasFnAttribute("sign-return-address")) + return false; + + StringRef Scope = F.getFnAttribute("sign-return-address").getValueAsString(); + if (Scope.equals("none")) + return false; + + if (Scope.equals("all")) + return true; + + assert(Scope.equals("non-leaf") && "Expected all, none or non-leaf"); + + for (const auto &Info : MF.getFrameInfo().getCalleeSavedInfo()) + if (Info.getReg() == AArch64::LR) + return true; + + return false; +} + void AArch64FrameLowering::emitCalleeSavedFrameMoves( MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) const { MachineFunction &MF = *MBB.getParent(); @@ -568,6 +594,11 @@ // to determine the end of the prologue. DebugLoc DL; + if (ShouldSignReturnAddress(MF)) { + BuildMI(MBB, MBBI, DL, TII->get(AArch64::PACIASP)) + .setMIFlag(MachineInstr::FrameSetup); + } + // All calls are tail calls in GHC calling conv, and functions have no // prologue/epilogue. if (MF.getFunction().getCallingConv() == CallingConv::GHC) @@ -832,6 +863,32 @@ } } +static void InsertReturnAddressAuth(MachineFunction &MF, + MachineBasicBlock &MBB) { + if (!ShouldSignReturnAddress(MF)) + return; + const AArch64Subtarget &Subtarget = MF.getSubtarget(); + const TargetInstrInfo *TII = Subtarget.getInstrInfo(); + + MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator(); + DebugLoc DL; + if (MBBI != MBB.end()) + DL = MBBI->getDebugLoc(); + + // The AUTIASP instruction assembles to a hint instruction before v8.3a so + // this instruction can safely used for any v8a architecture. + // From v8.3a onwards there are optimised authenticate LR and return + // instructions, namely RETA{A,B}, that can be used instead. + if (Subtarget.hasV8_3aOps() && MBBI != MBB.end() && + MBBI->getOpcode() == AArch64::RET_ReallyLR) { + BuildMI(MBB, MBBI, DL, TII->get(AArch64::RETAA)).copyImplicitOps(*MBBI); + MBB.erase(MBBI); + } else { + BuildMI(MBB, MBBI, DL, TII->get(AArch64::AUTIASP)) + .setMIFlag(MachineInstr::FrameDestroy); + } +} + void AArch64FrameLowering::emitEpilogue(MachineFunction &MF, MachineBasicBlock &MBB) const { MachineBasicBlock::iterator MBBI = MBB.getLastNonDebugInstr(); @@ -899,6 +956,8 @@ // AArch64TargetLowering::LowerCall figures out ArgumentPopSize and keeps // it as the 2nd argument of AArch64ISD::TC_RETURN. + auto Cleanup = make_scope_exit([&] { InsertReturnAddressAuth(MF, MBB); }); + bool IsWin64 = Subtarget.isCallingConvWin64(MF.getFunction().getCallingConv()); unsigned FixedObject = IsWin64 ? alignTo(AFI->getVarArgsGPRSize(), 16) : 0; Index: test/CodeGen/AArch64/sign-return-address.ll =================================================================== --- /dev/null +++ test/CodeGen/AArch64/sign-return-address.ll @@ -0,0 +1,86 @@ +; RUN: llc -mtriple=aarch64-none-eabi < %s | FileCheck %s + +; CHECK-LABEL: @leaf +; CHECK-NOT: paci{{[a,b]}}sp +; CHECK-NOT: auti{{[a,b]}}sp +define i32 @leaf(i32 %x) { + ret i32 %x +} + +; CHECK-LABEL: @leaf_sign_none +; CHECK-NOT: paci{{[a,b]}}sp +; CHECK-NOT: auti{{[a,b]}}sp +define i32 @leaf_sign_none(i32 %x) "sign-return-address"="none" { + ret i32 %x +} + +; CHECK-LABEL: @leaf_sign_non_leaf +; CHECK-NOT: paci{{[a,b]}}sp +; CHECK-NOT: auti{{[a,b]}}sp +define i32 @leaf_sign_non_leaf(i32 %x) "sign-return-address"="non-leaf" { + ret i32 %x +} + +; CHECK-LABEL: @leaf_sign_all +; CHECK: paciasp +; CHECK: autiasp +; CHECK-NEXT: ret +define i32 @leaf_sign_all(i32 %x) "sign-return-address"="all" { + ret i32 %x +} + +; CHECK: @leaf_clobbers_lr +; CHECK: paciasp +; CHECK-NEXT: str x30, [sp, #-16]! +; CHECK: ldr x30, [sp], #16 +; CHECK-NEXT: autiasp +; CHECK-NEXT: ret +define i64 @leaf_clobbers_lr(i64 %x) "sign-return-address"="non-leaf" { + call void asm sideeffect "mov x30, $0", "r,~{lr}"(i64 %x) #1 + ret i64 %x +} + +declare i32 @foo(i32) + +; CHECK: @non_leaf_sign_all +; CHECK: paciasp +; CHECK: autiasp +; CHECK-NEXT: ret +define i32 @non_leaf_sign_all(i32 %x) "sign-return-address"="all" { + %call = call i32 @foo(i32 %x) + ret i32 %call +} + +; CHECK: @non_leaf_sign_non_leaf +; CHECK: paciasp +; CHECK-NEXT: str x30, [sp, #-16]! +; CHECK: ldr x30, [sp], #16 +; CHECK-NEXT: autiasp +; CHECK-NEXT: ret +define i32 @non_leaf_sign_non_leaf(i32 %x) "sign-return-address"="non-leaf" { + %call = call i32 @foo(i32 %x) + ret i32 %call +} + +; CHECK-LABEL: @leaf_sign_all_v83 +; CHECK: paciasp +; CHECK-NOT: ret +; CHECK-NEXT: retaa +; CHECK-NOT: ret +define i32 @leaf_sign_all_v83(i32 %x) "sign-return-address"="all" "target-features"="+v8.3a" { + ret i32 %x +} + +declare fastcc i64 @bar(i64) + +; CHECK-LABEL: @spill_lr_and_tail_call +; CHECK: paciasp +; CHECK-NEXT: str x30, [sp, #-16]! +; CHECK: ldr x30, [sp], #16 +; CHECK-NEXT: autiasp +; CHECK-NEXT: b bar +define fastcc void @spill_lr_and_tail_call(i64 %x) "sign-return-address"="all" { + call void asm sideeffect "mov x30, $0", "r,~{lr}"(i64 %x) #1 + tail call fastcc i64 @bar(i64 %x) + ret void +}