Improve coerce logic and add coerce testusers/pcc/spr/add-pointer-field-protection-feature

Created using spr 1.3.6-beta.1

8 files changed, 167 insertions, 34 deletions

diff --git a/clang/include/clang/AST/ASTContext.h b/clang/include/clang/AST/ASTContext.h
index 7c04a39..b1864ca 100644
--- a/clang/include/clang/AST/ASTContext.h
+++ b/clang/include/clang/AST/ASTContext.h
@@ -186,6 +186,7 @@ struct TypeInfoChars {
 struct PFPField {
   CharUnits offset;
   FieldDecl *field;
+  bool isWithinUnion;
 };
 
 /// Holds long-lived AST nodes (such as types and decls) that can be
@@ -3727,7 +3728,8 @@ public:
 
   bool isPFPStruct(const RecordDecl *rec) const;
   void findPFPFields(QualType Ty, CharUnits Offset,
-                     std::vector<PFPField> &Fields, bool IncludeVBases) const;
+                     std::vector<PFPField> &Fields, bool IncludeVBases,
+                     bool IsWithinUnion = false) const;
   bool hasPFPFields(QualType ty) const;
   bool isPFPField(const FieldDecl *field) const;
 
diff --git a/clang/lib/AST/ASTContext.cpp b/clang/lib/AST/ASTContext.cpp
index 1dc5b6b..5c9d481 100644
--- a/clang/lib/AST/ASTContext.cpp
+++ b/clang/lib/AST/ASTContext.cpp
@@ -15183,6 +15183,10 @@ bool ASTContext::useAbbreviatedThunkName(GlobalDecl VirtualMethodDecl,
 }
 
 bool ASTContext::arePFPFieldsTriviallyRelocatable(const RecordDecl *RD) const {
+  bool IsPAuthSupported =
+      getTargetInfo().getTriple().getArch() == llvm::Triple::aarch64;
+  if (!IsPAuthSupported)
+    return true;
   if (getLangOpts().getPointerFieldProtection() ==
       LangOptions::PointerFieldProtectionKind::Tagged)
     return !isa<CXXRecordDecl>(RD) ||
@@ -15200,7 +15204,7 @@ bool ASTContext::isPFPStruct(const RecordDecl *rec) const {
 
 void ASTContext::findPFPFields(QualType Ty, CharUnits Offset,
                                std::vector<PFPField> &Fields,
-                               bool IncludeVBases) const {
+                               bool IncludeVBases, bool IsWithinUnion) const {
   if (auto *AT = getAsConstantArrayType(Ty)) {
     if (auto *ElemDecl = AT->getElementType()->getAsCXXRecordDecl()) {
       const ASTRecordLayout &ElemRL = getASTRecordLayout(ElemDecl);
@@ -15213,26 +15217,27 @@ void ASTContext::findPFPFields(QualType Ty, CharUnits Offset,
   auto *Decl = Ty->getAsCXXRecordDecl();
   if (!Decl)
     return;
+  IsWithinUnion |= Decl->isUnion();
   const ASTRecordLayout &RL = getASTRecordLayout(Decl);
   for (FieldDecl *field : Decl->fields()) {
     CharUnits fieldOffset =
         Offset + toCharUnitsFromBits(RL.getFieldOffset(field->getFieldIndex()));
     if (isPFPField(field))
-      Fields.push_back({fieldOffset, field});
-    findPFPFields(field->getType(), fieldOffset, Fields, true);
+      Fields.push_back({fieldOffset, field, IsWithinUnion});
+    findPFPFields(field->getType(), fieldOffset, Fields, true, IsWithinUnion);
   }
   for (auto &Base : Decl->bases()) {
     if (Base.isVirtual())
       continue;
     CharUnits BaseOffset =
         Offset + RL.getBaseClassOffset(Base.getType()->getAsCXXRecordDecl());
-    findPFPFields(Base.getType(), BaseOffset, Fields, false);
+    findPFPFields(Base.getType(), BaseOffset, Fields, false, IsWithinUnion);
   }
   if (IncludeVBases) {
     for (auto &Base : Decl->vbases()) {
       CharUnits BaseOffset =
           Offset + RL.getVBaseClassOffset(Base.getType()->getAsCXXRecordDecl());
-      findPFPFields(Base.getType(), BaseOffset, Fields, false);
+      findPFPFields(Base.getType(), BaseOffset, Fields, false, IsWithinUnion);
     }
   }
 }
diff --git a/clang/lib/CodeGen/CGBuiltin.cpp b/clang/lib/CodeGen/CGBuiltin.cpp
index d11726f..7be70ee 100644
--- a/clang/lib/CodeGen/CGBuiltin.cpp
+++ b/clang/lib/CodeGen/CGBuiltin.cpp
@@ -4499,14 +4499,13 @@ RValue CodeGenFunction::EmitBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
     Address Dest = EmitPointerWithAlignment(E->getArg(0));
     Address Src = EmitPointerWithAlignment(E->getArg(1));
     Value *SizeVal = EmitScalarExpr(E->getArg(2));
+    Value *TypeSize = ConstantInt::get(
+        SizeVal->getType(),
+        getContext()
+            .getTypeSizeInChars(E->getArg(0)->getType()->getPointeeType())
+            .getQuantity());
     if (BuiltinIDIfNoAsmLabel == Builtin::BI__builtin_trivially_relocate)
-      SizeVal = Builder.CreateMul(
-          SizeVal,
-          ConstantInt::get(
-              SizeVal->getType(),
-              getContext()
-                  .getTypeSizeInChars(E->getArg(0)->getType()->getPointeeType())
-                  .getQuantity()));
+      SizeVal = Builder.CreateMul(SizeVal, TypeSize);
     EmitArgCheck(TCK_Store, Dest, E->getArg(0), 0);
     EmitArgCheck(TCK_Load, Src, E->getArg(1), 1);
     auto *I = Builder.CreateMemMove(Dest, Src, SizeVal, false);
@@ -4515,13 +4514,38 @@ RValue CodeGenFunction::EmitBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
       std::vector<PFPField> PFPFields;
       getContext().findPFPFields(E->getArg(0)->getType()->getPointeeType(),
                                  CharUnits::Zero(), PFPFields, true);
-      for (auto &Field : PFPFields) {
-        if (getContext().arePFPFieldsTriviallyRelocatable(
-                Field.field->getParent()))
-          continue;
-        auto DestFieldPtr = EmitAddressOfPFPField(Dest, Field);
-        auto SrcFieldPtr = EmitAddressOfPFPField(Src, Field);
-        Builder.CreateStore(Builder.CreateLoad(SrcFieldPtr), DestFieldPtr);
+      if (!PFPFields.empty()) {
+        BasicBlock *Entry = Builder.GetInsertBlock();
+        BasicBlock *Loop = createBasicBlock("loop");
+        BasicBlock *LoopEnd = createBasicBlock("loop.end");
+        Builder.CreateCondBr(
+            Builder.CreateICmpEQ(SizeVal,
+                                 ConstantInt::get(SizeVal->getType(), 0)),
+            LoopEnd, Loop);
+
+        EmitBlock(Loop);
+        PHINode *Offset = Builder.CreatePHI(SizeVal->getType(), 2);
+        Offset->addIncoming(ConstantInt::get(SizeVal->getType(), 0), Entry);
+        Address DestRec = Dest.withPointer(
+            Builder.CreateInBoundsGEP(Int8Ty, Dest.getBasePointer(), {Offset}),
+            KnownNonNull);
+        Address SrcRec = Src.withPointer(
+            Builder.CreateInBoundsGEP(Int8Ty, Src.getBasePointer(), {Offset}),
+            KnownNonNull);
+        for (auto &Field : PFPFields) {
+          if (getContext().arePFPFieldsTriviallyRelocatable(
+                  Field.field->getParent()))
+            continue;
+          auto DestFieldPtr = EmitAddressOfPFPField(DestRec, Field);
+          auto SrcFieldPtr = EmitAddressOfPFPField(SrcRec, Field);
+          Builder.CreateStore(Builder.CreateLoad(SrcFieldPtr), DestFieldPtr);
+        }
+
+        Value *NextOffset = Builder.CreateAdd(Offset, TypeSize);
+        Offset->addIncoming(NextOffset, Loop);
+        Builder.CreateCondBr(Builder.CreateICmpEQ(NextOffset, SizeVal), LoopEnd, Loop);
+
+        EmitBlock(LoopEnd);
       }
     }
     return RValue::get(Dest, *this);
diff --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp
index 3457ec0..cd86818 100644
--- a/clang/lib/CodeGen/CGExpr.cpp
+++ b/clang/lib/CodeGen/CGExpr.cpp
@@ -5112,6 +5112,7 @@ static Address emitRawAddrOfFieldStorage(CodeGenFunction &CGF, Address base,
     return emitAddrOfZeroSizeField(CGF, base, field, IsInBounds);
 
   const RecordDecl *rec = field->getParent();
+
   unsigned idx =
     CGF.CGM.getTypes().getCGRecordLayout(rec).getLLVMFieldNo(field);
 
diff --git a/clang/lib/Sema/SemaTypeTraits.cpp b/clang/lib/Sema/SemaTypeTraits.cpp
index c06f530..c171b70 100644
--- a/clang/lib/Sema/SemaTypeTraits.cpp
+++ b/clang/lib/Sema/SemaTypeTraits.cpp
@@ -235,16 +235,16 @@ static bool IsEligibleForReplacement(Sema &SemaRef, const CXXRecordDecl *D) {
 
 static bool IsImplementationDefinedNonRelocatable(Sema &SemaRef,
                                                   const CXXRecordDecl *D) {
-  // FIXME: Should also check for polymorphic union members here if PAuth ABI is
-  // enabled.
-
-  // FIXME: PFP should not affect trivial relocatability except in cases where a
-  // PFP field is a member of a union, instead it should affect the
-  // implementation of std::trivially_relocate. See:
-  // https://discourse.llvm.org/t/rfc-structure-protection-a-family-of-uaf-mitigation-techniques/85555/16?u=pcc
-  if (!SemaRef.Context.arePFPFieldsTriviallyRelocatable(D) &&
-      SemaRef.Context.hasPFPFields(QualType(D->getTypeForDecl(), 0)))
-    return true;
+  // The implementation-defined carveout only exists for polymorphic types.
+  if (!D->isPolymorphic())
+    return false;
+
+  std::vector<PFPField> pfpFields;
+  SemaRef.Context.findPFPFields(QualType(D->getTypeForDecl(), 0),
+                                CharUnits::Zero(), pfpFields, true);
+  for (PFPField f : pfpFields)
+    if (f.isWithinUnion)
+      return true;
 
   return false;
 }
diff --git a/clang/test/CodeGenCXX/pfp-coerce.cpp b/clang/test/CodeGenCXX/pfp-coerce.cpp
index db53f2b..636d61c 100644
--- a/clang/test/CodeGenCXX/pfp-coerce.cpp
+++ b/clang/test/CodeGenCXX/pfp-coerce.cpp
@@ -186,7 +186,7 @@ void pass_trivial_abi_pointer(TrivialAbiPointer p, TrivialAbiPointer *pp) {
   // X86_64: store ptr %p.coerce, ptr %1, align 8
   // X86_64: store ptr %pp, ptr %pp.addr, align 8
   // X86_64: %2 = load ptr, ptr %pp.addr, align 8
-  // X86_64: %call = call noundef nonnull align 8 dereferenceable(8) ptr @_ZN17TrivialAbiPointeraSERKS_(ptr noundef nonnull align 8 dereferenceable(8) %2, ptr noundef nonnull align 8 dereferenceable(8) %p)
+  // X86_64: call void @llvm.memcpy.p0.p0.i64(ptr align 8 %2, ptr align 8 %p, i64 8, i1 false)
   // X86_64: call void @_ZN17TrivialAbiPointerD1Ev(ptr noundef nonnull align 8 dereferenceable(8) %p)
   *pp = p;
 }
@@ -210,7 +210,7 @@ TrivialAbiPointer return_trivial_abi_pointer(TrivialAbiPointer *pp) {
   // X86_64: %pp.addr = alloca ptr, align 8
   // X86_64: store ptr %pp, ptr %pp.addr, align 8
   // X86_64: %0 = load ptr, ptr %pp.addr, align 8
-  // X86_64: call void @_ZN17TrivialAbiPointerC1ERKS_(ptr noundef nonnull align 8 dereferenceable(8) %retval, ptr noundef nonnull align 8 dereferenceable(8) %0)
+  // X86_64: call void @llvm.memcpy.p0.p0.i64(ptr align 8 %retval, ptr align 8 %0, i64 8, i1 false)
   // X86_64: %1 = getelementptr inbounds i8, ptr %retval, i64 0
   // X86_64: %2 = call ptr @llvm.protected.field.ptr(ptr %1, i64 33, i1 false) [ "deactivation-symbol"(ptr @__pfp_ds__ZTS17TrivialAbiPointer.ptr) ]
   // X86_64: %3 = load ptr, ptr %2, align 8
diff --git a/clang/test/CodeGenCXX/pfp-memcpy.cpp b/clang/test/CodeGenCXX/pfp-memcpy.cpp
index 3e92db6..9df6355 100644
--- a/clang/test/CodeGenCXX/pfp-memcpy.cpp
+++ b/clang/test/CodeGenCXX/pfp-memcpy.cpp
@@ -1,4 +1,4 @@
-// RUN: %clang_cc1  -fexperimental-pointer-field-protection=tagged -emit-llvm -o - %s | FileCheck %s
+// RUN: %clang_cc1 -triple aarch64-linux -fexperimental-pointer-field-protection=tagged -emit-llvm -o - %s | FileCheck %s
 
 struct ClassWithTrivialCopy {
   ClassWithTrivialCopy();
@@ -16,4 +16,4 @@ void make_trivial_copy(ClassWithTrivialCopy *s1, ClassWithTrivialCopy *s2) {
 
 // CHECK-LABEL: define{{.*}} void @_Z17make_trivial_copyP20ClassWithTrivialCopyS0_
 // CHECK-NOT: memcpy
-// CHECK: ret void
\ No newline at end of file
+// CHECK: ret void
diff --git a/clang/test/CodeGenCXX/pfp-trivially-relocatable.cpp b/clang/test/CodeGenCXX/pfp-trivially-relocatable.cpp
new file mode 100644
index 0000000..e540d26
--- /dev/null
+++ b/clang/test/CodeGenCXX/pfp-trivially-relocatable.cpp
@@ -0,0 +1,101 @@
+// RUN: %clang_cc1 -std=c++26 -triple x86_64-linux-gnu -emit-llvm -fexperimental-pointer-field-protection=untagged -o - %s | FileCheck --check-prefix=RELOC %s
+// RUN: %clang_cc1 -std=c++26 -triple x86_64-linux-gnu -emit-llvm -fexperimental-pointer-field-protection=tagged -o - %s | FileCheck --check-prefix=RELOC %s
+// RUN: %clang_cc1 -std=c++26 -triple aarch64-linux-gnu -emit-llvm -fexperimental-pointer-field-protection=untagged -o - %s | FileCheck --check-prefix=RELOC %s
+// RUN: %clang_cc1 -std=c++26 -triple aarch64-linux-gnu -emit-llvm -fexperimental-pointer-field-protection=tagged -o - %s | FileCheck --check-prefix=NONRELOC %s
+
+typedef __SIZE_TYPE__ size_t;
+
+struct S trivially_relocatable_if_eligible {
+    S(const S&);
+    ~S();
+    int* a;
+private:
+    int* b;
+};
+
+// CHECK: define dso_local void @_Z5test1P1SS0_(
+void test1(S* source, S* dest) {
+  // RELOC:       %0 = load ptr, ptr %dest.addr, align 8
+  // RELOC-NEXT:  %1 = load ptr, ptr %source.addr, align 8
+  // RELOC-NEXT:  call void @llvm.memmove.p0.p0.i64(ptr align 8 %0, ptr align 8 %1, i64 16, i1 false)
+  // RELOC-NOT: @llvm.protected.field.ptr
+
+  // NONRELOC:        %0 = load ptr, ptr %dest.addr, align 8
+  // NONRELOC-NEXT:   %1 = load ptr, ptr %source.addr, align 8
+  // NONRELOC-NEXT:   call void @llvm.memmove.p0.p0.i64(ptr align 8 %0, ptr align 8 %1, i64 16, i1 false)
+  // NONRELOC-NEXT:   br i1 false, label %loop.end, label %loop
+
+  // NONRELOC:      loop:
+  // NONRELOC-NEXT:   %2 = phi i64 [ 0, %entry ], [ %19, %loop ]
+  // NONRELOC-NEXT:   %3 = getelementptr inbounds i8, ptr %0, i64 %2
+  // NONRELOC-NEXT:   %4 = getelementptr inbounds i8, ptr %1, i64 %2
+  // NONRELOC-NEXT:   %5 = getelementptr inbounds i8, ptr %3, i64 0
+  // NONRELOC-NEXT:   %6 = ptrtoint ptr %3 to i64
+  // NONRELOC-NEXT:   %7 = call ptr @llvm.protected.field.ptr(ptr %5, i64 %6, i1 true) [ "deactivation-symbol"(ptr @__pfp_ds__ZTS1S.a) ]
+  // NONRELOC-NEXT:   %8 = getelementptr inbounds i8, ptr %4, i64 0
+  // NONRELOC-NEXT:   %9 = ptrtoint ptr %4 to i64
+  // NONRELOC-NEXT:   %10 = call ptr @llvm.protected.field.ptr(ptr %8, i64 %9, i1 true) [ "deactivation-symbol"(ptr @__pfp_ds__ZTS1S.a) ]
+  // NONRELOC-NEXT:   %11 = load ptr, ptr %10, align 8
+  // NONRELOC-NEXT:   store ptr %11, ptr %7, align 8
+  // NONRELOC-NEXT:   %12 = getelementptr inbounds i8, ptr %3, i64 8
+  // NONRELOC-NEXT:   %13 = ptrtoint ptr %3 to i64
+  // NONRELOC-NEXT:   %14 = call ptr @llvm.protected.field.ptr(ptr %12, i64 %13, i1 true) [ "deactivation-symbol"(ptr @__pfp_ds__ZTS1S.b) ]
+  // NONRELOC-NEXT:   %15 = getelementptr inbounds i8, ptr %4, i64 8
+  // NONRELOC-NEXT:   %16 = ptrtoint ptr %4 to i64
+  // NONRELOC-NEXT:   %17 = call ptr @llvm.protected.field.ptr(ptr %15, i64 %16, i1 true) [ "deactivation-symbol"(ptr @__pfp_ds__ZTS1S.b) ]
+  // NONRELOC-NEXT:   %18 = load ptr, ptr %17, align 8
+  // NONRELOC-NEXT:   store ptr %18, ptr %14, align 8
+  // NONRELOC-NEXT:   %19 = add i64 %2, 16
+  // NONRELOC-NEXT:   %20 = icmp eq i64 %19, 16
+  // NONRELOC-NEXT:   br i1 %20, label %loop.end, label %loop
+
+  // NONRELOC:      loop.end:
+  // NONRELOC-NEXT:   ret void
+  __builtin_trivially_relocate(dest, source, 1);
+}
+
+// CHECK: define dso_local void @_Z5testNP1SS0_m(
+void testN(S* source, S* dest, size_t count) {
+  // RELOC:       %0 = load ptr, ptr %dest.addr, align 8
+  // RELOC-NEXT:  %1 = load ptr, ptr %source.addr, align 8
+  // RELOC-NEXT:  %2 = load i64, ptr %count.addr, align 8
+  // RELOC-NEXT:  %3 = mul i64 %2, 16
+  // RELOC-NEXT:  call void @llvm.memmove.p0.p0.i64(ptr align 8 %0, ptr align 8 %1, i64 %3, i1 false)
+  // RELOC-NOT: @llvm.protected.field.ptr
+ 
+  // NONRELOC:        %0 = load ptr, ptr %dest.addr, align 8
+  // NONRELOC-NEXT:   %1 = load ptr, ptr %source.addr, align 8
+  // NONRELOC-NEXT:   %2 = load i64, ptr %count.addr, align 8
+  // NONRELOC-NEXT:   %3 = mul i64 %2, 16
+  // NONRELOC-NEXT:   call void @llvm.memmove.p0.p0.i64(ptr align 8 %0, ptr align 8 %1, i64 %3, i1 false)
+  // NONRELOC-NEXT:   %4 = icmp eq i64 %3, 0
+  // NONRELOC-NEXT:   br i1 %4, label %loop.end, label %loop
+
+  // NONRELOC:      loop:
+  // NONRELOC-NEXT:   %5 = phi i64 [ 0, %entry ], [ %22, %loop ]
+  // NONRELOC-NEXT:   %6 = getelementptr inbounds i8, ptr %0, i64 %5
+  // NONRELOC-NEXT:   %7 = getelementptr inbounds i8, ptr %1, i64 %5
+  // NONRELOC-NEXT:   %8 = getelementptr inbounds i8, ptr %6, i64 0
+  // NONRELOC-NEXT:   %9 = ptrtoint ptr %6 to i64
+  // NONRELOC-NEXT:   %10 = call ptr @llvm.protected.field.ptr(ptr %8, i64 %9, i1 true) [ "deactivation-symbol"(ptr @__pfp_ds__ZTS1S.a) ]
+  // NONRELOC-NEXT:   %11 = getelementptr inbounds i8, ptr %7, i64 0
+  // NONRELOC-NEXT:   %12 = ptrtoint ptr %7 to i64
+  // NONRELOC-NEXT:   %13 = call ptr @llvm.protected.field.ptr(ptr %11, i64 %12, i1 true) [ "deactivation-symbol"(ptr @__pfp_ds__ZTS1S.a) ]
+  // NONRELOC-NEXT:   %14 = load ptr, ptr %13, align 8
+  // NONRELOC-NEXT:   store ptr %14, ptr %10, align 8
+  // NONRELOC-NEXT:   %15 = getelementptr inbounds i8, ptr %6, i64 8
+  // NONRELOC-NEXT:   %16 = ptrtoint ptr %6 to i64
+  // NONRELOC-NEXT:   %17 = call ptr @llvm.protected.field.ptr(ptr %15, i64 %16, i1 true) [ "deactivation-symbol"(ptr @__pfp_ds__ZTS1S.b) ]
+  // NONRELOC-NEXT:   %18 = getelementptr inbounds i8, ptr %7, i64 8
+  // NONRELOC-NEXT:   %19 = ptrtoint ptr %7 to i64
+  // NONRELOC-NEXT:   %20 = call ptr @llvm.protected.field.ptr(ptr %18, i64 %19, i1 true) [ "deactivation-symbol"(ptr @__pfp_ds__ZTS1S.b) ]
+  // NONRELOC-NEXT:   %21 = load ptr, ptr %20, align 8
+  // NONRELOC-NEXT:   store ptr %21, ptr %17, align 8
+  // NONRELOC-NEXT:   %22 = add i64 %5, 16
+  // NONRELOC-NEXT:   %23 = icmp eq i64 %22, %3
+  // NONRELOC-NEXT:   br i1 %23, label %loop.end, label %loop
+
+  // NONRELOC:      loop.end:
+  // NONRELOC-NEXT:   ret void
+  __builtin_trivially_relocate(dest, source, count);
+};