[CodeGen][AArch64] ptrauth intrinsic to safely construct relative pointer for swift coroutines

AbhayKanhere · AbhayKanhere · commit 2b431c0cd3ad · 2025-07-08T14:11:27.000-07:00
A ptrauth intrinsic for swift co-routine support that allows creation of signed pointer
 from offset stored at address relative to the pointer.

 Following C-like pseudo code (ignoring keys,discriminators) explains its operation:
      let rawptr = PACauth(inputptr);
      return PACsign(  rawptr +  *(int32*)(rawptr+addend) )

 What: Authenticate a signed pointer, load a 32bit value at offset 'addend' from pointer,
       add this value to pointer, sign this new pointer.
 builtin: __builtin_ptrauth_auth_load_relative_and_sign
 intrinsic: ptrauth_auth_resign_load_relative
diff --git a/clang/include/clang/Basic/Builtins.td b/clang/include/clang/Basic/Builtins.td
@@ -4643,6 +4643,12 @@ def PtrauthAuthAndResign : Builtin {
   let Prototype = "void*(void*,int,void*,int,void*)";
 }
 
+def PtrauthAuthLoadRelativeAndSign : Builtin {
+  let Spellings = ["__builtin_ptrauth_auth_load_relative_and_sign"];
+  let Attributes = [CustomTypeChecking, NoThrow];
+  let Prototype = "void*(void*,int,void*,int,void*,ptrdiff_t)";
+}
+
 def PtrauthAuth : Builtin {
   let Spellings = ["__builtin_ptrauth_auth"];
   let Attributes = [CustomTypeChecking, NoThrow];
diff --git a/clang/lib/CodeGen/CGBuiltin.cpp b/clang/lib/CodeGen/CGBuiltin.cpp
@@ -5584,12 +5584,13 @@ RValue CodeGenFunction::EmitBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
 
   case Builtin::BI__builtin_ptrauth_auth:
   case Builtin::BI__builtin_ptrauth_auth_and_resign:
+  case Builtin::BI__builtin_ptrauth_auth_load_relative_and_sign:
   case Builtin::BI__builtin_ptrauth_blend_discriminator:
   case Builtin::BI__builtin_ptrauth_sign_generic_data:
   case Builtin::BI__builtin_ptrauth_sign_unauthenticated:
   case Builtin::BI__builtin_ptrauth_strip: {
     // Emit the arguments.
-    SmallVector<llvm::Value *, 5> Args;
+    SmallVector<llvm::Value *, 6> Args;
     for (auto argExpr : E->arguments())
       Args.push_back(EmitScalarExpr(argExpr));
 
@@ -5600,6 +5601,7 @@ RValue CodeGenFunction::EmitBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
 
     switch (BuiltinID) {
     case Builtin::BI__builtin_ptrauth_auth_and_resign:
+    case Builtin::BI__builtin_ptrauth_auth_load_relative_and_sign:
       if (Args[4]->getType()->isPointerTy())
         Args[4] = Builder.CreatePtrToInt(Args[4], IntPtrTy);
       [[fallthrough]];
@@ -5627,6 +5629,8 @@ RValue CodeGenFunction::EmitBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
         return Intrinsic::ptrauth_auth;
       case Builtin::BI__builtin_ptrauth_auth_and_resign:
         return Intrinsic::ptrauth_resign;
+      case Builtin::BI__builtin_ptrauth_auth_load_relative_and_sign:
+        return Intrinsic::ptrauth_resign_load_relative;
       case Builtin::BI__builtin_ptrauth_blend_discriminator:
         return Intrinsic::ptrauth_blend;
       case Builtin::BI__builtin_ptrauth_sign_generic_data:
diff --git a/clang/lib/Headers/ptrauth.h b/clang/lib/Headers/ptrauth.h
@@ -165,6 +165,31 @@ typedef __UINTPTR_TYPE__ ptrauth_generic_signature_t;
   __builtin_ptrauth_auth_and_resign(__value, __old_key, __old_data, __new_key, \
                                     __new_data)
 
+/* Authenticate a pointer using one scheme, load 32bit value at offset addend
+   from the pointer, and add this value to the pointer, sign using specified
+   scheme.
+
+   If the result is subsequently authenticated using the new scheme, that
+   authentication is guaranteed to fail if and only if the initial
+   authentication failed.
+
+   The value must be an expression of pointer type.
+   The key must be a constant expression of type ptrauth_key.
+   The extra data must be an expression of pointer or integer type;
+   if an integer, it will be coerced to ptrauth_extra_data_t.
+   The addend must be an immediate ptrdiff_t value.
+   The result will have the same type as the original value.
+
+   This operation is guaranteed to not leave the intermediate value
+   available for attack before it is re-signed.
+
+   Do not pass a null pointer to this function. A null pointer
+   will not successfully authenticate. */
+#define ptrauth_auth_load_relative_and_sign(__value, __old_key, __old_data,    \
+                                            __new_key, __new_data, __addend)   \
+  __builtin_ptrauth_auth_load_relative_and_sign(                               \
+      __value, __old_key, __old_data, __new_key, __new_data, __addend)
+
 /* Authenticate a pointer using one scheme and resign it as a C
    function pointer.
 
@@ -327,7 +352,6 @@ typedef __UINTPTR_TYPE__ ptrauth_generic_signature_t;
     ((ptrauth_generic_signature_t)0);                                          \
   })
 
-
 #define ptrauth_cxx_vtable_pointer(key, address_discrimination,                \
                                    extra_discrimination...)
 
diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
@@ -2818,6 +2818,19 @@ def int_ptrauth_resign : Intrinsic<[llvm_i64_ty],
                                    [IntrNoMem, ImmArg<ArgIndex<1>>,
                                     ImmArg<ArgIndex<3>>]>;
 
+// Authenticate a signed pointer, load 32bit value at offset from pointer, add
+// both, and sign it. The second (key) and third (discriminator) arguments
+// specify the signing schema used for authenticating. The fourth and fifth
+// arguments specify the schema used for signing. The sixth argument is addend
+// added to pointer to load the relative offset. The signature must be valid.
+// This is a combined form of int_ptrauth_resign for relative pointers
+def int_ptrauth_resign_load_relative
+    : Intrinsic<[llvm_i64_ty],
+                [llvm_i64_ty, llvm_i32_ty, llvm_i64_ty, llvm_i32_ty,
+                 llvm_i64_ty, llvm_i64_ty],
+                [IntrReadMem, ImmArg<ArgIndex<1>>, ImmArg<ArgIndex<3>>,
+                 ImmArg<ArgIndex<5>>]>;
+
 // Strip the embedded signature out of a signed pointer.
 // The second argument specifies the key.
 // This behaves like @llvm.ptrauth.auth, but doesn't require the signature to
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -168,7 +168,7 @@ class AArch64AsmPrinter : public AsmPrinter {
   // Check authenticated LR before tail calling.
   void emitPtrauthTailCallHardening(const MachineInstr *TC);
 
-  // Emit the sequence for AUT or AUTPAC.
+  // Emit the sequence for AUT, AUTPAC, or AUTRELLOADPAC.
   void emitPtrauthAuthResign(const MachineInstr *MI);
 
   // Emit the sequence to compute the discriminator.
@@ -2066,8 +2066,9 @@ void AArch64AsmPrinter::emitPtrauthTailCallHardening(const MachineInstr *TC) {
 }
 
 void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
-  const bool IsAUTPAC = MI->getOpcode() == AArch64::AUTPAC;
-
+  const bool IsAUTPAC = MI->getOpcode() == AArch64::AUTPAC ||
+                        MI->getOpcode() == AArch64::AUTRELLOADPAC;
+  const bool HasLoad = MI->getOpcode() == AArch64::AUTRELLOADPAC;
   // We expand AUT/AUTPAC into a sequence of the form
   //
   //      ; authenticate x16
@@ -2142,11 +2143,75 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(const MachineInstr *MI) {
   }
 
   // We already emitted unchecked and checked-but-non-trapping AUTs.
-  // That left us with trapping AUTs, and AUTPACs.
+  // That left us with trapping AUTs, and AUTPA/AUTRELLOADPACs.
   // Trapping AUTs don't need PAC: we're done.
   if (!IsAUTPAC)
     return;
 
+  if (HasLoad) {
+    int64_t Addend = MI->getOperand(6).getImm();
+    // incoming rawpointer in X16, X17 is not live at this point.
+    //   LDSRWpre x17, x16, simm9  ; note: x16+simm9 used later.
+    if (isInt<9>(Addend)) {
+      EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::LDRSWpre)
+                                       .addReg(AArch64::X16)
+                                       .addReg(AArch64::X17)
+                                       .addReg(AArch64::X16)
+                                       .addImm(/*simm9:*/ Addend));
+    } else {
+      //   x16 = x16 + Addend computation has 2 variants
+      if (isUInt<24>(Addend)) {
+        // variant 1: add x16, x16, Addend >> shift12 ls shift12
+        // This can take upto 2 instructions.
+        for (int BitPos = 0; BitPos != 24 && (Addend >> BitPos); BitPos += 12) {
+          EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::ADDXri)
+                                           .addReg(AArch64::X16)
+                                           .addReg(AArch64::X16)
+                                           .addImm((Addend >> BitPos) & 0xfff)
+                                           .addImm(AArch64_AM::getShifterImm(
+                                               AArch64_AM::LSL, BitPos)));
+        }
+      } else {
+        // variant 2: accumulate constant in X17 16 bits at a time, and add to
+        // X16 This can take 2-5 instructions.
+        EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::MOVZXi)
+                                         .addReg(AArch64::X17)
+                                         .addImm(Addend & 0xffff)
+                                         .addImm(AArch64_AM::getShifterImm(
+                                             AArch64_AM::LSL, 0)));
+
+        for (int Offset = 16; Offset < 64; Offset += 16) {
+          uint16_t Fragment = static_cast<uint16_t>(Addend >> Offset);
+          if (!Fragment)
+            continue;
+          EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::MOVKXi)
+                                           .addReg(AArch64::X17)
+                                           .addReg(AArch64::X17)
+                                           .addImm(Fragment)
+                                           .addImm(/*shift:*/ Offset));
+        }
+        // addx x16, x16, x17
+        EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::ADDXrs)
+                                         .addReg(AArch64::X16)
+                                         .addReg(AArch64::X16)
+                                         .addReg(AArch64::X17)
+                                         .addImm(0));
+      }
+      // ldrsw x17,x16(0)
+      EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::LDRSWui)
+                                       .addReg(AArch64::X17)
+                                       .addReg(AArch64::X16)
+                                       .addImm(0));
+    }
+    // addx x16, x16, x17
+    EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::ADDXrs)
+                                     .addReg(AArch64::X16)
+                                     .addReg(AArch64::X16)
+                                     .addReg(AArch64::X17)
+                                     .addImm(0));
+
+  } /* HasLoad == true */
+
   auto PACKey = (AArch64PACKey::ID)MI->getOperand(3).getImm();
   uint64_t PACDisc = MI->getOperand(4).getImm();
   unsigned PACAddrDisc = MI->getOperand(5).getReg();
@@ -2864,6 +2929,7 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
 
   case AArch64::AUT:
   case AArch64::AUTPAC:
+  case AArch64::AUTRELLOADPAC:
     emitPtrauthAuthResign(MI);
     return;
 
diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -1544,12 +1544,15 @@ void AArch64DAGToDAGISel::SelectPtrauthAuth(SDNode *N) {
 
 void AArch64DAGToDAGISel::SelectPtrauthResign(SDNode *N) {
   SDLoc DL(N);
-  // IntrinsicID is operand #0
-  SDValue Val = N->getOperand(1);
-  SDValue AUTKey = N->getOperand(2);
-  SDValue AUTDisc = N->getOperand(3);
-  SDValue PACKey = N->getOperand(4);
-  SDValue PACDisc = N->getOperand(5);
+  // IntrinsicID is operand #0, if W_CHAIN it is #1
+  int OffsetBase = N->getOpcode() == ISD::INTRINSIC_W_CHAIN ? 1 : 0;
+  SDValue Val = N->getOperand(OffsetBase + 1);
+  SDValue AUTKey = N->getOperand(OffsetBase + 2);
+  SDValue AUTDisc = N->getOperand(OffsetBase + 3);
+  SDValue PACKey = N->getOperand(OffsetBase + 4);
+  SDValue PACDisc = N->getOperand(OffsetBase + 5);
+  uint32_t IntNum = N->getConstantOperandVal(OffsetBase + 0);
+  bool HasLoad = IntNum == Intrinsic::ptrauth_resign_load_relative;
 
   unsigned AUTKeyC = cast<ConstantSDNode>(AUTKey)->getZExtValue();
   unsigned PACKeyC = cast<ConstantSDNode>(PACKey)->getZExtValue();
@@ -1568,11 +1571,22 @@ void AArch64DAGToDAGISel::SelectPtrauthResign(SDNode *N) {
   SDValue X16Copy = CurDAG->getCopyToReg(CurDAG->getEntryNode(), DL,
                                          AArch64::X16, Val, SDValue());
 
-  SDValue Ops[] = {AUTKey,       AUTConstDisc, AUTAddrDisc,        PACKey,
-                   PACConstDisc, PACAddrDisc,  X16Copy.getValue(1)};
+  if (HasLoad) {
+    SDValue Addend = N->getOperand(OffsetBase + 6);
+    SDValue Ops[] = {AUTKey, AUTConstDisc,       AUTAddrDisc,
+                     PACKey, PACConstDisc,       PACAddrDisc,
+                     Addend, X16Copy.getValue(1)};
 
-  SDNode *AUTPAC = CurDAG->getMachineNode(AArch64::AUTPAC, DL, MVT::i64, Ops);
-  ReplaceNode(N, AUTPAC);
+    SDNode *AUTRELLOADPAC = CurDAG->getMachineNode(AArch64::AUTRELLOADPAC, DL,
+                                                   MVT::i64, MVT::Other, Ops);
+    ReplaceNode(N, AUTRELLOADPAC);
+  } else {
+    SDValue Ops[] = {AUTKey,       AUTConstDisc, AUTAddrDisc,        PACKey,
+                     PACConstDisc, PACAddrDisc,  X16Copy.getValue(1)};
+
+    SDNode *AUTPAC = CurDAG->getMachineNode(AArch64::AUTPAC, DL, MVT::i64, Ops);
+    ReplaceNode(N, AUTPAC);
+  }
 }
 
 bool AArch64DAGToDAGISel::tryIndexedLoad(SDNode *N) {
@@ -5769,6 +5783,9 @@ void AArch64DAGToDAGISel::Select(SDNode *Node) {
               {AArch64::BF2CVT_2ZZ_BtoH, AArch64::F2CVT_2ZZ_BtoH}))
         SelectCVTIntrinsicFP8(Node, 2, Opc);
       return;
+    case Intrinsic::ptrauth_resign_load_relative:
+      SelectPtrauthResign(Node);
+      return;
     }
   } break;
   case ISD::INTRINSIC_WO_CHAIN: {
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -2159,6 +2159,26 @@ let Predicates = [HasPAuth] in {
     let Uses = [X16];
   }
 
+  // Similiar to AUTPAC, except a 32bit value is loaded at Addend offset from
+  // pointer and this value is added to the pointer before signing. This
+  // directly manipulates x16/x17, which are the only registers the OS
+  // guarantees are safe to use for sensitive operations.
+  def AUTRELLOADPAC
+      : Pseudo<(outs),
+               (ins i32imm:$AUTKey, i64imm:$AUTDisc, GPR64:$AUTAddrDisc,
+                   i32imm:$PACKey, i64imm:$PACDisc, GPR64noip:$PACAddrDisc,
+                   i64imm:$Addend),
+               []>,
+        Sched<[WriteI, ReadI]> {
+    let isCodeGenOnly = 1;
+    let hasSideEffects = 1;
+    let mayStore = 0;
+    let mayLoad = 1;
+    let Size = 84;
+    let Defs = [X16, X17, NZCV];
+    let Uses = [X16];
+  }
+
   // Materialize a signed global address, with adrp+add and PAC.
   def MOVaddrPAC : Pseudo<(outs),
                           (ins i64imm:$Addr, i32imm:$Key,
diff --git a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
@@ -6665,6 +6665,45 @@ bool AArch64InstructionSelector::selectIntrinsicWithSideEffects(
     constrainSelectedInstRegOperands(*Memset, TII, TRI, RBI);
     break;
   }
+  case Intrinsic::ptrauth_resign_load_relative: {
+    Register DstReg = I.getOperand(0).getReg();
+    Register ValReg = I.getOperand(2).getReg();
+    uint64_t AUTKey = I.getOperand(3).getImm();
+    Register AUTDisc = I.getOperand(4).getReg();
+    uint64_t PACKey = I.getOperand(5).getImm();
+    Register PACDisc = I.getOperand(6).getReg();
+    int64_t Addend = I.getOperand(7).getImm();
+
+    Register AUTAddrDisc = AUTDisc;
+    uint16_t AUTConstDiscC = 0;
+    std::tie(AUTConstDiscC, AUTAddrDisc) =
+        extractPtrauthBlendDiscriminators(AUTDisc, MRI);
+
+    Register PACAddrDisc = PACDisc;
+    uint16_t PACConstDiscC = 0;
+    std::tie(PACConstDiscC, PACAddrDisc) =
+        extractPtrauthBlendDiscriminators(PACDisc, MRI);
+
+    MIB.buildCopy({AArch64::X16}, {ValReg});
+
+    // we dont support AUTRELLOADPACunchecked
+    unsigned Opc = AArch64::AUTRELLOADPAC;
+
+    MIB.buildInstr(Opc)
+        .addImm(AUTKey)
+        .addImm(AUTConstDiscC)
+        .addUse(AUTAddrDisc)
+        .addImm(PACKey)
+        .addImm(PACConstDiscC)
+        .addUse(PACAddrDisc)
+        .addImm(Addend)
+        .constrainAllUses(TII, TRI, RBI);
+    MIB.buildCopy({DstReg}, Register(AArch64::X16));
+
+    RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
+    I.eraseFromParent();
+    return true;
+  }
   }
 
   I.eraseFromParent();
diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
@@ -461,6 +461,7 @@ bool llvm::wouldInstructionBeTriviallyDead(const Instruction *I,
     case Intrinsic::wasm_trunc_unsigned:
     case Intrinsic::ptrauth_auth:
     case Intrinsic::ptrauth_resign:
+    case Intrinsic::ptrauth_resign_load_relative:
       return true;
     default:
       return false;
diff --git a/llvm/test/CodeGen/AArch64/ptrauth-intrinsic-auth-resign-relative-load.ll b/llvm/test/CodeGen/AArch64/ptrauth-intrinsic-auth-resign-relative-load.ll
