implement support for __rust_alloc_error_handler

10 files changed, 207 insertions, 7 deletions

diff --git a/src/tools/miri/src/shims/extern_static.rs b/src/tools/miri/src/shims/extern_static.rs
index 0284e5b606c..7c4a54fb461 100644
--- a/src/tools/miri/src/shims/extern_static.rs
+++ b/src/tools/miri/src/shims/extern_static.rs
@@ -32,9 +32,14 @@ impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
     /// Sets up the "extern statics" for this machine.
     pub fn init_extern_statics(this: &mut MiriInterpCx<'mir, 'tcx>) -> InterpResult<'tcx> {
         // "__rust_no_alloc_shim_is_unstable"
-        let val = ImmTy::from_int(0, this.machine.layouts.u8);
+        let val = ImmTy::from_int(0, this.machine.layouts.u8); // always 0, value does not matter
         Self::alloc_extern_static(this, "__rust_no_alloc_shim_is_unstable", val)?;
 
+        // "__rust_alloc_error_handler_should_panic"
+        let val = this.tcx.sess.opts.unstable_opts.oom.should_panic();
+        let val = ImmTy::from_int(val, this.machine.layouts.u8);
+        Self::alloc_extern_static(this, "__rust_alloc_error_handler_should_panic", val)?;
+
         match this.tcx.sess.target.os.as_ref() {
             "linux" => {
                 Self::null_ptr_extern_statics(
diff --git a/src/tools/miri/src/shims/foreign_items.rs b/src/tools/miri/src/shims/foreign_items.rs
index 6b0797f6da1..e6fc29a5ae0 100644
--- a/src/tools/miri/src/shims/foreign_items.rs
+++ b/src/tools/miri/src/shims/foreign_items.rs
@@ -1,7 +1,7 @@
 use std::{collections::hash_map::Entry, io::Write, iter, path::Path};
 
 use rustc_apfloat::Float;
-use rustc_ast::expand::allocator::AllocatorKind;
+use rustc_ast::expand::allocator::{alloc_error_handler_name, AllocatorKind};
 use rustc_hir::{def::DefKind, def_id::CrateNum};
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir;
@@ -80,6 +80,20 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                             panic_impl_instance,
                         )));
                     }
+                    "__rust_alloc_error_handler" => {
+                        // Forward to the right symbol that implements this function.
+                        let Some(handler_kind) = this.tcx.alloc_error_handler_kind(()) else {
+                            // in real code, this symbol does not exist without an allocator
+                            throw_unsup_format!(
+                                "`__rust_alloc_error_handler` cannot be called when no alloc error handler is set"
+                            );
+                        };
+                        let name = alloc_error_handler_name(handler_kind);
+                        let handler = this
+                            .lookup_exported_symbol(Symbol::intern(name))?
+                            .expect("missing alloc error handler symbol");
+                        return Ok(Some(handler));
+                    }
                     #[rustfmt::skip]
                     | "exit"
                     | "ExitProcess"
diff --git a/src/tools/miri/tests/fail/alloc/alloc_error_handler.rs b/src/tools/miri/tests/fail/alloc/alloc_error_handler.rs
new file mode 100644
index 00000000000..dc8e8c73800
--- /dev/null
+++ b/src/tools/miri/tests/fail/alloc/alloc_error_handler.rs
@@ -0,0 +1,25 @@
+//@error-in-other-file: aborted
+//@normalize-stderr-test: "unsafe \{ libc::abort\(\) \}|crate::intrinsics::abort\(\);" -> "ABORT();"
+//@normalize-stderr-test: "\| +\^+" -> "| ^"
+#![feature(allocator_api)]
+
+use std::alloc::*;
+use std::ptr::NonNull;
+
+struct BadAlloc;
+
+// Create a failing allocator; Miri's native allocator never fails so this is the only way to
+// actually call the alloc error handler.
+unsafe impl Allocator for BadAlloc {
+    fn allocate(&self, _l: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        Err(AllocError)
+    }
+
+    unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {
+        unreachable!();
+    }
+}
+
+fn main() {
+    let _b = Box::new_in(0, BadAlloc);
+}
diff --git a/src/tools/miri/tests/fail/alloc/alloc_error_handler.stderr b/src/tools/miri/tests/fail/alloc/alloc_error_handler.stderr
new file mode 100644
index 00000000000..f9d8e80c0fa
--- /dev/null
+++ b/src/tools/miri/tests/fail/alloc/alloc_error_handler.stderr
@@ -0,0 +1,26 @@
+memory allocation of 4 bytes failed
+error: abnormal termination: the program aborted execution
+  --> RUSTLIB/std/src/sys/pal/PLATFORM/mod.rs:LL:CC
+   |
+LL |     ABORT();
+   | ^ the program aborted execution
+   |
+   = note: BACKTRACE:
+   = note: inside `std::sys::pal::PLATFORM::abort_internal` at RUSTLIB/std/src/sys/pal/PLATFORM/mod.rs:LL:CC
+   = note: inside `std::process::abort` at RUSTLIB/std/src/process.rs:LL:CC
+   = note: inside `std::alloc::rust_oom` at RUSTLIB/std/src/alloc.rs:LL:CC
+   = note: inside `std::alloc::_::__rg_oom` at RUSTLIB/std/src/alloc.rs:LL:CC
+   = note: inside `std::alloc::handle_alloc_error::rt_error` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `std::alloc::handle_alloc_error` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `std::boxed::Box::<i32, BadAlloc>::new_uninit_in` at RUSTLIB/alloc/src/boxed.rs:LL:CC
+   = note: inside `std::boxed::Box::<i32, BadAlloc>::new_in` at RUSTLIB/alloc/src/boxed.rs:LL:CC
+note: inside `main`
+  --> $DIR/alloc_error_handler.rs:LL:CC
+   |
+LL |     let _b = Box::new_in(0, BadAlloc);
+   | ^
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to 1 previous error
+
diff --git a/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.rs b/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.rs
new file mode 100644
index 00000000000..8103296f47b
--- /dev/null
+++ b/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.rs
@@ -0,0 +1,65 @@
+//@compile-flags: -Cpanic=abort
+#![feature(start, core_intrinsics)]
+#![feature(alloc_error_handler)]
+#![feature(allocator_api)]
+#![no_std]
+
+extern crate alloc;
+
+use alloc::alloc::*;
+use alloc::boxed::Box;
+use core::ptr::NonNull;
+
+struct BadAlloc;
+
+// Create a failing allocator; that is the only way to actually call the alloc error handler.
+unsafe impl Allocator for BadAlloc {
+    fn allocate(&self, _l: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        Err(AllocError)
+    }
+
+    unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {
+        unreachable!();
+    }
+}
+
+#[alloc_error_handler]
+fn alloc_error_handler(_: Layout) -> ! {
+    extern "Rust" {
+        fn miri_write_to_stderr(bytes: &[u8]);
+    }
+    let msg = "custom alloc error handler called!\n";
+    unsafe { miri_write_to_stderr(msg.as_bytes()) };
+    core::intrinsics::abort(); //~ERROR: aborted
+}
+
+// rustc requires us to provide some more things that aren't actually used by this test
+mod plumbing {
+    use super::*;
+
+    #[panic_handler]
+    fn panic_handler(_: &core::panic::PanicInfo) -> ! {
+        core::intrinsics::abort();
+    }
+
+    struct NoAlloc;
+
+    unsafe impl GlobalAlloc for NoAlloc {
+        unsafe fn alloc(&self, _: Layout) -> *mut u8 {
+            unreachable!();
+        }
+
+        unsafe fn dealloc(&self, _: *mut u8, _: Layout) {
+            unreachable!();
+        }
+    }
+
+    #[global_allocator]
+    static GLOBAL: NoAlloc = NoAlloc;
+}
+
+#[start]
+fn start(_: isize, _: *const *const u8) -> isize {
+    let _b = Box::new_in(0, BadAlloc);
+    0
+}
diff --git a/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.stderr b/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.stderr
new file mode 100644
index 00000000000..b40ffb70121
--- /dev/null
+++ b/src/tools/miri/tests/fail/alloc/alloc_error_handler_no_std.stderr
@@ -0,0 +1,29 @@
+custom alloc error handler called!
+error: abnormal termination: the program aborted execution
+  --> $DIR/alloc_error_handler_no_std.rs:LL:CC
+   |
+LL |     core::intrinsics::abort();
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^ the program aborted execution
+   |
+   = note: BACKTRACE:
+   = note: inside `alloc_error_handler` at $DIR/alloc_error_handler_no_std.rs:LL:CC
+note: inside `_::__rg_oom`
+  --> $DIR/alloc_error_handler_no_std.rs:LL:CC
+   |
+LL | #[alloc_error_handler]
+   | ---------------------- in this procedural macro expansion
+LL | fn alloc_error_handler(_: Layout) -> ! {
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   = note: inside `alloc::alloc::handle_alloc_error::rt_error` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `alloc::alloc::handle_alloc_error` at RUSTLIB/alloc/src/alloc.rs:LL:CC
+   = note: inside `alloc::boxed::Box::<i32, BadAlloc>::new_uninit_in` at RUSTLIB/alloc/src/boxed.rs:LL:CC
+   = note: inside `alloc::boxed::Box::<i32, BadAlloc>::new_in` at RUSTLIB/alloc/src/boxed.rs:LL:CC
+note: inside `start`
+  --> $DIR/alloc_error_handler_no_std.rs:LL:CC
+   |
+LL |     let _b = Box::new_in(0, BadAlloc);
+   |              ^^^^^^^^^^^^^^^^^^^^^^^^
+   = note: this error originates in the attribute macro `alloc_error_handler` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error: aborting due to 1 previous error
+
diff --git a/src/tools/miri/tests/panic/alloc_error_handler_panic.rs b/src/tools/miri/tests/panic/alloc_error_handler_panic.rs
new file mode 100644
index 00000000000..186c9667a97
--- /dev/null
+++ b/src/tools/miri/tests/panic/alloc_error_handler_panic.rs
@@ -0,0 +1,32 @@
+//@compile-flags: -Zoom=panic
+#![feature(allocator_api)]
+
+use std::alloc::*;
+use std::ptr::NonNull;
+
+struct BadAlloc;
+
+// Create a failing allocator; Miri's native allocator never fails so this is the only way to
+// actually call the alloc error handler.
+unsafe impl Allocator for BadAlloc {
+    fn allocate(&self, _l: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        Err(AllocError)
+    }
+
+    unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {
+        unreachable!();
+    }
+}
+
+struct Bomb;
+impl Drop for Bomb {
+    fn drop(&mut self) {
+        eprintln!("yes we are unwinding!");
+    }
+}
+
+fn main() {
+    let bomb = Bomb;
+    let _b = Box::new_in(0, BadAlloc);
+    std::mem::forget(bomb); // defuse unwinding bomb
+}
diff --git a/src/tools/miri/tests/panic/alloc_error_handler_panic.stderr b/src/tools/miri/tests/panic/alloc_error_handler_panic.stderr
new file mode 100644
index 00000000000..202325468bf
--- /dev/null
+++ b/src/tools/miri/tests/panic/alloc_error_handler_panic.stderr
@@ -0,0 +1,4 @@
+thread 'main' panicked at RUSTLIB/std/src/alloc.rs:LL:CC:
+memory allocation of 4 bytes failed
+note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
+yes we are unwinding!
diff --git a/src/tools/miri/tests/pass/alloc-access-tracking.rs b/src/tools/miri/tests/pass/alloc-access-tracking.rs
index 5c782fca2df..29c1ee2f7b7 100644
--- a/src/tools/miri/tests/pass/alloc-access-tracking.rs
+++ b/src/tools/miri/tests/pass/alloc-access-tracking.rs
@@ -1,6 +1,6 @@
 #![feature(start)]
 #![no_std]
-//@compile-flags: -Zmiri-track-alloc-id=17 -Zmiri-track-alloc-accesses -Cpanic=abort
+//@compile-flags: -Zmiri-track-alloc-id=18 -Zmiri-track-alloc-accesses -Cpanic=abort
 //@only-target-linux: alloc IDs differ between OSes for some reason
 
 extern "Rust" {
diff --git a/src/tools/miri/tests/pass/alloc-access-tracking.stderr b/src/tools/miri/tests/pass/alloc-access-tracking.stderr
index 5e219fa1bed..bef13701ea2 100644
--- a/src/tools/miri/tests/pass/alloc-access-tracking.stderr
+++ b/src/tools/miri/tests/pass/alloc-access-tracking.stderr
@@ -2,7 +2,7 @@ note: tracking was triggered
   --> $DIR/alloc-access-tracking.rs:LL:CC
    |
 LL |         let ptr = miri_alloc(123, 1);
-   |                   ^^^^^^^^^^^^^^^^^^ created Miri bare-metal heap allocation of 123 bytes (alignment ALIGN bytes) with id 17
+   |                   ^^^^^^^^^^^^^^^^^^ created Miri bare-metal heap allocation of 123 bytes (alignment ALIGN bytes) with id 18
    |
    = note: BACKTRACE:
    = note: inside `start` at $DIR/alloc-access-tracking.rs:LL:CC
@@ -11,7 +11,7 @@ note: tracking was triggered
   --> $DIR/alloc-access-tracking.rs:LL:CC
    |
 LL |         *ptr = 42; // Crucially, only a write is printed here, no read!
-   |         ^^^^^^^^^ write access to allocation with id 17
+   |         ^^^^^^^^^ write access to allocation with id 18
    |
    = note: BACKTRACE:
    = note: inside `start` at $DIR/alloc-access-tracking.rs:LL:CC
@@ -20,7 +20,7 @@ note: tracking was triggered
   --> $DIR/alloc-access-tracking.rs:LL:CC
    |
 LL |         assert_eq!(*ptr, 42);
-   |         ^^^^^^^^^^^^^^^^^^^^ read access to allocation with id 17
+   |         ^^^^^^^^^^^^^^^^^^^^ read access to allocation with id 18
    |
    = note: BACKTRACE:
    = note: inside `start` at RUSTLIB/core/src/macros/mod.rs:LL:CC
@@ -30,7 +30,7 @@ note: tracking was triggered
   --> $DIR/alloc-access-tracking.rs:LL:CC
    |
 LL |         miri_dealloc(ptr, 123, 1);
-   |         ^^^^^^^^^^^^^^^^^^^^^^^^^ freed allocation with id 17
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^^ freed allocation with id 18
    |
    = note: BACKTRACE:
    = note: inside `start` at $DIR/alloc-access-tracking.rs:LL:CC