Add pointer field protection feature.

Pointer field protection is a use-after-free vulnerability
mitigation that works by changing how data structures' pointer
fields are stored in memory. For more information, see the RFC:
https://discourse.llvm.org/t/rfc-structure-protection-a-family-of-uaf-mitigation-techniques/85555

TODO:
- Fix test failure.
- Add more tests.
- Add documentation.

Pull Request: llvm#133538

Author: pcc

clang/include/clang/AST/ASTContext.h

@@ -183,6 +183,12 @@ struct TypeInfoChars {
   }
 };
 
+struct PFPField {
+  CharUnits structOffset;
+  CharUnits offset;
+  FieldDecl *field;
+};
+
 /// Holds long-lived AST nodes (such as types and decls) that can be
 /// referred to throughout the semantic analysis of a file.
 class ASTContext : public RefCountedBase<ASTContext> {
@@ -3668,6 +3674,22 @@ OPT_LIST(V)
 
   StringRef getCUIDHash() const;
 
+  bool isPFPStruct(const RecordDecl *rec) const;
+  void findPFPFields(QualType Ty, CharUnits Offset,
+                     std::vector<PFPField> &Fields, bool IncludeVBases) const;
+  bool hasPFPFields(QualType ty) const;
+  bool isPFPField(const FieldDecl *field) const;
+
+  /// Returns whether this record's PFP fields (if any) are trivially
+  /// relocatable (i.e. may be memcpy'd). This may also return true if the
+  /// record does not have any PFP fields, so it may be necessary for the caller
+  /// to check for PFP fields, e.g. by calling hasPFPFields().
+  bool arePFPFieldsTriviallyRelocatable(const RecordDecl *RD) const;
+
+  llvm::SetVector<const FieldDecl *> PFPFieldsWithEvaluatedOffset;
+  void recordMemberDataPointerEvaluation(const ValueDecl *VD);
+  void recordOffsetOfEvaluation(const OffsetOfExpr *E);
+
 private:
   /// All OMPTraitInfo objects live in this collection, one per
   /// `pragma omp [begin] declare variant` directive.

clang/include/clang/Basic/Attr.td

@@ -2524,6 +2524,12 @@ def CountedByOrNull : DeclOrTypeAttr {
   let LangOpts = [COnly];
 }
 
+def NoPointerFieldProtection : DeclOrTypeAttr {
+  let Spellings = [Clang<"no_field_protection">];
+  let Subjects = SubjectList<[Field], ErrorDiag>;
+  let Documentation = [Undocumented];
+}
+
 def SizedBy : DeclOrTypeAttr {
   let Spellings = [Clang<"sized_by">];
   let Subjects = SubjectList<[Field], ErrorDiag>;

clang/include/clang/Basic/Features.def

@@ -265,6 +265,9 @@ FEATURE(shadow_call_stack,
 FEATURE(tls, PP.getTargetInfo().isTLSSupported())
 FEATURE(underlying_type, LangOpts.CPlusPlus)
 FEATURE(experimental_library, LangOpts.ExperimentalLibrary)
+FEATURE(pointer_field_protection,
+        LangOpts.getPointerFieldProtection() !=
+        LangOptions::PointerFieldProtectionKind::None)
 
 // C11 features supported by other languages as extensions.
 EXTENSION(c_alignas, true)

clang/include/clang/Basic/LangOptions.def

@@ -500,6 +500,9 @@ LANGOPT(RelativeCXXABIVTables, 1, 0,
 LANGOPT(OmitVTableRTTI, 1, 0,
         "Use an ABI-incompatible v-table layout that omits the RTTI component")
 
+ENUM_LANGOPT(PointerFieldProtection, PointerFieldProtectionKind, 2, PointerFieldProtectionKind::None,
+             "Encode struct pointer fields to protect against UAF vulnerabilities")
+
 LANGOPT(VScaleMin, 32, 0, "Minimum vscale value")
 LANGOPT(VScaleMax, 32, 0, "Maximum vscale value")
 

clang/include/clang/Basic/LangOptions.h

@@ -362,6 +362,17 @@ class LangOptionsBase {
     BKey
   };
 
+  enum class PointerFieldProtectionKind {
+    /// Pointer field protection disabled
+    None,
+    /// Pointer field protection enabled, allocator does not tag heap
+    /// allocations.
+    Untagged,
+    /// Pointer field protection enabled, allocator is expected to tag heap
+    /// allocations.
+    Tagged,
+  };
+
   enum class ThreadModelKind {
     /// POSIX Threads.
     POSIX,

clang/include/clang/Driver/Options.td

@@ -3011,6 +3011,12 @@ defm experimental_omit_vtable_rtti : BoolFOption<"experimental-omit-vtable-rtti"
   NegFlag<SetFalse, [], [CC1Option], "Do not omit">,
   BothFlags<[], [CC1Option], " the RTTI component from virtual tables">>;
 
+def experimental_pointer_field_protection_EQ : Joined<["-"], "fexperimental-pointer-field-protection=">, Group<f_Group>,
+  Visibility<[ClangOption, CC1Option]>,
+  Values<"none,untagged,tagged">, NormalizedValuesScope<"LangOptions::PointerFieldProtectionKind">,
+  NormalizedValues<["None", "Untagged", "Tagged"]>,
+  MarshallingInfoEnum<LangOpts<"PointerFieldProtection">, "None">;
+
 def fcxx_abi_EQ : Joined<["-"], "fc++-abi=">,
                   Group<f_clang_Group>, Visibility<[ClangOption, CC1Option]>,
                   HelpText<"C++ ABI to use. This will override the target C++ ABI.">;

clang/lib/AST/ASTContext.cpp

@@ -80,6 +80,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/Support/Capacity.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MD5.h"
@@ -15113,3 +15114,97 @@ bool ASTContext::useAbbreviatedThunkName(GlobalDecl VirtualMethodDecl,
   ThunksToBeAbbreviated[VirtualMethodDecl] = std::move(SimplifiedThunkNames);
   return Result;
 }
+
+bool ASTContext::arePFPFieldsTriviallyRelocatable(const RecordDecl *RD) const {
+  if (getLangOpts().getPointerFieldProtection() ==
+      LangOptions::PointerFieldProtectionKind::Tagged)
+    return !isa<CXXRecordDecl>(RD) ||
+           cast<CXXRecordDecl>(RD)->hasTrivialDestructor();
+  return true;
+}
+
+bool ASTContext::isPFPStruct(const RecordDecl *rec) const {
+  if (getLangOpts().getPointerFieldProtection() !=
+      LangOptions::PointerFieldProtectionKind::None)
+    if (auto *cxxRec = dyn_cast<CXXRecordDecl>(rec))
+      return !cxxRec->isStandardLayout();
+  return false;
+}
+
+void ASTContext::findPFPFields(QualType Ty, CharUnits Offset,
+                               std::vector<PFPField> &Fields,
+                               bool IncludeVBases) const {
+  if (auto *AT = getAsConstantArrayType(Ty)) {
+    if (auto *ElemDecl = AT->getElementType()->getAsCXXRecordDecl()) {
+      const ASTRecordLayout &ElemRL = getASTRecordLayout(ElemDecl);
+      for (unsigned i = 0; i != AT->getSize(); ++i) {
+        findPFPFields(AT->getElementType(), Offset + i * ElemRL.getSize(),
+                      Fields, true);
+      }
+    }
+  }
+  auto *Decl = Ty->getAsCXXRecordDecl();
+  if (!Decl)
+    return;
+  const ASTRecordLayout &RL = getASTRecordLayout(Decl);
+  for (FieldDecl *field : Decl->fields()) {
+    CharUnits fieldOffset =
+        Offset + toCharUnitsFromBits(RL.getFieldOffset(field->getFieldIndex()));
+    if (isPFPField(field))
+      Fields.push_back({Offset, fieldOffset, field});
+    findPFPFields(field->getType(), fieldOffset, Fields, true);
+  }
+  for (auto &Base : Decl->bases()) {
+    if (Base.isVirtual())
+      continue;
+    CharUnits BaseOffset =
+        Offset + RL.getBaseClassOffset(Base.getType()->getAsCXXRecordDecl());
+    findPFPFields(Base.getType(), BaseOffset, Fields, false);
+  }
+  if (IncludeVBases) {
+    for (auto &Base : Decl->vbases()) {
+      CharUnits BaseOffset =
+          Offset + RL.getVBaseClassOffset(Base.getType()->getAsCXXRecordDecl());
+      findPFPFields(Base.getType(), BaseOffset, Fields, false);
+    }
+  }
+}
+
+bool ASTContext::hasPFPFields(QualType ty) const {
+  std::vector<PFPField> pfpFields;
+  findPFPFields(ty, CharUnits::Zero(), pfpFields, true);
+  return !pfpFields.empty();
+}
+
+bool ASTContext::isPFPField(const FieldDecl *field) const {
+  if (!isPFPStruct(field->getParent()))
+    return false;
+  return field->getType()->isPointerType() &&
+         !field->hasAttr<NoPointerFieldProtectionAttr>();
+}
+
+void ASTContext::recordMemberDataPointerEvaluation(const ValueDecl *VD) {
+  if (getLangOpts().getPointerFieldProtection() ==
+      LangOptions::PointerFieldProtectionKind::None)
+    return;
+  auto *FD = dyn_cast<FieldDecl>(VD);
+  if (!FD)
+    FD = cast<FieldDecl>(cast<IndirectFieldDecl>(VD)->chain().back());
+  if (!isPFPField(FD))
+    return;
+  PFPFieldsWithEvaluatedOffset.insert(FD);
+}
+
+void ASTContext::recordOffsetOfEvaluation(const OffsetOfExpr *E) {
+  if (getLangOpts().getPointerFieldProtection() ==
+          LangOptions::PointerFieldProtectionKind::None ||
+      E->getNumComponents() == 0)
+    return;
+  OffsetOfNode Comp = E->getComponent(E->getNumComponents() - 1);
+  if (Comp.getKind() != OffsetOfNode::Field)
+    return;
+  FieldDecl *FD = Comp.getField();
+  if (!isPFPField(FD))
+    return;
+  PFPFieldsWithEvaluatedOffset.insert(FD);
+}

clang/lib/AST/ExprConstant.cpp

@@ -14979,6 +14979,7 @@ bool IntExprEvaluator::VisitUnaryExprOrTypeTraitExpr(
 }
 
 bool IntExprEvaluator::VisitOffsetOfExpr(const OffsetOfExpr *OOE) {
+  Info.Ctx.recordOffsetOfEvaluation(OOE);
   CharUnits Result;
   unsigned n = OOE->getNumComponents();
   if (n == 0)

clang/lib/AST/TypePrinter.cpp

@@ -2098,6 +2098,9 @@ void TypePrinter::printAttributedAfter(const AttributedType *T,
   case attr::ExtVectorType:
     OS << "ext_vector_type";
     break;
+  case attr::NoPointerFieldProtection:
+    OS << "no_field_protection";
+    break;
   }
   OS << "))";
 }

clang/lib/CodeGen/CGBuiltin.cpp

@@ -4464,6 +4464,19 @@ RValue CodeGenFunction::EmitBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
     EmitArgCheck(TCK_Store, Dest, E->getArg(0), 0);
     EmitArgCheck(TCK_Load, Src, E->getArg(1), 1);
     Builder.CreateMemMove(Dest, Src, SizeVal, false);
+    if (BuiltinIDIfNoAsmLabel == Builtin::BI__builtin_trivially_relocate) {
+      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);
+      }
+    }
     return RValue::get(Dest, *this);
   }
   case Builtin::BImemset: