[clang] Implement the ptrauth_struct type attribute.

Author: ahmedbougacha

clang/include/clang/AST/ASTContext.h

@@ -1239,6 +1239,15 @@ class ASTContext : public RefCountedBase<ASTContext> {
   uint16_t
   getPointerAuthVTablePointerDiscriminator(const CXXRecordDecl *record);
 
+  bool typeContainsAuthenticatedNull(QualType) const;
+  bool typeContainsAuthenticatedNull(const Type *) const;
+  std::optional<bool> tryTypeContainsAuthenticatedNull(QualType) const;
+
+  /// Return the key of attribute ptrauth_struct on the record. If the attribute
+  /// isn't on the record, return the none key. If the key argument is value
+  /// dependent, set the boolean flag to false.
+  std::pair<bool, int> getPointerAuthStructKey(const RecordDecl *RD) const;
+
   /// Return the discriminator of attribute ptrauth_struct on the record. If the
   /// key argument is value dependent, set the boolean flag to false.
   std::pair<bool, unsigned>
@@ -1253,12 +1262,19 @@ class ASTContext : public RefCountedBase<ASTContext> {
     if (PointeeType->isFunctionType())
       return false;
 
+    // Disallow the qualifer on classes annotated with attribute ptrauth_struct
+    // with a key that isn't the none key.
+    if (auto *RD = PointeeType->getAsRecordDecl()) {
+      std::pair<bool, int> P = getPointerAuthStructKey(RD);
+      if (P.first && P.second != PointerAuthKeyNone)
+        return false;
+    }
+
     return true;
   }
 
-  bool typeContainsAuthenticatedNull(QualType) const;
-  bool typeContainsAuthenticatedNull(const Type *) const;
-  std::optional<bool> tryTypeContainsAuthenticatedNull(QualType) const;
+  bool hasPointerAuthStructMismatch(const RecordDecl *RD0,
+                                    const RecordDecl *RD1) const;
 
   /// Apply Objective-C protocol qualifiers to the given type.
   /// \param allowOnPointerType specifies if we can apply protocol

clang/include/clang/AST/CXXRecordDeclDefinitionBits.def

@@ -249,4 +249,8 @@ FIELD(HasDeclaredCopyAssignmentWithConstParam, 1, MERGE_OR)
 /// base classes or fields have a no-return destructor
 FIELD(IsAnyDestructorNoReturn, 1, NO_MERGE)
 
+/// Whether this class or any of its direct or indirect base classes are
+/// annotated with `ptrauth_struct` that uses a key that isn't the none key.
+FIELD(HasSignedClassInHierarchy, 1, NO_MERGE)
+
 #undef FIELD

clang/include/clang/AST/DeclCXX.h

@@ -1005,6 +1005,14 @@ class CXXRecordDecl : public RecordDecl {
     return data().NeedOverloadResolutionForDestructor;
   }
 
+  bool hasSignedClassInHierarchy() const {
+    return data().HasSignedClassInHierarchy;
+  }
+
+  void setHasSignedClassInHierarchy() {
+    data().HasSignedClassInHierarchy = true;
+  }
+
   /// Determine whether this class describes a lambda function object.
   bool isLambda() const {
     // An update record can't turn a non-lambda into a lambda.

clang/include/clang/Basic/Attr.td

@@ -2994,6 +2994,13 @@ def PointerAuth : TypeAttr {
   let Documentation = [PtrAuthDocs];
 }
 
+def PointerAuthStruct : InheritableAttr {
+  let Spellings = [Clang<"ptrauth_struct">];
+  let Args = [ExprArgument<"Key">, ExprArgument<"Discriminator">];
+  let Subjects = SubjectList<[Record]>;
+  let Documentation = [PointerAuthStructDocs];
+}
+
 def Unused : InheritableAttr {
   let Spellings = [CXX11<"", "maybe_unused", 201603>, GCC<"unused">,
                    C23<"", "maybe_unused", 202106>];

clang/include/clang/Basic/AttrDocs.td

@@ -1887,6 +1887,47 @@ features that implicitly use pointer authentication.
   }];
 }
 
+def PointerAuthStructDocs : Documentation {
+  let Category = DocCatDecl;
+  let Content = [{
+The ``ptrauth_struct`` attribute allows programmers to opt in to using signed
+pointers for all pointers to a particular struct, union, or C++ class type. In
+C++, this also includes references to the type.
+
+The attribute takes two arguments:
+
+- The first argument is the abstract signing key, which should be a constant
+  from ``<ptrauth.h>``.
+- The second argument is a constant discriminator.
+
+These arguments are identical to the arguments to the ``__ptrauth`` qualifier
+except that ``ptrauth_struct`` does not take an argument for enabling address
+diversity.
+
+.. code-block:: c++
+
+  // key=ptrauth_key_process_dependent_data,discriminator=100
+  struct __attribute__((ptrauth_struct(ptrauth_key_process_dependent_data,100))) S0 {
+    int f0[4];
+  };
+
+If a particular pointer to the type is qualified with a ``__ptrauth`` qualifier,
+that qualifier takes precedence over ptrauth_struct:
+
+.. code-block:: c++
+
+  struct __attribute__((ptrauth_struct(ptrauth_key_process_dependent_data,100))) S0 {
+    int f0[4];
+  };
+
+  void test(S0 *s0) {
+    // `p` is signed using `__ptrauth(ptrauth_key_process_dependent_data, 1, 200)`.
+    S0 * __ptrauth(ptrauth_key_process_dependent_data, 1, 200) p = s0;
+  }
+
+  }];
+}
+
 def ExternalSourceSymbolDocs : Documentation {
   let Category = DocCatDecl;
   let Content = [{

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -904,6 +904,15 @@ def note_ptrauth_virtual_function_incomplete_arg_ret_type :
 def err_ptrauth_type_disc_undiscriminated : Error<
   "cannot pass undiscriminated type %0 to "
   "'__builtin_ptrauth_type_discriminator'">;
+def err_ptrauth_struct_key_disc_not_record : Error<
+  "argument to '__builtin_ptrauth_type_discriminator' isn't a record type">;
+def err_ptrauth_struct_signing_mismatch : Error<
+  "%0 is signed differently from the previous declaration">;
+def note_previous_ptrauth_struct_declaration : Note<
+  "previous declaration of %0 is here">;
+def err_ptrauth_invalid_struct_attr_dynamic_class : Error<
+  "attribute ptrauth_struct cannot be used on class %0 or its subclasses "
+  "because it is a dynamic class">;
 def err_ptrauth_non_string_authentication_option : Error<
   "__ptrauth options must be a string of comma separated flags, found %0">;
 def err_ptrauth_unknown_authentication_option : Error<
@@ -939,11 +948,12 @@ def err_ptrauth_qualifier_signed_pointer_type : Error<
 def err_ptrauth_qualifier_bad_arg_count : Error<
   "%0 qualifier must take between 1 and 4 arguments">;
 def err_ptrauth_arg_not_ice : Error<
-  "argument to __ptrauth must be an integer constant expression">;
+  "argument to %select{__ptrauth|ptrauth_struct}0 must be an integer constant "
+  "expression">;
 def err_ptrauth_address_discrimination_invalid : Error<
   "address discrimination flag for __ptrauth must be 0 or 1; value is %0">;
 def err_ptrauth_extra_discriminator_invalid : Error<
-  "extra discriminator for __ptrauth must be between "
+  "extra discriminator for %select{__ptrauth|ptrauth_struct}2 must be between "
   "0 and %1; value is %0">;
 
 /// main()

clang/include/clang/Basic/TokenKinds.def

@@ -576,6 +576,8 @@ KEYWORD(__private_extern__          , KEYALL)
 KEYWORD(__module_private__          , KEYALL)
 
 UNARY_EXPR_OR_TYPE_TRAIT(__builtin_ptrauth_type_discriminator, PtrAuthTypeDiscriminator, KEYALL)
+UNARY_EXPR_OR_TYPE_TRAIT(__builtin_ptrauth_struct_key, PtrAuthStructKey, KEYALL)
+UNARY_EXPR_OR_TYPE_TRAIT(__builtin_ptrauth_struct_disc, PtrAuthStructDiscriminator, KEYALL)
 
 // Extension that will be enabled for Microsoft, Borland and PS4, but can be
 // disabled via '-fno-declspec'.

clang/include/clang/Parse/Parser.h

@@ -3628,6 +3628,7 @@ class Parser : public CodeCompletionHandler {
   ExprResult ParseExpressionTrait();
 
   ExprResult ParseBuiltinPtrauthTypeDiscriminator();
+  ExprResult ParseBuiltinPtrauthStructKeyOrDisc(bool ParseKey);
 
   //===--------------------------------------------------------------------===//
   // Preprocessor code-completion pass-through

clang/include/clang/Sema/Sema.h

@@ -2936,6 +2936,9 @@ class Sema final {
 
     // Extra discriminator argument of __ptrauth.
     PADAK_ExtraDiscPtrAuth,
+
+    // Type discriminator argument of ptrauth_struct.
+    PADAK_TypeDiscPtrAuthStruct,
   };
 
   bool checkPointerAuthDiscriminatorArg(Expr *arg, PointerAuthDiscArgKind kind,
@@ -4766,6 +4769,9 @@ class Sema final {
   bool checkCommonAttributeFeatures(const Stmt *S, const ParsedAttr &A,
                                     bool SkipArgCountCheck = false);
 
+  void addPointerAuthStructAttr(Decl *D, const AttributeCommonInfo &CI,
+                                Expr *Key, Expr *Disc);
+
   /// Determine if type T is a valid subject for a nonnull and similar
   /// attributes. By default, we look through references (the behavior used by
   /// nonnull), but if the second parameter is true, then we treat a reference

clang/lib/AST/ASTContext.cpp

@@ -3404,6 +3404,62 @@ uint16_t ASTContext::getPointerAuthVTablePointerDiscriminator(
   return getPointerAuthStringDiscriminator(*this, Str.c_str());
 }
 
+std::pair<bool, int>
+ASTContext::getPointerAuthStructKey(const RecordDecl *RD) const {
+  if (auto *Attr = RD->getAttr<PointerAuthStructAttr>()) {
+    Expr *E = Attr->getKey();
+    if (E->isValueDependent())
+      return {false, 0};
+    std::optional<llvm::APSInt> V = E->getIntegerConstantExpr(*this);
+    return {true, V->getExtValue()};
+  }
+
+  return {true, PointerAuthKeyNone};
+}
+
+std::pair<bool, unsigned>
+ASTContext::getPointerAuthStructDisc(const RecordDecl *RD) const {
+  if (auto *Attr = RD->getAttr<PointerAuthStructAttr>()) {
+    Expr *E = Attr->getDiscriminator();
+    if (E->isValueDependent())
+      return {false, 0};
+    std::optional<llvm::APSInt> V = E->getIntegerConstantExpr(*this);
+    return {true, V->getExtValue()};
+  }
+
+  return {true, 0};
+}
+
+bool ASTContext::hasPointerAuthStructMismatch(const RecordDecl *RD0,
+                                              const RecordDecl *RD1) const {
+  if (RD0 == RD1)
+    return true;
+
+  int Key0, Key1;
+  unsigned Disc0, Disc1;
+  bool Known0, Known1;
+  std::tie(Known0, Key0) = getPointerAuthStructKey(RD0);
+  std::tie(Known1, Key1) = getPointerAuthStructKey(RD1);
+
+  if (!Known0 || !Known1)
+    return false;
+
+  if (Key0 != Key1)
+    return true;
+
+  // If the key is none, skip checking the discriminators.
+  if (Key0 == PointerAuthKeyNone)
+    return false;
+
+  std::tie(Known0, Disc0) = getPointerAuthStructDisc(RD0);
+  std::tie(Known1, Disc1) = getPointerAuthStructDisc(RD1);
+
+  if (!Known0 || !Known1)
+    return false;
+
+  return Disc0 != Disc1;
+}
+
 bool ASTContext::typeContainsAuthenticatedNull(const Type *type) const {
   return typeContainsAuthenticatedNull(QualType(type, 0));
 }