[AllocToken, Clang] Implement TypeHashPointerSplit mode

Implement the TypeHashPointerSplit mode: This mode assigns a token ID
based on the hash of the allocated type's name, where the top half
ID-space is reserved for types that contain pointers and the bottom half
for types that do not contain pointers.

This mode with max tokens of 2 (`-falloc-token-max=2`) may also
be valuable for heap hardening strategies that simply separate pointer
types from non-pointer types.

Make it the new default mode.

Link: https://discourse.llvm.org/t/rfc-a-framework-for-allocator-partitioning-hints/87434

Pull Request: llvm#156840

Author: melver

clang/docs/AllocToken.rst

@@ -31,13 +31,18 @@ Token Assignment Mode
 
 The default mode to calculate tokens is:
 
-* ``typehash``: This mode assigns a token ID based on the hash of the allocated
-  type's name.
+* ``typehashpointersplit``: This mode assigns a token ID based on the hash of
+  the allocated type's name, where the top half ID-space is reserved for types
+  that contain pointers and the bottom half for types that do not contain
+  pointers.
 
 Other token ID assignment modes are supported, but they may be subject to
 change or removal. These may (experimentally) be selected with ``-mllvm
 -alloc-token-mode=<mode>``:
 
+* ``typehash``: This mode assigns a token ID based on the hash of the allocated
+  type's name.
+
 * ``random``: This mode assigns a statically-determined random token ID to each
   allocation site.
 

clang/lib/CodeGen/CGExpr.cpp

@@ -1272,20 +1272,84 @@ void CodeGenFunction::EmitBoundsCheckImpl(const Expr *E, llvm::Value *Bound,
   EmitCheck(std::make_pair(Check, CheckKind), CheckHandler, StaticData, Index);
 }
 
+static bool
+typeContainsPointer(QualType T,
+                    llvm::SmallPtrSet<const RecordDecl *, 4> &VisitedRD,
+                    bool &IncompleteType) {
+  QualType CanonicalType = T.getCanonicalType();
+  if (CanonicalType->isPointerType())
+    return true; // base case
+
+  // Look through typedef chain to check for special types.
+  for (QualType CurrentT = T; const auto *TT = CurrentT->getAs<TypedefType>();
+       CurrentT = TT->getDecl()->getUnderlyingType()) {
+    const IdentifierInfo *II = TT->getDecl()->getIdentifier();
+    // Special Case: Syntactically uintptr_t is not a pointer; semantically,
+    // however, very likely used as such. Therefore, classify uintptr_t as a
+    // pointer, too.
+    if (II && II->isStr("uintptr_t"))
+      return true;
+  }
+
+  // The type is an array; check the element type.
+  if (const ArrayType *AT = dyn_cast<ArrayType>(CanonicalType))
+    return typeContainsPointer(AT->getElementType(), VisitedRD, IncompleteType);
+  // The type is a struct, class, or union.
+  if (const RecordDecl *RD = CanonicalType->getAsRecordDecl()) {
+    if (!RD->isCompleteDefinition()) {
+      IncompleteType = true;
+      return false;
+    }
+    if (!VisitedRD.insert(RD).second)
+      return false; // already visited
+    // Check all fields.
+    for (const FieldDecl *Field : RD->fields()) {
+      if (typeContainsPointer(Field->getType(), VisitedRD, IncompleteType))
+        return true;
+    }
+    // For C++ classes, also check base classes.
+    if (const CXXRecordDecl *CXXRD = dyn_cast<CXXRecordDecl>(RD)) {
+      // Polymorphic types require a vptr.
+      if (CXXRD->isDynamicClass())
+        return true;
+      for (const CXXBaseSpecifier &Base : CXXRD->bases()) {
+        if (typeContainsPointer(Base.getType(), VisitedRD, IncompleteType))
+          return true;
+      }
+    }
+  }
+  return false;
+}
+
 void CodeGenFunction::EmitAllocToken(llvm::CallBase *CB, QualType AllocType) {
   assert(SanOpts.has(SanitizerKind::AllocToken) &&
          "Only needed with -fsanitize=alloc-token");
 
+  llvm::MDBuilder MDB(getLLVMContext());
+
+  // Get unique type name.
   PrintingPolicy Policy(CGM.getContext().getLangOpts());
   Policy.SuppressTagKeyword = true;
   Policy.FullyQualifiedName = true;
   SmallString<64> TypeName;
   llvm::raw_svector_ostream TypeNameOS(TypeName);
   AllocType.getCanonicalType().print(TypeNameOS, Policy);
-  auto *TypeMDS = llvm::MDString::get(CGM.getLLVMContext(), TypeNameOS.str());
+  auto *TypeNameMD = MDB.createString(TypeNameOS.str());
+
+  // Check if QualType contains a pointer. Implements a simple DFS to
+  // recursively check if a type contains a pointer type.
+  llvm::SmallPtrSet<const RecordDecl *, 4> VisitedRD;
+  bool IncompleteType = false;
+  const bool ContainsPtr =
+      typeContainsPointer(AllocType, VisitedRD, IncompleteType);
+  if (!ContainsPtr && IncompleteType)
+    return;
+  auto *ContainsPtrC = Builder.getInt1(ContainsPtr);
+  auto *ContainsPtrMD = MDB.createConstant(ContainsPtrC);
 
-  // Format: !{<type-name>}
-  auto *MDN = llvm::MDNode::get(CGM.getLLVMContext(), {TypeMDS});
+  // Format: !{<type-name>, <contains-pointer>}
+  auto *MDN =
+      llvm::MDNode::get(CGM.getLLVMContext(), {TypeNameMD, ContainsPtrMD});
   CB->setMetadata(llvm::LLVMContext::MD_alloc_token, MDN);
 }
 

clang/test/CodeGenCXX/alloc-token-pointer.cpp

@@ -0,0 +1,175 @@
+// RUN: %clang_cc1 -fsanitize=alloc-token -triple x86_64-linux-gnu -std=c++20 -emit-llvm -disable-llvm-passes %s -o - | FileCheck %s
+
+#include "../Analysis/Inputs/system-header-simulator-cxx.h"
+
+typedef __UINTPTR_TYPE__ uintptr_t;
+
+extern "C" {
+void *malloc(size_t size);
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z15test_malloc_intv(
+// CHECK: call ptr @malloc(i64 noundef 4)
+void *test_malloc_int() {
+  int *a = (int *)malloc(sizeof(int));
+  *a = 42;
+  return a;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z15test_malloc_ptrv(
+// CHECK: call ptr @malloc(i64 noundef 8)
+int **test_malloc_ptr() {
+  int **a = (int **)malloc(sizeof(int*));
+  *a = nullptr;
+  return a;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z12test_new_intv(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 4){{.*}} !alloc_token [[META_INT:![0-9]+]]
+int *test_new_int() {
+  return new int;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z20test_new_ulong_arrayv(
+// CHECK: call noalias noundef nonnull ptr @_Znam(i64 noundef 80){{.*}} !alloc_token [[META_ULONG:![0-9]+]]
+unsigned long *test_new_ulong_array() {
+  return new unsigned long[10];
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z12test_new_ptrv(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 8){{.*}} !alloc_token [[META_INTPTR:![0-9]+]]
+int **test_new_ptr() {
+  return new int*;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z18test_new_ptr_arrayv(
+// CHECK: call noalias noundef nonnull ptr @_Znam(i64 noundef 80){{.*}} !alloc_token [[META_INTPTR]]
+int **test_new_ptr_array() {
+  return new int*[10];
+}
+
+struct ContainsPtr {
+  int a;
+  char *buf;
+};
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z27test_malloc_struct_with_ptrv(
+// CHECK: call ptr @malloc(i64 noundef 16)
+ContainsPtr *test_malloc_struct_with_ptr() {
+  ContainsPtr *c = (ContainsPtr *)malloc(sizeof(ContainsPtr));
+  return c;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z33test_malloc_struct_array_with_ptrv(
+// CHECK: call ptr @malloc(i64 noundef 160)
+ContainsPtr *test_malloc_struct_array_with_ptr() {
+  ContainsPtr *c = (ContainsPtr *)malloc(10 * sizeof(ContainsPtr));
+  return c;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z32test_operatornew_struct_with_ptrv(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 16)
+ContainsPtr *test_operatornew_struct_with_ptr() {
+  ContainsPtr *c = (ContainsPtr *)__builtin_operator_new(sizeof(ContainsPtr));
+  return c;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z38test_operatornew_struct_array_with_ptrv(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 160)
+ContainsPtr *test_operatornew_struct_array_with_ptr() {
+  ContainsPtr *c = (ContainsPtr *)__builtin_operator_new(10 * sizeof(ContainsPtr));
+  return c;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z33test_operatornew_struct_with_ptr2v(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 16)
+ContainsPtr *test_operatornew_struct_with_ptr2() {
+  ContainsPtr *c = (ContainsPtr *)__builtin_operator_new(sizeof(*c));
+  return c;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z39test_operatornew_struct_array_with_ptr2v(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 160)
+ContainsPtr *test_operatornew_struct_array_with_ptr2() {
+  ContainsPtr *c = (ContainsPtr *)__builtin_operator_new(10 * sizeof(*c));
+  return c;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z24test_new_struct_with_ptrv(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 16){{.*}} !alloc_token [[META_CONTAINSPTR:![0-9]+]]
+ContainsPtr *test_new_struct_with_ptr() {
+  return new ContainsPtr;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z30test_new_struct_array_with_ptrv(
+// CHECK: call noalias noundef nonnull ptr @_Znam(i64 noundef 160){{.*}} !alloc_token [[META_CONTAINSPTR]]
+ContainsPtr *test_new_struct_array_with_ptr() {
+  return new ContainsPtr[10];
+}
+
+class TestClass {
+public:
+  void Foo();
+  ~TestClass();
+  int data[16];
+};
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z14test_new_classv(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 64){{.*}} !alloc_token [[META_TESTCLASS:![0-9]+]]
+TestClass *test_new_class() {
+  return new TestClass();
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z20test_new_class_arrayv(
+// CHECK: call noalias noundef nonnull ptr @_Znam(i64 noundef 648){{.*}} !alloc_token [[META_TESTCLASS]]
+TestClass *test_new_class_array() {
+  return new TestClass[10];
+}
+
+// Test that we detect that virtual classes have implicit vtable pointer.
+class VirtualTestClass {
+public:
+  virtual void Foo();
+  virtual ~VirtualTestClass();
+  int data[16];
+};
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z22test_new_virtual_classv(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 72){{.*}} !alloc_token [[META_VIRTUALTESTCLASS:![0-9]+]]
+VirtualTestClass *test_new_virtual_class() {
+  return new VirtualTestClass();
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z28test_new_virtual_class_arrayv(
+// CHECK: call noalias noundef nonnull ptr @_Znam(i64 noundef 728){{.*}} !alloc_token [[META_VIRTUALTESTCLASS]]
+VirtualTestClass *test_new_virtual_class_array() {
+  return new VirtualTestClass[10];
+}
+
+// uintptr_t is treated as a pointer.
+struct MyStructUintptr {
+  int a;
+  uintptr_t ptr;
+};
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z18test_uintptr_isptrv(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 16){{.*}} !alloc_token [[META_MYSTRUCTUINTPTR:![0-9]+]]
+MyStructUintptr *test_uintptr_isptr() {
+  return new MyStructUintptr;
+}
+
+using uptr = uintptr_t;
+// CHECK-LABEL: define dso_local noundef ptr @_Z19test_uintptr_isptr2v(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 8){{.*}} !alloc_token [[META_UINTPTR:![0-9]+]]
+uptr *test_uintptr_isptr2() {
+  return new uptr;
+}
+
+// CHECK: [[META_INT]] = !{!"int", i1 false}
+// CHECK: [[META_ULONG]] = !{!"unsigned long", i1 false}
+// CHECK: [[META_INTPTR]] = !{!"int *", i1 true}
+// CHECK: [[META_CONTAINSPTR]] = !{!"ContainsPtr", i1 true}
+// CHECK: [[META_TESTCLASS]] = !{!"TestClass", i1 false}
+// CHECK: [[META_VIRTUALTESTCLASS]] = !{!"VirtualTestClass", i1 true}
+// CHECK: [[META_MYSTRUCTUINTPTR]] = !{!"MyStructUintptr", i1 true}
+// CHECK: [[META_UINTPTR]] = !{!"unsigned long", i1 true}

clang/test/CodeGenCXX/alloc-token.cpp

@@ -137,5 +137,5 @@ TestClass *test_new_class_array() {
   return arr;
 }
 
-// CHECK: [[META_INT]] = !{!"int"}
-// CHECK: [[META_TESTCLASS]] = !{!"TestClass"}
+// CHECK: [[META_INT]] = !{!"int", i1 false}
+// CHECK: [[META_TESTCLASS]] = !{!"TestClass", i1 true}

llvm/docs/LangRef.rst

@@ -8588,13 +8588,14 @@ functions, and contains richer semantic information about the type of the
 allocation. This information is consumed by the ``alloc-token`` pass to
 instrument such calls with allocation token IDs.
 
-The metadata contains a string with the type of an allocation.
+The metadata contains: string with the type of an allocation, and a boolean
+denoting if the type contains a pointer.
 
 .. code-block:: none
 
   call ptr @malloc(i64 64), !alloc_token !0
 
-  !0 = !{!"<type-name>"}
+  !0 = !{!"<type-name>", i1 <contains-pointer>}
 
 Module Flags Metadata
 =====================

llvm/lib/IR/Verifier.cpp

@@ -5398,8 +5398,10 @@ void Verifier::visitCapturesMetadata(Instruction &I, const MDNode *Captures) {
 
 void Verifier::visitAllocTokenMetadata(Instruction &I, MDNode *MD) {
   Check(isa<CallBase>(I), "!alloc_token should only exist on calls", &I);
-  Check(MD->getNumOperands() == 1, "!alloc_token must have 1 operand", MD);
+  Check(MD->getNumOperands() == 2, "!alloc_token must have 2 operands", MD);
   Check(isa<MDString>(MD->getOperand(0)), "expected string", MD);
+  Check(mdconst::dyn_extract_or_null<ConstantInt>(MD->getOperand(1)),
+        "expected integer constant", MD);
 }
 
 /// verifyInstruction - Verify that an instruction is well formed.