Support for custom allocator in miri

[jacob-hughes]

If I understand correctly, Miri replaces the allocator, even when the #[global_allocator] attribute is used to specify a custom allocator. Unfortunately, this can lead to Miri thinking that a program invokes UB when it does not.

In the following example, I implement a custom allocator where alloc reserves the first machine word in each block in order to store a header, it then returns a pointer to the next address. Since Miri replaces this with its own allocator, dereferencing a pointer to the header's offset will be out-of-bounds.

#![feature(alloc_layout_extra)]
use std::alloc::{GlobalAlloc, Layout, System};
use std::mem::size_of;

#[global_allocator]
static ALLOCATOR: Allocator = Allocator {};

const HEADER: usize = 42;

struct Allocator {}

unsafe impl GlobalAlloc for Allocator {
    // Allocate a block with a usize sized header, and return a pointer to the
    // next address.
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let (new_layout, _) = Layout::new::<usize>().extend(layout).unwrap();
        let base_ptr = System.alloc(new_layout) as *mut usize;
        ::std::ptr::write(base_ptr, HEADER);
        base_ptr.add(1) as *mut u8
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let (new_layout, _) = Layout::new::<usize>().extend(layout).unwrap();
        let base_ptr = (ptr as *mut usize).wrapping_sub(1) as *mut u8;
        System.dealloc(base_ptr, new_layout);
    }
}

fn main() {
    let block = Box::into_raw(Box::new(123_usize));

    // Get a pointer to the block header using integer arithmetic.
    let header_ptr = (block as usize - size_of::<usize>()) as *mut usize;
    assert_eq!(unsafe { *header_ptr }, HEADER); // Miri error: dangling pointer dereference

    // Get a pointer to the block header using wrapping_sub.
    let header_ptr = block.wrapping_sub(1);
    assert_eq!(unsafe { *header_ptr }, HEADER); // Miri error: overflowing in-bounds pointer arithmetic

    unsafe { Box::from_raw(block) };
}

Running this under miri causes an evaluation error upon the first misuse of header_ptr:

   Compiling playground v0.0.1 (/playground)
error: Miri evaluation error: dangling pointer was dereferenced
  --> src/main.rs:34:25
   |
34 |     assert_eq!(unsafe { *header_ptr }, HEADER); // Miri error: dangling pointer dereference
   |                         ^^^^^^^^^^^ dangling pointer was dereferenced
   |

Playground link

Is this kind of support possible? Alternatively, perhaps if suppression support landed this may be an ideal use-case. [issue #788]

[RalfJung]

Yeah, right now your custom allocator will just not be called. (You can make it panic! if you want to verify that.)

Is this kind of support possible?

Possible, definitely. The __rust_alloc shim would have to see if there is a global_allocator attribute, and in that case dispatch to that instead of its own implementation. Same for the other allocator methods. See __rust_start_panic for how that dispatching-to-implementation can work.

I could mentor you on this if you want, except I am not sure how one would figure out if there is a global_allocator somewhere in the crate tree and what it is attached to.

[jacob-hughes]

Nice! I'll try and get a PR up for this then. One thing I'm unclear on though is the following:

The __rust_alloc shim would have to see if there is a global_allocator attribute, and in that case dispatch to that instead of its own implementation.

Surely, at some point though, we will need to dispatch to Miri's implementation? In the example I showed in this issue, my custom allocator dispatches the System allocator. I'm guessing this would break under Miri because System would need replacing with Miri's allocator implementation?

I could mentor you on this if you want, except I am not sure how one would figure out if there is a global_allocator somewhere in the crate tree and what it is attached to.

Appreciate this, but I think I'll be okay for now. I have some minor experience with rustc but will probably ping you if I get stuck :)

[RalfJung]

In the example I showed in this issue, my custom allocator dispatches the System allocator. I'm guessing this would break under Miri because System would need replacing with Miri's allocator implementation?

The System allocator is implemented in libstd (so for Miri that is just normal user code), and if your target is a Linux target, it will call malloc under the hood which is shimmed here.

Basically, Miri currently separately shims the Global and System allocators because they bottom out at different foreign function calls (__rust_* for the former, malloc/free/HeapAlloc/HeapFree for the latter). In a real Rust application, the Global calls are actually linked back to Rust code and most of the time are just tiny wrappers around the System allocator, but Miri has to specifically handle any call to a "foreign" item even if that item ends up being implemented in Rust.

So, __rust_alloc should always dispatch to the user-defined allocator if there is one.
(This means that with a user-defined allocator, Miri can no longer detect issues such as allocating with Global.alloc and then freeing with System.dealloc, but arguably that is permitted behavior if the user-defined allocator really uses the system allocator internally. But detecting such issues is why I think Miri should keep using the shim if the user did not explicitly pick an allocator. We should probably have a test ensuring this problem gets caught!)

[jacob-hughes]

Basically, Miri currently separately shims the Global and System allocators because they bottom out at different foreign function calls

Ah, I understand now.

But detecting such issues is why I think Miri should keep using the shim if the user did not explicitly pick an allocator

Yeah, agreed.

We should probably have a test ensuring this problem gets caught

No problem. I'll add one.

[RalfJung]

@jacob-hughes do you still plan to work on this? Let me know if you need any help. :)

[jacob-hughes]

Hi Ralf. Unfortunately the events over the last few weeks have made be a bit unproductive. I hope to have something in the next few days.

[RalfJung]

This has been implemented by #1885.