[PAC] Fix address discrimination for type info vtable pointers

In llvm#99726, `-fptrauth-type-info-vtable-pointer-discrimination` was introduced, which is intended to enable type and address discrimination for type_info vtable pointers. However, some codegen logic for actually enabling address discrimination was missing. This patch addresses the issue. Fixes llvm#101716
kovdan01 · Aug 6, 2024 · 0c20bcd · 0c20bcd
1 parent 642259a
commit 0c20bcd
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 10 deletions.
diff --git a/clang/lib/CodeGen/ItaniumCXXABI.cpp b/clang/lib/CodeGen/ItaniumCXXABI.cpp
@@ -3955,9 +3955,23 @@ void ItaniumRTTIBuilder::BuildVTablePointer(const Type *Ty) {
                                                           VTable, Two);
   }
 
-  if (auto &Schema = CGM.getCodeGenOpts().PointerAuth.CXXTypeInfoVTablePointer)
-    VTable = CGM.getConstantSignedPointer(VTable, Schema, nullptr, GlobalDecl(),
-                                          QualType(Ty, 0));
+  if (const auto &Schema =
+          CGM.getCodeGenOpts().PointerAuth.CXXTypeInfoVTablePointer) {
+    llvm::PointerType *PtrTy = llvm::PointerType::get(
+        CGM.getLLVMContext(),
+        CGM.getModule().getDataLayout().getProgramAddressSpace());
+    llvm::Constant *StorageAddress =
+        (Schema.isAddressDiscriminated()
+             ? llvm::ConstantExpr::getIntToPtr(
+                   llvm::ConstantInt::get(
+                       CGM.IntPtrTy,
+                       llvm::ConstantPtrAuth::
+                           AddrDiscriminator_CXXTypeInfoVTablePointer),
+                   PtrTy)
+             : nullptr);
+    VTable = CGM.getConstantSignedPointer(VTable, Schema, StorageAddress,
+                                          GlobalDecl(), QualType(Ty, 0));
+  }
 
   Fields.push_back(VTable);
 }

diff --git a/clang/test/CodeGenCXX/ptrauth-type-info-vtable.cpp b/clang/test/CodeGenCXX/ptrauth-type-info-vtable.cpp
@@ -65,7 +65,7 @@ extern "C" int disc_std_type_info = __builtin_ptrauth_string_discriminator("_ZTV
 
 // NODISC: @_ZTI10TestStruct = constant { ptr, ptr } { ptr ptrauth (ptr getelementptr inbounds (ptr, ptr @_ZTVN10__cxxabiv117__class_type_infoE, i64 2), i32 2), ptr @_ZTS10TestStruct }, align 8
 
-// DISC: @_ZTI10TestStruct = constant { ptr, ptr } { ptr ptrauth (ptr getelementptr inbounds (ptr, ptr @_ZTVN10__cxxabiv117__class_type_infoE, i64 2), i32 2, i64 [[STDTYPEINFO_DISC]]), ptr @_ZTS10TestStruct }, align 8
+// DISC: @_ZTI10TestStruct = constant { ptr, ptr } { ptr ptrauth (ptr getelementptr inbounds (ptr, ptr @_ZTVN10__cxxabiv117__class_type_infoE, i64 2), i32 2, i64 [[STDTYPEINFO_DISC]], ptr inttoptr (i64 1 to ptr)), ptr @_ZTS10TestStruct }, align 8
 
 struct TestStruct {
   virtual ~TestStruct();

diff --git a/llvm/include/llvm/IR/Constants.h b/llvm/include/llvm/IR/Constants.h
@@ -1056,12 +1056,18 @@ class ConstantPtrAuth final : public Constant {
     return !getAddrDiscriminator()->isNullValue();
   }
 
-  /// A constant value for the address discriminator which has special
-  /// significance to ctors/dtors lowering. Regular address discrimination can't
-  /// be applied for them since uses of llvm.global_{c|d}tors are disallowed
-  /// (see Verifier::visitGlobalVariable) and we can't emit getelementptr
-  /// expressions referencing these special arrays.
-  enum { AddrDiscriminator_CtorsDtors = 1 };
+  /// Constant values for the address discriminator which have special
+  /// significance to lowering in some contexts.
+  /// - For ctors/dtors, regular address discrimination can't
+  ///   be applied for them since uses of llvm.global_{c|d}tors are disallowed
+  ///   (see Verifier::visitGlobalVariable) and we can't emit getelementptr
+  ///   expressions referencing these special arrays.
+  /// - For vtable pointers of std::type_info and classes derived from it,
+  ///   we do not know the storage address when emitting ptrauth constant.
+  enum {
+    AddrDiscriminator_CtorsDtors = 1,
+    AddrDiscriminator_CXXTypeInfoVTablePointer = 1
+  };
 
   /// Whether the address uses a special address discriminator.
   /// These discriminators can't be used in real pointer-auth values; they

diff --git a/llvm/test/CodeGen/AArch64/ptrauth-type-info-vptr-discr.ll b/llvm/test/CodeGen/AArch64/ptrauth-type-info-vptr-discr.ll
@@ -0,0 +1,20 @@
+; RUN: llc -mtriple aarch64-linux-gnu    -mattr=+pauth -filetype=asm -o - %s | FileCheck --check-prefix=ELF %s
+; RUN: llc -mtriple aarch64-apple-darwin -mattr=+pauth -filetype=asm -o - %s | FileCheck --check-prefix=MACHO %s
+
+; ELF-LABEL:   _ZTI10Disc:
+; ELF-NEXT:      .xword  (_ZTVN10__cxxabiv117__class_type_infoE+16)@AUTH(da,45546,addr)
+; ELF-LABEL:   _ZTI10NoDisc:
+; ELF-NEXT:      .xword  (_ZTVN10__cxxabiv117__class_type_infoE+16)@AUTH(da,45546)
+
+; MACHO-LABEL: __ZTI10Disc:
+; MACHO-NEXT:    .quad   (__ZTVN10__cxxabiv117__class_type_infoE+16)@AUTH(da,45546,addr)
+; MACHO-LABEL: __ZTI10NoDisc:
+; MACHO-NEXT:    .quad   (__ZTVN10__cxxabiv117__class_type_infoE+16)@AUTH(da,45546)
+
+@_ZTVN10__cxxabiv117__class_type_infoE = external global [0 x ptr]
+
+@_ZTS10Disc   = constant [4 x i8] c"Disc", align 1
+@_ZTI10Disc   = constant { ptr, ptr } { ptr ptrauth (ptr getelementptr inbounds (ptr, ptr @_ZTVN10__cxxabiv117__class_type_infoE, i64 2), i32 2, i64 45546, ptr inttoptr (i64 1 to ptr)), ptr @_ZTS10Disc }, align 8
+
+@_ZTS10NoDisc = constant [6 x i8] c"NoDisc", align 1
+@_ZTI10NoDisc = constant { ptr, ptr } { ptr ptrauth (ptr getelementptr inbounds (ptr, ptr @_ZTVN10__cxxabiv117__class_type_infoE, i64 2), i32 2, i64 45546), ptr @_ZTS10NoDisc }, align 8