interpret/allocation: expose init + write_wildcards on a range

Author: nia-e

compiler/rustc_const_eval/src/interpret/memory.rs

@@ -1472,7 +1472,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             dest_alloc
                 .write_uninit(&tcx, dest_range)
                 .map_err(|e| e.to_interp_error(dest_alloc_id))?;
-            // We can forget about the provenance, this is all not initialized anyway.
+            // `write_uninit` also resets the provenance, so we are done.
             return interp_ok(());
         }
 

compiler/rustc_middle/src/mir/interpret/allocation.rs

@@ -101,6 +101,8 @@ pub struct Allocation<Prov: Provenance = CtfeProvenance, Extra = (), Bytes = Box
     /// at the given offset.
     provenance: ProvenanceMap<Prov>,
     /// Denotes which part of this allocation is initialized.
+    ///
+    /// Invariant: the uninitialized parts are 0 (in `bytes`) and have no provenance.
     init_mask: InitMask,
     /// The alignment of the allocation to detect unaligned reads.
     /// (`Align` guarantees that this is a power of two.)
@@ -792,28 +794,25 @@ impl<Prov: Provenance, Extra, Bytes: AllocBytes> Allocation<Prov, Extra, Bytes>
     /// Write "uninit" to the given memory range.
     pub fn write_uninit(&mut self, cx: &impl HasDataLayout, range: AllocRange) -> AllocResult {
         self.mark_init(range, false);
+        // Restore the invariant that the uninit parts are 0 and have no provenance.
+        self.bytes[range.start.bytes_usize()..range.end().bytes_usize()].fill(0);
         self.provenance.clear(range, cx)?;
         Ok(())
     }
 
-    /// Initialize all previously uninitialized bytes in the entire allocation, and set
-    /// provenance of everything to `Wildcard`. Before calling this, make sure all
-    /// provenance in this allocation is exposed!
-    pub fn prepare_for_native_access(&mut self) {
-        let full_range = AllocRange { start: Size::ZERO, size: Size::from_bytes(self.len()) };
-        // Overwrite uninitialized bytes with 0, to ensure we don't leak whatever their value happens to be.
-        for chunk in self.init_mask.range_as_init_chunks(full_range) {
-            if !chunk.is_init() {
-                let uninit_bytes = &mut self.bytes
-                    [chunk.range().start.bytes_usize()..chunk.range().end.bytes_usize()];
-                uninit_bytes.fill(0);
-            }
-        }
-        // Mark everything as initialized now.
-        self.mark_init(full_range, true);
-
-        // Set provenance of all bytes to wildcard.
-        self.provenance.write_wildcards(self.len());
+    /// Mark all bytes in the given range as initialised and reset the provenance
+    /// to wildcards. This entirely breaks the normal mechanisms for tracking
+    /// initialisation and is only provided for Miri operating in native-lib
+    /// mode. UB will be missed if the underlying bytes were not actually written to.
+    ///
+    /// If `range` is `None`, defaults to performing this on the whole allocation.
+    pub fn process_native_write(&mut self, cx: &impl HasDataLayout, range: Option<AllocRange>) {
+        let range = range.unwrap_or_else(|| AllocRange {
+            start: Size::ZERO,
+            size: Size::from_bytes(self.len()),
+        });
+        self.mark_init(range, true);
+        self.provenance.write_wildcards(cx, range);
     }
 
     /// Remove all provenance in the given memory range.

compiler/rustc_middle/src/mir/interpret/allocation/provenance_map.rs

@@ -212,21 +212,37 @@ impl<Prov: Provenance> ProvenanceMap<Prov> {
         Ok(())
     }
 
-    /// Overwrites all provenance in the allocation with wildcard provenance.
+    /// Overwrites all provenance in the given range with wildcard provenance.
+    /// Pointers partially overwritten will have their provenances preserved
+    /// bytewise on their remaining bytes.
     ///
     /// Provided for usage in Miri and panics otherwise.
-    pub fn write_wildcards(&mut self, alloc_size: usize) {
+    pub fn write_wildcards(&mut self, cx: &impl HasDataLayout, range: AllocRange) {
         assert!(
             Prov::OFFSET_IS_ADDR,
             "writing wildcard provenance is not supported when `OFFSET_IS_ADDR` is false"
         );
         let wildcard = Prov::WILDCARD.unwrap();
 
-        // Remove all pointer provenances, then write wildcards into the whole byte range.
-        self.ptrs.clear();
-        let last = Size::from_bytes(alloc_size);
         let bytes = self.bytes.get_or_insert_with(Box::default);
-        for offset in Size::ZERO..last {
+
+        // Remove pointer provenances that overlap with the range, then readd the edge ones bytewise.
+        let ptr_range = Self::adjusted_range_ptrs(range, cx);
+        let ptrs = self.ptrs.range(ptr_range.clone());
+        if let Some((offset, prov)) = ptrs.first() {
+            for byte_ofs in *offset..range.start {
+                bytes.insert(byte_ofs, *prov);
+            }
+        }
+        if let Some((offset, prov)) = ptrs.last() {
+            for byte_ofs in range.end()..*offset + cx.data_layout().pointer_size() {
+                bytes.insert(byte_ofs, *prov);
+            }
+        }
+        self.ptrs.remove_range(ptr_range);
+
+        // Overwrite bytewise provenance.
+        for offset in range.start..range.end() {
             bytes.insert(offset, wildcard);
         }
     }

src/tools/miri/src/shims/native_lib/mod.rs

@@ -251,8 +251,8 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
 
             // Prepare for possible write from native code if mutable.
             if info.mutbl.is_mut() {
-                let alloc = &mut this.get_alloc_raw_mut(alloc_id)?.0;
-                alloc.prepare_for_native_access();
+                let (alloc, cx) = this.get_alloc_raw_mut(alloc_id)?;
+                alloc.process_native_write(&cx.tcx, None);
                 // Also expose *mutable* provenance for the interpreter-level allocation.
                 std::hint::black_box(alloc.get_bytes_unchecked_raw_mut().expose_provenance());
             }