CFI: Introduce CFI shims

Indirect calls through vtables (trait objects or drop_in_place) expect
to have a type based on `dyn Trait` at the call-site. The actual
implementations have types based on `MyImplType`. These shims allow the
insertion of an explicit cast at the beginning of any instance, allowing
a different type to be assigned. These shims function for both CFI and
KCFI, as they have a single principal type.

Author: maurer

compiler/rustc_const_eval/src/interpret/terminator.rs

@@ -547,6 +547,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

@@ -939,6 +939,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

@@ -526,6 +526,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,14 +135,27 @@ 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.
+    ///
+    /// For `target_instance` which are generatable shims, this is done via inserting a cast at the
+    /// beginning of the shim generated by that instance.
+    /// For `Item`s, we instead build a direct call shim with that cast inserted.
+    CfiShim { target_instance: &'tcx InstanceDef<'tcx>, invoke_ty: Ty<'tcx> },
 }
 
 impl<'tcx> Instance<'tcx> {
     /// Returns the `Ty` corresponding to this `Instance`, with generic instantiations applied and
     /// lifetimes erased, allowing a `ParamEnv` to be specified for use during normalization.
     pub fn ty(&self, tcx: TyCtxt<'tcx>, param_env: ty::ParamEnv<'tcx>) -> Ty<'tcx> {
-        let ty = tcx.type_of(self.def.def_id());
-        tcx.instantiate_and_normalize_erasing_regions(self.args, param_env, ty)
+        let args = if let InstanceDef::CfiShim { invoke_ty, .. } = self.def {
+            tcx.mk_args_trait(invoke_ty, (*self.args).into_iter().skip(1))
+        } else {
+            self.args
+        };
+        tcx.instantiate_and_normalize_erasing_regions(args, param_env, tcx.type_of(self.def_id()))
     }
 
     /// Finds a crate that contains a monomorphization of this instance that
@@ -198,6 +211,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 +223,7 @@ impl<'tcx> InstanceDef<'tcx> {
                 Some(def_id)
             }
             InstanceDef::VTableShim(..)
+            | InstanceDef::CfiShim { .. }
             | InstanceDef::ReifyShim(..)
             | InstanceDef::FnPtrShim(..)
             | InstanceDef::Virtual(..)
@@ -319,6 +334,9 @@ impl<'tcx> InstanceDef<'tcx> {
             | InstanceDef::ReifyShim(..)
             | InstanceDef::Virtual(..)
             | InstanceDef::VTableShim(..) => true,
+            InstanceDef::CfiShim { target_instance, .. } => {
+                target_instance.has_polymorphic_mir_body()
+            }
         }
     }
 
@@ -360,6 +378,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})")
+        }
     }
 }
 
@@ -600,6 +622,75 @@ 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 tcx.is_closure_like(self.def.def_id()) {
+                // Closures don't have a fn_sig and can't be called directly.
+                // Adjust it into a call through
+                // `Fn`/`FnMut`/`FnOnce`/`AsyncFn`/`AsyncFnMut`/`AsyncFnOnce`/`Coroutine`
+                // based on its receiver.
+                let ty::TyKind::Dynamic(pep, _, _) = invoke_ty.kind() else {
+                    bug!("Closure-like with non-dynamic invoke_ty {invoke_ty}")
+                };
+                let Some(fn_trait) = pep.principal_def_id() else {
+                    bug!("Closure-like with no principal trait on invoke_ty {invoke_ty}")
+                };
+                let call = tcx
+                    .associated_items(fn_trait)
+                    .in_definition_order()
+                    .find(|it| it.kind == ty::AssocKind::Fn)
+                    .expect("No call-family function on closure-like principal trait?")
+                    .def_id;
+
+                let self_ty = self.ty(tcx, ty::ParamEnv::reveal_all());
+                let tupled_inputs_ty =
+                    self.args.as_closure().sig().map_bound(|sig| sig.inputs()[0]);
+                let tupled_inputs_ty = tcx.instantiate_bound_regions_with_erased(tupled_inputs_ty);
+                self.args = tcx.mk_args_trait(self_ty, [tupled_inputs_ty.into()]);
+                self.def = InstanceDef::Item(call);
+            } else if let Some(impl_id) = tcx.impl_of_method(self.def.def_id()) {
+                // Trait methods will have a Self polymorphic parameter, where the concreteized
+                // implementatation will not. We need to walk back to the more general trait method
+                // so that we can swap out Self when generating a type signature.
+                let Some(trait_ref) = tcx.impl_trait_ref(impl_id) else {
+                    bug!("When trying to rewrite the type on {self}, found inherent impl method")
+                };
+                let trait_ref = trait_ref.instantiate(tcx, self.args);
+
+                let method_id = tcx
+                    .impl_item_implementor_ids(impl_id)
+                    .items()
+                    .filter_map(|(trait_method, impl_method)| {
+                        (*impl_method == self.def.def_id()).then_some(*trait_method)
+                    })
+                    .min()
+                    .unwrap();
+                self.def = InstanceDef::Item(method_id);
+                self.args = trait_ref.args;
+            }
+            // At this point, we should have gauranteed that the first item in the args list is
+            // the dispatch type. We can't check for Self, because `drop_in_place` takes `T`.
+            self.def = InstanceDef::CfiShim {
+                target_instance: (&self.def).lift_to_tcx(tcx).expect("Could not lift for shimming"),
+                invoke_ty,
+            };
+        }
+        self
+    }
+
+    pub fn force_thin_self(&self) -> bool {
+        matches!(self.def, InstanceDef::Virtual(..) | InstanceDef::CfiShim { .. })
+    }
+
     #[instrument(level = "debug", skip(tcx), ret)]
     pub fn fn_once_adapter_instance(
         tcx: TyCtxt<'tcx>,

compiler/rustc_middle/src/ty/mod.rs

@@ -1722,7 +1722,8 @@ impl<'tcx> TyCtxt<'tcx> {
             | ty::InstanceDef::DropGlue(..)
             | ty::InstanceDef::CloneShim(..)
             | ty::InstanceDef::ThreadLocalShim(..)
-            | ty::InstanceDef::FnPtrAddrShim(..) => self.mir_shims(instance),
+            | ty::InstanceDef::FnPtrAddrShim(..)
+            | ty::InstanceDef::CfiShim { .. } => self.mir_shims(instance),
         }
     }
 

compiler/rustc_middle/src/ty/structural_impls.rs

@@ -478,6 +478,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.
 
@@ -490,6 +498,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) -> ControlFlow<V::BreakTy> {
+        (*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/lib.rs

@@ -101,6 +101,7 @@ mod remove_unneeded_drops;
 mod remove_zsts;
 mod required_consts;
 mod reveal_all;
+mod rewrite_receiver;
 mod shim;
 mod ssa;
 // This pass is public to allow external drivers to perform MIR cleanup