Index: lib/Target/AArch64/AArch64FrameLowering.cpp =================================================================== --- lib/Target/AArch64/AArch64FrameLowering.cpp +++ lib/Target/AArch64/AArch64FrameLowering.cpp @@ -572,6 +572,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(); @@ -595,7 +606,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); unsigned CFIIndex = @@ -888,10 +901,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 =================================================================== --- test/CodeGen/MIR/AArch64/return-address-signing.mir +++ test/CodeGen/MIR/AArch64/return-address-signing.mir @@ -9,7 +9,14 @@ 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} @@ -58,3 +65,46 @@ 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: frame-setup CFI_INSTRUCTION negate_ra_sign_state +#CHECK: RETAB +body: | + bb.0.entry: + $w0 = MOVi32imm 2 + RET_ReallyLR implicit killed $w0 + +... +