CFI: Introduce CFI shims

Indirect calls through vtables (trait objects or drop_in_place) expect
to have an alias set based on `dyn Trait` at the call-site. The actual
implementations have aslias sets based on `MyImplType`. These shims create a
separate `InstanceDef`, allowing a different type to be assigned. These
function for both CFI and KCFI, as they have a single principal type.

Author: maurer

compiler/rustc_const_eval/src/interpret/terminator.rs

@@ -558,6 +558,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             | ty::InstanceDef::CloneShim(..)
             | ty::InstanceDef::FnPtrAddrShim(..)
             | ty::InstanceDef::ThreadLocalShim(..)
+            | ty::InstanceDef::CfiShim { .. }
             | ty::InstanceDef::Item(_) => {
                 // We need MIR for this fn
                 let Some((body, instance)) = M::find_mir_or_eval_fn(

compiler/rustc_middle/src/mir/mono.rs

@@ -406,7 +406,8 @@ impl<'tcx> CodegenUnit<'tcx> {
                             | InstanceDef::DropGlue(..)
                             | InstanceDef::CloneShim(..)
                             | InstanceDef::ThreadLocalShim(..)
-                            | InstanceDef::FnPtrAddrShim(..) => None,
+                            | InstanceDef::FnPtrAddrShim(..)
+                            | InstanceDef::CfiShim { .. } => None,
                         }
                     }
                     MonoItem::Static(def_id) => def_id.as_local().map(Idx::index),

compiler/rustc_middle/src/mir/visit.rs

@@ -940,6 +940,21 @@ macro_rules! make_mir_visitor {
                         // FIXME(eddyb) use a better `TyContext` here.
                         self.visit_ty($(& $mutability *)? ty, TyContext::Location(location));
                     }
+                    ty::InstanceDef::CfiShim { target_instance, invoke_ty } => {
+                        self.visit_ty($(& $mutability *)? invoke_ty, TyContext::Location(location));
+                        let $($mutability)? local_target_instance = {
+                            $(
+                                let $mutability _unused = ();
+                                *
+                            )?
+                            *target_instance
+                        };
+                        self.visit_instance_def($(& $mutability)? local_target_instance);
+                        $(
+                            *target_instance = self.tcx().arena.alloc(local_target_instance);
+                            let $mutability _unused = ();
+                        )?
+                    }
                 }
             }
 

compiler/rustc_middle/src/ty/codec.rs

@@ -536,6 +536,7 @@ impl_arena_copy_decoder! {<'tcx>
     rustc_span::def_id::LocalDefId,
     (rustc_middle::middle::exported_symbols::ExportedSymbol<'tcx>, rustc_middle::middle::exported_symbols::SymbolExportInfo),
     ty::DeducedParamAttrs,
+    ty::InstanceDef<'tcx>,
 }
 
 #[macro_export]

compiler/rustc_middle/src/ty/instance.rs

@@ -1,7 +1,7 @@
 use crate::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use crate::ty::print::{FmtPrinter, Printer};
 use crate::ty::{self, Ty, TyCtxt, TypeFoldable, TypeSuperFoldable};
-use crate::ty::{EarlyBinder, GenericArgs, GenericArgsRef, TypeVisitableExt};
+use crate::ty::{EarlyBinder, GenericArgs, GenericArgsRef, Lift, TypeVisitableExt};
 use rustc_errors::ErrorGuaranteed;
 use rustc_hir as hir;
 use rustc_hir::def::Namespace;
@@ -135,6 +135,14 @@ pub enum InstanceDef<'tcx> {
     ///
     /// The `DefId` is for `FnPtr::addr`, the `Ty` is the type `T`.
     FnPtrAddrShim(DefId, Ty<'tcx>),
+
+    /// Typecast shim which replaces the `Self` type with the provided type.
+    /// This is used in vtable calls, where the type of `Self` is abstract as of the time of
+    /// the call.
+    ///
+    /// `target_instance` will be an instantiable `InstanceDef`, either an `Item` for which
+    /// we have MIR available or a generatable shim.
+    CfiShim { target_instance: &'tcx InstanceDef<'tcx>, invoke_ty: Ty<'tcx> },
 }
 
 impl<'tcx> Instance<'tcx> {
@@ -198,6 +206,7 @@ impl<'tcx> InstanceDef<'tcx> {
             | InstanceDef::DropGlue(def_id, _)
             | InstanceDef::CloneShim(def_id, _)
             | InstanceDef::FnPtrAddrShim(def_id, _) => def_id,
+            InstanceDef::CfiShim { target_instance, .. } => target_instance.def_id(),
         }
     }
 
@@ -209,6 +218,7 @@ impl<'tcx> InstanceDef<'tcx> {
                 Some(def_id)
             }
             InstanceDef::VTableShim(..)
+            | InstanceDef::CfiShim { .. }
             | InstanceDef::ReifyShim(..)
             | InstanceDef::FnPtrShim(..)
             | InstanceDef::Virtual(..)
@@ -319,6 +329,17 @@ impl<'tcx> InstanceDef<'tcx> {
             | InstanceDef::ReifyShim(..)
             | InstanceDef::Virtual(..)
             | InstanceDef::VTableShim(..) => true,
+            InstanceDef::CfiShim { target_instance, .. } => {
+                target_instance.has_polymorphic_mir_body()
+            }
+        }
+    }
+
+    pub fn is_vtable_shim(&self) -> bool {
+        match self {
+            InstanceDef::VTableShim(..) => true,
+            InstanceDef::CfiShim { target_instance, .. } => target_instance.is_vtable_shim(),
+            _ => false,
         }
     }
 }
@@ -339,6 +360,10 @@ fn fmt_instance_def(f: &mut fmt::Formatter<'_>, instance_def: &InstanceDef<'_>)
         InstanceDef::DropGlue(_, Some(ty)) => write!(f, " - shim(Some({ty}))"),
         InstanceDef::CloneShim(_, ty) => write!(f, " - shim({ty})"),
         InstanceDef::FnPtrAddrShim(_, ty) => write!(f, " - shim({ty})"),
+        InstanceDef::CfiShim { invoke_ty, target_instance } => {
+            fmt_instance_def(f, target_instance)?;
+            write!(f, " - cfi-shim({invoke_ty})")
+        }
     }
 }
 
@@ -579,6 +604,34 @@ impl<'tcx> Instance<'tcx> {
         Instance::expect_resolve(tcx, ty::ParamEnv::reveal_all(), def_id, args)
     }
 
+    pub fn cfi_shim(
+        mut self,
+        tcx: TyCtxt<'tcx>,
+        invoke_trait: Option<ty::PolyTraitRef<'tcx>>,
+    ) -> ty::Instance<'tcx> {
+        if tcx.sess.cfi_shims() {
+            let invoke_ty = if let Some(poly_trait_ref) = invoke_trait {
+                tcx.trait_object_ty(poly_trait_ref)
+            } else {
+                Ty::new_dynamic(tcx, ty::List::empty(), tcx.lifetimes.re_erased, ty::Dyn)
+            };
+            // If we're an Item and the `def_id` is not local, we may not have MIR available.
+            // If it's a closure, we can't ReifyShim it, just use it directly.
+            if let InstanceDef::Item(def_id) = self.def
+                && !def_id.is_local()
+                && !tcx.is_closure_like(def_id)
+            {
+                self.def = InstanceDef::ReifyShim(def_id);
+            }
+
+            self.def = InstanceDef::CfiShim {
+                target_instance: (&self.def).lift_to_tcx(tcx).expect("Could not lift for shimming"),
+                invoke_ty,
+            };
+        }
+        self
+    }
+
     #[instrument(level = "debug", skip(tcx), ret)]
     pub fn fn_once_adapter_instance(
         tcx: TyCtxt<'tcx>,

compiler/rustc_middle/src/ty/mod.rs

@@ -1730,6 +1730,7 @@ impl<'tcx> TyCtxt<'tcx> {
             | ty::InstanceDef::CloneShim(..)
             | ty::InstanceDef::ThreadLocalShim(..)
             | ty::InstanceDef::FnPtrAddrShim(..) => self.mir_shims(instance),
+            ty::InstanceDef::CfiShim { target_instance, .. } => self.instance_mir(*target_instance),
         }
     }
 

compiler/rustc_middle/src/ty/structural_impls.rs

@@ -479,6 +479,14 @@ impl<'a, 'tcx> Lift<'tcx> for Term<'a> {
     }
 }
 
+impl<'a, 'b, 'tcx> Lift<'tcx> for &'b ty::InstanceDef<'a> {
+    type Lifted = &'tcx ty::InstanceDef<'tcx>;
+    fn lift_to_tcx(self, tcx: TyCtxt<'tcx>) -> Option<Self::Lifted> {
+        let lifted: ty::InstanceDef<'tcx> = (*self).lift_to_tcx(tcx)?;
+        Some(tcx.arena.alloc(lifted))
+    }
+}
+
 ///////////////////////////////////////////////////////////////////////////
 // Traversal implementations.
 
@@ -488,6 +496,22 @@ impl<'tcx> TypeVisitable<TyCtxt<'tcx>> for ty::AdtDef<'tcx> {
     }
 }
 
+impl<'tcx> TypeVisitable<TyCtxt<'tcx>> for &ty::InstanceDef<'tcx> {
+    fn visit_with<V: TypeVisitor<TyCtxt<'tcx>>>(&self, visitor: &mut V) -> V::Result {
+        (*self).visit_with(visitor)
+    }
+}
+
+impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for &'tcx ty::InstanceDef<'tcx> {
+    fn try_fold_with<F: FallibleTypeFolder<TyCtxt<'tcx>>>(
+        self,
+        folder: &mut F,
+    ) -> Result<Self, F::Error> {
+        let folded: ty::InstanceDef<'tcx> = (*self).try_fold_with(folder)?;
+        Ok(folder.interner().arena.alloc(folded))
+    }
+}
+
 impl<'tcx, T: TypeFoldable<TyCtxt<'tcx>>> TypeFoldable<TyCtxt<'tcx>> for ty::Binder<'tcx, T> {
     fn try_fold_with<F: FallibleTypeFolder<TyCtxt<'tcx>>>(
         self,

compiler/rustc_mir_transform/src/inline.rs

@@ -331,7 +331,8 @@ impl<'tcx> Inliner<'tcx> {
             | InstanceDef::DropGlue(..)
             | InstanceDef::CloneShim(..)
             | InstanceDef::ThreadLocalShim(..)
-            | InstanceDef::FnPtrAddrShim(..) => return Ok(()),
+            | InstanceDef::FnPtrAddrShim(..)
+            | InstanceDef::CfiShim { .. } => return Ok(()),
         }
 
         if self.tcx.is_constructor(callee_def_id) {

compiler/rustc_mir_transform/src/inline/cycle.rs

@@ -90,7 +90,8 @@ pub(crate) fn mir_callgraph_reachable<'tcx>(
                 | InstanceDef::ConstructCoroutineInClosureShim { .. }
                 | InstanceDef::CoroutineKindShim { .. }
                 | InstanceDef::ThreadLocalShim { .. }
-                | InstanceDef::CloneShim(..) => {}
+                | InstanceDef::CloneShim(..)
+                | InstanceDef::CfiShim { .. } => {}
 
                 // This shim does not call any other functions, thus there can be no recursion.
                 InstanceDef::FnPtrAddrShim(..) => continue,

compiler/rustc_mir_transform/src/shim.rs

@@ -31,6 +31,7 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
 
     let mut result = match instance {
         ty::InstanceDef::Item(..) => bug!("item {:?} passed to make_shim", instance),
+        ty::InstanceDef::CfiShim { .. } => bug!("cfi shim {instance:?} passed to make_shim"),
         ty::InstanceDef::VTableShim(def_id) => {
             let adjustment = Adjustment::Deref { source: DerefSource::MutPtr };
             build_call_shim(tcx, instance, Some(adjustment), CallKind::Direct(def_id))