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

[rcvalle]

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 rust-lang/rust@f43ca59 (or rust-lang/rust#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 rust-lang/rust#160298 draft PR, MCP available in the rust-lang/compiler-team#1032 issue, and the design is described in detail in the design document in the tracking issue rust-lang/rust#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.

[ChayimFriedman2]

Rust explicitly has untyped memory, unlike C and C++. That means it is impossible to determine at allocation time if a memory will contain pointers. While exploit mitigations that rule out valid unsafe code were already implemented in the past (e.g. CFI prohibiting ABI-compatible function pointer transmutes), this feels like a big deviation that will rule out a large percentage of code.

Also, how is the (generic, e.g. Vec) caller code supposed to determine if the type contains pointers? If some API like fn contain_pointers<T>() -> bool is warranted, this probably require involvement of T-lang as well.

[RalfJung]

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).

So can I imagine this as the Rust wrapper being written as follows?

if token == 0 {
  __alloc_token_0_malloc(...)
} else {
  __alloc_token_1_malloc(..)
}

Implicit-only (i.e., just the AllocTokenPass without API):

What's the AllocTokenPass?

Alternatives

Would be good to summarize the discussion that led to this name (*_with_token), since it is a pretty obscure term that we are considering solely because there is value in following the (IMO bad) precedent set by clang.

token: usize,

Given that there will be a magic intrinsic to generate those tokens (why is that not in the ACP?), I think this should have a newtype, not a plain usize.

[RalfJung]

Rust explicitly has untyped memory, unlike C and C++. That means it is impossible to determine at allocation time if a memory will contain pointers. While exploit mitigations that rule out valid unsafe code were already implemented in the past (e.g. CFI prohibiting ABI-compatible function pointer transmutes), this feels like a big deviation that will rule out a large percentage of code.

The ACP does not do a good job at explaining this feature to Rust people. To summarize what we learned in prior discussions: This isn't typed memory, it's just a hint to the allocator for where to place the allocation. Nothing goes wrong if you put a pointer into a "no-pointer allocation", except that the mitigation becomes less effective.

Also, how is the (generic, e.g. Vec) caller code supposed to determine if the type contains pointers? If some API like fn contain_pointers() -> bool is warranted, this probably require involvement of T-lang as well.

There's a magic intrinsic that takes the type as generic parameter and returns a token, based on how the sanitizer is configured. And yes that's a new language feature and thus needs t-lang involvement.

[matthieu-m]

Rust explicitly has untyped memory, unlike C and C++. That means it is impossible to determine at allocation time if a memory will contain pointers. While exploit mitigations that rule out valid unsafe code were already implemented in the past (e.g. CFI prohibiting ABI-compatible function pointer transmutes), this feels like a big deviation that will rule out a large percentage of code.

In particular, the nightly Vec::recycle is specifically meant to help with reusing a memory block while changing its type.

It appears that this usage is supported by the above API, as I do not see any hint of a token being provided for deallocation, only for new allocations, but... :/

---

Adding a whole 4 methods to Allocator for just 1 specific partitioning strategy seems like it would scale poorly.

Are there any other partitioning strategies in use today, or in the works, which would also require "duplicating" the API.

---

The Linux Kernel also uses a contingent of extra flags when allocating memory (see kmalloc). How do they mix flags & allocation tokens?

Should we expect Allocator::allocate_with_flags and Allocator::allocate_with_token_flags or could flags & partition token be combined in a single integer?

---

Does the requirement of working with EII & dyn Allocator truly precludes a "polymorphic" token, which could be reused in various situations?

Would it possible to instrument all Allocator methods with an arbitrary foreign argument:

trait Allocator {
    #[inline(always)]
    fn allocate_with<F>(&self, layout: Layout, foreign: F) -> Result<NonNull<u8>, AllocError>
    where
         Self: Sized,
    {
        self.allocate_with_dyn(layout, &foreign)
    }

    #[inline(always)]
    fn allocate_with_dyn(&self, layout: Layout, _foreign: &dyn Any) -> Result<NonNull<u8>, AllocError> {
        self.allocate(layout)
    }

Implementers would then use std::any::try_as_dyn (or something) to check whether a specific trait of interest (PartitionToken or KernelFlags) is implemented, and extract the values they care about out of it. With the non-dynamic foreign, the try_as_dyn should be resolved at compile-time, so would have 0-overhead. With the dynamic foreign, well, I'm not a magician.

(std::any::try_as_dyn does lack delegation, that is the argument cannot itself delegate the call to any of it fields, which means that a compiler pass inserting tokens would need to rewrap existing foreigns in a type which implement any trait the foreign implements... or something...)

This would mean Allocator gets triple its API surface... once and for all. No matter what other scheme is cooked up which requires extra arguments, they can just be passed as foreign, from those who care to those who care.

[nia-e]

Is there value in having this be an extension trait of Allocator, with the extra guarantee that a type that impls said extension trait must actually respect the token? I'm a bit concerned about adding a bunch of methods to a dyn-compatible trait and about API explosion, since this runs into many of the same issues - notably not including the need for an associated type - as flag passing. If maybe-bounds are likely anywhere in the near future I'd prefer going in that direction personally

[ChayimFriedman2]

@nia-e How do you call the extension trait's methods without specialization?

[nia-e]

That's the proposal for T: maybe Trait bounds - i don't see a clear place where it's explained, but the tl;dr for anyone in this thread who's unfamiliar is that it's effectively the same compilation strategy as a const generic bool for whether each maybe-trait is implemented that's propagated up the callstack and determined at whatever callsite gives a concrete type, and users of that bound can branch on whether their T impls Trait (and call those methods on it if so). this avoids all the soundness problems with specialisation and is basically just impl + syntax bikeshed work as I understand it. thus, methods on e.g. Vec could have a bound of A: Allocator + maybe TokenPartitioning or the likes and branch on whether that trait s implemented when they want to call allocator methods

[ChayimFriedman2]

That is still essentially a form of specialization (and one proposed form for it, there are other possible ways to make specialization sound). Blocking this on any form of specialization means blocking this for a long time (even though today specialization might actually be in sight). Which we might want to do, but this is a trade-off to consider.

[nia-e]

yep, that's fair. i was under the impression that maybe bounds are coming much sooner than full specialisation (arguably they're quite different and mostly orthogonal - one is impl specialisation, this is callsite specialisation). if this would also take a long time, though, i'm less inclined to block such a stabilisation