CFI: Fix types that implement Fn, FnMut, or FnOnce

Fix type transformation for types that implement Fn, FnMut, or FnOnce traits
(such as Box<dyn Fn()>, Box<dyn FnMut()>, etc.) by transforming self into a
trait object of the trait that defines the method.

Author: rcvalle

compiler/rustc_sanitizers/src/cfi/typeid/itanium_cxx_abi/transform.rs

@@ -453,6 +453,35 @@ pub(crate) fn transform_instance<'tcx>(
             instance.def = ty::InstanceKind::Virtual(call, 0);
             instance.args = abstract_args;
         }
+        let fn_traits = [
+            (LangItem::Fn, sym::call),
+            (LangItem::FnMut, sym::call_mut),
+            (LangItem::FnOnce, sym::call_once),
+        ];
+        for (lang_item, method_sym) in fn_traits {
+            if let Some(trait_id) = tcx.lang_items().get(lang_item) {
+                let items = tcx.associated_items(trait_id);
+                if let Some(call_method) =
+                    items.in_definition_order().find(|item| item.name() == method_sym)
+                {
+                    if instance.def_id() == call_method.def_id {
+                        // This is a call to a method of Fn, FnMut, or FnOnce. Transform self into a
+                        // trait object of the trait that defines the method.
+                        let self_ty = trait_object_ty(
+                            tcx,
+                            ty::Binder::dummy(ty::TraitRef::from_method(
+                                tcx,
+                                trait_id,
+                                instance.args,
+                            )),
+                        );
+                        instance.args =
+                            tcx.mk_args_trait(self_ty, instance.args.into_iter().skip(1));
+                        break;
+                    }
+                }
+            }
+        }
     }
 
     instance

tests/ui/sanitizer/cfi/fn-trait-objects.rs

@@ -0,0 +1,30 @@
+// Verifies that types that implement the Fn, FnMut, or FnOnce traits can be
+// called through their trait methods.
+//
+//@ needs-sanitizer-cfi
+//@ compile-flags: -Ctarget-feature=-crt-static -Ccodegen-units=1 -Clto -Cprefer-dynamic=off -Copt-level=0 -Zsanitizer=cfi --test
+//@ run-pass
+
+#![feature(fn_traits)]
+#![feature(unboxed_closures)]
+
+fn foo(_a: u32) {}
+
+#[test]
+fn test_fn_trait() {
+    let f: Box<dyn Fn(u32)> = Box::new(foo);
+    Fn::call(&f, (0,));
+}
+
+#[test]
+fn test_fnmut_trait() {
+    let mut a = 0;
+    let mut f: Box<dyn FnMut(u32)> = Box::new(|x| a += x);
+    FnMut::call_mut(&mut f, (1,));
+}
+
+#[test]
+fn test_fnonce_trait() {
+    let f: Box<dyn FnOnce(u32)> = Box::new(foo);
+    FnOnce::call_once(f, (2,));
+}

tests/ui/sanitizer/kcfi/fn-trait-objects.rs

@@ -0,0 +1,30 @@
+// Verifies that types that implement the Fn, FnMut, or FnOnce traits can be
+// called through their trait methods.
+//
+//@ needs-sanitizer-cfi
+//@ compile-flags: -Ctarget-feature=-crt-static -Zpanic_abort_tests -Cpanic=abort -Cprefer-dynamic=off -Copt-level=0 -Zsanitizer=kcfi --test
+//@ run-pass
+
+#![feature(fn_traits)]
+#![feature(unboxed_closures)]
+
+fn foo(_a: u32) {}
+
+#[test]
+fn test_fn_trait() {
+    let f: Box<dyn Fn(u32)> = Box::new(foo);
+    Fn::call(&f, (0,));
+}
+
+#[test]
+fn test_fnmut_trait() {
+    let mut a = 0;
+    let mut f: Box<dyn FnMut(u32)> = Box::new(|x| a += x);
+    FnMut::call_mut(&mut f, (1,));
+}
+
+#[test]
+fn test_fnonce_trait() {
+    let f: Box<dyn FnOnce(u32)> = Box::new(foo);
+    FnOnce::call_once(f, (2,));
+}