Auto merge of #113569 - RalfJung:miri, r=oli-obk

miri: protect Move() function arguments during the call

This gives `Move` operands a meaning specific to function calls:
- for the duration of the call, the place the operand comes from is protected, making all read and write accesses insta-UB.
- the contents of that place are reset to `Uninit`, so looking at them again after the function returns, we cannot observe their contents

Turns out we can replace the existing "retag return place" hack with the exact same sort of protection on the return place, which is nicely symmetric.

Fixes rust-lang/rust#112564
Fixes rust-lang#2927

This starts with a Miri rustc-push, since we'd otherwise conflict with a PR that recently landed in Miri.
(The "miri tree borrows" commit is an unrelated cleanup I noticed while doing the PR. I can remove it if you prefer.)
r? `@oli-obk`

rust-version

@@ -1 +1 @@
-d4096e0412ac5de785d739a0aa2b1c1c7b9d3b7d
+743333f3dd90721461c09387ec73d09c080d5f5f

src/borrow_tracker/mod.rs

@@ -302,12 +302,12 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
     }
 
-    fn retag_return_place(&mut self) -> InterpResult<'tcx> {
+    fn protect_place(&mut self, place: &MPlaceTy<'tcx, Provenance>) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
         let method = this.machine.borrow_tracker.as_ref().unwrap().borrow().borrow_tracker_method;
         match method {
-            BorrowTrackerMethod::StackedBorrows => this.sb_retag_return_place(),
-            BorrowTrackerMethod::TreeBorrows => this.tb_retag_return_place(),
+            BorrowTrackerMethod::StackedBorrows => this.sb_protect_place(place),
+            BorrowTrackerMethod::TreeBorrows => this.tb_protect_place(place),
         }
     }
 

src/borrow_tracker/stacked_borrows/diagnostics.rs

@@ -189,7 +189,7 @@ struct RetagOp {
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub enum RetagCause {
     Normal,
-    FnReturnPlace,
+    InPlaceFnPassing,
     FnEntry,
     TwoPhase,
 }
@@ -501,7 +501,7 @@ impl RetagCause {
         match self {
             RetagCause::Normal => "retag",
             RetagCause::FnEntry => "function-entry retag",
-            RetagCause::FnReturnPlace => "return-place retag",
+            RetagCause::InPlaceFnPassing => "in-place function argument/return passing protection",
             RetagCause::TwoPhase => "two-phase retag",
         }
         .to_string()

src/borrow_tracker/stacked_borrows/mod.rs

@@ -994,35 +994,25 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
     }
 
-    /// After a stack frame got pushed, retag the return place so that we are sure
-    /// it does not alias with anything.
-    ///
-    /// This is a HACK because there is nothing in MIR that would make the retag
-    /// explicit. Also see <https://github.com/rust-lang/rust/issues/71117>.
-    fn sb_retag_return_place(&mut self) -> InterpResult<'tcx> {
+    /// Protect a place so that it cannot be used any more for the duration of the current function
+    /// call.
+    /// 
+    /// This is used to ensure soundness of in-place function argument/return passing.
+    fn sb_protect_place(&mut self, place: &MPlaceTy<'tcx, Provenance>) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
-        let return_place = &this.frame().return_place;
-        if return_place.layout.is_zst() {
-            // There may not be any memory here, nothing to do.
-            return Ok(());
-        }
-        // We need this to be in-memory to use tagged pointers.
-        let return_place = this.force_allocation(&return_place.clone())?;
 
-        // We have to turn the place into a pointer to use the existing code.
+        // We have to turn the place into a pointer to use the usual retagging logic.
         // (The pointer type does not matter, so we use a raw pointer.)
-        let ptr_layout = this.layout_of(Ty::new_mut_ptr(this.tcx.tcx, return_place.layout.ty))?;
-        let val = ImmTy::from_immediate(return_place.to_ref(this), ptr_layout);
-        // Reborrow it. With protection! That is part of the point.
+        let ptr_layout = this.layout_of(Ty::new_mut_ptr(this.tcx.tcx, place.layout.ty))?;
+        let ptr = ImmTy::from_immediate(place.to_ref(this), ptr_layout);
+        // Reborrow it. With protection! That is the entire point.
         let new_perm = NewPermission::Uniform {
             perm: Permission::Unique,
             access: Some(AccessKind::Write),
             protector: Some(ProtectorKind::StrongProtector),
         };
-        let val = this.sb_retag_reference(&val, new_perm, RetagCause::FnReturnPlace)?;
-        // And use reborrowed pointer for return place.
-        let return_place = this.ref_to_mplace(&val)?;
-        this.frame_mut().return_place = return_place.into();
+        let _new_ptr = this.sb_retag_reference(&ptr, new_perm, RetagCause::InPlaceFnPassing)?;
+        // We just throw away `new_ptr`, so nobody can access this memory while it is protected.
 
         Ok(())
     }

src/borrow_tracker/tree_borrows/mod.rs

@@ -178,7 +178,7 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
         &mut self,
         place: &MPlaceTy<'tcx, Provenance>, // parent tag extracted from here
         ptr_size: Size,
-        new_perm: Option<NewPermission>,
+        new_perm: NewPermission,
         new_tag: BorTag,
     ) -> InterpResult<'tcx, Option<(AllocId, BorTag)>> {
         let this = self.eval_context_mut();
@@ -256,10 +256,6 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
             ptr_size.bytes()
         );
 
-        let Some(new_perm) = new_perm else {
-            return Ok(Some((alloc_id, orig_tag)));
-        };
-
         if let Some(protect) = new_perm.protector {
             // We register the protection in two different places.
             // This makes creating a protector slower, but checking whether a tag
@@ -305,7 +301,7 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
     fn tb_retag_reference(
         &mut self,
         val: &ImmTy<'tcx, Provenance>,
-        new_perm: Option<NewPermission>,
+        new_perm: NewPermission,
     ) -> InterpResult<'tcx, ImmTy<'tcx, Provenance>> {
         let this = self.eval_context_mut();
         // We want a place for where the ptr *points to*, so we get one.
@@ -317,7 +313,7 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
         // - if the pointer is not reborrowed (raw pointer) or if `zero_size` is set
         // then we override the size to do a zero-length reborrow.
         let reborrow_size = match new_perm {
-            Some(NewPermission { zero_size: false, .. }) =>
+            NewPermission { zero_size: false, .. } =>
                 this.size_and_align_of_mplace(&place)?
                     .map(|(size, _)| size)
                     .unwrap_or(place.layout.size),
@@ -374,7 +370,11 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 NewPermission::from_ref_ty(pointee, mutability, kind, this),
             _ => None,
         };
-        this.tb_retag_reference(val, new_perm)
+        if let Some(new_perm) = new_perm {
+            this.tb_retag_reference(val, new_perm)
+        } else {
+            Ok(val.clone())
+        }
     }
 
     /// Retag all pointers that are stored in this place.
@@ -405,9 +405,11 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 place: &PlaceTy<'tcx, Provenance>,
                 new_perm: Option<NewPermission>,
             ) -> InterpResult<'tcx> {
-                let val = self.ecx.read_immediate(&self.ecx.place_to_op(place)?)?;
-                let val = self.ecx.tb_retag_reference(&val, new_perm)?;
-                self.ecx.write_immediate(*val, place)?;
+                if let Some(new_perm) = new_perm {
+                    let val = self.ecx.read_immediate(&self.ecx.place_to_op(place)?)?;
+                    let val = self.ecx.tb_retag_reference(&val, new_perm)?;
+                    self.ecx.write_immediate(*val, place)?;
+                }
                 Ok(())
             }
         }
@@ -493,37 +495,25 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
     }
 
-    /// After a stack frame got pushed, retag the return place so that we are sure
-    /// it does not alias with anything.
-    ///
-    /// This is a HACK because there is nothing in MIR that would make the retag
-    /// explicit. Also see <https://github.com/rust-lang/rust/issues/71117>.
-    fn tb_retag_return_place(&mut self) -> InterpResult<'tcx> {
+    /// Protect a place so that it cannot be used any more for the duration of the current function
+    /// call.
+    /// 
+    /// This is used to ensure soundness of in-place function argument/return passing.
+    fn tb_protect_place(&mut self, place: &MPlaceTy<'tcx, Provenance>) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
-        //this.debug_hint_location();
-        let return_place = &this.frame().return_place;
-        if return_place.layout.is_zst() {
-            // There may not be any memory here, nothing to do.
-            return Ok(());
-        }
-        // We need this to be in-memory to use tagged pointers.
-        let return_place = this.force_allocation(&return_place.clone())?;
 
-        // We have to turn the place into a pointer to use the existing code.
+        // We have to turn the place into a pointer to use the usual retagging logic.
         // (The pointer type does not matter, so we use a raw pointer.)
-        let ptr_layout = this.layout_of(Ty::new_mut_ptr(this.tcx.tcx, return_place.layout.ty))?;
-        let val = ImmTy::from_immediate(return_place.to_ref(this), ptr_layout);
-        // Reborrow it. With protection! That is part of the point.
-        // FIXME: do we truly want a 2phase borrow here?
-        let new_perm = Some(NewPermission {
-            initial_state: Permission::new_unique_2phase(/*freeze*/ false),
+        let ptr_layout = this.layout_of(Ty::new_mut_ptr(this.tcx.tcx, place.layout.ty))?;
+        let ptr = ImmTy::from_immediate(place.to_ref(this), ptr_layout);
+        // Reborrow it. With protection! That is the entire point.
+        let new_perm = NewPermission {
+            initial_state: Permission::new_active(),
             zero_size: false,
             protector: Some(ProtectorKind::StrongProtector),
-        });
-        let val = this.tb_retag_reference(&val, new_perm)?;
-        // And use reborrowed pointer for return place.
-        let return_place = this.ref_to_mplace(&val)?;
-        this.frame_mut().return_place = return_place.into();
+        };
+        let _new_ptr = this.tb_retag_reference(&ptr, new_perm)?;
+        // We just throw away `new_ptr`, so nobody can access this memory while it is protected.
 
         Ok(())
     }

src/borrow_tracker/tree_borrows/perms.rs

@@ -138,6 +138,11 @@ impl Permission {
         Self { inner: Reserved { ty_is_freeze } }
     }
 
+    /// Default initial permission for return place.
+    pub fn new_active() -> Self {
+        Self { inner: Active }
+    }
+
     /// Default initial permission of a reborrowed shared reference
     pub fn new_frozen() -> Self {
         Self { inner: Frozen }

src/lib.rs

@@ -55,6 +55,7 @@ extern crate rustc_index;
 extern crate rustc_session;
 extern crate rustc_span;
 extern crate rustc_target;
+extern crate either; // the one from rustc
 
 // Necessary to pull in object code as the rest of the rustc crates are shipped only as rmeta
 // files.

src/machine.rs

@@ -7,6 +7,7 @@ use std::fmt;
 use std::path::Path;
 use std::process;
 
+use either::Either;
 use rand::rngs::StdRng;
 use rand::SeedableRng;
 
@@ -533,15 +534,16 @@ impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
             let target = &tcx.sess.target;
             match target.arch.as_ref() {
                 "wasm32" | "wasm64" => 64 * 1024, // https://webassembly.github.io/spec/core/exec/runtime.html#memory-instances
-                "aarch64" =>
+                "aarch64" => {
                     if target.options.vendor.as_ref() == "apple" {
                         // No "definitive" source, but see:
                         // https://www.wwdcnotes.com/notes/wwdc20/10214/
                         // https://github.com/ziglang/zig/issues/11308 etc.
                         16 * 1024
                     } else {
                         4 * 1024
-                    },
+                    }
+                }
                 _ => 4 * 1024,
             }
         };
@@ -892,7 +894,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         ecx: &mut MiriInterpCx<'mir, 'tcx>,
         instance: ty::Instance<'tcx>,
         abi: Abi,
-        args: &[OpTy<'tcx, Provenance>],
+        args: &[FnArg<'tcx, Provenance>],
         dest: &PlaceTy<'tcx, Provenance>,
         ret: Option<mir::BasicBlock>,
         unwind: mir::UnwindAction,
@@ -905,12 +907,13 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         ecx: &mut MiriInterpCx<'mir, 'tcx>,
         fn_val: Dlsym,
         abi: Abi,
-        args: &[OpTy<'tcx, Provenance>],
+        args: &[FnArg<'tcx, Provenance>],
         dest: &PlaceTy<'tcx, Provenance>,
         ret: Option<mir::BasicBlock>,
         _unwind: mir::UnwindAction,
     ) -> InterpResult<'tcx> {
-        ecx.call_dlsym(fn_val, abi, args, dest, ret)
+        let args = ecx.copy_fn_args(args)?; // FIXME: Should `InPlace` arguments be reset to uninit?
+        ecx.call_dlsym(fn_val, abi, &args, dest, ret)
     }
 
     #[inline(always)]
@@ -1094,8 +1097,9 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         ptr: Pointer<Self::Provenance>,
     ) -> InterpResult<'tcx> {
         match ptr.provenance {
-            Provenance::Concrete { alloc_id, tag } =>
-                intptrcast::GlobalStateInner::expose_ptr(ecx, alloc_id, tag),
+            Provenance::Concrete { alloc_id, tag } => {
+                intptrcast::GlobalStateInner::expose_ptr(ecx, alloc_id, tag)
+            }
             Provenance::Wildcard => {
                 // No need to do anything for wildcard pointers as
                 // their provenances have already been previously exposed.
@@ -1206,6 +1210,25 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         Ok(())
     }
 
+    fn protect_in_place_function_argument(
+        ecx: &mut InterpCx<'mir, 'tcx, Self>,
+        place: &PlaceTy<'tcx, Provenance>,
+    ) -> InterpResult<'tcx> {
+        // We do need to write `uninit` so that even after the call ends, the former contents of
+        // this place cannot be observed any more.
+        ecx.write_uninit(place)?;
+        // If we have a borrow tracker, we also have it set up protection so that all reads *and
+        // writes* during this call are insta-UB.
+        if ecx.machine.borrow_tracker.is_some() {
+            if let Either::Left(place) = place.as_mplace_or_local() {
+                ecx.protect_place(&place)?;
+            } else {
+                // Locals that don't have their address taken are as protected as they can ever be.
+            }
+        }
+        Ok(())
+    }
+
     #[inline(always)]
     fn init_frame_extra(
         ecx: &mut InterpCx<'mir, 'tcx, Self>,
@@ -1288,8 +1311,17 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
             let stack_len = ecx.active_thread_stack().len();
             ecx.active_thread_mut().set_top_user_relevant_frame(stack_len - 1);
         }
-        if ecx.machine.borrow_tracker.is_some() {
-            ecx.retag_return_place()?;
+        Ok(())
+    }
+
+    fn before_stack_pop(
+        ecx: &InterpCx<'mir, 'tcx, Self>,
+        frame: &Frame<'mir, 'tcx, Self::Provenance, Self::FrameExtra>,
+    ) -> InterpResult<'tcx> {
+        // We want this *before* the return value copy, because the return place itself is protected
+        // until we do `end_call` here.
+        if let Some(borrow_tracker) = &ecx.machine.borrow_tracker {
+            borrow_tracker.borrow_mut().end_call(&frame.extra);
         }
         Ok(())
     }
@@ -1308,9 +1340,6 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
             ecx.active_thread_mut().recompute_top_user_relevant_frame();
         }
         let timing = frame.extra.timing.take();
-        if let Some(borrow_tracker) = &ecx.machine.borrow_tracker {
-            borrow_tracker.borrow_mut().end_call(&frame.extra);
-        }
         let res = ecx.handle_stack_pop_unwind(frame.extra, unwinding);
         if let Some(profiler) = ecx.machine.profiler.as_ref() {
             profiler.finish_recording_interval_event(timing.unwrap());

src/shims/mod.rs

@@ -31,7 +31,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         &mut self,
         instance: ty::Instance<'tcx>,
         abi: Abi,
-        args: &[OpTy<'tcx, Provenance>],
+        args: &[FnArg<'tcx, Provenance>],
         dest: &PlaceTy<'tcx, Provenance>,
         ret: Option<mir::BasicBlock>,
         unwind: mir::UnwindAction,
@@ -41,7 +41,8 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
 
         // There are some more lang items we want to hook that CTFE does not hook (yet).
         if this.tcx.lang_items().align_offset_fn() == Some(instance.def.def_id()) {
-            let [ptr, align] = check_arg_count(args)?;
+            let args = this.copy_fn_args(args)?;
+            let [ptr, align] = check_arg_count(&args)?;
             if this.align_offset(ptr, align, dest, ret, unwind)? {
                 return Ok(None);
             }
@@ -55,7 +56,8 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             // to run extra MIR), and Ok(Some(body)) if we found MIR to run for the
             // foreign function
             // Any needed call to `goto_block` will be performed by `emulate_foreign_item`.
-            return this.emulate_foreign_item(instance.def_id(), abi, args, dest, ret, unwind);
+            let args = this.copy_fn_args(args)?; // FIXME: Should `InPlace` arguments be reset to uninit?
+            return this.emulate_foreign_item(instance.def_id(), abi, &args, dest, ret, unwind);
         }
 
         // Otherwise, load the MIR.