safe transmute: forbid reference lifetime extension

Modifies `BikeshedIntrinsicFrom` to forbid lifetime extensions on
references. This static check can be opted out of with the
`Assume::lifetimes` flag.

Fixes #129097

Author: jswrenn

compiler/rustc_trait_selection/src/traits/select/confirmation.rs

@@ -17,8 +17,7 @@ use rustc_infer::infer::{DefineOpaqueTypes, HigherRankedType, InferOk};
 use rustc_infer::traits::ObligationCauseCode;
 use rustc_middle::traits::{BuiltinImplSource, SignatureMismatchData};
 use rustc_middle::ty::{
-    self, GenericArgs, GenericArgsRef, GenericParamDefKind, ToPolyTraitRef, TraitPredicate, Ty,
-    TyCtxt, Upcast,
+    self, GenericArgs, GenericArgsRef, GenericParamDefKind, ToPolyTraitRef, Ty, TyCtxt, Upcast,
 };
 use rustc_middle::{bug, span_bug};
 use rustc_span::def_id::DefId;
@@ -292,90 +291,120 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
         &mut self,
         obligation: &PolyTraitObligation<'tcx>,
     ) -> Result<Vec<PredicateObligation<'tcx>>, SelectionError<'tcx>> {
-        use rustc_transmute::{Answer, Condition};
-        #[instrument(level = "debug", skip(tcx, obligation, predicate))]
+        use rustc_transmute::{Answer, Assume, Condition};
+
+        /// Generate sub-obligations for reference-to-reference transmutations.
+        fn reference_obligations<'tcx>(
+            tcx: TyCtxt<'tcx>,
+            obligation: &PolyTraitObligation<'tcx>,
+            (src_lifetime, src_ty, src_mut): (ty::Region<'tcx>, Ty<'tcx>, Mutability),
+            (dst_lifetime, dst_ty, dst_mut): (ty::Region<'tcx>, Ty<'tcx>, Mutability),
+            assume: Assume,
+        ) -> Vec<PredicateObligation<'tcx>> {
+            let make_transmute_obl = |src, dst| {
+                let transmute_trait = obligation.predicate.def_id();
+                let assume = obligation.predicate.skip_binder().trait_ref.args.const_at(2);
+                let trait_ref = ty::TraitRef::new(
+                    tcx,
+                    transmute_trait,
+                    [
+                        ty::GenericArg::from(dst),
+                        ty::GenericArg::from(src),
+                        ty::GenericArg::from(assume),
+                    ],
+                );
+                Obligation::with_depth(
+                    tcx,
+                    obligation.cause.clone(),
+                    obligation.recursion_depth + 1,
+                    obligation.param_env,
+                    obligation.predicate.rebind(trait_ref),
+                )
+            };
+
+            let make_freeze_obl = |ty| {
+                let trait_ref = ty::TraitRef::new(
+                    tcx,
+                    tcx.require_lang_item(LangItem::Freeze, None),
+                    [ty::GenericArg::from(ty)],
+                );
+                Obligation::with_depth(
+                    tcx,
+                    obligation.cause.clone(),
+                    obligation.recursion_depth + 1,
+                    obligation.param_env,
+                    trait_ref,
+                )
+            };
+
+            let make_outlives_obl = |target, region| {
+                let outlives = ty::OutlivesPredicate(target, region);
+                Obligation::with_depth(
+                    tcx,
+                    obligation.cause.clone(),
+                    obligation.recursion_depth + 1,
+                    obligation.param_env,
+                    obligation.predicate.rebind(outlives),
+                )
+            };
+
+            // Given a transmutation from `&'a (mut) Src` and `&'dst (mut) Dst`,
+            // it is always the case that `Src` must be transmutable into `Dst`,
+            // and that that `'src` must outlive `'dst`.
+            let mut obls = vec![make_transmute_obl(src_ty, dst_ty)];
+            if !assume.lifetimes {
+                obls.push(make_outlives_obl(src_lifetime, dst_lifetime));
+            }
+
+            // Given a transmutation from `&Src`, both `Src` and `Dst` must be
+            // `Freeze`, otherwise, using the transmuted value could lead to
+            // data races.
+            if src_mut == Mutability::Not {
+                obls.extend([make_freeze_obl(src_ty), make_freeze_obl(dst_ty)])
+            }
+
+            // Given a transmutation into `&'dst mut Dst`, it also must be the
+            // case that `Dst` is transmutable into `Src`. For example,
+            // transmuting bool -> u8 is OK as long as you can't update that u8
+            // to be > 1, because you could later transmute the u8 back to a
+            // bool and get undefined behavior. It also must be the case that
+            // `'dst` lives exactly as long as `'src`.
+            if dst_mut == Mutability::Mut {
+                obls.push(make_transmute_obl(dst_ty, src_ty));
+                if !assume.lifetimes {
+                    obls.push(make_outlives_obl(dst_lifetime, src_lifetime));
+                }
+            }
+
+            obls
+        }
+
+        /// Flatten the `Condition` tree into a conjunction of obligations.
+        #[instrument(level = "debug", skip(tcx, obligation))]
         fn flatten_answer_tree<'tcx>(
             tcx: TyCtxt<'tcx>,
             obligation: &PolyTraitObligation<'tcx>,
-            predicate: TraitPredicate<'tcx>,
             cond: Condition<rustc_transmute::layout::rustc::Ref<'tcx>>,
+            assume: Assume,
         ) -> Vec<PredicateObligation<'tcx>> {
             match cond {
                 // FIXME(bryangarza): Add separate `IfAny` case, instead of treating as `IfAll`
                 // Not possible until the trait solver supports disjunctions of obligations
                 Condition::IfAll(conds) | Condition::IfAny(conds) => conds
                     .into_iter()
-                    .flat_map(|cond| flatten_answer_tree(tcx, obligation, predicate, cond))
+                    .flat_map(|cond| flatten_answer_tree(tcx, obligation, cond, assume))
                     .collect(),
-                Condition::IfTransmutable { src, dst } => {
-                    let transmute_trait = obligation.predicate.def_id();
-                    let assume_const = predicate.trait_ref.args.const_at(2);
-                    let make_transmute_obl = |from_ty, to_ty| {
-                        let trait_ref = ty::TraitRef::new(
-                            tcx,
-                            transmute_trait,
-                            [
-                                ty::GenericArg::from(to_ty),
-                                ty::GenericArg::from(from_ty),
-                                ty::GenericArg::from(assume_const),
-                            ],
-                        );
-                        Obligation::with_depth(
-                            tcx,
-                            obligation.cause.clone(),
-                            obligation.recursion_depth + 1,
-                            obligation.param_env,
-                            trait_ref,
-                        )
-                    };
-
-                    let make_freeze_obl = |ty| {
-                        let trait_ref = ty::TraitRef::new(
-                            tcx,
-                            tcx.require_lang_item(LangItem::Freeze, None),
-                            [ty::GenericArg::from(ty)],
-                        );
-                        Obligation::with_depth(
-                            tcx,
-                            obligation.cause.clone(),
-                            obligation.recursion_depth + 1,
-                            obligation.param_env,
-                            trait_ref,
-                        )
-                    };
-
-                    let mut obls = vec![];
-
-                    // If the source is a shared reference, it must be `Freeze`;
-                    // otherwise, transmuting could lead to data races.
-                    if src.mutability == Mutability::Not {
-                        obls.extend([make_freeze_obl(src.ty), make_freeze_obl(dst.ty)])
-                    }
-
-                    // If Dst is mutable, check bidirectionally.
-                    // For example, transmuting bool -> u8 is OK as long as you can't update that u8
-                    // to be > 1, because you could later transmute the u8 back to a bool and get UB.
-                    match dst.mutability {
-                        Mutability::Not => obls.push(make_transmute_obl(src.ty, dst.ty)),
-                        Mutability::Mut => obls.extend([
-                            make_transmute_obl(src.ty, dst.ty),
-                            make_transmute_obl(dst.ty, src.ty),
-                        ]),
-                    }
-
-                    obls
-                }
+                Condition::IfTransmutable { src, dst } => reference_obligations(
+                    tcx,
+                    obligation,
+                    (src.lifetime, src.ty, src.mutability),
+                    (dst.lifetime, dst.ty, dst.mutability),
+                    assume,
+                ),
             }
         }
 
-        // We erase regions here because transmutability calls layout queries,
-        // which does not handle inference regions and doesn't particularly
-        // care about other regions. Erasing late-bound regions is equivalent
-        // to instantiating the binder with placeholders then erasing those
-        // placeholder regions.
-        let predicate = self
-            .tcx()
-            .erase_regions(self.tcx().instantiate_bound_regions_with_erased(obligation.predicate));
+        let predicate = obligation.predicate.skip_binder();
 
         let Some(assume) = rustc_transmute::Assume::from_const(
             self.infcx.tcx,
@@ -387,6 +416,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
 
         let dst = predicate.trait_ref.args.type_at(0);
         let src = predicate.trait_ref.args.type_at(1);
+
         debug!(?src, ?dst);
         let mut transmute_env = rustc_transmute::TransmuteTypeEnv::new(self.infcx);
         let maybe_transmutable = transmute_env.is_transmutable(
@@ -397,7 +427,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
 
         let fully_flattened = match maybe_transmutable {
             Answer::No(_) => Err(Unimplemented)?,
-            Answer::If(cond) => flatten_answer_tree(self.tcx(), obligation, predicate, cond),
+            Answer::If(cond) => flatten_answer_tree(self.tcx(), obligation, cond, assume),
             Answer::Yes => vec![],
         };
 

compiler/rustc_transmute/src/layout/mod.rs

@@ -63,7 +63,9 @@ pub mod rustc {
     use std::fmt::{self, Write};
 
     use rustc_middle::mir::Mutability;
-    use rustc_middle::ty::{self, Ty};
+    use rustc_middle::ty::layout::{LayoutCx, LayoutError};
+    use rustc_middle::ty::{self, Ty, TyCtxt};
+    use rustc_target::abi::Layout;
 
     /// A reference in the layout.
     #[derive(Debug, Hash, Eq, PartialEq, Clone, Copy)]
@@ -120,4 +122,12 @@ pub mod rustc {
             self != &Self::Primitive
         }
     }
+
+    pub(crate) fn layout_of<'tcx>(
+        cx: LayoutCx<'tcx, TyCtxt<'tcx>>,
+        ty: Ty<'tcx>,
+    ) -> Result<Layout<'tcx>, &'tcx LayoutError<'tcx>> {
+        use rustc_middle::ty::layout::LayoutOf;
+        cx.layout_of(ty).map(|tl| tl.layout)
+    }
 }