don't panic when allocator symbols called without allocator

Author: beepster4096

src/shims/foreign_items.rs

@@ -298,6 +298,33 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
         Ok(None)
     }
 
+    /// Emulates calling the internal __rust_* allocator functions
+    fn emulate_allocator(
+        &mut self,
+        symbol: Symbol,
+        default: impl FnOnce(&mut MiriEvalContext<'mir, 'tcx>) -> InterpResult<'tcx>,
+    ) -> InterpResult<'tcx, EmulateByNameResult<'mir, 'tcx>> {
+        let this = self.eval_context_mut();
+
+        let allocator_kind = if let Some(allocator_kind) = this.tcx.allocator_kind(()) {
+            allocator_kind
+        } else {
+            // in real code, this symbol does not exist without an allocator
+            return Ok(EmulateByNameResult::NotSupported);
+        };
+
+        match allocator_kind {
+            AllocatorKind::Global => {
+                let body = this
+                    .lookup_exported_symbol(symbol)?
+                    .expect("symbol should be present if there is a global allocator");
+
+                Ok(EmulateByNameResult::MirBody(body))
+            }
+            AllocatorKind::Default => default(this).map(|()| EmulateByNameResult::NeedsJumping),
+        }
+    }
+
     /// Emulates calling a foreign item using its name.
     fn emulate_foreign_item_by_name(
         &mut self,
@@ -372,77 +399,48 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
                 let align = this.read_scalar(align)?.to_machine_usize(this)?;
                 Self::check_alloc_request(size, align)?;
 
-                let allocator_kind = this.tcx.allocator_kind(());
-
-                match allocator_kind.expect("allocator should be present if used") {
-                    AllocatorKind::Global => {
-                        let body = this.lookup_exported_symbol(Symbol::intern("__rg_alloc"))?
-                            .expect("symbol should be present if there is a global allocator");
-
-                        return Ok(EmulateByNameResult::MirBody(body))
-                    },
-                    AllocatorKind::Default => {
-                        let ptr = this.memory.allocate(
-                            Size::from_bytes(size),
-                            Align::from_bytes(align).unwrap(),
-                            MiriMemoryKind::Rust.into(),
-                        )?;
-                        this.write_pointer(ptr, dest)?;
-                    }
-                };
+                return this.emulate_allocator(Symbol::intern("__rg_alloc"), |this| {
+                    let ptr = this.memory.allocate(
+                        Size::from_bytes(size),
+                        Align::from_bytes(align).unwrap(),
+                        MiriMemoryKind::Rust.into(),
+                    )?;
+
+                    this.write_pointer(ptr, dest)
+                });
             }
             "__rust_alloc_zeroed" => {
                 let &[ref size, ref align] = this.check_shim(abi, Abi::Rust, link_name, args)?;
                 let size = this.read_scalar(size)?.to_machine_usize(this)?;
                 let align = this.read_scalar(align)?.to_machine_usize(this)?;
                 Self::check_alloc_request(size, align)?;
 
-                let allocator_kind = this.tcx.allocator_kind(());
-
-                match allocator_kind.expect("allocator should be present if used") {
-                    AllocatorKind::Global => {
-                        let body = this.lookup_exported_symbol(Symbol::intern("__rg_alloc_zeroed"))?
-                            .expect("symbol should be present if there is a global allocator");
-
-                        return Ok(EmulateByNameResult::MirBody(body))
-                    },
-                    AllocatorKind::Default => {
-                        let ptr = this.memory.allocate(
-                            Size::from_bytes(size),
-                            Align::from_bytes(align).unwrap(),
-                            MiriMemoryKind::Rust.into(),
-                        )?;
-
-                        // We just allocated this, the access is definitely in-bounds.
-                        this.memory.write_bytes(ptr.into(), iter::repeat(0u8).take(usize::try_from(size).unwrap())).unwrap();
-                        this.write_pointer(ptr, dest)?;
-                    }
-                }
+                return this.emulate_allocator(Symbol::intern("__rg_alloc_zeroed"), |this| {
+                    let ptr = this.memory.allocate(
+                        Size::from_bytes(size),
+                        Align::from_bytes(align).unwrap(),
+                        MiriMemoryKind::Rust.into(),
+                    )?;
+
+                    // We just allocated this, the access is definitely in-bounds.
+                    this.memory.write_bytes(ptr.into(), iter::repeat(0u8).take(usize::try_from(size).unwrap())).unwrap();
+                    this.write_pointer(ptr, dest)
+                });
             }
             "__rust_dealloc" => {
                 let &[ref ptr, ref old_size, ref align] = this.check_shim(abi, Abi::Rust, link_name, args)?;
                 let ptr = this.read_pointer(ptr)?;
                 let old_size = this.read_scalar(old_size)?.to_machine_usize(this)?;
                 let align = this.read_scalar(align)?.to_machine_usize(this)?;
 
-                let allocator_kind = this.tcx.allocator_kind(());
-
-                match allocator_kind.expect("allocator should be present if used") {
-                    AllocatorKind::Global => {
-                        let body = this.lookup_exported_symbol(Symbol::intern("__rg_dealloc"))?
-                            .expect("symbol should be present if there is a global allocator");
-
-                        return Ok(EmulateByNameResult::MirBody(body))
-                    },
-                    AllocatorKind::Default => {
-                        // No need to check old_size/align; we anyway check that they match the allocation.
-                        this.memory.deallocate(
-                            ptr,
-                            Some((Size::from_bytes(old_size), Align::from_bytes(align).unwrap())),
-                            MiriMemoryKind::Rust.into(),
-                        )?;
-                    }
-                }
+                return this.emulate_allocator(Symbol::intern("__rg_dealloc"), |this| {
+                    // No need to check old_size/align; we anyway check that they match the allocation.
+                    this.memory.deallocate(
+                        ptr,
+                        Some((Size::from_bytes(old_size), Align::from_bytes(align).unwrap())),
+                        MiriMemoryKind::Rust.into(),
+                    )
+                });
             }
             "__rust_realloc" => {
                 let &[ref ptr, ref old_size, ref align, ref new_size] = this.check_shim(abi, Abi::Rust, link_name, args)?;
@@ -453,27 +451,17 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
                 Self::check_alloc_request(new_size, align)?;
                 // No need to check old_size; we anyway check that they match the allocation.
 
-                let allocator_kind = this.tcx.allocator_kind(());
-
-                match allocator_kind.expect("allocator should be present if used") {
-                    AllocatorKind::Global => {
-                        let body = this.lookup_exported_symbol(Symbol::intern("__rg_realloc"))?
-                            .expect("symbol should be present if there is a global allocator");
-
-                        return Ok(EmulateByNameResult::MirBody(body))
-                    },
-                    AllocatorKind::Default => {
-                        let align = Align::from_bytes(align).unwrap();
-                        let new_ptr = this.memory.reallocate(
-                            ptr,
-                            Some((Size::from_bytes(old_size), align)),
-                            Size::from_bytes(new_size),
-                            align,
-                            MiriMemoryKind::Rust.into(),
-                        )?;
-                        this.write_pointer(new_ptr, dest)?;
-                    }
-                }
+                return this.emulate_allocator(Symbol::intern("__rg_realloc"), |this| {
+                    let align = Align::from_bytes(align).unwrap();
+                    let new_ptr = this.memory.reallocate(
+                        ptr,
+                        Some((Size::from_bytes(old_size), align)),
+                        Size::from_bytes(new_size),
+                        align,
+                        MiriMemoryKind::Rust.into(),
+                    )?;
+                    this.write_pointer(new_ptr, dest)
+                });
             }
 
             // C memory handling functions