[clang] Implement the ptrauth_struct type attribute.

FIXME: the ptrauth-struct-attr.cpp/.m tests regressed
after 7eca38c, which uses CreateGEP (that eagerly
emits a GEP instr to put a nuw flag on it, which in turn
causes the null-checked resign to be emitted as the raw address),
whereas it previously used CreateInBoundsGEP (which lazily
emits the GEP when it's needed, by accumulating the offset
into the Address, thus allowing the later dereference use
to mark the Address as known-non-null).

I have hacks to fix it, and we can obviously just update the test
to include the null-checks in resigns, but we should come up with
a long term solution.

Author: ahmedbougacha

clang/include/clang/AST/ASTContext.h

@@ -1347,6 +1347,16 @@ class ASTContext : public RefCountedBase<ASTContext> {
   /// Return the "other" type-specific discriminator for the given type.
   uint16_t getPointerAuthTypeDiscriminator(QualType T);
 
+  /// 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, unsigned> 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>
+  getPointerAuthStructDisc(const RecordDecl *RD) const;
+
   // Determine whether the type can have qualifier __ptrauth with key
   // ptrauth_key_none.
   bool canQualifyWithPtrAuthKeyNone(QualType T) const {
@@ -1359,9 +1369,20 @@ 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, unsigned> P = getPointerAuthStructKey(RD);
+      if (P.first && P.second != PointerAuthKeyNone)
+        return false;
+    }
+
     return true;
   }
 
+  bool hasPointerAuthStructMismatch(const RecordDecl *RD0,
+                                    const RecordDecl *RD1) const;
+
   bool typeContainsAuthenticatedNull(QualType) const;
   bool typeContainsAuthenticatedNull(const Type *) const;
   std::optional<bool> tryTypeContainsAuthenticatedNull(QualType) const;

clang/include/clang/AST/CXXRecordDeclDefinitionBits.def

@@ -253,4 +253,8 @@ FIELD(IsAnyDestructorNoReturn, 1, NO_MERGE)
 /// type that is intangible). HLSL only.
 FIELD(IsHLSLIntangible, 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

@@ -1026,6 +1026,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

@@ -3550,6 +3550,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

@@ -2339,6 +2339,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

@@ -1023,11 +1023,12 @@ def err_ptrauth_qualifier_redundant : 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">;
 
 def err_ptrauth_restricted_intptr_qualifier_pointer : Error<
@@ -1056,6 +1057,16 @@ def err_ptrauth_invalid_authenticated_null_global : Error<
   "%select{static locals|globals}0 with authenticated null values are currently unsupported"
 >;
 
+// ptrauth_struct attribute.
+def err_ptrauth_struct_key_disc_not_record : Error<
+  "attribute ptrauth_struct can only be used on 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">;
 
 /// main()
 // static main() is not an error in C, just in C++.

clang/include/clang/Basic/TokenKinds.def

@@ -602,6 +602,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

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

clang/include/clang/Sema/Sema.h

@@ -3544,6 +3544,9 @@ class Sema final : public SemaBase {
 
     // Extra discriminator argument of __ptrauth.
     PADAK_ExtraDiscPtrAuth,
+
+    // Type discriminator argument of ptrauth_struct.
+    PADAK_TypeDiscPtrAuthStruct,
   };
 
   bool checkPointerAuthDiscriminatorArg(Expr *Arg, PointerAuthDiscArgKind Kind,
@@ -4588,6 +4591,9 @@ class Sema final : public SemaBase {
                                       StringRef &Str,
                                       SourceLocation *ArgLocation = nullptr);
 
+  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. Dependent types are considered valid so they can be checked
   /// during instantiation time. By default, we look through references (the

clang/lib/AST/ASTContext.cpp

@@ -3523,6 +3523,62 @@ uint16_t ASTContext::getPointerAuthTypeDiscriminator(QualType T) {
   return llvm::getPointerAuthStableSipHash(Str);
 }
 
+std::pair<bool, unsigned>
+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, static_cast<unsigned>(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;
+
+  unsigned 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;
+}
+
 std::optional<bool>
 ASTContext::tryTypeContainsAuthenticatedNull(QualType T) const {
   if (!LangOpts.PointerAuthIntrinsics)