From 9fe0fe15c9d6b9d94438d5706edb3501732903ae Mon Sep 17 00:00:00 2001 From: Matthew Maurer Date: Tue, 26 Mar 2024 17:00:57 +0000 Subject: [PATCH] CFI: Rewrite closure and coroutine instances to their trait method Similar to methods on a trait object, the most common way to indirectly call a closure or coroutine is through the vtable on the appropriate trait. This uses the same approach as we use for trait methods, after backing out the trait arguments from the type. --- .../src/typeid/typeid_itanium_cxx_abi.rs | 32 +++++++++++++ tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs | 18 -------- tests/ui/sanitizer/cfi-closures.rs | 45 +++++++++++++++++++ tests/ui/sanitizer/cfi-coroutine.rs | 23 ++++++++++ 4 files changed, 100 insertions(+), 18 deletions(-) delete mode 100644 tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs create mode 100644 tests/ui/sanitizer/cfi-closures.rs create mode 100644 tests/ui/sanitizer/cfi-coroutine.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..ecfb020848817 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 @@ -10,6 +10,7 @@ use rustc_data_structures::base_n; use rustc_data_structures::fx::FxHashMap; use rustc_hir as hir; +use rustc_hir::lang_items::LangItem; use rustc_middle::ty::layout::IntegerExt; use rustc_middle::ty::TypeVisitableExt; use rustc_middle::ty::{ @@ -1150,6 +1151,37 @@ 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()); + let (trait_id, tuple_args) = match closure_ty.kind() { + ty::Closure(_def_id, args) => { + let closure_args = ty::ClosureArgs { args }; + let trait_id = tcx.fn_trait_kind_to_def_id(closure_args.kind()).unwrap(); + let tuple_args = closure_args.sig().inputs().no_bound_vars().unwrap()[0]; + (trait_id, tuple_args) + } + ty::Coroutine(_def_id, args) => ( + tcx.require_lang_item(LangItem::Coroutine, None), + ty::CoroutineArgs { args }.resume_ty(), + ), + x => bug!("Unexpected type kind for closure-like: {x:?}"), + }; + 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; } let fn_abi = tcx diff --git a/tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs b/tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs deleted file mode 100644 index a46a3afd7343a..0000000000000 --- a/tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs +++ /dev/null @@ -1,18 +0,0 @@ -// Tests that converting a closure to a function pointer works -// The notable thing being tested here is that when the closure does not capture anything, -// the call method from its Fn trait takes a ZST representing its environment. The compiler then -// uses the assumption that the ZST is non-passed to reify this into a function pointer. -// -// This checks that the reified function pointer will have the expected alias set at its call-site. - -//@ needs-sanitizer-cfi -// FIXME(#122848) Remove only-linux once OSX CFI binaries work -//@ only-linux -//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi -//@ compile-flags: -C target-feature=-crt-static -C codegen-units=1 -C opt-level=0 -//@ run-pass - -pub fn main() { - let f: &fn() = &((|| ()) as _); - f(); -} diff --git a/tests/ui/sanitizer/cfi-closures.rs b/tests/ui/sanitizer/cfi-closures.rs new file mode 100644 index 0000000000000..073589f68ea4c --- /dev/null +++ b/tests/ui/sanitizer/cfi-closures.rs @@ -0,0 +1,45 @@ +// Check various forms of dynamic closure calls +// +// 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 +//@ compile-flags: --test +//@ run-pass + +#![feature(fn_traits)] + +fn foo<'a, T>() -> Box &'a T> { + Box::new(|x| x) +} + +#[test] +fn dyn_fn_with_params() { + let x = 3; + let f = foo(); + f(&x); + // FIXME remove once drops are working. + std::mem::forget(f); +} + +#[test] +fn call_fn_trait() { + let f: &(dyn Fn()) = &(|| {}) as _; + f.call(()); +} + +#[test] +fn fn_ptr_cast() { + let f: &fn() = &((|| ()) as _); + f(); +} + +fn use_fnmut(mut f: F) { + f() +} + +#[test] +fn fn_to_fnmut() { + let f: &(dyn Fn()) = &(|| {}) as _; + use_fnmut(f); +} diff --git a/tests/ui/sanitizer/cfi-coroutine.rs b/tests/ui/sanitizer/cfi-coroutine.rs new file mode 100644 index 0000000000000..0971bf88a025e --- /dev/null +++ b/tests/ui/sanitizer/cfi-coroutine.rs @@ -0,0 +1,23 @@ +// Verifies that we can call dynamic coroutines +// +// 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(coroutines)] +#![feature(coroutine_trait)] + +use std::ops::{Coroutine, CoroutineState}; +use std::pin::{pin, Pin}; + +fn main() { + let mut coro = |x: i32| { + yield x; + "done" + }; + let mut abstract_coro: Pin<&mut dyn Coroutine> = pin!(coro); + assert_eq!(abstract_coro.as_mut().resume(2), CoroutineState::Yielded(2)); + assert_eq!(abstract_coro.as_mut().resume(0), CoroutineState::Complete("done")); +}