ACP: Add token-enabled allocator interface support for LLVM AllocToken and heap partitioning

Proposal

Problem statement

Overview

Heap partitioning is an exploit mitigation that protects programs from having their heap layout controlled for exploitation by separating allocations into isolated heap partitions. It is implemented by modern memory allocators and operating systems, such as the Apple XNU kernel kalloc_type allocator, the Chromium PartitionAlloc allocator, the GrapheneOS hardened_malloc allocator, and the Linux kernel randomized kmalloc caches. LLVM AllocToken (available in LLVM 22 or higher, with Clang support) assigns a token identifier (i.e., an integer) to each allocation, and passes it to the allocator by rewriting allocation calls to token-enabled versions of the allocation functions.

The Rust compiler currently does not support heap partitioning when

• building Rust-compiled code only programs.
• linking foreign C- or C++-compiled code into a program written in Rust.
• linking foreign Rust-compiled code into a program written in C or C++.

Therefore, the absence of support for heap partitioning in the Rust compiler is a security concern when gradually migrating from C and C++ to Rust, and when C or C++ and Rust-compiled code share the same virtual address space and memory allocator.

Standard library and allocator support

Rust heap allocations are performed through type-erased interfaces (i.e., GlobalAlloc::alloc, Allocator::allocate, and alloc::alloc::alloc are passed a core::alloc::Layout), so a token identifier can not be passed to a token-enabled memory allocator, leaving allocator implementers with no way to implement the token-enabled allocator interface in Rust.

Motivating examples or use cases

Heap partitioning is implemented by modern memory allocators and operating systems, such as the Apple XNU kernel kalloc_type allocator, the Chromium PartitionAlloc allocator, the GrapheneOS hardened_malloc allocator, and the Linux kernel randomized kmalloc caches. LLVM AllocToken (available in LLVM 22 or higher, with Clang support) assigns a token identifier (i.e., an integer) to each allocation, and passes it to the allocator by rewriting allocation calls to token-enabled versions of the allocation functions.

TCMalloc implements the token-enabled allocator interface using the fast ABI (i.e., the token identifier is encoded in the allocation function name, such as __alloc_token_0_malloc and __alloc_token_1_malloc, instead of being appended as an argument) with the maximum number of tokens set to two (i.e., the pointer-split heap partitioning scheme).

$ git clone https://github.com/google/tcmalloc
$ cd tcmalloc
$ bazel build --copt=-fsanitize=alloc-token --copt=-fsanitize-alloc-token-fast-abi --copt=-falloc-token-max=2 --copt=-DALLOC_TOKEN_MAX=2 //tcmalloc:tcmalloc

Fig. 1. Build of TCMalloc with allocation token instrumentation enabled.

#include <stdlib.h>

struct buffer {
    char data[4096];
};

struct node {
    struct node *next;
    unsigned long value;
};

void *
alloc_buffer(void)
{
    /* Allocations of types not containing pointers (i.e., `struct buffer`) are
       placed in partition 0. */
    return malloc(sizeof(struct buffer));
}

void *
alloc_node(void)
{
    /* Allocations of types containing pointers (i.e., `struct node`) are placed
       in partition 1. */
    return malloc(sizeof(struct node));
}

Fig. 2. Example C library.

pub struct Buffer {
    data: [u8; 4096],
}

pub struct Node {
    next: *mut Node,
    value: u64,
}

#[link(name = "foo")]
extern "C" {
    fn alloc_buffer() -> *mut u8;
    fn alloc_node() -> *mut u8;
}

fn main() {
    // Allocations of types not containing pointers (i.e., `Buffer`) are placed
    // in partition 0.
    let buffer = Box::new(Buffer { data: [1; 4096] });

    // Allocations of types containing pointers (i.e., `Node`) are placed in
    // partition 1.
    let node = Box::new(Node {
        next: core::ptr::null_mut(),
        value: 1,
    });

    let c_buffer = unsafe { alloc_buffer() };
    let c_node = unsafe { alloc_node() };

    println!("Rust Buffer allocation: {:p}", &*buffer);
    println!("Rust Node allocation:   {:p}", &*node);
    println!("C buffer allocation:    {:p}", c_buffer);
    println!("C node allocation:      {:p}", c_node);

    std::hint::black_box((buffer, node, c_buffer, c_node));
}

Fig. 3. Example allocations of types not containing pointers (i.e., Buffer and struct buffer) and types containing pointers (i.e., Node and struct node) by both C- and Rust-compiled code in a mixed binary.

$ make
mkdir -p target/release
clang -I. -Isrc -Wall -fsanitize=alloc-token -fsanitize-alloc-token-fast-abi -falloc-token-max=2 --target=x86_64-unknown-linux-gnu -c src/foo.c -o target/release/libfoo.o
llvm-ar rcs target/release/libfoo.a target/release/libfoo.o
RUSTFLAGS="-L./target/release -L/opt/tcmalloc -Clink-arg=-ltcmalloc -Clink-arg=-lstdc++ -Clink-arg=-lm -Clinker=clang -Clink-arg=-fuse-ld=lld -Zsanitizer=alloc-token -Zsanitizer-alloc-token-fast-abi" cargo build -Zbuild-std --release --target x86_64-unknown-linux-gnu
   ...
   Compiling rust-heap-partitioning-3 v0.1.0 (/workdir/rust-heap-partitioning-3)
   ...
    Finished `release` profile [optimized] target(s) in 1m 02s
$ TCMALLOC_HEAP_PARTITIONING=1 ./target/x86_64-unknown-linux-gnu/release/rust-heap-partitioning-3
Rust Buffer allocation: 0x32ad7fe00000
Rust Node allocation:   0x3b293fc000c0
C buffer allocation:    0x32ad7fe01000
C node allocation:      0x3b293fc000b0

Fig. 4. Build and execution of Figs. 2 and 3 with LLVM AllocToken enabled.

When LLVM AllocToken is enabled and a single token-enabled memory allocator serves all compiled code in the mixed binary, allocations of types not containing pointers (i.e., Buffer and struct buffer) and allocations of
types containing pointers (i.e., Node and struct node) are placed in separate heap partitions (i.e., the printed allocation addresses are in separate memory regions), for both the C- and Rust-compiled code (see Fig. 4).

See more examples in the https://github.com/rcvalle/rust-heap-partitioning-examples repository.

Solution sketch

Four unstable methods on core::alloc::Allocator, mirroring the existing methods (i.e., allocate, allocate_zeroed, grow, shrink):

    /// Behaves like `allocate`, but also passes a token identifier, for LLVM AllocToken and heap
    /// partitioning support (i.e., for token-enabled memory allocators that use token
    /// identifiers to separate allocations into partitions).
    ///
    /// The default implementation ignores `token` and calls `allocate`, for backward
    /// compatibility with existing allocators.
    #[unstable(feature = "alloc_with_token", issue = "159111")]
    fn allocate_with_token(
        &self,
        layout: Layout,
        token: usize,
    ) -> Result<NonNull<[u8]>, AllocError> {
        let _ = token;
        self.allocate(layout)
    }

    /// Behaves like `allocate_zeroed`, but also passes a token identifier, for LLVM AllocToken
    /// and heap partitioning support (i.e., for token-enabled memory allocators that use token
    /// identifiers to separate allocations into partitions).
    ///
    /// The default implementation ignores `token` and calls `allocate_zeroed`, for backward
    /// compatibility with existing allocators.
    #[unstable(feature = "alloc_with_token", issue = "159111")]
    fn allocate_zeroed_with_token(
        &self,
        layout: Layout,
        token: usize,
    ) -> Result<NonNull<[u8]>, AllocError> {
        let _ = token;
        self.allocate_zeroed(layout)
    }

    /// Behaves like `grow`, but also passes a token identifier, for LLVM AllocToken and heap
    /// partitioning support (i.e., for token-enabled memory allocators that use token
    /// identifiers to separate allocations into partitions).
    ///
    /// The default implementation ignores `token` and calls `grow`, for backward compatibility
    /// with existing allocators.
    ///
    /// # Safety
    ///
    /// Same as `grow`.
    #[unstable(feature = "alloc_with_token", issue = "159111")]
    unsafe fn grow_with_token(
        &self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
        token: usize,
    ) -> Result<NonNull<[u8]>, AllocError> {
        let _ = token;
        // SAFETY: the safety contract for `grow` must be upheld by the caller.
        unsafe { self.grow(ptr, old_layout, new_layout) }
    }

    /// Behaves like `shrink`, but also passes a token identifier, for LLVM AllocToken and heap
    /// partitioning support (i.e., for token-enabled memory allocators that use token
    /// identifiers to separate allocations into partitions).
    ///
    /// The default implementation ignores `token` and calls `shrink`, for backward
    /// compatibility with existing allocators.
    ///
    /// # Safety
    ///
    /// Same as `shrink`.
    #[unstable(feature = "alloc_with_token", issue = "159111")]
    unsafe fn shrink_with_token(
        &self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
        token: usize,
    ) -> Result<NonNull<[u8]>, AllocError> {
        let _ = token;
        // SAFETY: the safety contract for `shrink` must be upheld by the caller.
        unsafe { self.shrink(ptr, old_layout, new_layout) }
    }

See the complete implementation in https://github.com/rust-lang/rust/pull/160298/changes/f43ca594a0c479b87eb6068aa231b585864b2468 (or https://github.com/rust-lang/rust/pull/160298, commit 6).

Alternatives

Each was considered or actually tried in the series:

• Token methods on GlobalAlloc: were in the original implementation and were rejected by T-libs feedback (re: no new API on the legacy trait).
• Free functions in the alloc module: were also rejected by T-libs feedback (re: no free functions).
• Inherent methods on Global/System only: would be insufficient because generic code and custom allocators couldn't participate, and future dispatch (e.g., an EII-based allocator interface generating calls on the registered static) would need a trait.
• Changed signatures on existing methods or a separate trait: would be a breaking change or cause fragmentation (as opposed to additions with default implementations).
• Implicit-only (i.e., just the AllocTokenPass without API): would cover std paths but would leave allocator implementers with no way to implement the token-enabled allocator interface in Rust.
• A crates.io crate: is not viable because it requires methods on the Allocator trait in core and the System impl in std, coordinated with compiler instrumentation.

Links and related work

Working examples are available in the https://github.com/rcvalle/rust-heap-partitioning-examples repository, the complete implementation is available in the https://github.com/rust-lang/rust/pull/160298 draft PR, MCP available in the https://github.com/rust-lang/compiler-team/issues/1032 issue, and the design is described in detail in the design document in the tracking issue https://github.com/rust-lang/rust/issues/159111.

What happens now?

This issue contains an API change proposal (or ACP) and is part of the libs-api team feature lifecycle. Once this issue is filed, the libs-api team will review open proposals as capability becomes available. Current response times do not have a clear estimate, but may be up to several months.

Possible responses

The libs team may respond in various different ways. First, the team will consider the problem (this doesn't require any concrete solution or alternatives to have been proposed):

• We think this problem seems worth solving, and the standard library might be the right place to solve it.
• We think that this probably doesn't belong in the standard library.

Second, if there's a concrete solution:

• We think this specific solution looks roughly right, approved, you or someone else should implement this. (Further review will still happen on the subsequent implementation PR.)
• We're not sure this is the right solution, and the alternatives or other materials don't give us enough information to be sure about that. Here are some questions we have that aren't answered, or rough ideas about alternatives we'd want to see discussed.

Start with the proposed methods on core::alloc::Allocator and compare them with the complete implementation in rust-lang/rust#160298, especially commit 6. Read the tracking issue rust-lang/rust#159111 and compiler-team issue #1032 for the design and instrumentation context. Done means the libs-api review reaches a decision on the proposed interface.