[AArch64] - return address signing

[LukeCheeseman]

Summary

Add a command line options -msign-return-address to enable return address signing

• Armv8.3a added instructions to sign the return address to help mitigate against ROP attacks
• This patch adds command line options to generate function attributes that signal to the back whether return address signing instructions should be added

[javed.absar]

Maybe you can provide some more context to this patch (why you need this, or point to some document), if possible.

include/clang/Frontend/CodeGenOptions.h

111 ↗

(On Diff #157241)

Please conform this to the code around (i.e. each option on a separate line with comments explaining the option).

test/CodeGen/aarch64-sign-return-address.c

5 ↗

(On Diff #157241)

Can the label check not be under one 'CHECK' prefix?

[kcc]

include/clang/Frontend/CodeGenOptions.h

114 ↗

(On Diff #157266)

what's the purpose of signing LR if it is not spilled?

[LukeCheeseman]

Change codegen option partial to non-leaf to match msign-return-address scope values

[LukeCheeseman]

include/clang/Frontend/CodeGenOptions.h

114 ↗

(On Diff #157266)

Assuming you are in a context where you have managed to gain control of the flow of execution. If you don't sign functions that spill LR then those functions become good candidates for finding gadgets as now execution can start from any point in that function.

[LukeCheeseman]

Move -msign-return-address argument handling into AArch64TargetArgs

[LukeCheeseman]

ping

[olista01]

LGTM, thanks!

[hctim]

This comment was removed by hctim.

llvm/trunk/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.

llvm/trunk/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
+}