CFI: Apply CFI shims to drops

Fixes: 118761

Author: maurer

compiler/rustc_middle/src/ty/vtable.rs

@@ -53,11 +53,11 @@ pub(super) fn vtable_allocation_provider<'tcx>(
 ) -> AllocId {
     let (ty, poly_trait_ref) = key;
 
-    let vtable_entries = if let Some(poly_trait_ref) = poly_trait_ref {
-        let trait_ref = poly_trait_ref.with_self_ty(tcx, ty);
-        let trait_ref = tcx.erase_regions(trait_ref);
+    let invoke_trait = poly_trait_ref
+        .map(|poly_trait_ref| tcx.erase_regions(poly_trait_ref.with_self_ty(tcx, ty)));
 
-        tcx.vtable_entries(trait_ref)
+    let vtable_entries = if let Some(invoke_trait) = invoke_trait {
+        tcx.vtable_entries(invoke_trait)
     } else {
         TyCtxt::COMMON_VTABLE_ENTRIES
     };
@@ -83,7 +83,8 @@ pub(super) fn vtable_allocation_provider<'tcx>(
         let idx: u64 = u64::try_from(idx).unwrap();
         let scalar = match entry {
             VtblEntry::MetadataDropInPlace => {
-                let instance = ty::Instance::resolve_drop_in_place(tcx, ty);
+                let instance =
+                    ty::Instance::resolve_drop_in_place(tcx, ty).cfi_shim(tcx, invoke_trait);
                 let fn_alloc_id = tcx.reserve_and_set_fn_alloc(instance);
                 let fn_ptr = Pointer::from(fn_alloc_id);
                 Scalar::from_pointer(fn_ptr, &tcx)

compiler/rustc_monomorphize/src/collector.rs

@@ -381,7 +381,7 @@ fn collect_items_rec<'tcx>(
             debug_assert!(should_codegen_locally(tcx, &instance));
 
             let ty = instance.ty(tcx, ty::ParamEnv::reveal_all());
-            visit_drop_use(tcx, ty, true, starting_item.span, &mut used_items);
+            visit_drop_use(tcx, ty, None, true, starting_item.span, &mut used_items);
 
             recursion_depth_reset = None;
 
@@ -855,7 +855,7 @@ impl<'a, 'tcx> MirVisitor<'tcx> for MirUsedCollector<'a, 'tcx> {
             mir::TerminatorKind::Drop { ref place, .. } => {
                 let ty = place.ty(self.body, self.tcx).ty;
                 let ty = self.monomorphize(ty);
-                visit_drop_use(self.tcx, ty, true, source, self.output);
+                visit_drop_use(self.tcx, ty, None, true, source, self.output);
             }
             mir::TerminatorKind::InlineAsm { ref operands, .. } => {
                 for op in operands {
@@ -917,11 +917,17 @@ impl<'a, 'tcx> MirVisitor<'tcx> for MirUsedCollector<'a, 'tcx> {
 fn visit_drop_use<'tcx>(
     tcx: TyCtxt<'tcx>,
     ty: Ty<'tcx>,
+    invoke_trait: Option<ty::PolyTraitRef<'tcx>>,
     is_direct_call: bool,
     source: Span,
     output: &mut MonoItems<'tcx>,
 ) {
-    let instance = Instance::resolve_drop_in_place(tcx, ty);
+    let mut instance = Instance::resolve_drop_in_place(tcx, ty);
+    if tcx.sess.cfi_shims() && !is_direct_call {
+        // The CFI shim may generate a direct call to the unshimmed drop
+        visit_drop_use(tcx, ty, None, true, source, output);
+        instance = instance.cfi_shim(tcx, invoke_trait);
+    }
     visit_instance_use(tcx, instance, is_direct_call, source, output);
 }
 
@@ -1188,8 +1194,10 @@ fn create_mono_items_for_vtable_methods<'tcx>(
     assert!(!trait_ty.has_escaping_bound_vars() && !impl_ty.has_escaping_bound_vars());
 
     if let ty::Dynamic(trait_ty, ..) = trait_ty.kind() {
-        if let Some(principal) = trait_ty.principal() {
-            let poly_trait_ref = principal.with_self_ty(tcx, impl_ty);
+        let invoke_trait =
+            trait_ty.principal().map(|principal| principal.with_self_ty(tcx, impl_ty));
+
+        if let Some(poly_trait_ref) = invoke_trait {
             assert!(!poly_trait_ref.has_escaping_bound_vars());
 
             // Walk all methods of the trait, including those of its supertraits
@@ -1211,10 +1219,10 @@ fn create_mono_items_for_vtable_methods<'tcx>(
                 })
                 .map(|item| create_fn_mono_item(tcx, item, source));
             output.extend(methods);
-        }
+        };
 
         // Also add the destructor.
-        visit_drop_use(tcx, impl_ty, false, source, output);
+        visit_drop_use(tcx, impl_ty, invoke_trait, false, source, output);
     }
 }
 
@@ -1239,7 +1247,7 @@ impl<'v> RootCollector<'_, 'v> {
                     debug!("RootCollector: ADT drop-glue for `{id:?}`",);
 
                     let ty = self.tcx.type_of(id.owner_id.to_def_id()).no_bound_vars().unwrap();
-                    visit_drop_use(self.tcx, ty, true, DUMMY_SP, self.output);
+                    visit_drop_use(self.tcx, ty, None, true, DUMMY_SP, self.output);
                 }
             }
             DefKind::GlobalAsm => {

tests/codegen/sanitizer/kcfi-emit-type-metadata-trait-objects.rs

@@ -12,7 +12,7 @@
 #![no_core]
 
 #[lang="sized"]
-trait Sized { }
+pub trait Sized { }
 #[lang="copy"]
 trait Copy { }
 #[lang="receiver"]
@@ -28,7 +28,9 @@ impl<'a, 'b: 'a, T: ?Sized + Unsize<U>, U: ?Sized> CoerceUnsized<&'a U> for &'b
 #[lang="freeze"]
 trait Freeze { }
 #[lang="drop_in_place"]
-fn drop_in_place_fn<T>() { }
+fn drop_in_place_fn<T: ?Sized>(_to_drop: *mut T) { }
+#[lang="unpin"]
+pub trait Unpin { }
 
 pub trait Trait1 {
     fn foo(&self);

tests/ui/sanitizer/cfi-drop-issue-118761.rs

@@ -0,0 +1,14 @@
+// Validate that objects that might have custom drop can be dropped with CFI on. See #118761
+
+//@ needs-sanitizer-cfi
+//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi -C codegen-units=1
+//@ compile-flags: -C opt-level=0
+//@ run-pass
+
+struct Bar;
+trait Fooable {}
+impl Fooable for Bar {}
+
+fn main() {
+   let _: Box<dyn Fooable> = Box::new(Bar);
+}

tests/ui/sanitizer/cfi-nontrivial-drop-glue.rs

@@ -0,0 +1,10 @@
+//@ needs-sanitizer-cfi
+//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi
+//@ compile-flags: -C codegen-units=1 -C opt-level=0
+//@ run-pass
+
+use std::env;
+
+fn main() {
+    env::current_exe().unwrap();
+}

tests/ui/sanitizer/cfi-trait-object.rs

@@ -0,0 +1,22 @@
+// Check trait objects run correctly
+
+//@ needs-sanitizer-cfi
+//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi
+//@ compile-flags: -C codegen-units=1 -C opt-level=0
+//@ run-pass
+
+struct Bar;
+trait Fooable {
+    fn foo(&self) -> i32;
+}
+
+impl Fooable for Bar {
+    fn foo(&self) -> i32 {
+        3
+    }
+}
+
+fn main() {
+   let bar: Box<dyn Fooable> = Box::new(Bar);
+   bar.foo();
+}