[clang] Implement pointer authentication for C++ virtual functions, v-tables, and VTTs.

This also implements the ptrauth_vtable_pointer attribute to allow
overriding the default ptrauth schema for vtable pointers.

Co-Authored-By: John McCall <[email protected]>

Author: ojhunt

clang/include/clang/AST/ASTContext.h

@@ -37,6 +37,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/Support/TypeSize.h"
 #include <optional>
@@ -1261,6 +1262,10 @@ class ASTContext : public RefCountedBase<ASTContext> {
   /// space.
   QualType removeAddrSpaceQualType(QualType T) const;
 
+  /// Return the "other" type-specific discriminator for the given type.
+  uint16_t
+  getPointerAuthVTablePointerDiscriminator(const CXXRecordDecl *record);
+
   /// Apply Objective-C protocol qualifiers to the given type.
   /// \param allowOnPointerType specifies if we can apply protocol
   /// qualifiers on ObjCObjectPointerType. It can be set to true when
@@ -3418,12 +3423,20 @@ OPT_LIST(V)
   /// Whether a C++ static variable or CUDA/HIP kernel should be externalized.
   bool shouldExternalize(const Decl *D) const;
 
+  /// Resolve the root record to be used to derive the vtable pointer
+  /// authentication policy for the specified record.
+  const CXXRecordDecl *baseForVTableAuthentication(const CXXRecordDecl *);
+  bool useAbbreviatedThunkName(GlobalDecl virtualMethodDecl,
+                               StringRef mangledName);
+
   StringRef getCUIDHash() const;
 
 private:
   /// All OMPTraitInfo objects live in this collection, one per
   /// `pragma omp [begin] declare variant` directive.
   SmallVector<std::unique_ptr<OMPTraitInfo>, 4> OMPTraitInfoVector;
+
+  llvm::DenseMap<GlobalDecl, llvm::StringSet<>> thunksToBeAbbreviated;
 };
 
 /// Insertion operator for diagnostics.

clang/include/clang/AST/Mangle.h

@@ -130,15 +130,15 @@ class MangleContext {
   // FIXME: consider replacing raw_ostream & with something like SmallString &.
   void mangleName(GlobalDecl GD, raw_ostream &);
   virtual void mangleCXXName(GlobalDecl GD, raw_ostream &) = 0;
-  virtual void mangleThunk(const CXXMethodDecl *MD,
-                          const ThunkInfo &Thunk,
-                          raw_ostream &) = 0;
+  virtual void mangleThunk(const CXXMethodDecl *MD, const ThunkInfo &Thunk,
+                           bool elideOverrideInfo, raw_ostream &) = 0;
   virtual void mangleCXXDtorThunk(const CXXDestructorDecl *DD, CXXDtorType Type,
-                                  const ThisAdjustment &ThisAdjustment,
-                                  raw_ostream &) = 0;
+                                  const ThunkInfo &Thunk,
+                                  bool elideOverrideInfo, raw_ostream &) = 0;
   virtual void mangleReferenceTemporary(const VarDecl *D,
                                         unsigned ManglingNumber,
                                         raw_ostream &) = 0;
+  virtual void mangleCXXVTable(const CXXRecordDecl *RD, raw_ostream &) = 0;
   virtual void mangleCXXRTTI(QualType T, raw_ostream &) = 0;
   virtual void mangleCXXRTTIName(QualType T, raw_ostream &,
                                  bool NormalizeIntegers = false) = 0;
@@ -192,7 +192,6 @@ class ItaniumMangleContext : public MangleContext {
                                 bool IsAux = false)
       : MangleContext(C, D, MK_Itanium, IsAux) {}
 
-  virtual void mangleCXXVTable(const CXXRecordDecl *RD, raw_ostream &) = 0;
   virtual void mangleCXXVTT(const CXXRecordDecl *RD, raw_ostream &) = 0;
   virtual void mangleCXXCtorVTable(const CXXRecordDecl *RD, int64_t Offset,
                                    const CXXRecordDecl *Type,

clang/include/clang/Basic/Attr.td

@@ -680,6 +680,10 @@ class Attr {
   bit PragmaAttributeSupport;
   // Set to true if this attribute accepts parameter pack expansion expressions.
   bit AcceptsExprPack = 0;
+  // To support multiple enum parameters to an attribute without breaking
+  // our existing general parsing we need to have a separate flag that
+  // opts an attribute into strict parsing of attribute parameters
+  bit StrictEnumParameters = 0;
   // Lists language options, one of which is required to be true for the
   // attribute to be applicable. If empty, no language options are required.
   list<LangOpt> LangOpts = [];
@@ -4546,6 +4550,31 @@ def NoRandomizeLayout : InheritableAttr {
 }
 def : MutualExclusions<[RandomizeLayout, NoRandomizeLayout]>;
 
+def VTablePointerAuthentication : InheritableAttr {
+  let Spellings = [Clang<"ptrauth_vtable_pointer">];
+  let Subjects = SubjectList<[CXXRecord]>;
+  let Documentation = [Undocumented];
+  let StrictEnumParameters = 1;
+  let Args = [EnumArgument<"Key", "VPtrAuthKeyType", /*is_string=*/ true,
+                ["default_key", "no_authentication", "process_dependent",
+                 "process_independent"],
+                ["DefaultKey", "NoKey", "ProcessDependent",
+                 "ProcessIndependent"]>,
+              EnumArgument<"AddressDiscrimination", "AddressDiscriminationMode",
+                /*is_string=*/ true,
+                ["default_address_discrimination", "no_address_discrimination",
+                 "address_discrimination"],
+                ["DefaultAddressDiscrimination", "NoAddressDiscrimination",
+                 "AddressDiscrimination"]>,
+              EnumArgument<"ExtraDiscrimination", "ExtraDiscrimination",
+                /*is_string=*/ true,
+                ["default_extra_discrimination", "no_extra_discrimination",
+                 "type_discrimination", "custom_discrimination"],
+                ["DefaultExtraDiscrimination", "NoExtraDiscrimination",
+                 "TypeDiscrimination", "CustomDiscrimination"]>,
+              IntArgument<"CustomDiscriminationValue", 1>];
+}
+
 def FunctionReturnThunks : InheritableAttr,
     TargetSpecificAttr<TargetAnyX86> {
   let Spellings = [GCC<"function_return">];

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -938,6 +938,13 @@ def warn_ptrauth_auth_null_pointer :
 def err_ptrauth_string_not_literal : Error<
   "argument must be a string literal%select{| of char type}0">;
 
+def note_ptrauth_virtual_function_pointer_incomplete_arg_ret :
+  Note<"cannot take an address of a virtual member function if its return or "
+       "argument types are incomplete">;
+def note_ptrauth_virtual_function_incomplete_arg_ret_type :
+  Note<"%0 is incomplete">;
+
+
 /// main()
 // static main() is not an error in C, just in C++.
 def warn_static_main : Warning<"'main' should not be declared static">,
@@ -12175,6 +12182,30 @@ def warn_cuda_maxclusterrank_sm_90 : Warning<
   "maxclusterrank requires sm_90 or higher, CUDA arch provided: %0, ignoring "
   "%1 attribute">, InGroup<IgnoredAttributes>;
 
+// VTable pointer authentication errors
+def err_non_polymorphic_vtable_pointer_auth : Error<
+  "cannot set vtable pointer authentication on monomorphic type %0">;
+def err_incomplete_type_vtable_pointer_auth : Error<
+  "cannot set vtable pointer authentication on an incomplete type %0">;
+def err_non_top_level_vtable_pointer_auth : Error<
+  "cannot set vtable pointer authentication on %0 which is a subclass of polymorphic type %1">;
+def err_duplicated_vtable_pointer_auth : Error<
+  "multiple vtable pointer authentication policies on %0">;
+def err_invalid_authentication_key : Error<
+  "invalid authentication key %0">;
+def err_invalid_address_discrimination : Error<
+  "invalid address discrimination mode %0">;
+def err_invalid_extra_discrimination : Error<
+  "invalid extra discrimination selection %0">;
+def err_invalid_custom_discrimination : Error<
+  "invalid custom discrimination">;
+def err_missing_custom_discrimination : Error<
+  "missing custom discrimination">;
+def err_no_default_vtable_pointer_auth : Error<
+  "cannot specify a default vtable pointer authentication "
+  "%select{key|address discrimination mode|discriminator}0 with no default set"
+>;
+
 def err_bit_int_bad_size : Error<"%select{signed|unsigned}0 _BitInt must "
                                  "have a bit size of at least %select{2|1}0">;
 def err_bit_int_max_size : Error<"%select{signed|unsigned}0 _BitInt of bit "

clang/include/clang/Basic/PointerAuthOptions.h

@@ -49,6 +49,12 @@ class PointerAuthSchema {
     /// No additional discrimination.
     None,
 
+    /// Include a hash of the entity's type.
+    Type,
+
+    /// Include a hash of the entity's identity.
+    Decl,
+
     /// Discriminate using a constant value.
     Constant,
   };
@@ -152,6 +158,25 @@ class PointerAuthSchema {
 struct PointerAuthOptions {
   /// The ABI for C function pointers.
   PointerAuthSchema FunctionPointers;
+
+  /// The ABI for C++ virtual table pointers (the pointer to the table
+  /// itself) as installed in an actual class instance.
+  PointerAuthSchema CXXVTablePointers;
+
+  /// TypeInfo has external ABI requirements and is emitted without
+  /// actually having parsed the libcxx definition, so we can't simply
+  /// perform a look up. The settings for this should match the exact
+  /// specification in type_info.h
+  PointerAuthSchema CXXTypeInfoVTablePointer;
+
+  /// The ABI for C++ virtual table pointers as installed in a VTT.
+  PointerAuthSchema CXXVTTVTablePointers;
+
+  /// The ABI for most C++ virtual function pointers, i.e. v-table entries.
+  PointerAuthSchema CXXVirtualFunctionPointers;
+
+  /// The ABI for variadic C++ virtual function pointers.
+  PointerAuthSchema CXXVirtualVariadicFunctionPointers;
 };
 
 } // end namespace clang

clang/include/clang/Basic/Thunk.h

@@ -169,16 +169,20 @@ struct ThunkInfo {
   /// CAUTION: In the unlikely event you need to sort ThunkInfos, consider using
   /// an ABI-specific comparator.
   const CXXMethodDecl *Method;
+  const Type *ThisType { nullptr };
 
   ThunkInfo() : Method(nullptr) {}
 
   ThunkInfo(const ThisAdjustment &This, const ReturnAdjustment &Return,
+            const Type *thisType,
             const CXXMethodDecl *Method = nullptr)
-      : This(This), Return(Return), Method(Method) {}
+      : This(This), Return(Return), Method(Method),
+        ThisType(thisType) {}
 
   friend bool operator==(const ThunkInfo &LHS, const ThunkInfo &RHS) {
     return LHS.This == RHS.This && LHS.Return == RHS.Return &&
-           LHS.Method == RHS.Method;
+           LHS.Method == RHS.Method &&
+           LHS.ThisType == RHS.ThisType;
   }
 
   bool isEmpty() const {

clang/include/clang/CodeGen/CodeGenABITypes.h

@@ -42,6 +42,7 @@ class CXXConstructorDecl;
 class CXXDestructorDecl;
 class CXXRecordDecl;
 class CXXMethodDecl;
+class GlobalDecl;
 class ObjCMethodDecl;
 class ObjCProtocolDecl;
 
@@ -104,6 +105,9 @@ llvm::Type *convertTypeForMemory(CodeGenModule &CGM, QualType T);
 unsigned getLLVMFieldNumber(CodeGenModule &CGM,
                             const RecordDecl *RD, const FieldDecl *FD);
 
+/// Return a declaration discriminator for the given global decl.
+uint16_t getPointerAuthDeclDiscriminator(CodeGenModule &CGM, GlobalDecl GD);
+
 /// Return a signed constant pointer.
 llvm::Constant *getConstantSignedPointer(CodeGenModule &CGM,
                                          llvm::Constant *pointer,

clang/include/clang/CodeGen/ConstantInitBuilder.h

@@ -25,8 +25,11 @@
 #include <vector>
 
 namespace clang {
-namespace CodeGen {
+class GlobalDecl;
+class PointerAuthSchema;
+class QualType;
 
+namespace CodeGen {
 class CodeGenModule;
 
 /// A convenience builder class for complex constant initializers,
@@ -199,6 +202,17 @@ class ConstantAggregateBuilderBase {
     add(llvm::ConstantInt::get(intTy, value, isSigned));
   }
 
+  /// Add a signed pointer using the given pointer authentication schema.
+  void addSignedPointer(llvm::Constant *pointer,
+                        const PointerAuthSchema &schema, GlobalDecl calleeDecl,
+                        QualType calleeType);
+
+  /// Add a signed pointer using the given pointer authentication schema.
+  void addSignedPointer(llvm::Constant *pointer,
+                        unsigned key,
+                        bool useAddressDiscrimination,
+                        llvm::Constant *otherDiscriminator);
+
   /// Add a null pointer of a specific type.
   void addNullPointer(llvm::PointerType *ptrTy) {
     add(llvm::ConstantPointerNull::get(ptrTy));

clang/include/clang/InstallAPI/Visitor.h

@@ -60,8 +60,8 @@ class InstallAPIVisitor final : public ASTConsumer,
   std::string getMangledName(const NamedDecl *D) const;
   std::string getBackendMangledName(llvm::Twine Name) const;
   std::string getMangledCXXVTableName(const CXXRecordDecl *D) const;
-  std::string getMangledCXXThunk(const GlobalDecl &D,
-                                 const ThunkInfo &Thunk) const;
+  std::string getMangledCXXThunk(const GlobalDecl &D, const ThunkInfo &Thunk,
+                                 bool ElideOverrideInfo) const;
   std::string getMangledCXXRTTI(const CXXRecordDecl *D) const;
   std::string getMangledCXXRTTIName(const CXXRecordDecl *D) const;
   std::string getMangledCtorDtor(const CXXMethodDecl *D, int Type) const;

clang/include/clang/Sema/Sema.h

@@ -4401,6 +4401,10 @@ class Sema final : public SemaBase {
   /// conditions that are needed for the attribute to have an effect.
   void checkIllFormedTrivialABIStruct(CXXRecordDecl &RD);
 
+  /// Check that VTable Pointer authentication is only being set on the first
+  /// first instantiation of the vtable
+  void checkIncorrectVTablePointerAuthenticationAttribute(CXXRecordDecl &RD);
+
   void ActOnFinishCXXMemberSpecification(Scope *S, SourceLocation RLoc,
                                          Decl *TagDecl, SourceLocation LBrac,
                                          SourceLocation RBrac,

clang/lib/AST/ASTContext.cpp

@@ -86,6 +86,7 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MD5.h"
 #include "llvm/Support/MathExtras.h"
+#include "llvm/Support/SipHash.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/TargetParser/Triple.h"
 #include <algorithm>
@@ -3088,6 +3089,17 @@ QualType ASTContext::removeAddrSpaceQualType(QualType T) const {
     return QualType(TypeNode, Quals.getFastQualifiers());
 }
 
+uint16_t ASTContext::getPointerAuthVTablePointerDiscriminator(
+    const CXXRecordDecl *record) {
+  assert(record->isPolymorphic() &&
+         "Attempted to get vtable pointer discriminator on a monomorphic type");
+  std::unique_ptr<MangleContext> MC(createMangleContext());
+  SmallString<256> Str;
+  llvm::raw_svector_ostream Out(Str);
+  MC->mangleCXXVTable(record, Out);
+  return llvm::getPointerAuthStableSipHash16(Str.c_str());
+}
+
 QualType ASTContext::getObjCGCQualType(QualType T,
                                        Qualifiers::GC GCAttr) const {
   QualType CanT = getCanonicalType(T);
@@ -13812,3 +13824,77 @@ StringRef ASTContext::getCUIDHash() const {
   CUIDHash = llvm::utohexstr(llvm::MD5Hash(LangOpts.CUID), /*LowerCase=*/true);
   return CUIDHash;
 }
+
+const CXXRecordDecl *
+ASTContext::baseForVTableAuthentication(const CXXRecordDecl *thisClass) {
+  assert(thisClass);
+  assert(thisClass->isPolymorphic());
+  const CXXRecordDecl *primaryBase = thisClass;
+  while (1) {
+    assert(primaryBase);
+    assert(primaryBase->isPolymorphic());
+    auto &layout = getASTRecordLayout(primaryBase);
+    auto base = layout.getPrimaryBase();
+    if (!base || base == primaryBase || !base->isPolymorphic())
+      break;
+    primaryBase = base;
+  }
+  return primaryBase;
+}
+
+bool ASTContext::useAbbreviatedThunkName(GlobalDecl virtualMethodDecl,
+                                         StringRef mangledName) {
+  auto method = cast<CXXMethodDecl>(virtualMethodDecl.getDecl());
+  assert(method->isVirtual());
+  bool defaultIncludesPointerAuth =
+      LangOpts.PointerAuthCalls || LangOpts.PointerAuthIntrinsics;
+
+  if (!defaultIncludesPointerAuth)
+    return true;
+
+  auto existing = thunksToBeAbbreviated.find(virtualMethodDecl);
+  if (existing != thunksToBeAbbreviated.end())
+    return existing->second.contains(mangledName.str());
+
+  std::unique_ptr<MangleContext> mangler(createMangleContext());
+  llvm::StringMap<llvm::SmallVector<std::string, 2>> thunks;
+  auto vtableContext = getVTableContext();
+  if (const auto *thunkInfos = vtableContext->getThunkInfo(virtualMethodDecl)) {
+    auto destructor = dyn_cast<CXXDestructorDecl>(method);
+    for (const auto &thunk : *thunkInfos) {
+      SmallString<256> elidedName;
+      llvm::raw_svector_ostream elidedNameStream(elidedName);
+      if (destructor) {
+        mangler->mangleCXXDtorThunk(destructor, virtualMethodDecl.getDtorType(),
+                                    thunk, /* elideOverrideInfo */ true,
+                                    elidedNameStream);
+      } else {
+        mangler->mangleThunk(method, thunk, /* elideOverrideInfo */ true,
+                             elidedNameStream);
+      }
+      SmallString<256> mangledName;
+      llvm::raw_svector_ostream mangledNameStream(mangledName);
+      if (destructor) {
+        mangler->mangleCXXDtorThunk(destructor, virtualMethodDecl.getDtorType(),
+                                    thunk, /* elideOverrideInfo */ false,
+                                    mangledNameStream);
+      } else {
+        mangler->mangleThunk(method, thunk, /* elideOverrideInfo */ false,
+                             mangledNameStream);
+      }
+
+      if (thunks.find(elidedName) == thunks.end()) {
+        thunks[elidedName] = {};
+      }
+      thunks[elidedName].push_back(std::string(mangledName));
+    }
+  }
+  llvm::StringSet<> simplifiedThunkNames;
+  for (auto &thunkList : thunks) {
+    llvm::sort(thunkList.second);
+    simplifiedThunkNames.insert(thunkList.second[0]);
+  }
+  bool result = simplifiedThunkNames.contains(mangledName);
+  thunksToBeAbbreviated[virtualMethodDecl] = std::move(simplifiedThunkNames);
+  return result;
+}