TheRock/compiler/amd-llvm/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp

//===-- AArch64PointerAuth.cpp -- Harden code using PAuth ------------------==//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "AArch64PointerAuth.h"

#include "AArch64.h"
#include "AArch64FrameLowering.h"
#include "AArch64InstrInfo.h"
#include "AArch64MachineFunctionInfo.h"
#include "AArch64Subtarget.h"
#include "llvm/CodeGen/CFIInstBuilder.h"
#include "llvm/CodeGen/MachineBasicBlock.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/MachineModuleInfo.h"

using namespace llvm;
using namespace llvm::AArch64PAuth;

#define AARCH64_POINTER_AUTH_NAME "AArch64 Pointer Authentication"

namespace {

class AArch64PointerAuthImpl {
public:
  bool run(MachineFunction &MF);

private:
  const AArch64Subtarget *Subtarget = nullptr;
  const AArch64InstrInfo *TII = nullptr;

  void signLR(MachineFunction &MF, MachineBasicBlock::iterator MBBI) const;

  void authenticateLR(MachineFunction &MF,
                      MachineBasicBlock::iterator MBBI) const;
};

class AArch64PointerAuthLegacy : public MachineFunctionPass {
public:
  static char ID;

  AArch64PointerAuthLegacy() : MachineFunctionPass(ID) {}

  bool runOnMachineFunction(MachineFunction &MF) override;

  StringRef getPassName() const override { return AARCH64_POINTER_AUTH_NAME; }
};

} // end anonymous namespace

INITIALIZE_PASS(AArch64PointerAuthLegacy, "aarch64-ptrauth",
                AARCH64_POINTER_AUTH_NAME, false, false)

FunctionPass *llvm::createAArch64PointerAuthPass() {
  return new AArch64PointerAuthLegacy();
}

char AArch64PointerAuthLegacy::ID = 0;

static void emitPACSymOffsetIntoReg(const TargetInstrInfo &TII,
                                    MachineBasicBlock &MBB,
                                    MachineBasicBlock::iterator I, DebugLoc DL,
                                    MCSymbol *PACSym, Register Reg) {
  BuildMI(MBB, I, DL, TII.get(AArch64::ADRP), Reg)
      .addSym(PACSym, AArch64II::MO_PAGE);
  BuildMI(MBB, I, DL, TII.get(AArch64::ADDXri), Reg)
      .addReg(Reg)
      .addSym(PACSym, AArch64II::MO_PAGEOFF | AArch64II::MO_NC)
      .addImm(0);
}

static void emitPACCFI(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
                       MachineInstr::MIFlag Flags, bool EmitCFI) {
  if (!EmitCFI)
    return;

  auto &MF = *MBB.getParent();
  auto &MFnI = *MF.getInfo<AArch64FunctionInfo>();

  CFIInstBuilder CFIBuilder(MBB, MBBI, Flags);
  MFnI.branchProtectionPAuthLR() ? CFIBuilder.buildNegateRAStateWithPC()
                                 : CFIBuilder.buildNegateRAState();
}

void AArch64PointerAuthImpl::signLR(MachineFunction &MF,
                                    MachineBasicBlock::iterator MBBI) const {
  auto &MFnI = *MF.getInfo<AArch64FunctionInfo>();
  bool UseBKey = MFnI.shouldSignWithBKey();
  bool EmitCFI = MFnI.needsDwarfUnwindInfo(MF);
  bool NeedsWinCFI = MF.hasWinCFI();

  MachineBasicBlock &MBB = *MBBI->getParent();

  // Debug location must be unknown, see AArch64FrameLowering::emitPrologue.
  DebugLoc DL;

  if (UseBKey && !MF.getTarget().getTargetTriple().isOSBinFormatMachO()) {
    BuildMI(MBB, MBBI, DL, TII->get(AArch64::EMITBKEY))
        .setMIFlag(MachineInstr::FrameSetup);
  }

  // PAuthLR authentication instructions need to know the value of PC at the
  // point of signing (PACI*).
  if (MFnI.branchProtectionPAuthLR()) {
    MCSymbol *PACSym = MF.getContext().createTempSymbol();
    MFnI.setSigningInstrLabel(PACSym);
  }

  // No SEH opcode for this one; it doesn't materialize into an
  // instruction on Windows.
  if (MFnI.branchProtectionPAuthLR() && Subtarget->hasPAuthLR()) {
    emitPACCFI(MBB, MBBI, MachineInstr::FrameSetup, EmitCFI);
    BuildMI(MBB, MBBI, DL,
            TII->get(UseBKey ? AArch64::PACIBSPPC : AArch64::PACIASPPC))
        .setMIFlag(MachineInstr::FrameSetup)
        ->setPreInstrSymbol(MF, MFnI.getSigningInstrLabel());
  } else {
    if (MFnI.branchProtectionPAuthLR()) {
      BuildMI(MBB, MBBI, DL, TII->get(AArch64::PACM))
          .setMIFlag(MachineInstr::FrameSetup);
      emitPACCFI(MBB, MBBI, MachineInstr::FrameSetup, EmitCFI);
    }
    BuildMI(MBB, MBBI, DL,
            TII->get(UseBKey ? AArch64::PACIBSP : AArch64::PACIASP))
        .setMIFlag(MachineInstr::FrameSetup)
        ->setPreInstrSymbol(MF, MFnI.getSigningInstrLabel());
    if (!MFnI.branchProtectionPAuthLR())
      emitPACCFI(MBB, MBBI, MachineInstr::FrameSetup, EmitCFI);
  }

  if (!EmitCFI && NeedsWinCFI) {
    BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_PACSignLR))
        .setMIFlag(MachineInstr::FrameSetup);
  }
}

void AArch64PointerAuthImpl::authenticateLR(
    MachineFunction &MF, MachineBasicBlock::iterator MBBI) const {
  const AArch64FunctionInfo *MFnI = MF.getInfo<AArch64FunctionInfo>();
  bool UseBKey = MFnI->shouldSignWithBKey();
  bool EmitAsyncCFI = MFnI->needsAsyncDwarfUnwindInfo(MF);
  bool NeedsWinCFI = MF.hasWinCFI();

  MachineBasicBlock &MBB = *MBBI->getParent();
  DebugLoc DL = MBBI->getDebugLoc();
  // MBBI points to a PAUTH_EPILOGUE instruction to be replaced and
  // TI points to a terminator instruction that may or may not be combined.
  // Note that inserting new instructions "before MBBI" and "before TI" is
  // not the same because if ShadowCallStack is enabled, its instructions
  // are placed between MBBI and TI.
  MachineBasicBlock::iterator TI = MBB.getFirstInstrTerminator();

  // 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. In this case the
  // DW_CFA_AARCH64_negate_ra_state can't be emitted.
  bool TerminatorIsCombinable =
      TI != MBB.end() && TI->getOpcode() == AArch64::RET;
  MCSymbol *PACSym = MFnI->getSigningInstrLabel();

  if (Subtarget->hasPAuth() && TerminatorIsCombinable && !NeedsWinCFI &&
      !MF.getFunction().hasFnAttribute(Attribute::ShadowCallStack)) {
    if (MFnI->branchProtectionPAuthLR() && Subtarget->hasPAuthLR()) {
      assert(PACSym && "No PAC instruction to refer to");
      BuildMI(MBB, TI, DL,
              TII->get(UseBKey ? AArch64::RETABSPPCi : AArch64::RETAASPPCi))
          .addSym(PACSym)
          .copyImplicitOps(*MBBI)
          .setMIFlag(MachineInstr::FrameDestroy);
    } else {
      if (MFnI->branchProtectionPAuthLR()) {
        emitPACSymOffsetIntoReg(*TII, MBB, MBBI, DL, PACSym, AArch64::X16);
        BuildMI(MBB, MBBI, DL, TII->get(AArch64::PACM))
            .setMIFlag(MachineInstr::FrameDestroy);
      }
      BuildMI(MBB, TI, DL, TII->get(UseBKey ? AArch64::RETAB : AArch64::RETAA))
          .copyImplicitOps(*MBBI)
          .setMIFlag(MachineInstr::FrameDestroy);
    }
    MBB.erase(TI);
    return;
  }

  // When FPDiff != 0 (tail call with callee-popped stack arg space), SP has
  // been adjusted and no longer matches the entry SP used as the signing
  // modifier. We must reconstruct entry SP for authentication.
  auto &AFL = *static_cast<const AArch64FrameLowering *>(
      MF.getSubtarget().getFrameLowering());
  if (int64_t FPDiff = AFL.getArgumentStackToRestore(MF, MBB)) {
    // Use AUTI[AB]1716 variants: x17=LR, x16=entry_SP.
    BuildMI(MBB, MBBI, DL, TII->get(AArch64::ORRXrs), AArch64::X17)
        .addReg(AArch64::XZR)
        .addReg(AArch64::LR)
        .addImm(0)
        .setMIFlag(MachineInstr::FrameDestroy);
    emitFrameOffset(MBB, MBBI, DL, AArch64::X16, AArch64::SP,
                    StackOffset::getFixed(-FPDiff), TII,
                    MachineInstr::FrameDestroy);

    if (MFnI->branchProtectionPAuthLR() && Subtarget->hasPAuthLR()) {
      assert(PACSym && "No PAC instruction to refer to");
      emitPACSymOffsetIntoReg(*TII, MBB, MBBI, DL, PACSym, AArch64::X15);

      emitPACCFI(MBB, MBBI, MachineInstr::FrameDestroy, EmitAsyncCFI);
      unsigned AutOpc = UseBKey ? AArch64::AUTIB171615 : AArch64::AUTIA171615;
      BuildMI(MBB, MBBI, DL, TII->get(AutOpc))
          .setMIFlag(MachineInstr::FrameDestroy);
    } else if (MFnI->branchProtectionPAuthLR()) {
      assert(PACSym && "No PAC instruction to refer to");
      emitPACSymOffsetIntoReg(*TII, MBB, MBBI, DL, PACSym, AArch64::X15);

      // The PACM hint-space instruction modifies the following AUTI[AB]1716
      // to optionally take x15 as an extra operand depending on the
      // presence of +pauth-lr at runtime. On machines without +pauth-lr, it
      // behaves as a nop, and the address of the PACI[AB]SP in x15 is
      // ignored.
      BuildMI(MBB, MBBI, DL, TII->get(AArch64::PACM))
          .setMIFlag(MachineInstr::FrameDestroy);

      emitPACCFI(MBB, MBBI, MachineInstr::FrameDestroy, EmitAsyncCFI);
      unsigned AutOpc = UseBKey ? AArch64::AUTIB1716 : AArch64::AUTIA1716;
      BuildMI(MBB, MBBI, DL, TII->get(AutOpc))
          .setMIFlag(MachineInstr::FrameDestroy);
    } else {
      unsigned AutOpc = UseBKey ? AArch64::AUTIB1716 : AArch64::AUTIA1716;
      BuildMI(MBB, MBBI, DL, TII->get(AutOpc))
          .setMIFlag(MachineInstr::FrameDestroy);
      emitPACCFI(MBB, MBBI, MachineInstr::FrameDestroy, EmitAsyncCFI);
    }

    BuildMI(MBB, MBBI, DL, TII->get(AArch64::ORRXrs), AArch64::LR)
        .addReg(AArch64::XZR)
        .addReg(AArch64::X17)
        .addImm(0)
        .setMIFlag(MachineInstr::FrameDestroy);
    return;
  }

  if (MFnI->branchProtectionPAuthLR() && Subtarget->hasPAuthLR()) {
    assert(PACSym && "No PAC instruction to refer to");
    emitPACCFI(MBB, MBBI, MachineInstr::FrameDestroy, EmitAsyncCFI);
    BuildMI(MBB, MBBI, DL,
            TII->get(UseBKey ? AArch64::AUTIBSPPCi : AArch64::AUTIASPPCi))
        .addSym(PACSym)
        .setMIFlag(MachineInstr::FrameDestroy);
  } else {
    if (MFnI->branchProtectionPAuthLR()) {
      emitPACSymOffsetIntoReg(*TII, MBB, MBBI, DL, PACSym, AArch64::X16);
      BuildMI(MBB, MBBI, DL, TII->get(AArch64::PACM))
          .setMIFlag(MachineInstr::FrameDestroy);
      emitPACCFI(MBB, MBBI, MachineInstr::FrameDestroy, EmitAsyncCFI);
    }
    BuildMI(MBB, MBBI, DL,
            TII->get(UseBKey ? AArch64::AUTIBSP : AArch64::AUTIASP))
        .setMIFlag(MachineInstr::FrameDestroy);
    if (!MFnI->branchProtectionPAuthLR())
      emitPACCFI(MBB, MBBI, MachineInstr::FrameDestroy, EmitAsyncCFI);
  }

  if (NeedsWinCFI) {
    BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_PACSignLR))
        .setMIFlag(MachineInstr::FrameDestroy);
  }
}

unsigned llvm::AArch64PAuth::getCheckerSizeInBytes(AuthCheckMethod Method) {
  switch (Method) {
  case AuthCheckMethod::None:
    return 0;
  case AuthCheckMethod::DummyLoad:
    return 4;
  case AuthCheckMethod::HighBitsNoTBI:
    return 12;
  case AuthCheckMethod::XPACHint:
  case AuthCheckMethod::XPAC:
    return 20;
  }
  llvm_unreachable("Unknown AuthCheckMethod enum");
}

bool AArch64PointerAuthImpl::run(MachineFunction &MF) {
  Subtarget = &MF.getSubtarget<AArch64Subtarget>();
  TII = Subtarget->getInstrInfo();

  SmallVector<MachineBasicBlock::instr_iterator> PAuthPseudoInstrs;

  bool Modified = false;

  for (auto &MBB : MF) {
    for (auto &MI : MBB) {
      switch (MI.getOpcode()) {
      default:
        break;
      case AArch64::PAUTH_PROLOGUE:
      case AArch64::PAUTH_EPILOGUE:
        PAuthPseudoInstrs.push_back(MI.getIterator());
        break;
      }
    }
  }

  for (auto It : PAuthPseudoInstrs) {
    switch (It->getOpcode()) {
    case AArch64::PAUTH_PROLOGUE:
      signLR(MF, It);
      break;
    case AArch64::PAUTH_EPILOGUE:
      authenticateLR(MF, It);
      break;
    default:
      llvm_unreachable("Unhandled opcode");
    }
    It->eraseFromParent();
    Modified = true;
  }

  return Modified;
}

bool AArch64PointerAuthLegacy::runOnMachineFunction(MachineFunction &MF) {
  return AArch64PointerAuthImpl().run(MF);
}

PreservedAnalyses
AArch64PointerAuthPass::run(MachineFunction &MF,
                            MachineFunctionAnalysisManager &MFAM) {
  const bool Changed = AArch64PointerAuthImpl().run(MF);
  if (!Changed)
    return PreservedAnalyses::all();
  PreservedAnalyses PA = getMachineFunctionPassPreservedAnalyses();
  PA.preserveSet<CFGAnalyses>();
  return PA;
}