From 1d2d8675cce7bc5ec6191d885a8f2c89d2a4b392 Mon Sep 17 00:00:00 2001
From: Matthew Maurer <mmaurer@google.com>
Date: Tue, 26 Mar 2024 17:00:57 +0000
Subject: [PATCH] CFI: Rewrite closure and coroutine instances to their trait
 method

---
 .../src/typeid/typeid_itanium_cxx_abi.rs      | 34 +++++++++++++++++++
 .../ui/sanitizer/cfi-closure-fn-trait-cast.rs | 13 +++++++
 tests/ui/sanitizer/cfi-closure-with-params.rs | 19 +++++++++++
 3 files changed, 66 insertions(+)
 create mode 100644 tests/ui/sanitizer/cfi-closure-fn-trait-cast.rs
 create mode 100644 tests/ui/sanitizer/cfi-closure-with-params.rs

diff --git a/compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs b/compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs
index 54f59ad7bc3ed..03c44ebdedd8a 100644
--- a/compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs
+++ b/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
diff --git a/tests/ui/sanitizer/cfi-closure-fn-trait-cast.rs b/tests/ui/sanitizer/cfi-closure-fn-trait-cast.rs
new file mode 100644
index 0000000000000..d78d74cb9d5ec
--- /dev/null
+++ b/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(());
+}
diff --git a/tests/ui/sanitizer/cfi-closure-with-params.rs b/tests/ui/sanitizer/cfi-closure-with-params.rs
new file mode 100644
index 0000000000000..75a33617ad96e
--- /dev/null
+++ b/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);
+}