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

       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

note: conflicts resolved llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp due to upstream change
in args for emitPtrauthAuthResign

Author: AbhayKanhere

clang/include/clang/Basic/Builtins.td

@@ -4655,6 +4655,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];

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -11097,6 +11097,8 @@ def err_builtin_requires_language : Error<"'%0' is only available in %1">;
 
 def err_constant_integer_arg_type : Error<
   "argument to %0 must be a constant integer">;
+def err_constant_integer_last_arg_type
+    : Error<"last argument to %0 must be a constant integer">;
 
 def ext_mixed_decls_code : Extension<
   "mixing declarations and code is a C99 extension">,

clang/lib/CodeGen/CGBuiltin.cpp

@@ -5601,12 +5601,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));
 
@@ -5617,6 +5618,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]];
@@ -5644,6 +5646,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:

clang/lib/Headers/ptrauth.h

@@ -178,6 +178,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.
 
@@ -351,6 +376,17 @@ typedef __UINTPTR_TYPE__ ptrauth_generic_signature_t;
     __value;                                                                   \
   })
 
+#define ptrauth_auth_load_relative_and_sign(__value, __old_key, __old_data,    \
+                                            __new_key, __new_data, __addend)   \
+  ({                                                                           \
+    (void)__old_key;                                                           \
+    (void)__old_data;                                                          \
+    (void)__new_key;                                                           \
+    (void)__new_data;                                                          \
+    const char *__value_tmp = (const char *)(__value);                         \
+    (void *)(__value_tmp + *(const int *)(__value_tmp + (__addend)));          \
+  })
+
 #define ptrauth_auth_function(__value, __old_key, __old_data)                  \
   ({                                                                           \
     (void)__old_key;                                                           \
@@ -380,7 +416,6 @@ typedef __UINTPTR_TYPE__ ptrauth_generic_signature_t;
     ((ptrauth_generic_signature_t)0);                                          \
   })
 
-
 #define ptrauth_cxx_vtable_pointer(key, address_discrimination,                \
                                    extra_discrimination...)
 

clang/lib/Sema/SemaChecking.cpp

@@ -1814,6 +1814,32 @@ static ExprResult PointerAuthAuthAndResign(Sema &S, CallExpr *Call) {
   return Call;
 }
 
+static ExprResult PointerAuthAuthLoadRelativeAndSign(Sema &S, CallExpr *Call) {
+  if (S.checkArgCount(Call, 6))
+    return ExprError();
+  if (checkPointerAuthEnabled(S, Call))
+    return ExprError();
+  const Expr *AddendExpr = Call->getArg(5);
+  bool AddendIsConstInt = AddendExpr->isIntegerConstantExpr(S.Context);
+  if (!AddendIsConstInt) {
+    const Expr *Arg = Call->getArg(5)->IgnoreParenImpCasts();
+    DeclRefExpr *DRE = cast<DeclRefExpr>(Call->getCallee()->IgnoreParenCasts());
+    FunctionDecl *FDecl = cast<FunctionDecl>(DRE->getDecl());
+    S.Diag(Arg->getBeginLoc(), diag::err_constant_integer_last_arg_type)
+        << FDecl->getDeclName() << Arg->getSourceRange();
+  }
+  if (checkPointerAuthValue(S, Call->getArgs()[0], PAO_Auth) ||
+      checkPointerAuthKey(S, Call->getArgs()[1]) ||
+      checkPointerAuthValue(S, Call->getArgs()[2], PAO_Discriminator) ||
+      checkPointerAuthKey(S, Call->getArgs()[3]) ||
+      checkPointerAuthValue(S, Call->getArgs()[4], PAO_Discriminator) ||
+      !AddendIsConstInt)
+    return ExprError();
+
+  Call->setType(Call->getArgs()[0]->getType());
+  return Call;
+}
+
 static ExprResult PointerAuthStringDiscriminator(Sema &S, CallExpr *Call) {
   if (checkPointerAuthEnabled(S, Call))
     return ExprError();
@@ -2866,6 +2892,8 @@ Sema::CheckBuiltinFunctionCall(FunctionDecl *FDecl, unsigned BuiltinID,
     return PointerAuthSignGenericData(*this, TheCall);
   case Builtin::BI__builtin_ptrauth_auth_and_resign:
     return PointerAuthAuthAndResign(*this, TheCall);
+  case Builtin::BI__builtin_ptrauth_auth_load_relative_and_sign:
+    return PointerAuthAuthLoadRelativeAndSign(*this, TheCall);
   case Builtin::BI__builtin_ptrauth_string_discriminator:
     return PointerAuthStringDiscriminator(*this, TheCall);
 

clang/test/CodeGen/ptrauth-intrinsics.c

@@ -55,6 +55,18 @@ void test_auth_and_resign() {
   fnptr = __builtin_ptrauth_auth_and_resign(fnptr, 0, ptr_discriminator, 3, 15);
 }
 
+// CHECK-LABEL: define {{.*}}void @test_auth_load_relative_and_sign()
+void test_auth_load_relative_and_sign() {
+  // CHECK:      [[PTR:%.*]] = load ptr, ptr @fnptr,
+  // CHECK-NEXT: [[DISC0:%.*]] = load ptr, ptr @ptr_discriminator,
+  // CHECK-NEXT: [[T0:%.*]] = ptrtoint ptr [[PTR]] to i64
+  // CHECK-NEXT: [[DISC:%.*]] = ptrtoint ptr [[DISC0]] to i64
+  // CHECK-NEXT: [[T1:%.*]] = call i64 @llvm.ptrauth.resign.load.relative(i64 [[T0]], i32 0, i64 [[DISC]], i32 3, i64 15, i64  16)
+  // CHECK-NEXT: [[RESULT:%.*]] = inttoptr  i64 [[T1]] to ptr
+  // CHECK-NEXT: store ptr [[RESULT]], ptr @fnptr,
+  fnptr = __builtin_ptrauth_auth_load_relative_and_sign(fnptr, 0, ptr_discriminator, 3, 15, 16L);
+}
+
 // CHECK-LABEL: define {{.*}}void @test_blend_discriminator()
 void test_blend_discriminator() {
   // CHECK:      [[PTR:%.*]] = load ptr, ptr @fnptr,

clang/test/Sema/ptrauth.c

@@ -121,6 +121,29 @@ void test_auth_and_resign(int *dp, int (*fp)(int)) {
 
   float *mismatch = __builtin_ptrauth_auth_and_resign(dp, VALID_DATA_KEY, 0, VALID_DATA_KEY, dp); // expected-warning {{incompatible pointer types initializing 'float *' with an expression of type 'int *'}}
 }
+void test_auth_load_relative_and_sign(int *dp, int (*fp)(int)) {
+  __builtin_ptrauth_auth_load_relative_and_sign(dp, VALID_DATA_KEY, 0, VALID_DATA_KEY, 0); // expected-error {{too few arguments}}
+  __builtin_ptrauth_auth_load_relative_and_sign(dp, VALID_DATA_KEY, dp, VALID_DATA_KEY, dp, 0, 0); // expected-error {{too many arguments}}
+  int n = *dp;
+  __builtin_ptrauth_auth_load_relative_and_sign(dp, VALID_DATA_KEY, dp, VALID_DATA_KEY, dp,n); // expected-error {{last argument to '__builtin_ptrauth_auth_load_relative_and_sign' must be a constant integer}}
+  __builtin_ptrauth_auth_load_relative_and_sign(mismatched_type, VALID_DATA_KEY, 0, VALID_DATA_KEY, dp, 0); // expected-error {{signed value must have pointer type; type here is 'struct A'}}
+  __builtin_ptrauth_auth_load_relative_and_sign(dp, mismatched_type, 0, VALID_DATA_KEY, dp, 0); // expected-error {{passing 'struct A' to parameter of incompatible type 'int'}}
+  __builtin_ptrauth_auth_load_relative_and_sign(dp, VALID_DATA_KEY, mismatched_type, VALID_DATA_KEY, dp, 0); // expected-error {{extra discriminator must have pointer or integer type; type here is 'struct A'}}
+  __builtin_ptrauth_auth_load_relative_and_sign(dp, VALID_DATA_KEY, 0, mismatched_type, dp, 0); // expected-error {{passing 'struct A' to parameter of incompatible type 'int'}}
+  __builtin_ptrauth_auth_load_relative_and_sign(dp, VALID_DATA_KEY, 0, VALID_DATA_KEY, mismatched_type, 0); // expected-error {{extra discriminator must have pointer or integer type; type here is 'struct A'}}
+
+  (void) __builtin_ptrauth_auth_and_resign(NULL, VALID_DATA_KEY, 0, VALID_DATA_KEY, dp); // expected-warning {{authenticating a null pointer will almost certainly trap}}
+
+  int *dr = __builtin_ptrauth_auth_load_relative_and_sign(dp, VALID_DATA_KEY, 0, VALID_DATA_KEY, dp, 10);
+  dr = __builtin_ptrauth_auth_load_relative_and_sign(dp, INVALID_KEY, 0, VALID_DATA_KEY, dp, 10); // expected-error {{does not identify a valid pointer authentication key for the current target}}
+  dr = __builtin_ptrauth_auth_load_relative_and_sign(dp, VALID_DATA_KEY, 0, INVALID_KEY, dp, 10); // expected-error {{does not identify a valid pointer authentication key for the current target}}
+
+  int (*fr)(int) = __builtin_ptrauth_auth_load_relative_and_sign(fp, VALID_CODE_KEY, 0, VALID_CODE_KEY, dp, 10);
+  fr = __builtin_ptrauth_auth_load_relative_and_sign(fp, INVALID_KEY, 0, VALID_CODE_KEY, dp, 10); // expected-error {{does not identify a valid pointer authentication key for the current target}}
+  fr = __builtin_ptrauth_auth_load_relative_and_sign(fp, VALID_CODE_KEY, 0, INVALID_KEY, dp, 10); // expected-error {{does not identify a valid pointer authentication key for the current target}}
+
+  float *mismatch = __builtin_ptrauth_auth_load_relative_and_sign(dp, VALID_DATA_KEY, 0, VALID_DATA_KEY, dp,0); // expected-warning {{incompatible pointer types initializing 'float *' with an expression of type 'int *'}}
+}
 
 void test_sign_generic_data(int *dp) {
   __builtin_ptrauth_sign_generic_data(dp); // expected-error {{too few arguments}}

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

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -168,13 +168,14 @@ 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 or AUTPAC. Addend if AUTRELLOADPAC
   void emitPtrauthAuthResign(Register AUTVal, AArch64PACKey::ID AUTKey,
                              uint64_t AUTDisc,
                              const MachineOperand *AUTAddrDisc,
                              Register Scratch,
                              std::optional<AArch64PACKey::ID> PACKey,
-                             uint64_t PACDisc, Register PACAddrDisc);
+                             uint64_t PACDisc, Register PACAddrDisc,
+                             std::optional<uint64_t> Addend);
 
   // Emit the sequence for PAC.
   void emitPtrauthSign(const MachineInstr *MI);
@@ -2078,9 +2079,9 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(
     Register AUTVal, AArch64PACKey::ID AUTKey, uint64_t AUTDisc,
     const MachineOperand *AUTAddrDisc, Register Scratch,
     std::optional<AArch64PACKey::ID> PACKey, uint64_t PACDisc,
-    Register PACAddrDisc) {
+    Register PACAddrDisc, std::optional<uint64_t> OptAddend) {
   const bool IsAUTPAC = PACKey.has_value();
-
+  const bool HasLoad = OptAddend.has_value();
   // We expand AUT/AUTPAC into a sequence of the form
   //
   //      ; authenticate x16
@@ -2151,12 +2152,76 @@ void AArch64AsmPrinter::emitPtrauthAuthResign(
   }
 
   // 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;
 
-  // Compute pac discriminator
+  if (HasLoad) {
+    int64_t Addend = *OptAddend;
+    // 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 */
+
+  // Compute pac discriminator into x17
   assert(isUInt<16>(PACDisc));
   Register PACDiscReg =
       emitPtrauthDiscriminator(PACDisc, PACAddrDisc, Scratch);
@@ -2906,22 +2971,31 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) {
     emitPtrauthAuthResign(AArch64::X16,
                           (AArch64PACKey::ID)MI->getOperand(0).getImm(),
                           MI->getOperand(1).getImm(), &MI->getOperand(2),
-                          AArch64::X17, std::nullopt, 0, 0);
+                          AArch64::X17, std::nullopt, 0, 0, std::nullopt);
     return;
 
   case AArch64::AUTxMxN:
     emitPtrauthAuthResign(MI->getOperand(0).getReg(),
                           (AArch64PACKey::ID)MI->getOperand(3).getImm(),
                           MI->getOperand(4).getImm(), &MI->getOperand(5),
-                          MI->getOperand(1).getReg(), std::nullopt, 0, 0);
+                          MI->getOperand(1).getReg(), std::nullopt, 0, 0,
+                          std::nullopt);
     return;
 
+  case AArch64::AUTRELLOADPAC:
+    emitPtrauthAuthResign(
+        AArch64::X16, (AArch64PACKey::ID)MI->getOperand(0).getImm(),
+        MI->getOperand(1).getImm(), &MI->getOperand(2), AArch64::X17,
+        (AArch64PACKey::ID)MI->getOperand(3).getImm(),
+        MI->getOperand(4).getImm(), MI->getOperand(5).getReg(),
+        MI->getOperand(6).getImm());
+    return;
   case AArch64::AUTPAC:
     emitPtrauthAuthResign(
         AArch64::X16, (AArch64PACKey::ID)MI->getOperand(0).getImm(),
         MI->getOperand(1).getImm(), &MI->getOperand(2), AArch64::X17,
         (AArch64PACKey::ID)MI->getOperand(3).getImm(),
-        MI->getOperand(4).getImm(), MI->getOperand(5).getReg());
+        MI->getOperand(4).getImm(), MI->getOperand(5).getReg(), std::nullopt);
     return;
 
   case AArch64::PAC: