[Clang] Wire up -fsanitize=alloc-token (#156839)

Wire up the `-fsanitize=alloc-token` command-line option, hooking up the `AllocToken` pass -- it provides allocation tokens to compatible runtime allocators, enabling different heap organization strategies, e.g. hardening schemes based on heap partitioning. The instrumentation rewrites standard allocation calls into variants that accept an additional `size_t token_id` argument. For example, calls to `malloc(size)` become `__alloc_token_malloc(size, token_id)`, and a C++ `new MyType` expression will call `__alloc_token__Znwm(size, token_id)`. Currently untyped allocation calls do not yet have `!alloc_token` metadata, and therefore receive the fallback token only. This will be fixed in subsequent changes through best-effort type-inference. One benefit of the instrumentation approach is that it can be applied transparently to large codebases, and scales in deployment as other sanitizers. Similarly to other sanitizers, instrumentation can selectively be controlled using `__attribute__((no_sanitize("alloc-token")))`. Support for sanitizer ignorelists to disable instrumentation for specific functions or source files is implemented. See clang/docs/AllocToken.rst for more usage instructions. Link: https://discourse.llvm.org/t/rfc-a-framework-for-allocator-partitioning-hints/87434 --- This change is part of the following series: 1. https://github.com/llvm/llvm-project/pull/160131 2. https://github.com/llvm/llvm-project/pull/156838 3. https://github.com/llvm/llvm-project/pull/162098 4. https://github.com/llvm/llvm-project/pull/162099 5. https://github.com/llvm/llvm-project/pull/156839 6. https://github.com/llvm/llvm-project/pull/156840 7. https://github.com/llvm/llvm-project/pull/156841 8. https://github.com/llvm/llvm-project/pull/156842
2025-10-08 20:59:24 +02:00 · 2025-10-08 20:59:24 +02:00 · 774ffe5cce
commit 774ffe5cce
parent 761be78dd7
20 changed files with 563 additions and 11 deletions
--- a/clang/docs/AllocToken.rst
+++ b/clang/docs/AllocToken.rst
@ -0,0 +1,173 @@
+=================
+Allocation Tokens
+=================
+
+.. contents::
+   :local:
+
+Introduction
+============
+
+Clang provides support for allocation tokens to enable allocator-level heap
+organization strategies. Clang assigns mode-dependent token IDs to allocation
+calls; the runtime behavior depends entirely on the implementation of a
+compatible memory allocator.
+
+Possible allocator strategies include:
+
+* **Security Hardening**: Placing allocations into separate, isolated heap
+  partitions. For example, separating pointer-containing types from raw data
+  can mitigate exploits that rely on overflowing a primitive buffer to corrupt
+  object metadata.
+
+* **Memory Layout Optimization**: Grouping related allocations to improve data
+  locality and cache utilization.
+
+* **Custom Allocation Policies**: Applying different management strategies to
+  different partitions.
+
+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.
+
+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>``:
+
+* ``random``: This mode assigns a statically-determined random token ID to each
+  allocation site.
+
+* ``increment``: This mode assigns a simple, incrementally increasing token ID
+  to each allocation site.
+
+Allocation Token Instrumentation
+================================
+
+To enable instrumentation of allocation functions, code can be compiled with
+the ``-fsanitize=alloc-token`` flag:
+
+.. code-block:: console
+
+    % clang++ -fsanitize=alloc-token example.cc
+
+The instrumentation transforms allocation calls to include a token ID. For
+example:
+
+.. code-block:: c
+
+    // Original:
+    ptr = malloc(size);
+
+    // Instrumented:
+    ptr = __alloc_token_malloc(size, <token id>);
+
+The following command-line options affect generated token IDs:
+
+* ``-falloc-token-max=<N>``
+    Configures the maximum number of tokens. No max by default (tokens bounded
+    by ``SIZE_MAX``).
+
+    .. code-block:: console
+
+        % clang++ -fsanitize=alloc-token -falloc-token-max=512 example.cc
+
+Runtime Interface
+-----------------
+
+A compatible runtime must be provided that implements the token-enabled
+allocation functions. The instrumentation generates calls to functions that
+take a final ``size_t token_id`` argument.
+
+.. code-block:: c
+
+    // C standard library functions
+    void *__alloc_token_malloc(size_t size, size_t token_id);
+    void *__alloc_token_calloc(size_t count, size_t size, size_t token_id);
+    void *__alloc_token_realloc(void *ptr, size_t size, size_t token_id);
+    // ...
+
+    // C++ operators (mangled names)
+    // operator new(size_t, size_t)
+    void *__alloc_token__Znwm(size_t size, size_t token_id);
+    // operator new[](size_t, size_t)
+    void *__alloc_token__Znam(size_t size, size_t token_id);
+    // ... other variants like nothrow, etc., are also instrumented.
+
+Fast ABI
+--------
+
+An alternative ABI can be enabled with ``-fsanitize-alloc-token-fast-abi``,
+which encodes the token ID hint in the allocation function name.
+
+.. code-block:: c
+
+    void *__alloc_token_0_malloc(size_t size);
+    void *__alloc_token_1_malloc(size_t size);
+    void *__alloc_token_2_malloc(size_t size);
+    ...
+    void *__alloc_token_0_Znwm(size_t size);
+    void *__alloc_token_1_Znwm(size_t size);
+    void *__alloc_token_2_Znwm(size_t size);
+    ...
+
+This ABI provides a more efficient alternative where
+``-falloc-token-max`` is small.
+
+Disabling Instrumentation
+-------------------------
+
+To exclude specific functions from instrumentation, you can use the
+``no_sanitize("alloc-token")`` attribute:
+
+.. code-block:: c
+
+    __attribute__((no_sanitize("alloc-token")))
+    void* custom_allocator(size_t size) {
+        return malloc(size);  // Uses original malloc
+    }
+
+Note: Independent of any given allocator support, the instrumentation aims to
+remain performance neutral. As such, ``no_sanitize("alloc-token")``
+functions may be inlined into instrumented functions and vice-versa. If
+correctness is affected, such functions should explicitly be marked
+``noinline``.
+
+The ``__attribute__((disable_sanitizer_instrumentation))`` is also supported to
+disable this and other sanitizer instrumentations.
+
+Suppressions File (Ignorelist)
+------------------------------
+
+AllocToken respects the ``src`` and ``fun`` entity types in the
+:doc:`SanitizerSpecialCaseList`, which can be used to omit specified source
+files or functions from instrumentation.
+
+.. code-block:: bash
+
+    [alloc-token]
+    # Exclude specific source files
+    src:third_party/allocator.c
+    # Exclude function name patterns
+    fun:*custom_malloc*
+    fun:LowLevel::*
+
+.. code-block:: console
+
+    % clang++ -fsanitize=alloc-token -fsanitize-ignorelist=my_ignorelist.txt example.cc
+
+Conditional Compilation with ``__SANITIZE_ALLOC_TOKEN__``
+-----------------------------------------------------------
+
+In some cases, one may need to execute different code depending on whether
+AllocToken instrumentation is enabled. The ``__SANITIZE_ALLOC_TOKEN__`` macro
+can be used for this purpose.
+
+.. code-block:: c
+
+    #ifdef __SANITIZE_ALLOC_TOKEN__
+    // Code specific to -fsanitize=alloc-token builds
+    #endif
--- a/clang/docs/ReleaseNotes.rst
+++ b/clang/docs/ReleaseNotes.rst
@ -257,10 +257,16 @@ Non-comprehensive list of changes in this release

 - Fixed a crash when the second argument to ``__builtin_assume_aligned`` was not constant (#GH161314)

+- Introduce support for :doc:`allocation tokens <AllocToken>` to enable
+  allocator-level heap organization strategies. A feature to instrument all
+  allocation functions with a token ID can be enabled via the
+  ``-fsanitize=alloc-token`` flag.
+
 New Compiler Flags
 ------------------
 - New option ``-fno-sanitize-debug-trap-reasons`` added to disable emitting trap reasons into the debug info when compiling with trapping UBSan (e.g. ``-fsanitize-trap=undefined``).
 - New option ``-fsanitize-debug-trap-reasons=`` added to control emitting trap reasons into the debug info when compiling with trapping UBSan (e.g. ``-fsanitize-trap=undefined``).
+- New options for enabling allocation token instrumentation: ``-fsanitize=alloc-token``, ``-falloc-token-max=``, ``-fsanitize-alloc-token-fast-abi``, ``-fsanitize-alloc-token-extended``.


 Lanai Support
--- a/clang/docs/UsersManual.rst
+++ b/clang/docs/UsersManual.rst
@ -2155,13 +2155,11 @@ are listed below.

 .. option:: -f[no-]sanitize=check1,check2,...

-   Turn on runtime checks for various forms of undefined or suspicious
-   behavior.
+   Turn on runtime checks or mitigations for various forms of undefined or
+   suspicious behavior. These are disabled by default.

-   This option controls whether Clang adds runtime checks for various
-   forms of undefined or suspicious behavior, and is disabled by
-   default. If a check fails, a diagnostic message is produced at
-   runtime explaining the problem. The main checks are:
+   The following options enable runtime checks for various forms of undefined
+   or suspicious behavior:

   -  .. _opt_fsanitize_address:

@ -2195,6 +2193,14 @@ are listed below.
   -  ``-fsanitize=realtime``: :doc:`RealtimeSanitizer`,
      a real-time safety checker.

+   The following options enable runtime mitigations for various forms of
+   undefined or suspicious behavior:
+
+   -  ``-fsanitize=alloc-token``: Enables :doc:`allocation tokens <AllocToken>`
+      for allocator-level heap organization strategies, such as for security
+      hardening. It passes type-derived token IDs to a compatible memory
+      allocator. Requires linking against a token-aware allocator.
+
   There are more fine-grained checks available: see
   the :ref:`list <ubsan-checks>` of specific kinds of
   undefined behavior that can be detected and the :ref:`list <cfi-schemes>`
--- a/clang/docs/index.rst
+++ b/clang/docs/index.rst
@ -40,6 +40,7 @@ Using Clang as a Compiler
   SanitizerCoverage
   SanitizerStats
   SanitizerSpecialCaseList
+   AllocToken
   BoundsSafety
   BoundsSafetyAdoptionGuide
   BoundsSafetyImplPlans
--- a/clang/include/clang/Basic/CodeGenOptions.def
+++ b/clang/include/clang/Basic/CodeGenOptions.def
@ -306,6 +306,8 @@ CODEGENOPT(SanitizeBinaryMetadataCovered, 1, 0, Benign) ///< Emit PCs for covere
 CODEGENOPT(SanitizeBinaryMetadataAtomics, 1, 0, Benign) ///< Emit PCs for atomic operations.
 CODEGENOPT(SanitizeBinaryMetadataUAR, 1, 0, Benign) ///< Emit PCs for start of functions
                                                    ///< that are subject for use-after-return checking.
+CODEGENOPT(SanitizeAllocTokenFastABI, 1, 0, Benign) ///< Use the AllocToken fast ABI.
+CODEGENOPT(SanitizeAllocTokenExtended, 1, 0, Benign) ///< Extend coverage to custom allocation functions.
 CODEGENOPT(SanitizeStats     , 1, 0, Benign) ///< Collect statistics for sanitizers.
 ENUM_CODEGENOPT(SanitizeDebugTrapReasons, SanitizeDebugTrapReasonKind, 2, SanitizeDebugTrapReasonKind::Detailed, Benign) ///< Control how "trap reasons" are emitted in debug info
 CODEGENOPT(SimplifyLibCalls  , 1, 1, Benign) ///< Set when -fbuiltin is enabled.
--- a/clang/include/clang/Basic/CodeGenOptions.h
+++ b/clang/include/clang/Basic/CodeGenOptions.h
@ -447,6 +447,10 @@ public:

  std::optional<double> AllowRuntimeCheckSkipHotCutoff;

+  /// Maximum number of allocation tokens (0 = no max), nullopt if none set (use
+  /// pass default).
+  std::optional<uint64_t> AllocTokenMax;
+
  /// List of backend command-line options for -fembed-bitcode.
  std::vector<uint8_t> CmdArgs;

--- a/clang/include/clang/Driver/Options.td
+++ b/clang/include/clang/Driver/Options.td
@ -2731,8 +2731,25 @@ def fsanitize_skip_hot_cutoff_EQ
          "(0.0 [default] = skip none; 1.0 = skip all). "
          "Argument format: <sanitizer1>=<value1>,<sanitizer2>=<value2>,...">;

+defm sanitize_alloc_token_fast_abi : BoolOption<"f", "sanitize-alloc-token-fast-abi",
+  CodeGenOpts<"SanitizeAllocTokenFastABI">, DefaultFalse,
+  PosFlag<SetTrue, [], [ClangOption], "Use the AllocToken fast ABI">,
+  NegFlag<SetFalse, [], [ClangOption], "Use the default AllocToken ABI">>,
+  Group<f_clang_Group>;
+defm sanitize_alloc_token_extended : BoolOption<"f", "sanitize-alloc-token-extended",
+  CodeGenOpts<"SanitizeAllocTokenExtended">, DefaultFalse,
+  PosFlag<SetTrue, [], [ClangOption], "Enable">,
+  NegFlag<SetFalse, [], [ClangOption], "Disable">,
+  BothFlags<[], [ClangOption], " extended coverage to custom allocation functions">>,
+  Group<f_clang_Group>;
+
 } // end -f[no-]sanitize* flags

+def falloc_token_max_EQ : Joined<["-"], "falloc-token-max=">,
+  Group<f_Group>, Visibility<[ClangOption, CC1Option]>,
+  MetaVarName<"<N>">,
+  HelpText<"Limit to maximum N allocation tokens (0 = no max)">;
+
 def fallow_runtime_check_skip_hot_cutoff_EQ
    : Joined<["-"], "fallow-runtime-check-skip-hot-cutoff=">,
      Group<f_clang_Group>,
--- a/clang/include/clang/Driver/SanitizerArgs.h
+++ b/clang/include/clang/Driver/SanitizerArgs.h
@ -75,6 +75,8 @@ class SanitizerArgs {
      llvm::AsanDetectStackUseAfterReturnMode::Invalid;

  std::string MemtagMode;
+  bool AllocTokenFastABI = false;
+  bool AllocTokenExtended = false;

 public:
  /// Parses the sanitizer arguments from an argument list.
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@ -60,11 +60,13 @@
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/HipStdPar/HipStdPar.h"
 #include "llvm/Transforms/IPO/EmbedBitcodePass.h"
+#include "llvm/Transforms/IPO/InferFunctionAttrs.h"
 #include "llvm/Transforms/IPO/LowerTypeTests.h"
 #include "llvm/Transforms/IPO/ThinLTOBitcodeWriter.h"
 #include "llvm/Transforms/InstCombine/InstCombine.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizer.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizerOptions.h"
+#include "llvm/Transforms/Instrumentation/AllocToken.h"
 #include "llvm/Transforms/Instrumentation/BoundsChecking.h"
 #include "llvm/Transforms/Instrumentation/DataFlowSanitizer.h"
 #include "llvm/Transforms/Instrumentation/GCOVProfiler.h"
@ -232,6 +234,14 @@ public:
 };
 } // namespace

+static AllocTokenOptions getAllocTokenOptions(const CodeGenOptions &CGOpts) {
+  AllocTokenOptions Opts;
+  Opts.MaxTokens = CGOpts.AllocTokenMax;
+  Opts.Extended = CGOpts.SanitizeAllocTokenExtended;
+  Opts.FastABI = CGOpts.SanitizeAllocTokenFastABI;
+  return Opts;
+}
+
 static SanitizerCoverageOptions
 getSancovOptsFromCGOpts(const CodeGenOptions &CGOpts) {
  SanitizerCoverageOptions Opts;
@ -789,6 +799,16 @@ static void addSanitizers(const Triple &TargetTriple,
      MPM.addPass(DataFlowSanitizerPass(LangOpts.NoSanitizeFiles,
                                        PB.getVirtualFileSystemPtr()));
    }
+
+    if (LangOpts.Sanitize.has(SanitizerKind::AllocToken)) {
+      if (Level == OptimizationLevel::O0) {
+        // The default pass builder only infers libcall function attrs when
+        // optimizing, so we insert it here because we need it for accurate
+        // memory allocation function detection.
+        MPM.addPass(InferFunctionAttrsPass());
+      }
+      MPM.addPass(AllocTokenPass(getAllocTokenOptions(CodeGenOpts)));
+    }
  };
  if (ClSanitizeOnOptimizerEarlyEP) {
    PB.registerOptimizerEarlyEPCallback(
--- a/clang/lib/Driver/SanitizerArgs.cpp
+++ b/clang/lib/Driver/SanitizerArgs.cpp
@ -61,8 +61,9 @@ static const SanitizerMask RecoverableByDefault =
    SanitizerKind::ImplicitConversion | SanitizerKind::Nullability |
    SanitizerKind::FloatDivideByZero | SanitizerKind::ObjCCast |
    SanitizerKind::Vptr;
-static const SanitizerMask Unrecoverable =
-    SanitizerKind::Unreachable | SanitizerKind::Return;
+static const SanitizerMask Unrecoverable = SanitizerKind::Unreachable |
+                                           SanitizerKind::Return |
+                                           SanitizerKind::AllocToken;
 static const SanitizerMask AlwaysRecoverable = SanitizerKind::KernelAddress |
                                               SanitizerKind::KernelHWAddress |
                                               SanitizerKind::KCFI;
@ -84,7 +85,8 @@ static const SanitizerMask CFIClasses =
 static const SanitizerMask CompatibleWithMinimalRuntime =
    TrappingSupported | SanitizerKind::Scudo | SanitizerKind::ShadowCallStack |
    SanitizerKind::MemtagStack | SanitizerKind::MemtagHeap |
-    SanitizerKind::MemtagGlobals | SanitizerKind::KCFI;
+    SanitizerKind::MemtagGlobals | SanitizerKind::KCFI |
+    SanitizerKind::AllocToken;

 enum CoverageFeature {
  CoverageFunc = 1 << 0,
@ -203,6 +205,7 @@ static void addDefaultIgnorelists(const Driver &D, SanitizerMask Kinds,
                     {"tysan_blacklist.txt", SanitizerKind::Type},
                     {"dfsan_abilist.txt", SanitizerKind::DataFlow},
                     {"cfi_ignorelist.txt", SanitizerKind::CFI},
+                     {"alloc_token_ignorelist.txt", SanitizerKind::AllocToken},
                     {"ubsan_ignorelist.txt",
                      SanitizerKind::Undefined | SanitizerKind::Vptr |
                          SanitizerKind::Integer | SanitizerKind::Nullability |
@ -650,7 +653,12 @@ SanitizerArgs::SanitizerArgs(const ToolChain &TC,
      std::make_pair(SanitizerKind::KCFI, SanitizerKind::Function),
      std::make_pair(SanitizerKind::Realtime,
                     SanitizerKind::Address | SanitizerKind::Thread |
-                         SanitizerKind::Undefined | SanitizerKind::Memory)};
+                         SanitizerKind::Undefined | SanitizerKind::Memory),
+      std::make_pair(SanitizerKind::AllocToken,
+                     SanitizerKind::Address | SanitizerKind::HWAddress |
+                         SanitizerKind::KernelAddress |
+                         SanitizerKind::KernelHWAddress |
+                         SanitizerKind::Memory)};

  // Enable toolchain specific default sanitizers if not explicitly disabled.
  SanitizerMask Default = TC.getDefaultSanitizers() & ~AllRemove;
@ -1159,6 +1167,15 @@ SanitizerArgs::SanitizerArgs(const ToolChain &TC,
        !TC.getTriple().isAndroid() && !TC.getTriple().isOSFuchsia();
  }

+  if (AllAddedKinds & SanitizerKind::AllocToken) {
+    AllocTokenFastABI = Args.hasFlag(
+        options::OPT_fsanitize_alloc_token_fast_abi,
+        options::OPT_fno_sanitize_alloc_token_fast_abi, AllocTokenFastABI);
+    AllocTokenExtended = Args.hasFlag(
+        options::OPT_fsanitize_alloc_token_extended,
+        options::OPT_fno_sanitize_alloc_token_extended, AllocTokenExtended);
+  }
+
  LinkRuntimes = Args.hasFlag(options::OPT_fsanitize_link_runtime,
                              options::OPT_fno_sanitize_link_runtime,
                              !Args.hasArg(options::OPT_r));
@ -1527,6 +1544,12 @@ void SanitizerArgs::addArgs(const ToolChain &TC, const llvm::opt::ArgList &Args,
      Sanitizers.has(SanitizerKind::Address))
    CmdArgs.push_back("-fno-assume-sane-operator-new");

+  // Flags for -fsanitize=alloc-token.
+  if (AllocTokenFastABI)
+    CmdArgs.push_back("-fsanitize-alloc-token-fast-abi");
+  if (AllocTokenExtended)
+    CmdArgs.push_back("-fsanitize-alloc-token-extended");
+
  // libFuzzer wants to intercept calls to certain library functions, so the
  // following -fno-builtin-* flags force the compiler to emit interposable
  // libcalls to these functions. Other sanitizers effectively do the same thing
--- a/clang/lib/Driver/ToolChain.cpp
+++ b/clang/lib/Driver/ToolChain.cpp
@ -1623,7 +1623,8 @@ SanitizerMask ToolChain::getSupportedSanitizers() const {
      SanitizerKind::CFICastStrict | SanitizerKind::FloatDivideByZero |
      SanitizerKind::KCFI | SanitizerKind::UnsignedIntegerOverflow |
      SanitizerKind::UnsignedShiftBase | SanitizerKind::ImplicitConversion |
-      SanitizerKind::Nullability | SanitizerKind::LocalBounds;
+      SanitizerKind::Nullability | SanitizerKind::LocalBounds |
+      SanitizerKind::AllocToken;
  if (getTriple().getArch() == llvm::Triple::x86 ||
      getTriple().getArch() == llvm::Triple::x86_64 ||
      getTriple().getArch() == llvm::Triple::arm ||
--- a/clang/lib/Driver/ToolChains/Clang.cpp
+++ b/clang/lib/Driver/ToolChains/Clang.cpp
@ -7618,6 +7618,8 @@ void Clang::ConstructJob(Compilation &C, const JobAction &JA,
  // features enabled through -Xclang -target-feature flags.
  SanitizeArgs.addArgs(TC, Args, CmdArgs, InputType);

+  Args.AddLastArg(CmdArgs, options::OPT_falloc_token_max_EQ);
+
 #if CLANG_ENABLE_CIR
  // Forward -mmlir arguments to to the MLIR option parser.
  for (const Arg *A : Args.filtered(options::OPT_mmlir)) {
--- a/clang/lib/Frontend/CompilerInvocation.cpp
+++ b/clang/lib/Frontend/CompilerInvocation.cpp
@ -1833,6 +1833,10 @@ void CompilerInvocationBase::GenerateCodeGenArgs(const CodeGenOptions &Opts,
       serializeSanitizerKinds(Opts.SanitizeAnnotateDebugInfo))
    GenerateArg(Consumer, OPT_fsanitize_annotate_debug_info_EQ, Sanitizer);

+  if (Opts.AllocTokenMax)
+    GenerateArg(Consumer, OPT_falloc_token_max_EQ,
+                std::to_string(*Opts.AllocTokenMax));
+
  if (!Opts.EmitVersionIdentMetadata)
    GenerateArg(Consumer, OPT_Qn);

@ -2346,6 +2350,15 @@ bool CompilerInvocation::ParseCodeGenArgs(CodeGenOptions &Opts, ArgList &Args,
    }
  }

+  if (const auto *Arg = Args.getLastArg(options::OPT_falloc_token_max_EQ)) {
+    StringRef S = Arg->getValue();
+    uint64_t Value = 0;
+    if (S.getAsInteger(0, Value))
+      Diags.Report(diag::err_drv_invalid_value) << Arg->getAsString(Args) << S;
+    else
+      Opts.AllocTokenMax = Value;
+  }
+
  Opts.EmitVersionIdentMetadata = Args.hasFlag(OPT_Qy, OPT_Qn, true);

  if (!LangOpts->CUDAIsDevice)
--- a/clang/lib/Frontend/InitPreprocessor.cpp
+++ b/clang/lib/Frontend/InitPreprocessor.cpp
@ -1530,6 +1530,8 @@ static void InitializePredefinedMacros(const TargetInfo &TI,
    Builder.defineMacro("__SANITIZE_HWADDRESS__");
  if (LangOpts.Sanitize.has(SanitizerKind::Thread))
    Builder.defineMacro("__SANITIZE_THREAD__");
+  if (LangOpts.Sanitize.has(SanitizerKind::AllocToken))
+    Builder.defineMacro("__SANITIZE_ALLOC_TOKEN__");

  // Target OS macro definitions.
  if (PPOpts.DefineTargetOSMacros) {
--- a/clang/test/CodeGen/alloc-token-ignorelist.c
+++ b/clang/test/CodeGen/alloc-token-ignorelist.c
@ -0,0 +1,27 @@
+// Test AllocToken respects ignorelist for functions and files.
+//
+// RUN: %clang_cc1 -fsanitize=alloc-token -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s --check-prefixes=CHECK,CHECK-ALLOW
+//
+// RUN: echo "fun:excluded_by_all" > %t.func.ignorelist
+// RUN: %clang_cc1 -fsanitize=alloc-token -fsanitize-ignorelist=%t.func.ignorelist -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s --check-prefixes=CHECK,CHECK-FUN
+//
+// RUN: echo "src:%s" | sed -e 's/\\/\\\\/g' > %t.file.ignorelist
+// RUN: %clang_cc1 -fsanitize=alloc-token -fsanitize-ignorelist=%t.file.ignorelist -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s --check-prefixes=CHECK,CHECK-SRC
+
+extern void* malloc(unsigned long size);
+
+// CHECK-LABEL: define{{.*}} @excluded_by_all(
+void* excluded_by_all(unsigned long size) {
+    // CHECK-ALLOW: call ptr @__alloc_token_malloc(
+    // CHECK-FUN: call ptr @malloc(
+    // CHECK-SRC: call ptr @malloc(
+    return malloc(size);
+}
+
+// CHECK-LABEL: define{{.*}} @excluded_by_src(
+void* excluded_by_src(unsigned long size) {
+    // CHECK-ALLOW: call ptr @__alloc_token_malloc(
+    // CHECK-FUN: call ptr @__alloc_token_malloc(
+    // CHECK-SRC: call ptr @malloc(
+    return malloc(size);
+}
--- a/clang/test/CodeGen/alloc-token-lower.c
+++ b/clang/test/CodeGen/alloc-token-lower.c
@ -0,0 +1,22 @@
+// Test optimization pipelines do not interfere with AllocToken lowering, and we
+// pass on function attributes correctly.
+//
+// RUN: %clang_cc1     -fsanitize=alloc-token -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s
+// RUN: %clang_cc1 -O1 -fsanitize=alloc-token -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s
+// RUN: %clang_cc1 -O2 -fsanitize=alloc-token -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s
+
+typedef __typeof(sizeof(int)) size_t;
+
+void *malloc(size_t size);
+
+// CHECK-LABEL: @test_malloc(
+// CHECK: call{{.*}} ptr @__alloc_token_malloc(i64 noundef 4, i64 0)
+void *test_malloc() {
+  return malloc(sizeof(int));
+}
+
+// CHECK-LABEL: @no_sanitize_malloc(
+// CHECK: call{{.*}} ptr @malloc(i64 noundef 4)
+void *no_sanitize_malloc(size_t size) __attribute__((no_sanitize("alloc-token"))) {
+  return malloc(sizeof(int));
+}
--- a/clang/test/CodeGen/alloc-token.c
+++ b/clang/test/CodeGen/alloc-token.c
@ -0,0 +1,37 @@
+// RUN: %clang_cc1 -fsanitize=alloc-token -triple x86_64-linux-gnu -emit-llvm -disable-llvm-passes %s -o - | FileCheck %s
+
+typedef __typeof(sizeof(int)) size_t;
+
+void *aligned_alloc(size_t alignment, size_t size);
+void *malloc(size_t size);
+void *calloc(size_t num, size_t size);
+void *realloc(void *ptr, size_t size);
+void *reallocarray(void *ptr, size_t nmemb, size_t size);
+void *memalign(size_t alignment, size_t size);
+void *valloc(size_t size);
+void *pvalloc(size_t size);
+int posix_memalign(void **memptr, size_t alignment, size_t size);
+
+void *sink;
+
+// CHECK-LABEL: define dso_local void @test_malloc_like(
+// CHECK: call ptr @malloc(i64 noundef 4)
+// CHECK: call ptr @calloc(i64 noundef 3, i64 noundef 4)
+// CHECK: call ptr @realloc(ptr noundef {{.*}}, i64 noundef 8)
+// CHECK: call ptr @reallocarray(ptr noundef {{.*}}, i64 noundef 5, i64 noundef 8)
+// CHECK: call align 128 ptr @aligned_alloc(i64 noundef 128, i64 noundef 1024)
+// CHECK: call align 16 ptr @memalign(i64 noundef 16, i64 noundef 256)
+// CHECK: call ptr @valloc(i64 noundef 4096)
+// CHECK: call ptr @pvalloc(i64 noundef 8192)
+// CHECK: call i32 @posix_memalign(ptr noundef @sink, i64 noundef 64, i64 noundef 4)
+void test_malloc_like() {
+  sink = malloc(sizeof(int));
+  sink = calloc(3, sizeof(int));
+  sink = realloc(sink, sizeof(long));
+  sink = reallocarray(sink, 5, sizeof(long));
+  sink = aligned_alloc(128, 1024);
+  sink = memalign(16, 256);
+  sink = valloc(4096);
+  sink = pvalloc(8192);
+  posix_memalign(&sink, 64, sizeof(int));
+}
--- a/clang/test/CodeGenCXX/alloc-token.cpp
+++ b/clang/test/CodeGenCXX/alloc-token.cpp
@ -0,0 +1,141 @@
+// RUN: %clang_cc1 -fsanitize=alloc-token -triple x86_64-linux-gnu -std=c++20 -fexceptions -fcxx-exceptions -emit-llvm -disable-llvm-passes %s -o - | FileCheck %s
+
+#include "../Analysis/Inputs/system-header-simulator-cxx.h"
+extern "C" {
+void *aligned_alloc(size_t alignment, size_t size);
+void *malloc(size_t size);
+void *calloc(size_t num, size_t size);
+void *realloc(void *ptr, size_t size);
+void *reallocarray(void *ptr, size_t nmemb, size_t size);
+void *memalign(size_t alignment, size_t size);
+void *valloc(size_t size);
+void *pvalloc(size_t size);
+int posix_memalign(void **memptr, size_t alignment, size_t size);
+
+struct __sized_ptr_t {
+  void *p;
+  size_t n;
+};
+enum class __hot_cold_t : uint8_t;
+__sized_ptr_t __size_returning_new(size_t size);
+__sized_ptr_t __size_returning_new_hot_cold(size_t, __hot_cold_t);
+__sized_ptr_t __size_returning_new_aligned(size_t, std::align_val_t);
+__sized_ptr_t __size_returning_new_aligned_hot_cold(size_t, std::align_val_t,  __hot_cold_t);
+}
+
+void *sink; // prevent optimizations from removing the calls
+
+// CHECK-LABEL: define dso_local void @_Z16test_malloc_likev(
+// CHECK: call ptr @malloc(i64 noundef 4)
+// CHECK: call ptr @calloc(i64 noundef 3, i64 noundef 4)
+// CHECK: call ptr @realloc(ptr noundef {{.*}}, i64 noundef 8)
+// CHECK: call ptr @reallocarray(ptr noundef {{.*}}, i64 noundef 5, i64 noundef 8)
+// CHECK: call align 128 ptr @aligned_alloc(i64 noundef 128, i64 noundef 1024)
+// CHECK: call ptr @memalign(i64 noundef 16, i64 noundef 256)
+// CHECK: call ptr @valloc(i64 noundef 4096)
+// CHECK: call ptr @pvalloc(i64 noundef 8192)
+// CHECK: call i32 @posix_memalign(ptr noundef @sink, i64 noundef 64, i64 noundef 4)
+void test_malloc_like() {
+  sink = malloc(sizeof(int));
+  sink = calloc(3, sizeof(int));
+  sink = realloc(sink, sizeof(long));
+  sink = reallocarray(sink, 5, sizeof(long));
+  sink = aligned_alloc(128, 1024);
+  sink = memalign(16, 256);
+  sink = valloc(4096);
+  sink = pvalloc(8192);
+  posix_memalign(&sink, 64, sizeof(int));
+}
+
+// CHECK-LABEL: define dso_local void @_Z17test_operator_newv(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 4)
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 4)
+void test_operator_new() {
+  sink = __builtin_operator_new(sizeof(int));
+  sink = ::operator new(sizeof(int));
+}
+
+// CHECK-LABEL: define dso_local void @_Z25test_operator_new_nothrowv(
+// CHECK: call noalias noundef ptr @_ZnwmRKSt9nothrow_t(i64 noundef 4, ptr noundef nonnull align 1 dereferenceable(1) @_ZSt7nothrow)
+// CHECK: call noalias noundef ptr @_ZnwmRKSt9nothrow_t(i64 noundef 4, ptr noundef nonnull align 1 dereferenceable(1) @_ZSt7nothrow)
+void test_operator_new_nothrow() {
+  sink = __builtin_operator_new(sizeof(int), std::nothrow);
+  sink = ::operator new(sizeof(int), std::nothrow);
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z8test_newv(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 4){{.*}} !alloc_token [[META_INT:![0-9]+]]
+int *test_new() {
+  return new int;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z14test_new_arrayv(
+// CHECK: call noalias noundef nonnull ptr @_Znam(i64 noundef 40){{.*}} !alloc_token [[META_INT]]
+int *test_new_array() {
+  return new int[10];
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z16test_new_nothrowv(
+// CHECK: call noalias noundef ptr @_ZnwmRKSt9nothrow_t(i64 noundef 4, ptr noundef nonnull align 1 dereferenceable(1) @_ZSt7nothrow){{.*}} !alloc_token [[META_INT]]
+int *test_new_nothrow() {
+  return new (std::nothrow) int;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z22test_new_array_nothrowv(
+// CHECK: call noalias noundef ptr @_ZnamRKSt9nothrow_t(i64 noundef 40, ptr noundef nonnull align 1 dereferenceable(1) @_ZSt7nothrow){{.*}} !alloc_token [[META_INT]]
+int *test_new_array_nothrow() {
+  return new (std::nothrow) int[10];
+}
+
+// CHECK-LABEL: define dso_local void @_Z23test_size_returning_newv(
+// CHECK: call { ptr, i64 } @__size_returning_new(i64 noundef 8)
+// CHECK: call { ptr, i64 } @__size_returning_new_hot_cold(i64 noundef 8, i8 noundef zeroext 1)
+// CHECK: call { ptr, i64 } @__size_returning_new_aligned(i64 noundef 8, i64 noundef 32)
+// CHECK: call { ptr, i64 } @__size_returning_new_aligned_hot_cold(i64 noundef 8, i64 noundef 32, i8 noundef zeroext 1)
+void test_size_returning_new() {
+  sink = __size_returning_new(sizeof(long)).p;
+  sink = __size_returning_new_hot_cold(sizeof(long), __hot_cold_t{1}).p;
+  sink = __size_returning_new_aligned(sizeof(long), std::align_val_t{32}).p;
+  sink = __size_returning_new_aligned_hot_cold(sizeof(long), std::align_val_t{32}, __hot_cold_t{1}).p;
+}
+
+class TestClass {
+public:
+  virtual void Foo();
+  virtual ~TestClass();
+  int data[16];
+};
+
+void may_throw();
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z27test_exception_handling_newv(
+// CHECK: invoke noalias noundef nonnull ptr @_Znwm(i64 noundef 72)
+// CHECK-NEXT: !alloc_token [[META_TESTCLASS:![0-9]+]]
+TestClass *test_exception_handling_new() {
+  try {
+    TestClass *obj = new TestClass();
+    may_throw();
+    return obj;
+  } catch (...) {
+    return nullptr;
+  }
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z14test_new_classv(
+// CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef 72){{.*}} !alloc_token [[META_TESTCLASS]]
+TestClass *test_new_class() {
+  TestClass *obj = new TestClass();
+  obj->data[0] = 42;
+  return obj;
+}
+
+// CHECK-LABEL: define dso_local noundef ptr @_Z20test_new_class_arrayv(
+// CHECK: call noalias noundef nonnull ptr @_Znam(i64 noundef 728){{.*}} !alloc_token [[META_TESTCLASS]]
+TestClass *test_new_class_array() {
+  TestClass* arr = new TestClass[10];
+  arr[0].data[0] = 123;
+  return arr;
+}
+
+// CHECK: [[META_INT]] = !{!"int"}
+// CHECK: [[META_TESTCLASS]] = !{!"TestClass"}
--- a/clang/test/Driver/fsanitize-alloc-token.c
+++ b/clang/test/Driver/fsanitize-alloc-token.c
@ -0,0 +1,43 @@
+// RUN: %clang --target=x86_64-linux-gnu -fsanitize=alloc-token %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-TOKEN-ALLOC
+// CHECK-TOKEN-ALLOC: "-fsanitize=alloc-token"
+
+// RUN: %clang --target=x86_64-linux-gnu -fsanitize=alloc-token -fno-sanitize=alloc-token %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-NO-TOKEN-ALLOC
+// CHECK-NO-TOKEN-ALLOC-NOT: "-fsanitize=alloc-token"
+
+// RUN: %clang --target=x86_64-linux-gnu -flto -fvisibility=hidden -fno-sanitize-ignorelist -fsanitize=alloc-token,undefined,cfi %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-COMPATIBLE
+// CHECK-COMPATIBLE: "-fsanitize={{.*}}alloc-token"
+
+// RUN: %clang --target=x86_64-linux-gnu -fsanitize=alloc-token -fsanitize-minimal-runtime %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-MINIMAL
+// CHECK-MINIMAL: "-fsanitize=alloc-token"
+// CHECK-MINIMAL: "-fsanitize-minimal-runtime"
+
+// RUN: %clang --target=arm-arm-non-eabi -fsanitize=alloc-token %s -### 2>&1 | FileCheck %s -check-prefix=CHECK-BAREMETAL
+// RUN: %clang --target=aarch64-none-elf -fsanitize=alloc-token %s -### 2>&1 | FileCheck %s -check-prefix=CHECK-BAREMETAL
+// CHECK-BAREMETAL: "-fsanitize=alloc-token"
+
+// RUN: not %clang --target=x86_64-linux-gnu -fsanitize=alloc-token,address %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-INCOMPATIBLE-ADDRESS
+// CHECK-INCOMPATIBLE-ADDRESS: error: invalid argument '-fsanitize=alloc-token' not allowed with '-fsanitize=address'
+
+// RUN: not %clang --target=x86_64-linux-gnu -fsanitize=alloc-token,memory %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-INCOMPATIBLE-MEMORY
+// CHECK-INCOMPATIBLE-MEMORY: error: invalid argument '-fsanitize=alloc-token' not allowed with '-fsanitize=memory'
+
+// RUN: not %clang --target=x86_64-linux-gnu -fsanitize=alloc-token -fsanitize-trap=alloc-token %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-INCOMPATIBLE-TRAP
+// CHECK-INCOMPATIBLE-TRAP: error: unsupported argument 'alloc-token' to option '-fsanitize-trap='
+
+// RUN: not %clang --target=x86_64-linux-gnu %s -fsanitize=alloc-token -fsanitize-recover=alloc-token -### 2>&1 | FileCheck %s --check-prefix=CHECK-INCOMPATIBLE-RECOVER
+// CHECK-INCOMPATIBLE-RECOVER: unsupported argument 'alloc-token' to option '-fsanitize-recover='
+
+// RUN: %clang --target=x86_64-linux-gnu -fsanitize=alloc-token -fsanitize-alloc-token-fast-abi %s -### 2>&1 | FileCheck -check-prefix=CHECK-FASTABI %s
+// CHECK-FASTABI: "-fsanitize-alloc-token-fast-abi"
+// RUN: %clang --target=x86_64-linux-gnu -fsanitize=alloc-token -fsanitize-alloc-token-fast-abi -fno-sanitize-alloc-token-fast-abi %s -### 2>&1 | FileCheck -check-prefix=CHECK-NOFASTABI %s
+// CHECK-NOFASTABI-NOT: "-fsanitize-alloc-token-fast-abi"
+
+// RUN: %clang --target=x86_64-linux-gnu -fsanitize=alloc-token -fsanitize-alloc-token-extended %s -### 2>&1 | FileCheck -check-prefix=CHECK-EXTENDED %s
+// CHECK-EXTENDED: "-fsanitize-alloc-token-extended"
+// RUN: %clang --target=x86_64-linux-gnu -fsanitize=alloc-token -fsanitize-alloc-token-extended -fno-sanitize-alloc-token-extended %s -### 2>&1 | FileCheck -check-prefix=CHECK-NOEXTENDED %s
+// CHECK-NOEXTENDED-NOT: "-fsanitize-alloc-token-extended"
+
+// RUN: %clang --target=x86_64-linux-gnu -falloc-token-max=0 -falloc-token-max=42 %s -### 2>&1 | FileCheck -check-prefix=CHECK-MAX %s
+// CHECK-MAX: "-falloc-token-max=42"
+// RUN: not %clang --target=x86_64-linux-gnu -fsanitize=alloc-token -falloc-token-max=-1 %s 2>&1 | FileCheck -check-prefix=CHECK-INVALID-MAX %s
+// CHECK-INVALID-MAX: error: invalid value
--- a/clang/test/Preprocessor/alloc_token.cpp
+++ b/clang/test/Preprocessor/alloc_token.cpp
@ -0,0 +1,10 @@
+// RUN: %clang_cc1 -E -fsanitize=alloc-token %s -o - | FileCheck --check-prefix=CHECK-SANITIZE %s
+// RUN: %clang_cc1 -E  %s -o - | FileCheck --check-prefix=CHECK-DEFAULT %s
+
+#if __SANITIZE_ALLOC_TOKEN__
+// CHECK-SANITIZE: has_sanitize_alloc_token
+int has_sanitize_alloc_token();
+#else
+// CHECK-DEFAULT: no_sanitize_alloc_token
+int no_sanitize_alloc_token();
+#endif