[AArch64] Return address signing B key support

[LukeCheeseman]

Summary

• Add support to generate AUTIBSP, PACIBSP, RETAB instructions for return address signing
• The key used to sign the function is controlled by the function attribute "sign-return-address-key"

[LukeCheeseman]

Added MIR test for B key support

[olista01]

LGTM

[olista01]

Hi Luke,

I've just noticed that this hasn't been committed yet, is there another patch which it depends on?

[LukeCheeseman]

• Removing CFI_negate_ra_state references as that patch hasn't gone in yet

lib/Target/AArch64/AArch64FrameLowering.cpp

static void adaptForLdStOpt(MachineBasicBlock &MBB,
   // add      sp, sp, #64
   //
   // and the load-store optimizer can merge the last two instructions into:
   //
   // ldp      x26, x25, [sp], #64
   //
 }
 
+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();
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   const Function &F = MF.getFunction();
   const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
   const AArch64RegisterInfo *RegInfo = Subtarget.getRegisterInfo();
   const TargetInstrInfo *TII = Subtarget.getInstrInfo();
   MachineModuleInfo &MMI = MF.getMMI();
   AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
   bool needsFrameMoves = MMI.hasDebugInfo() || F.needsUnwindTableEntry();
   bool HasFP = hasFP(MF);
 
   // At this point, we're going to decide whether or not the function uses a
   // redzone. In most cases, the function doesn't have a redzone so let's
   // assume that's false and set it to true in the case that there's a redzone.
   AFI->setHasRedZone(false);
 
   // Debug location must be unknown since the first debug location is used
   // to determine the end of the prologue.
   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 =
         MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
     BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
         .addCFIIndex(CFIIndex)
         .setMIFlags(MachineInstr::FrameSetup);
   }
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);
+    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);
 
     unsigned CFIIndex =
         MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
     BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
         .addCFIIndex(CFIIndex)
         .setMIFlags(MachineInstr::FrameSetup);
   }

test/CodeGen/AArch64/sign-return-address.ll

 ; CHECK: ldr  x30, [sp], #16
 ; CHECK: autiasp
 ; CHECK: 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
 }
+
+; 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
+}