[AArch64][PAC] Refactor aarch64-ptrauth pass (#70446)

Refactor Pointer Authentication pass in preparation for adding more
PAUTH_* pseudo instructions:
* dropped early return from runOnMachineFunction() as other PAUTH_*
  instructions need expansion even when pac-ret is disabled
* refactored runOnMachineFunction() to first collect all the
  instructions of interest without modifying anything and then performing
  changes in the later loops. There are two types of relevant
  instructions: PAUTH_* pseudos that should definitely be replaced by this
  pass and tail call instructions that may require attention if pac-ret is
  enabled
* made the loop iterating over all of the instructions handle
  instruction bundles by itself: even though this pass still does not
  support bundled TCRETURN* instructions (such as produced by KCFI) it
  does not crash anymore when no support is actually required
This commit is contained in:
Anatoly Trosinenko
2023-11-14 15:14:57 +03:00
committed by GitHub
parent 63537872ae
commit 9bc142a023

View File

@@ -297,52 +297,67 @@ bool AArch64PointerAuth::checkAuthenticatedLR(
bool AArch64PointerAuth::runOnMachineFunction(MachineFunction &MF) {
const auto *MFnI = MF.getInfo<AArch64FunctionInfo>();
if (!MFnI->shouldSignReturnAddress(true))
return false;
Subtarget = &MF.getSubtarget<AArch64Subtarget>();
TII = Subtarget->getInstrInfo();
TRI = Subtarget->getRegisterInfo();
SmallVector<MachineBasicBlock::iterator> DeletedInstrs;
SmallVector<MachineBasicBlock::iterator> TailCallInstrs;
SmallVector<MachineBasicBlock::instr_iterator> PAuthPseudoInstrs;
SmallVector<MachineBasicBlock::instr_iterator> TailCallInstrs;
bool Modified = false;
bool HasAuthenticationInstrs = false;
for (auto &MBB : MF) {
for (auto &MI : MBB) {
auto It = MI.getIterator();
// Using instr_iterator to catch unsupported bundled TCRETURN* instructions
// instead of just skipping them.
for (auto &MI : MBB.instrs()) {
switch (MI.getOpcode()) {
default:
// Bundled TCRETURN* instructions (such as created by KCFI)
// are not supported yet, but no support is required if no
// PAUTH_EPILOGUE instructions exist in the same function.
// Skip the BUNDLE instruction itself (actual bundled instructions
// follow it in the instruction list).
if (MI.isBundle())
continue;
if (AArch64InstrInfo::isTailCallReturnInst(MI))
TailCallInstrs.push_back(It);
TailCallInstrs.push_back(MI.getIterator());
break;
case AArch64::PAUTH_PROLOGUE:
signLR(MF, It);
DeletedInstrs.push_back(It);
Modified = true;
break;
case AArch64::PAUTH_EPILOGUE:
authenticateLR(MF, It);
DeletedInstrs.push_back(It);
Modified = true;
HasAuthenticationInstrs = true;
assert(!MI.isBundled());
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);
HasAuthenticationInstrs = true;
break;
default:
llvm_unreachable("Unhandled opcode");
}
It->eraseFromParent();
Modified = true;
}
// FIXME Do we need to emit any PAuth-related epilogue code at all
// when SCS is enabled?
if (HasAuthenticationInstrs &&
!MFnI->needsShadowCallStackPrologueEpilogue(MF)) {
for (auto TailCall : TailCallInstrs)
for (auto TailCall : TailCallInstrs) {
assert(!TailCall->isBundled() && "Not yet supported");
Modified |= checkAuthenticatedLR(TailCall);
}
}
for (auto MI : DeletedInstrs)
MI->eraseFromParent();
return Modified;
}