Merge pull request #2670 from TNorthover/swifttail-pacibsp

arm64e: authenticate LR against the correct SP when tail calling.

Author: TNorthover

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -1597,6 +1597,132 @@ static bool isFuncletReturnInstr(const MachineInstr &MI) {
   }
 }
 
+// Check if *II is a register update that can be merged into STGloop that ends
+// at (Reg + Size). RemainingOffset is the required adjustment to Reg after the
+// end of the loop.
+static bool canMergeRegUpdate(MachineBasicBlock::iterator II, unsigned Reg,
+                       int64_t Size, int64_t *TotalOffset) {
+  MachineInstr &MI = *II;
+  if ((MI.getOpcode() == AArch64::ADDXri ||
+       MI.getOpcode() == AArch64::SUBXri) &&
+      MI.getOperand(0).getReg() == Reg && MI.getOperand(1).getReg() == Reg) {
+    unsigned Shift = AArch64_AM::getShiftValue(MI.getOperand(3).getImm());
+    int64_t Offset = MI.getOperand(2).getImm() << Shift;
+    if (MI.getOpcode() == AArch64::SUBXri)
+      Offset = -Offset;
+    int64_t AbsPostOffset = std::abs(Offset - Size);
+    const int64_t kMaxOffset =
+        0xFFF; // Max encoding for unshifted ADDXri / SUBXri
+    if (AbsPostOffset <= kMaxOffset && AbsPostOffset % 16 == 0) {
+      *TotalOffset = Offset;
+      return true;
+    }
+  }
+  return false;
+}
+
+// If we're restoring LR, authenticate it before returning.
+void AArch64FrameLowering::insertAuthLR(MachineBasicBlock &MBB,
+                                        int64_t ArgumentStackToRestore,
+                                        DebugLoc DL) const {
+  MachineFunction &MF = *MBB.getParent();
+  const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
+  MachineFrameInfo &MFI = MF.getFrameInfo();
+
+  if (!shouldAuthenticateLR(MF))
+    return;
+
+  bool IsLRVulnerable = false;
+  for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo()) {
+    if (Info.getReg() != AArch64::LR)
+      continue;
+    IsLRVulnerable = true;
+    break;
+  }
+
+  if(!IsLRVulnerable)
+    return;
+
+  if (LLVM_UNLIKELY(!Subtarget.hasPA()))
+    report_fatal_error("arm64e LR authentication requires ptrauth");
+
+  const TargetInstrInfo *TII = Subtarget.getInstrInfo();
+  MachineBasicBlock::iterator TI = MBB.getFirstTerminator();
+  if (TI != MBB.end() && TI->getOpcode() == AArch64::RET_ReallyLR &&
+      ArgumentStackToRestore == 0) {
+    // If there is a terminator and it's a RET, we can fold AUTH into it.
+    // Be careful to keep the implicitly returned registers.
+    // By now, we don't need the ReallyLR pseudo, since it's only there
+    // to make it possible for LR to be used for non-RET purposes, and
+    // that happens in RA and PEI.
+    BuildMI(MBB, TI, DL, TII->get(AArch64::RETAB))
+        .copyImplicitOps(*TI)
+        .setMIFlag(MachineInstr::FrameDestroy);
+    MBB.erase(TI);
+    return;
+  }
+
+  // Otherwise, we could be in a shrink-wrapped or tail-calling block.
+  if (ArgumentStackToRestore >= 0) {
+    // We can safely move sp to where it was on entry, execute autibsp to
+    // authenticate LR, and deallocate the extra incoming argument stack.
+    int64_t Offset = 0;
+    MachineBasicBlock::iterator FrameI =
+        TI == MBB.end() ? MBB.end() : std::prev(TI);
+    if (FrameI != MBB.end() &&
+        canMergeRegUpdate(FrameI, AArch64::SP, ArgumentStackToRestore,
+                          &Offset)) {
+      Offset -= ArgumentStackToRestore;
+      if (Offset == 0)
+        MBB.erase(FrameI);
+      else {
+        FrameI->setDesc(
+            TII->get(Offset >= 0 ? AArch64::ADDXri : AArch64::SUBXri));
+        FrameI->getOperand(2).setImm(std::abs(Offset));
+        FrameI->getOperand(3).setImm(0);
+      }
+    } else {
+      emitFrameOffset(MBB, TI, DL, AArch64::SP, AArch64::SP,
+                      StackOffset::getFixed(-ArgumentStackToRestore), TII,
+                      MachineInstr::FrameDestroy);
+    }
+    BuildMI(MBB, TI, DL, TII->get(AArch64::AUTIBSP))
+        .setMIFlag(MachineInstr::FrameDestroy);
+    emitFrameOffset(MBB, TI, DL, AArch64::SP, AArch64::SP,
+                    StackOffset::getFixed(ArgumentStackToRestore), TII,
+                    MachineInstr::FrameDestroy);
+    return;
+  }
+
+  // SP is going to be below where it was on entry, so trying to use AUTIBSP
+  // risks leaving live arguments outside the redzone. We need to find a spare
+  // register for the discriminator and execute an AUTIB instead.
+  RegScavenger RS;
+  RS.enterBasicBlockEnd(MBB);
+  RS.backward(TI);
+
+  // Prefer x16 or x17 since if we get interrupted they have better protection
+  // in the kernel.
+  Register TmpReg;
+  if (!RS.isRegUsed(AArch64::X16))
+    TmpReg = AArch64::X16;
+  else if (!RS.isRegUsed(AArch64::X17))
+    TmpReg = AArch64::X17;
+  else
+    TmpReg = RS.scavengeRegisterBackwards(AArch64::GPR64commonRegClass, TI,
+                                          false, 0, false);
+  if (TmpReg == AArch64::NoRegister)
+    report_fatal_error("unable to claim register to authenticate LR");
+
+  emitFrameOffset(MBB, TI, DL, TmpReg, AArch64::SP,
+                  StackOffset::getFixed(-ArgumentStackToRestore), TII,
+                  MachineInstr::FrameDestroy);
+  BuildMI(MBB, TI, DL, TII->get(AArch64::AUTIB), AArch64::LR)
+    .addUse(AArch64::LR)
+    .addUse(TmpReg)
+    .setMIFlag(MachineInstr::FrameDestroy);
+}
+
 void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
                                         MachineBasicBlock &MBB) const {
   MachineBasicBlock::iterator MBBI = MBB.getLastNonDebugInstr();
@@ -1623,35 +1749,12 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
   if (MF.getFunction().getCallingConv() == CallingConv::GHC)
     return;
 
-  // If we're restoring LR, authenticate it before returning.
-  // Use scope_exit to ensure we do that last on all return paths.
-  auto InsertAuthLROnExit = make_scope_exit([&]() {
-    if (shouldAuthenticateLR(MF)) {
-      if (LLVM_UNLIKELY(!Subtarget.hasPA()))
-        report_fatal_error("arm64e LR authentication requires ptrauth");
-      for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo()) {
-        if (Info.getReg() != AArch64::LR)
-          continue;
-        MachineBasicBlock::iterator TI = MBB.getFirstTerminator();
-        if (TI != MBB.end() && TI->getOpcode() == AArch64::RET_ReallyLR) {
-          // If there is a terminator and it's a RET, we can fold AUTH into it.
-          // Be careful to keep the implicitly returned registers.
-          // By now, we don't need the ReallyLR pseudo, since it's only there
-          // to make it possible for LR to be used for non-RET purposes, and
-          // that happens in RA and PEI.
-          BuildMI(MBB, TI, DL, TII->get(AArch64::RETAB)).copyImplicitOps(*TI);
-          MBB.erase(TI);
-        } else {
-          // Otherwise, we could be in a shrink-wrapped or tail-calling block.
-          BuildMI(MBB, TI, DL, TII->get(AArch64::AUTIBSP));
-        }
-      }
-    }
-  });
-
   // How much of the stack used by incoming arguments this function is expected
   // to restore in this particular epilogue.
-  int64_t ArgumentStackToRestore = getArgumentStackToRestore(MF, MBB);
+  const int64_t ArgumentStackToRestore = getArgumentStackToRestore(MF, MBB);
+
+  auto InsertAuthLROnExit =
+      make_scope_exit([&]() { insertAuthLR(MBB, ArgumentStackToRestore, DL); });
 
   // The stack frame should be like below,
   //
@@ -3184,30 +3287,6 @@ void TagStoreEdit::emitLoop(MachineBasicBlock::iterator InsertI) {
   }
 }
 
-// Check if *II is a register update that can be merged into STGloop that ends
-// at (Reg + Size). RemainingOffset is the required adjustment to Reg after the
-// end of the loop.
-bool canMergeRegUpdate(MachineBasicBlock::iterator II, unsigned Reg,
-                       int64_t Size, int64_t *TotalOffset) {
-  MachineInstr &MI = *II;
-  if ((MI.getOpcode() == AArch64::ADDXri ||
-       MI.getOpcode() == AArch64::SUBXri) &&
-      MI.getOperand(0).getReg() == Reg && MI.getOperand(1).getReg() == Reg) {
-    unsigned Shift = AArch64_AM::getShiftValue(MI.getOperand(3).getImm());
-    int64_t Offset = MI.getOperand(2).getImm() << Shift;
-    if (MI.getOpcode() == AArch64::SUBXri)
-      Offset = -Offset;
-    int64_t AbsPostOffset = std::abs(Offset - Size);
-    const int64_t kMaxOffset =
-        0xFFF; // Max encoding for unshifted ADDXri / SUBXri
-    if (AbsPostOffset <= kMaxOffset && AbsPostOffset % 16 == 0) {
-      *TotalOffset = Offset;
-      return true;
-    }
-  }
-  return false;
-}
-
 void mergeMemRefs(const SmallVectorImpl<TagStoreInstr> &TSE,
                   SmallVectorImpl<MachineMemOperand *> &MemRefs) {
   MemRefs.clear();

llvm/lib/Target/AArch64/AArch64FrameLowering.h

@@ -61,6 +61,10 @@ class AArch64FrameLowering : public TargetFrameLowering {
                               MutableArrayRef<CalleeSavedInfo> CSI,
                               const TargetRegisterInfo *TRI) const override;
 
+  void insertAuthLR(MachineBasicBlock &MBB,
+                    int64_t ArgumentStackToRestore,
+                    DebugLoc DL) const;
+
   /// Can this function use the red zone for local allocations.
   bool canUseRedZone(const MachineFunction &MF) const;
 

llvm/test/CodeGen/AArch64/swifttail-ptrauth.ll

@@ -0,0 +1,49 @@
+; RUN: llc -verify-machineinstrs < %s -mtriple=arm64e-apple-macosx | FileCheck %s
+
+declare swifttailcc void @callee_stack0()
+declare swifttailcc void @callee_stack8([8 x i64], i64)
+declare swifttailcc void @callee_stack16([8 x i64], i64, i64)
+declare extern_weak swifttailcc void @callee_weak()
+
+define swifttailcc void @caller_to0_from0() "ptrauth-returns" "frame-pointer"="all" nounwind {
+; CHECK-LABEL: caller_to0_from0:
+; CHECK: stp x29, x30, [sp, #-16]!
+; [...]
+; CHECK: ldp x29, x30, [sp], #16
+; CHECK-NEXT: autibsp
+; CHECK-NOT: add sp
+; CHECK-NOT: sub sp
+  musttail call swifttailcc void @callee_stack0()
+  ret void
+
+}
+
+define swifttailcc void @caller_to0_from8([8 x i64], i64) "ptrauth-returns" "frame-pointer"="all" {
+; CHECK-LABEL: caller_to0_from8:
+; CHECK: stp x29, x30, [sp, #-16]!
+; [...]
+; CHECK: ldp x29, x30, [sp], #16
+; CHECK-NEXT: autibsp
+; CHECK: add sp, sp, #16
+
+  musttail call swifttailcc void @callee_stack0()
+  ret void
+
+}
+
+define swifttailcc void @caller_to8_from0() "ptrauth-returns" "frame-pointer"="all" {
+; CHECK-LABEL: caller_to8_from0:
+; CHECK: stp x29, x30, [sp, #-32]!
+; [...]
+; CHECK: ldp x29, x30, [sp], #16
+; CHECK-NEXT: add x16, sp, #16
+; CHECK-NEXT: autib x30, x16
+; CHECK-NOT: add sp
+; CHECK-NOT: sub sp
+
+; Key point is that we don't move sp then autibsp because that leaves live
+; arguments below sp, potentially outside the redzone.
+  musttail call swifttailcc void @callee_stack8([8 x i64] undef, i64 42)
+  ret void
+
+}