Index: lib/Target/AArch64/AArch64FrameLowering.cpp =================================================================== --- lib/Target/AArch64/AArch64FrameLowering.cpp +++ lib/Target/AArch64/AArch64FrameLowering.cpp @@ -597,6 +597,17 @@ // } +static bool ShouldSignWithAKey(MachineFunction &MF) { + const Function &F = MF.getFunction(); + if (!F.hasFnAttribute("sign-return-address-key")) + return true; + + const StringRef Key = + F.getFnAttribute("sign-return-address-key").getValueAsString(); + assert(Key.equals_lower("a_key") || Key.equals_lower("b_key")); + return Key.equals_lower("a_key"); +} + void AArch64FrameLowering::emitPrologue(MachineFunction &MF, MachineBasicBlock &MBB) const { MachineBasicBlock::iterator MBBI = MBB.begin(); @@ -620,7 +631,9 @@ DebugLoc DL; if (ShouldSignReturnAddress(MF)) { - BuildMI(MBB, MBBI, DL, TII->get(AArch64::PACIASP)) + BuildMI( + MBB, MBBI, DL, + TII->get(ShouldSignWithAKey(MF) ? AArch64::PACIASP : AArch64::PACIBSP)) .setMIFlag(MachineInstr::FrameSetup); } @@ -907,10 +920,14 @@ // 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); + BuildMI(MBB, MBBI, DL, + TII->get(ShouldSignWithAKey(MF) ? AArch64::RETAA : AArch64::RETAB)) + .copyImplicitOps(*MBBI); MBB.erase(MBBI); } else { - BuildMI(MBB, MBBI, DL, TII->get(AArch64::AUTIASP)) + BuildMI( + MBB, MBBI, DL, + TII->get(ShouldSignWithAKey(MF) ? AArch64::AUTIASP : AArch64::AUTIBSP)) .setMIFlag(MachineInstr::FrameDestroy); } } Index: test/CodeGen/AArch64/sign-return-address.ll =================================================================== --- test/CodeGen/AArch64/sign-return-address.ll +++ test/CodeGen/AArch64/sign-return-address.ll @@ -84,3 +84,26 @@ tail call fastcc i64 @bar(i64 %x) ret void } + +; CHECK-LABEL: @leaf_sign_all_a_key +; CHECK: paciasp +; CHECK: autiasp +define i32 @leaf_sign_all_a_key(i32 %x) "sign-return-address"="all" "sign-return-address-key"="a_key" { + ret i32 %x +} + +; CHECK-LABEL: @leaf_sign_all_b_key +; CHECK: pacibsp +; CHECK: autibsp +define i32 @leaf_sign_all_b_key(i32 %x) "sign-return-address"="all" "sign-return-address-key"="b_key" { + ret i32 %x +} + +; CHECK-LABEL: @leaf_sign_all_v83_b_key +; CHECK: pacibsp +; CHECK-NOT: ret +; CHECK: retab +; CHECK-NOT: ret +define i32 @leaf_sign_all_v83_b_key(i32 %x) "sign-return-address"="all" "target-features"="+v8.3a" "sign-return-address-key"="b_key" { + ret i32 %x +} Index: test/CodeGen/MIR/AArch64/return-address-signing.mir =================================================================== --- /dev/null +++ test/CodeGen/MIR/AArch64/return-address-signing.mir @@ -0,0 +1,109 @@ + +# RUN: llc -mtriple=aarch64-arm-none-eabi -run-pass=prologepilog -o - %s 2>&1 | FileCheck %s +--- | + target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128" + target triple = "aarch64-arm-none-eabi" + + ; Function Attrs: noinline nounwind optnone + define dso_local i32 @foo() #0 { + entry: + ret i32 2 + } + + ; Function Attrs: noinline nounwind optnone + define dso_local i32 @bar() #1 { + entry: + ret i32 2 + } + + attributes #0 = { noinline nounwind optnone "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "less-precise-fpmad"="false" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "sign-return-address"="all" "sign-return-address-key"="a_key" "stack-protector-buffer-size"="8" "target-cpu"="generic" "target-features"="+neon,+v8.3a" "unsafe-fp-math"="false" "use-soft-float"="false" } + attributes #1 = { noinline nounwind optnone "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "less-precise-fpmad"="false" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "sign-return-address"="all" "sign-return-address-key"="b_key" "stack-protector-buffer-size"="8" "target-cpu"="generic" "target-features"="+neon,+v8.3a" "unsafe-fp-math"="false" "use-soft-float"="false" } + + !llvm.module.flags = !{!0} + + !0 = !{i32 1, !"wchar_size", i32 4} + +... +--- +#CHECK: foo +name: foo +alignment: 2 +exposesReturnsTwice: false +legalized: false +regBankSelected: false +selected: false +failedISel: false +tracksRegLiveness: true +registers: +liveins: +frameInfo: + isFrameAddressTaken: false + isReturnAddressTaken: false + hasStackMap: false + hasPatchPoint: false + stackSize: 0 + offsetAdjustment: 0 + maxAlignment: 0 + adjustsStack: false + hasCalls: false + stackProtector: '' + maxCallFrameSize: 0 + hasOpaqueSPAdjustment: false + hasVAStart: false + hasMustTailInVarArgFunc: false + localFrameSize: 0 + savePoint: '' + restorePoint: '' +fixedStack: +stack: +constants: +#CHECK: frame-setup PACIASP implicit-def $lr, implicit $lr, implicit $sp +#CHECK: RETAA +body: | + bb.0.entry: + $w0 = MOVi32imm 2 + RET_ReallyLR implicit killed $w0 + +... +--- +#CHECK: bar +name: bar +alignment: 2 +exposesReturnsTwice: false +legalized: false +regBankSelected: false +selected: false +failedISel: false +tracksRegLiveness: true +registers: +liveins: +frameInfo: + isFrameAddressTaken: false + isReturnAddressTaken: false + hasStackMap: false + hasPatchPoint: false + stackSize: 0 + offsetAdjustment: 0 + maxAlignment: 0 + adjustsStack: false + hasCalls: false + stackProtector: '' + maxCallFrameSize: 0 + hasOpaqueSPAdjustment: false + hasVAStart: false + hasMustTailInVarArgFunc: false + localFrameSize: 0 + savePoint: '' + restorePoint: '' +fixedStack: +stack: +constants: +#CHECK: frame-setup PACIBSP implicit-def $lr, implicit $lr, implicit $sp +#CHECK: RETAB +body: | + bb.0.entry: + $w0 = MOVi32imm 2 + RET_ReallyLR implicit killed $w0 + +... +