[clang] Ensure correct copying of records with authenticated fields

clang/include/clang/AST/NonTrivialTypeVisitor.h

@@ -95,6 +95,8 @@ struct CopiedTypeVisitor {
       return asDerived().visitARCWeak(FT, std::forward<Ts>(Args)...);
     case QualType::PCK_Struct:
       return asDerived().visitStruct(FT, std::forward<Ts>(Args)...);
+    case QualType::PCK_PtrAuth:
+      return asDerived().visitPtrAuth(FT, std::forward<Ts>(Args)...);
     case QualType::PCK_Trivial:
       return asDerived().visitTrivial(FT, std::forward<Ts>(Args)...);
     case QualType::PCK_VolatileTrivial:

clang/include/clang/AST/Type.h

@@ -1507,6 +1507,9 @@ class QualType {
     /// with the ARC __weak qualifier.
     PCK_ARCWeak,
 
+    /// The type is an address-discriminated signed pointer type.
+    PCK_PtrAuth,
+
     /// The type is a struct containing a field whose type is neither
     /// PCK_Trivial nor PCK_VolatileTrivial.
     /// Note that a C++ struct type does not necessarily match this; C++ copying

clang/lib/AST/ASTContext.cpp

@@ -8372,6 +8372,9 @@ bool ASTContext::BlockRequiresCopying(QualType Ty,
     return true;
   }
 
+  if (Ty.hasAddressDiscriminatedPointerAuth())
+    return true;
+
   // The block needs copy/destroy helpers if Ty is non-trivial to destructively
   // move or destroy.
   if (Ty.isNonTrivialToPrimitiveDestructiveMove() || Ty.isDestructedType())

clang/lib/AST/Type.cpp

@@ -2801,6 +2801,10 @@ static bool isTriviallyCopyableTypeImpl(const QualType &type,
   if (type.hasNonTrivialObjCLifetime())
     return false;
 
+  QualType::PrimitiveCopyKind PCK = type.isNonTrivialToPrimitiveCopy();
+  if (PCK != QualType::PCK_Trivial && PCK != QualType::PCK_VolatileTrivial)
+    return false;
+
   // C++11 [basic.types]p9 - See Core 2094
   //   Scalar types, trivially copyable class types, arrays of such types, and
   //   cv-qualified versions of these types are collectively
@@ -2968,6 +2972,8 @@ QualType::PrimitiveCopyKind QualType::isNonTrivialToPrimitiveCopy() const {
   case Qualifiers::OCL_Weak:
     return PCK_ARCWeak;
   default:
+    if (hasAddressDiscriminatedPointerAuth())
+      return PCK_PtrAuth;
     return Qs.hasVolatile() ? PCK_VolatileTrivial : PCK_Trivial;
   }
 }

clang/lib/CodeGen/CGBlocks.cpp

@@ -1591,6 +1591,10 @@ computeCopyInfoForBlockCapture(const BlockDecl::Capture &CI, QualType T,
     return std::make_pair(BlockCaptureEntityKind::BlockObject, Flags);
   }
 
+  if (T.hasAddressDiscriminatedPointerAuth())
+    return std::make_pair(
+        BlockCaptureEntityKind::AddressDiscriminatedPointerAuth, Flags);
+
   Flags = BLOCK_FIELD_IS_OBJECT;
   bool isBlockPointer = T->isBlockPointerType();
   if (isBlockPointer)
@@ -1611,6 +1615,10 @@ computeCopyInfoForBlockCapture(const BlockDecl::Capture &CI, QualType T,
     return std::make_pair(!isBlockPointer ? BlockCaptureEntityKind::ARCStrong
                                           : BlockCaptureEntityKind::BlockObject,
                           Flags);
+  case QualType::PCK_PtrAuth:
+    return std::make_pair(
+        BlockCaptureEntityKind::AddressDiscriminatedPointerAuth,
+        BlockFieldFlags());
   case QualType::PCK_Trivial:
   case QualType::PCK_VolatileTrivial: {
     if (!T->isObjCRetainableType())
@@ -1713,6 +1721,13 @@ static std::string getBlockCaptureStr(const CGBlockInfo::Capture &Cap,
   case BlockCaptureEntityKind::ARCStrong:
     Str += "s";
     break;
+  case BlockCaptureEntityKind::AddressDiscriminatedPointerAuth: {
+    auto PtrAuth = CaptureTy.getPointerAuth();
+    assert(PtrAuth && PtrAuth.isAddressDiscriminated());
+    Str += "p" + llvm::to_string(PtrAuth.getKey()) + "d" +
+           llvm::to_string(PtrAuth.getExtraDiscriminator());
+    break;
+  }
   case BlockCaptureEntityKind::BlockObject: {
     const VarDecl *Var = CI.getVariable();
     unsigned F = Flags.getBitMask();
@@ -1829,6 +1844,7 @@ static void pushCaptureCleanup(BlockCaptureEntityKind CaptureKind,
     }
     break;
   }
+  case BlockCaptureEntityKind::AddressDiscriminatedPointerAuth:
   case BlockCaptureEntityKind::None:
     break;
   }
@@ -1925,6 +1941,14 @@ CodeGenFunction::GenerateCopyHelperFunction(const CGBlockInfo &blockInfo) {
     case BlockCaptureEntityKind::ARCWeak:
       EmitARCCopyWeak(dstField, srcField);
       break;
+    case BlockCaptureEntityKind::AddressDiscriminatedPointerAuth: {
+      QualType Type = CI.getVariable()->getType();
+      PointerAuthQualifier PointerAuth = Type.getPointerAuth();
+      assert(PointerAuth && PointerAuth.isAddressDiscriminated());
+      EmitPointerAuthCopy(PointerAuth, Type, dstField, srcField);
+      // We don't need to push cleanups for ptrauth types.
+      continue;
+    }
     case BlockCaptureEntityKind::NonTrivialCStruct: {
       // If this is a C struct that requires non-trivial copy construction,
       // emit a call to its copy constructor.
@@ -2261,6 +2285,33 @@ class CXXByrefHelpers final : public BlockByrefHelpers {
   }
 };
 
+/// Emits the copy/dispose helpers for a __block variable with
+/// address-discriminated pointer authentication.
+class AddressDiscriminatedByrefHelpers final : public BlockByrefHelpers {
+  QualType VarType;
+
+public:
+  AddressDiscriminatedByrefHelpers(CharUnits Alignment, QualType Type)
+      : BlockByrefHelpers(Alignment), VarType(Type) {
+    assert(Type.hasAddressDiscriminatedPointerAuth());
+  }
+
+  void emitCopy(CodeGenFunction &CGF, Address DestField,
+                Address SrcField) override {
+    CGF.EmitPointerAuthCopy(VarType.getPointerAuth(), VarType, DestField,
+                            SrcField);
+  }
+
+  bool needsDispose() const override { return false; }
+  void emitDispose(CodeGenFunction &CGF, Address Field) override {
+    llvm_unreachable("should never be called");
+  }
+
+  void profileImpl(llvm::FoldingSetNodeID &ID) const override {
+    ID.AddPointer(VarType.getCanonicalType().getAsOpaquePtr());
+  }
+};
+
 /// Emits the copy/dispose helpers for a __block variable that is a non-trivial
 /// C struct.
 class NonTrivialCStructByrefHelpers final : public BlockByrefHelpers {
@@ -2462,7 +2513,10 @@ CodeGenFunction::buildByrefHelpers(llvm::StructType &byrefType,
     return ::buildByrefHelpers(
         CGM, byrefInfo, CXXByrefHelpers(valueAlignment, type, copyExpr));
   }
-
+  if (type.hasAddressDiscriminatedPointerAuth()) {
+    return ::buildByrefHelpers(
+        CGM, byrefInfo, AddressDiscriminatedByrefHelpers(valueAlignment, type));
+  }
   // If type is a non-trivial C struct type that is non-trivial to
   // destructly move or destroy, build the copy and dispose helpers.
   if (type.isNonTrivialToPrimitiveDestructiveMove() == QualType::PCK_Struct ||

clang/lib/CodeGen/CGBlocks.h

@@ -146,6 +146,7 @@ class BlockByrefInfo {
 enum class BlockCaptureEntityKind {
   None,
   CXXRecord, // Copy or destroy
+  AddressDiscriminatedPointerAuth,
   ARCWeak,
   ARCStrong,
   NonTrivialCStruct,

clang/lib/CodeGen/CGCall.cpp

@@ -4857,7 +4857,7 @@ void CodeGenFunction::EmitCallArg(CallArgList &args, const Expr *E,
 
   if (HasAggregateEvalKind && isa<ImplicitCastExpr>(E) &&
       cast<CastExpr>(E)->getCastKind() == CK_LValueToRValue &&
-      !type->isArrayParameterType()) {
+      !type->isArrayParameterType() && !type.isNonTrivialToPrimitiveCopy()) {
     LValue L = EmitLValue(cast<CastExpr>(E)->getSubExpr());
     assert(L.isSimple());
     args.addUncopiedAggregate(L, type);

clang/lib/CodeGen/CGNonTrivialStruct.cpp

@@ -266,6 +266,18 @@ struct GenBinaryFuncName : CopyStructVisitor<GenBinaryFuncName<IsMove>, IsMove>,
     this->appendStr("_tv" + llvm::to_string(OffsetInBits) + "w" +
                     llvm::to_string(getFieldSize(FD, FT, this->Ctx)));
   }
+
+  void visitPtrAuth(QualType FT, const FieldDecl *FD,
+                    CharUnits CurStructOffset) {
+    this->appendStr("_pa");
+    PointerAuthQualifier PtrAuth = FT.getPointerAuth().withoutKeyNone();
+    this->appendStr(llvm::to_string(PtrAuth.getKey()) + "_");
+    this->appendStr(llvm::to_string(PtrAuth.getExtraDiscriminator()) + "_");
+    if (PtrAuth.authenticatesNullValues())
+      this->appendStr("anv_");
+    CharUnits FieldOffset = CurStructOffset + this->getFieldOffset(FD);
+    this->appendStr(llvm::to_string(FieldOffset.getQuantity()));
+  }
 };
 
 struct GenDefaultInitializeFuncName
@@ -568,6 +580,13 @@ struct GenBinaryFunc : CopyStructVisitor<Derived, IsMove>,
     RValue SrcVal = this->CGF->EmitLoadOfLValue(SrcLV, SourceLocation());
     this->CGF->EmitStoreThroughLValue(SrcVal, DstLV);
   }
+  void visitPtrAuth(QualType FT, const FieldDecl *FD, CharUnits CurStackOffset,
+                    std::array<Address, 2> Addrs) {
+    PointerAuthQualifier PtrAuth = FT.getPointerAuth().withoutKeyNone();
+    Addrs[DstIdx] = this->getAddrWithOffset(Addrs[DstIdx], CurStackOffset, FD);
+    Addrs[SrcIdx] = this->getAddrWithOffset(Addrs[SrcIdx], CurStackOffset, FD);
+    this->CGF->EmitPointerAuthCopy(PtrAuth, FT, Addrs[DstIdx], Addrs[SrcIdx]);
+  }
 };
 
 // These classes that emit the special functions for a non-trivial struct.

clang/lib/Sema/SemaChecking.cpp

@@ -9443,6 +9443,9 @@ struct SearchNonTrivialToCopyField
   void visitARCWeak(QualType FT, SourceLocation SL) {
     S.DiagRuntimeBehavior(SL, E, S.PDiag(diag::note_nontrivial_field) << 0);
   }
+  void visitPtrAuth(QualType FT, SourceLocation SL) {
+    S.DiagRuntimeBehavior(SL, E, S.PDiag(diag::note_nontrivial_field) << 0);
+  }
   void visitStruct(QualType FT, SourceLocation SL) {
     for (const FieldDecl *FD : FT->castAs<RecordType>()->getDecl()->fields())
       visit(FD->getType(), FD->getLocation());

clang/lib/Sema/SemaDecl.cpp

@@ -13395,6 +13395,12 @@ struct DiagNonTrivalCUnionCopyVisitor
         asDerived().visit(FD->getType(), FD, InNonTrivialUnion);
   }
 
+  void visitPtrAuth(QualType QT, const FieldDecl *FD, bool InNonTrivialUnion) {
+    if (InNonTrivialUnion)
+      S.Diag(FD->getLocation(), diag::note_non_trivial_c_union)
+          << 1 << 2 << QT << FD->getName();
+  }
+
   void preVisit(QualType::PrimitiveCopyKind PCK, QualType QT,
                 const FieldDecl *FD, bool InNonTrivialUnion) {}
   void visitTrivial(QualType QT, const FieldDecl *FD, bool InNonTrivialUnion) {}

clang/test/CodeGen/ptrauth-in-c-struct.c

@@ -0,0 +1,172 @@
+// RUN: %clang_cc1 -triple arm64-apple-ios -fblocks -fptrauth-calls -fptrauth-returns -fptrauth-intrinsics -emit-llvm -o - %s | FileCheck %s
+
+#define AQ1_50 __ptrauth(1,1,50)
+#define AQ2_30 __ptrauth(2,1,30)
+#define IQ __ptrauth(1,0,50)
+
+typedef void (^BlockTy)(void);
+
+// CHECK: %[[STRUCT_SA:.*]] = type { i32, ptr }
+// CHECK: %[[STRUCT_SA2:.*]] = type { i32, ptr }
+// CHECK: %[[STRUCT_SI:.*]] = type { ptr }
+
+typedef struct {
+  int f0;
+  int * AQ1_50 f1; // Signed using address discrimination.
+} SA;
+
+typedef struct {
+  int f0;
+  int * AQ2_30 f1; // Signed using address discrimination.
+} SA2;
+
+typedef struct {
+  int * IQ f; // No address discrimination.
+} SI;
+
+SA getSA(void);
+void calleeSA(SA);
+
+int g0;
+
+// CHECK: define void @test_copy_constructor_SA(ptr noundef %{{.*}})
+// CHECK: call void @__copy_constructor_8_8_t0w4_pa1_50_8(
+
+// CHECK: define linkonce_odr hidden void @__copy_constructor_8_8_t0w4_pa1_50_8(ptr noundef %[[DST:.*]], ptr noundef %[[SRC:.*]])
+// CHECK: %[[DST_ADDR:.*]] = alloca ptr, align 8
+// CHECK: %[[SRC_ADDR:.*]] = alloca ptr, align 8
+// CHECK: store ptr %[[DST]], ptr %[[DST_ADDR]], align 8
+// CHECK: store ptr %[[SRC]], ptr %[[SRC_ADDR]], align 8
+// CHECK: %[[V0:.*]] = load ptr, ptr %[[DST_ADDR]], align 8
+// CHECK: %[[V1:.*]] = load ptr, ptr %[[SRC_ADDR]], align 8
+// CHECK: %[[V6:.*]] = getelementptr inbounds i8, ptr %[[V0]], i64 8
+// CHECK: %[[V9:.*]] = getelementptr inbounds i8, ptr %[[V1]], i64 8
+// CHECK: %[[V11:.*]] = load ptr, ptr %[[V9]], align 8
+// CHECK: %[[V12:.*]] = ptrtoint ptr %[[V9]] to i64
+// CHECK: %[[V13:.*]] = call i64 @llvm.ptrauth.blend(i64 %[[V12]], i64 50)
+// CHECK: %[[V14:.*]] = ptrtoint ptr %[[V6]] to i64
+// CHECK: %[[V15:.*]] = call i64 @llvm.ptrauth.blend(i64 %[[V14]], i64 50)
+// CHECK: %[[V17:.*]] = ptrtoint ptr %[[V11]] to i64
+// CHECK: %[[V18:.*]] = call i64 @llvm.ptrauth.resign(i64 %[[V17]], i32 1, i64 %[[V13]], i32 1, i64 %[[V15]])
+
+void test_copy_constructor_SA(SA *s) {
+  SA t = *s;
+}
+
+// CHECK: define void @test_copy_constructor_SA2(ptr noundef %{{.*}})
+// CHECK: call void @__copy_constructor_8_8_t0w4_pa2_30_8(
+
+// CHECK: define linkonce_odr hidden void @__copy_constructor_8_8_t0w4_pa2_30_8(ptr noundef %[[DST:.*]], ptr noundef %[[SRC:.*]])
+// CHECK: %[[DST_ADDR:.*]] = alloca ptr, align 8
+// CHECK: %[[SRC_ADDR:.*]] = alloca ptr, align 8
+// CHECK: store ptr %[[DST]], ptr %[[DST_ADDR]], align 8
+// CHECK: store ptr %[[SRC]], ptr %[[SRC_ADDR]], align 8
+// CHECK: %[[V0:.*]] = load ptr, ptr %[[DST_ADDR]], align 8
+// CHECK: %[[V1:.*]] = load ptr, ptr %[[SRC_ADDR]], align 8
+// CHECK: %[[V6:.*]] = getelementptr inbounds i8, ptr %[[V0]], i64 8
+// CHECK: %[[V9:.*]] = getelementptr inbounds i8, ptr %[[V1]], i64 8
+// CHECK: %[[V11:.*]] = load ptr, ptr %[[V9]], align 8
+// CHECK: %[[V12:.*]] = ptrtoint ptr %[[V9]] to i64
+// CHECK: %[[V13:.*]] = call i64 @llvm.ptrauth.blend(i64 %[[V12]], i64 30)
+// CHECK: %[[V14:.*]] = ptrtoint ptr %[[V6]] to i64
+// CHECK: %[[V15:.*]] = call i64 @llvm.ptrauth.blend(i64 %[[V14]], i64 30)
+// CHECK: %[[V17:.*]] = ptrtoint ptr %[[V11]] to i64
+// CHECK: %[[V18:.*]] = call i64 @llvm.ptrauth.resign(i64 %[[V17]], i32 2, i64 %[[V13]], i32 2, i64 %[[V15]])
+
+void test_copy_constructor_SA2(SA2 *s) {
+  SA2 t = *s;
+}
+
+// CHECK: define void @test_copy_assignment_SA(
+// CHECK: call void @__copy_assignment_8_8_t0w4_pa1_50_8(
+
+// CHECK: define linkonce_odr hidden void @__copy_assignment_8_8_t0w4_pa1_50_8(
+
+void test_copy_assignment_SA(SA *d, SA *s) {
+  *d = *s;
+}
+
+// CHECK: define void @test_move_constructor_SA(
+// CHECK: define internal void @__Block_byref_object_copy_(
+// CHECK: define linkonce_odr hidden void @__move_constructor_8_8_t0w4_pa1_50_8(
+
+void test_move_constructor_SA(void) {
+  __block SA t;
+  BlockTy b = ^{ (void)t; };
+}
+
+// CHECK: define void @test_move_assignment_SA(
+// CHECK: call void @__move_assignment_8_8_t0w4_pa1_50_8(
+// CHECK: define linkonce_odr hidden void @__move_assignment_8_8_t0w4_pa1_50_8(
+
+void test_move_assignment_SA(SA *p) {
+  *p = getSA();
+}
+
+// CHECK: define void @test_parameter_SA(ptr noundef %{{.*}})
+// CHECK-NOT: call
+// CHECK: ret void
+
+void test_parameter_SA(SA a) {
+}
+
+// CHECK: define void @test_argument_SA(ptr noundef %[[A:.*]])
+// CHECK: %[[A_ADDR:.*]] = alloca ptr, align 8
+// CHECK: %[[AGG_TMP:.*]] = alloca %[[STRUCT_SA]], align 8
+// CHECK: store ptr %[[A]], ptr %[[A_ADDR]], align 8
+// CHECK: %[[V0:.*]] = load ptr, ptr %[[A_ADDR]], align 8
+// CHECK: call void @__copy_constructor_8_8_t0w4_pa1_50_8(ptr %[[AGG_TMP]], ptr %[[V0]])
+// CHECK: call void @calleeSA(ptr noundef %[[AGG_TMP]])
+// CHECK-NOT: call
+// CHECK: ret void
+
+void test_argument_SA(SA *a) {
+  calleeSA(*a);
+}
+
+// CHECK: define void @test_return_SA(ptr dead_on_unwind noalias writable sret(%struct.SA) align 8 %[[AGG_RESULT:.*]], ptr noundef %[[A:.*]])
+// CHECK: %[[A_ADDR:.*]] = alloca ptr, align 8
+// CHECK: store ptr %[[A]], ptr %[[A_ADDR]], align 8
+// CHECK: %[[V0:.*]] = load ptr, ptr %[[A_ADDR]], align 8
+// CHECK: call void @__copy_constructor_8_8_t0w4_pa1_50_8(ptr %[[AGG_RESULT]], ptr %[[V0]])
+// CHECK-NOT: call
+// CHECK: ret void
+
+SA test_return_SA(SA *a) {
+  return *a;
+}
+
+// CHECK: define void @test_copy_constructor_SI(
+// CHECK-NOT: call
+// CHECK: call void @llvm.memcpy.p0.p0.i64(
+// CHECK-NOT: call
+// CHECK: ret void
+
+void test_copy_constructor_SI(SI *s) {
+  SI t = *s;
+}
+
+// CHECK: define void @test_parameter_SI(i64 %{{.*}})
+// CHECK-NOT: call
+// CHECK: ret void
+
+void test_parameter_SI(SI a) {
+}
+
+// CHECK-LABEL: define void @test_array(
+// CHECK: %[[F1:.*]] = getelementptr inbounds nuw %[[STRUCT_SA]], ptr %{{.*}}, i32 0, i32 1
+// CHECK: %[[V0:.*]] = ptrtoint ptr %[[F1]] to i64
+// CHECK: %[[V1:.*]] = call i64 @llvm.ptrauth.blend(i64 %[[V0]], i64 50)
+// CHECK: %[[V2:.*]] = call i64 @llvm.ptrauth.sign(i64 ptrtoint (ptr @g0 to i64), i32 1, i64 %[[V1]])
+// CHECK: %[[V3:.*]] = inttoptr i64 %[[V2]] to ptr
+// CHECK: store ptr %[[V3]], ptr %[[F1]], align 8
+// CHECK: %[[F12:.*]] = getelementptr inbounds nuw %[[STRUCT_SA]], ptr %{{.*}}, i32 0, i32 1
+// CHECK: %[[V4:.*]] = ptrtoint ptr %[[F12]] to i64
+// CHECK: %[[V5:.*]] = call i64 @llvm.ptrauth.blend(i64 %[[V4]], i64 50)
+// CHECK: %[[V6:.*]] = call i64 @llvm.ptrauth.sign(i64 ptrtoint (ptr @g0 to i64), i32 1, i64 %[[V5]])
+// CHECK: %[[V7:.*]] = inttoptr i64 %[[V6]] to ptr
+// CHECK: store ptr %[[V7]], ptr %[[F12]], align 8
+
+void test_array(void) {
+  const SA a[] = {{0, &g0}, {1, &g0}};
+}