Check both register operands of AUTH_TCRETURN*

Author: atrosinenko

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -157,6 +157,9 @@ class AArch64AsmPrinter : public AsmPrinter {
                                           bool ShouldTrap,
                                           const MCSymbol *OnFailure);
 
+  // Check authenticated LR before tail calling.
+  void emitPtrauthTailCallHardening(const MachineInstr *TC);
+
   // Emit the sequence for AUT or AUTPAC.
   void emitPtrauthAuthResign(const MachineInstr *MI);
 
@@ -1870,6 +1873,30 @@ void AArch64AsmPrinter::emitPtrauthCheckAuthenticatedValue(
   OutStreamer->emitLabel(SuccessSym);
 }
 
+// With Pointer Authentication, it may be needed to explicitly check the
+// authenticated value in LR before performing a tail call.
+// Otherwise, the callee may re-sign the invalid return address,
+// introducing a signing oracle.
+void AArch64AsmPrinter::emitPtrauthTailCallHardening(const MachineInstr *TC) {
+  if (!AArch64FI->shouldSignReturnAddress(*MF))
+    return;
+
+  auto LRCheckMethod = STI->getAuthenticatedLRCheckMethod(*MF);
+  if (LRCheckMethod == AArch64PAuth::AuthCheckMethod::None)
+    return;
+
+  const AArch64RegisterInfo *TRI = STI->getRegisterInfo();
+  Register ScratchReg =
+      TC->readsRegister(AArch64::X16, TRI) ? AArch64::X17 : AArch64::X16;
+  assert(!TC->readsRegister(ScratchReg, TRI) &&
+         "Neither x16 nor x17 is available as a scratch register");
+  AArch64PACKey::ID Key =
+      AArch64FI->shouldSignWithBKey() ? AArch64PACKey::IB : AArch64PACKey::IA;
+  emitPtrauthCheckAuthenticatedValue(
+      AArch64::LR, ScratchReg, Key, LRCheckMethod,
+      /*ShouldTrap=*/true, /*OnFailure=*/nullptr);
+}
+
 void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
   const bool IsAUTPAC = MI->getOpcode() == AArch64::AUTPAC;
 
@@ -2312,27 +2339,6 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
     OutStreamer->emitLabel(LOHLabel);
   }
 
-  // With Pointer Authentication, it may be needed to explicitly check the
-  // authenticated value in LR when performing a tail call.
-  // Otherwise, the callee may re-sign the invalid return address,
-  // introducing a signing oracle.
-  auto CheckLRInTailCall = [this](Register CallDestinationReg) {
-    if (!AArch64FI->shouldSignReturnAddress(*MF))
-      return;
-
-    auto LRCheckMethod = STI->getAuthenticatedLRCheckMethod(*MF);
-    if (LRCheckMethod == AArch64PAuth::AuthCheckMethod::None)
-      return;
-
-    Register ScratchReg =
-        CallDestinationReg == AArch64::X16 ? AArch64::X17 : AArch64::X16;
-    AArch64PACKey::ID Key =
-        AArch64FI->shouldSignWithBKey() ? AArch64PACKey::IB : AArch64PACKey::IA;
-    emitPtrauthCheckAuthenticatedValue(
-        AArch64::LR, ScratchReg, Key, LRCheckMethod,
-        /*ShouldTrap=*/true, /*OnFailure=*/nullptr);
-  };
-
   AArch64TargetStreamer *TS =
     static_cast<AArch64TargetStreamer *>(OutStreamer->getTargetStreamer());
   // Do any manual lowerings.
@@ -2479,7 +2485,7 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
                               ? AArch64::X17
                               : AArch64::X16;
 
-    CheckLRInTailCall(MI->getOperand(0).getReg());
+    emitPtrauthTailCallHardening(MI);
 
     unsigned DiscReg = AddrDisc;
     if (Disc) {
@@ -2511,7 +2517,7 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
   case AArch64::TCRETURNrix17:
   case AArch64::TCRETURNrinotx16:
   case AArch64::TCRETURNriALL: {
-    CheckLRInTailCall(MI->getOperand(0).getReg());
+    emitPtrauthTailCallHardening(MI);
 
     MCInst TmpInst;
     TmpInst.setOpcode(AArch64::BR);
@@ -2520,7 +2526,7 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
     return;
   }
   case AArch64::TCRETURNdi: {
-    CheckLRInTailCall(AArch64::NoRegister);
+    emitPtrauthTailCallHardening(MI);
 
     MCOperand Dest;
     MCInstLowering.lowerOperand(MI->getOperand(0), Dest);

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -1905,15 +1905,19 @@ let Predicates = [HasPAuth] in {
   // Size 16: 4 fixed + 8 variable, to compute discriminator.
   // The size returned by getInstSizeInBytes() is incremented according
   // to the variant of LR check.
+  // As the check requires either x16 or x17 as a scratch register and
+  // authenticated tail call instructions have two register operands,
+  // make sure at least one register is usable as a scratch one - for that
+  // purpose, use tcGPRnotx16x17 register class for the second operand.
   let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1, Size = 16,
       Uses = [SP] in {
     def AUTH_TCRETURN
       : Pseudo<(outs), (ins tcGPR64:$dst, i32imm:$FPDiff, i32imm:$Key,
-                            i64imm:$Disc, tcGPR64:$AddrDisc),
+                            i64imm:$Disc, tcGPRnotx16x17:$AddrDisc),
                []>, Sched<[WriteBrReg]>;
     def AUTH_TCRETURN_BTI
       : Pseudo<(outs), (ins tcGPRx16x17:$dst, i32imm:$FPDiff, i32imm:$Key,
-                            i64imm:$Disc, tcGPR64:$AddrDisc),
+                            i64imm:$Disc, tcGPRnotx16x17:$AddrDisc),
                []>, Sched<[WriteBrReg]>;
   }
 

llvm/lib/Target/AArch64/AArch64RegisterInfo.td

@@ -238,6 +238,10 @@ def tcGPR64 : RegisterClass<"AArch64", [i64], 64, (sub GPR64common, X19, X20, X2
 def tcGPRx17 : RegisterClass<"AArch64", [i64], 64, (add X17)>;
 def tcGPRx16x17 : RegisterClass<"AArch64", [i64], 64, (add X16, X17)>;
 def tcGPRnotx16 : RegisterClass<"AArch64", [i64], 64, (sub tcGPR64, X16)>;
+// LR checking code expects either x16 or x17 to be available as a scratch
+// register - for that reason restrict one of two register operands of
+// AUTH_TCRETURN* pseudos.
+def tcGPRnotx16x17 : RegisterClass<"AArch64", [i64], 64, (sub tcGPR64, X16, X17)>;
 
 // Register set that excludes registers that are reserved for procedure calls.
 // This is used for pseudo-instructions that are actually implemented using a

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

@@ -173,9 +173,9 @@ define void @test_tailcall_omit_mov_x16_x16(ptr %objptr) #0 {
 ; CHECK:         mov     x17, x0
 ; CHECK:         movk    x17, #6503, lsl #48
 ; CHECK:         autda   x16, x17
-; CHECK:         ldr     x1, [x16]
+; CHECK:         ldr     x2, [x16]
 ; CHECK:         movk    x16, #54167, lsl #48
-; CHECK:         braa    x1, x16
+; CHECK:         braa    x2, x16
   %vtable.signed = load ptr, ptr %objptr, align 8
   %objptr.int = ptrtoint ptr %objptr to i64
   %vtable.discr = tail call i64 @llvm.ptrauth.blend(i64 %objptr.int, i64 6503)