[AArch64] - Generate pointer authentication instructions

[LukeCheeseman]

Summary

Generate pointer authentication instructions

• The functions instrumented depend on function attribtues: all (all functions instrumentent), non-leaf (only those that spill LR), none
• Function epilogues sign the LR before spilling to the stack and authenticate the LR once restored
• If the target is v8.3a or greater than can use the combined authenticate and return instruction

[javed.absar]

lib/Target/AArch64/AArch64FrameLowering.cpp

299

Please add reason string to the assert e.g. asset(... && "expected one of none|all|partial");

[olista01]

lib/Target/AArch64/AArch64FrameLowering.cpp

286

This comment and the gcc/clang command line options use "non-leaf", but you check for "partial" in the function attribute. I think it would be better to use "non-leaf" everywhere, unless there's a good reason to switch to "partial".

293

I can't find any other places where we do case-insensitive checks of function attributes, so I think we should stick to doing an exact string-compare here.

880

I think this could do with a comment explaining the different pre-v8.3 and post-v8.3 code.

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

7

Most of the code in these tests can be removed (e.g. this one could be "ret i32 %x", or just "ret void").

69

It would be clearer to add an external function to call, rather than having the tests call each other.

113

Putting these attributes directly on the function definitions (in place of the #<n> reference) would make the tests a bit easier to read.

[eugenis]

What happens if a function with signed LR tail calls a function with unsigned LR, or vice versa?

[LukeCheeseman]

• Address reviewer points
• Add a test for when a tail call occurs in a function that signs the LR

[LukeCheeseman]

In D49791#1175839, @eugenis wrote:

What happens if a function with signed LR tail calls a function with unsigned LR, or vice versa?

I've added a test for what I think is meant here, see spill_lr_and_tail call.

The LR is authenticated before the tail call happens.

[ab]

A few minor comments inline, and one question: it sounds like this is based on a GCC feature; what's the expected behavior for __builtin_returnaddress ? Shouldn't it xpaci the result?

Also, that's one of the reasons why I'm not sure it's wise to emit the instructions always, even if we're not explicitly targeting v8.3a. You run the risk of having some parts of the code (__builtin_returnaddress, backtrace, unwind, ...) be built for v8, not knowing that they need to strip, even though some other part might have enabled this LR signing. I'm not sure there's much to be done about that, though ;)

lib/Target/AArch64/AArch64FrameLowering.cpp

297

Is this actually necessary? "all" is still only functions that actually save LR, no? Why sign it if it's not saved?

1007

One trick to make sure this covers all 'return' paths (even when the code changes) is to use 'make_scope_exit' to guarantee the code is run.

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

6

To expand on Oliver's comments, you can also remove the basic block names, and 'dso_local' (unless its semantics matter, which doesn't seem to be the case here)

Also, another nit: having the check lines away from the check-label is kinda confusing, maybe put it all before the IR?

43

The ordering of the instructions is critical here, no? Whether the PAC/AUT occur before (resp. after) SP updates (or, for that matter, the LR save) is crucial to the feature. Maybe have explicit CHECK-NEXT for the rest of the prologue/epilogue?

[LukeCheeseman]

Update patch based on comments

[LukeCheeseman]

lib/Target/AArch64/AArch64FrameLowering.cpp

297

Yes, it is necessary. "all" is all functions, even those that don't save the LR. Assume you are in a context where you have managed to gain control over the flow of execution, if you don't sign the LR in any function that doesn't save the LR then these functions becomes the ideal place to find gadgets.

[vlad.tsyrklevich]

lib/Target/AArch64/AArch64FrameLowering.cpp

289

nit: clang-format this change

[LukeCheeseman]

Ran through clang-format

[LukeCheeseman]

When the basic block has no terminator (as it will fall through to the next block), don't bail out of inserting an authentication instruction.

[LukeCheeseman]

ping

[olista01]

LGTM, thanks!

[LukeCheeseman]

This was closed by r340018 (I attached the wrong differential revision to the commit and so this was not automatically closed)

lib/Target/AArch64/AArch64FrameLowering.cpp

 
 #include "AArch64FrameLowering.h"
 #include "AArch64InstrInfo.h"
 #include "AArch64MachineFunctionInfo.h"
 #include "AArch64RegisterInfo.h"
 #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"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
if (!TFI->hasReservedCallFrame(MF)) {
     // stack, we want to add it back if we have a reserved call frame.
     assert(CalleePopAmount < 0xffffff && "call frame too large");
     emitFrameOffset(MBB, I, DL, AArch64::SP, AArch64::SP, -CalleePopAmount,
                     TII);
   }
   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();
   MachineFrameInfo &MFI = MF.getFrameInfo();
   const TargetSubtargetInfo &STI = MF.getSubtarget();
   const MCRegisterInfo *MRI = STI.getRegisterInfo();
   const TargetInstrInfo *TII = STI.getInstrInfo();
   DebugLoc DL = MBB.findDebugLoc(MBBI);
void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
   // 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))
+        .setMIFlag(MachineInstr::FrameSetup);
+  }
+
   // All calls are tail calls in GHC calling conv, and functions have no
   // prologue/epilogue.
   if (MF.getFunction().getCallingConv() == CallingConv::GHC)
     return;
 
   int NumBytes = (int)MFI.getStackSize();
   if (!AFI->hasStackFrame() && !windowsRequiresStackProbe(MF, NumBytes)) {
     assert(!HasFP && "unexpected function without stack frame but with FP");
if (needsFrameMoves) {
     }
 
     // Now emit the moves for whatever callee saved regs we have (including FP,
     // LR if those are saved).
     emitCalleeSavedFrameMoves(MBB, MBBI);
   }
 }
 
+static void InsertReturnAddressAuth(MachineFunction &MF,
+                                    MachineBasicBlock &MBB) {
+  if (!ShouldSignReturnAddress(MF))
+    return;
+  const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
+  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();
   MachineFrameInfo &MFI = MF.getFrameInfo();
   const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
   const TargetInstrInfo *TII = Subtarget.getInstrInfo();
   DebugLoc DL;
   bool IsTailCallReturn = false;
void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
   //      ----------------------        ---          ---
   //
   // So NumBytes = StackSize + BytesInStackArgArea - CalleeArgStackSize
   //             = StackSize + ArgumentPopSize
   //
   // 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;
 
   uint64_t AfterCSRPopSize = ArgumentPopSize;
   auto PrologueSaveSize = AFI->getCalleeSavedStackSize() + FixedObject;
   bool CombineSPBump = shouldCombineCSRLocalStackBump(MF, NumBytes);
   // Assume we can't combine the last pop with the sp restore.
if (!LastPopI->getFlag(MachineInstr::FrameDestroy)) {
       fixupCalleeSaveRestoreStackOffset(*LastPopI, AFI->getLocalStackSize());
   }
 
   // If there is a single SP update, insert it before the ret and we're done.
   if (CombineSPBump) {
     emitFrameOffset(MBB, MBB.getFirstTerminator(), DL, AArch64::SP, AArch64::SP,
                     NumBytes + AfterCSRPopSize, TII,
                     MachineInstr::FrameDestroy);
     return;
   }
 
   NumBytes -= PrologueSaveSize;
   assert(NumBytes >= 0 && "Negative stack allocation size!?");
 
   if (!hasFP(MF)) {
     bool RedZone = canUseRedZone(MF);
     // If this was a redzone leaf function, we don't need to restore the

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

+; 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
+}