CFI: Rewrite closure and coroutine instances to their trait method

Author: maurer

compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs

@@ -1150,6 +1150,40 @@ pub fn typeid_for_instance<'tcx>(
                 tcx.mk_args_trait(invoke_ty, trait_ref.args.into_iter().skip(1));
             instance.args = instance.args.rebase_onto(tcx, impl_id, abstract_trait_args);
         }
+    } else if tcx.is_closure_like(instance.def_id()) {
+        // We're either a closure or a coroutine. Our goal is to find the trait we're defined on,
+        // instantiate it, and take the type of its only method as our own.
+        let closure_ty = instance.ty(tcx, ty::ParamEnv::reveal_all());
+        match closure_ty.kind() {
+            ty::Closure(_def_id, args) => {
+                let ca = ty::ClosureArgs { args };
+                let trait_id = tcx
+                    .fn_trait_kind_to_def_id(ca.kind())
+                    .expect("Resolving abstract closure typeid for undefined Fn trait?");
+                let tuple_args = ca
+                    .sig()
+                    .inputs()
+                    .no_bound_vars()
+                    .expect("We are at codegen, this instance should be fully instantiated.")[0];
+                let trait_ref = ty::TraitRef::new(tcx, trait_id, [closure_ty, tuple_args]);
+                let invoke_ty = trait_object_ty(tcx, ty::Binder::dummy(trait_ref));
+                let abstract_args =
+                    tcx.mk_args_trait(invoke_ty, trait_ref.args.into_iter().skip(1));
+                // There should be exactly one method on this trait, and it should be the one we're
+                // defining.
+                let call = tcx
+                    .associated_items(trait_id)
+                    .in_definition_order()
+                    .find(|it| it.kind == ty::AssocKind::Fn)
+                    .expect("No call-family function on closure-like Fn trait?")
+                    .def_id;
+
+                instance.def = ty::InstanceDef::Virtual(call, 0);
+                instance.args = abstract_args;
+            }
+            ty::CoroutineClosure(_def_id, _args) => unimplemented!(),
+            x => bug!("Unexpected type kind for closure-like: {x:?}"),
+        }
     }
 
     let fn_abi = tcx

tests/ui/sanitizer/cfi-closure-fn-trait-cast.rs

@@ -0,0 +1,13 @@
+// Verifies that casting a closure to a Fn trait object works.
+//
+// FIXME(#122848): Remove only-linux when fixed.
+//@ only-linux
+//@ needs-sanitizer-cfi
+//@ compile-flags: -Clto -Copt-level=0 -Cprefer-dynamic=off -Ctarget-feature=-crt-static -Zsanitizer=cfi
+//@ run-pass
+
+#![feature(fn_traits)]
+fn main() {
+   let f: &(dyn Fn()) = &(|| {}) as _;
+   f.call(());
+}

tests/ui/sanitizer/cfi-closure-with-params.rs

@@ -0,0 +1,19 @@
+// Verifies that casting a closure to a Fn trait object works.
+//
+// FIXME(#122848): Remove only-linux when fixed.
+//@ only-linux
+//@ needs-sanitizer-cfi
+//@ compile-flags: -Clto -Copt-level=0 -Cprefer-dynamic=off -Ctarget-feature=-crt-static -Zsanitizer=cfi
+//@ run-pass
+
+fn foo<'a, T>() -> Box<dyn Fn(&'a T) -> &'a T> {
+    Box::new(|x| x)
+}
+
+fn main() {
+    let x = 3;
+    let f = foo();
+    f(&x);
+    // FIXME remove once drops are working.
+    std::mem::forget(f);
+}