From e974570c42f4460b30fce11003b7891435931cbd Mon Sep 17 00:00:00 2001 From: Matthew Maurer Date: Wed, 27 Mar 2024 15:50:53 +0000 Subject: [PATCH 1/2] CFI: Only encode Coroutine Parent Args Fixes #122705 --- .../src/typeid/typeid_itanium_cxx_abi.rs | 16 ++++++++-- tests/ui/sanitizer/cfi-async-closures.rs | 30 +++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) create mode 100644 tests/ui/sanitizer/cfi-async-closures.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 aba3212b4cea9..82ceccba66a24 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 @@ -641,9 +641,7 @@ fn encode_ty<'tcx>( } // Function types - ty::FnDef(def_id, args) - | ty::Closure(def_id, args) - | ty::CoroutineClosure(def_id, args) => { + ty::FnDef(def_id, args) | ty::Closure(def_id, args) => { // u[IE], where is , // as vendor extended type. let mut s = String::new(); @@ -654,6 +652,18 @@ fn encode_ty<'tcx>( typeid.push_str(&s); } + ty::CoroutineClosure(def_id, args) => { + // u[IE], where is , + // as vendor extended type. + let mut s = String::new(); + let name = encode_ty_name(tcx, *def_id); + let _ = write!(s, "u{}{}", name.len(), &name); + let parent_args = tcx.mk_args(args.as_coroutine_closure().parent_args()); + s.push_str(&encode_args(tcx, parent_args, dict, options)); + compress(dict, DictKey::Ty(ty, TyQ::None), &mut s); + typeid.push_str(&s); + } + ty::Coroutine(def_id, args, ..) => { // u[IE], where is , // as vendor extended type. diff --git a/tests/ui/sanitizer/cfi-async-closures.rs b/tests/ui/sanitizer/cfi-async-closures.rs new file mode 100644 index 0000000000000..50edeb852bd6e --- /dev/null +++ b/tests/ui/sanitizer/cfi-async-closures.rs @@ -0,0 +1,30 @@ +// Check various forms of dynamic closure calls + +//@ edition: 2021 +//@ revisions: cfi kcfi +// FIXME(#122848) Remove only-linux once OSX CFI binaries work +//@ only-linux +//@ [cfi] needs-sanitizer-cfi +//@ [kcfi] needs-sanitizer-kcfi +//@ compile-flags: -C target-feature=-crt-static +//@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0 +//@ [cfi] compile-flags: -Z sanitizer=cfi +//@ [kcfi] compile-flags: -Z sanitizer=kcfi +//@ run-pass + +#![feature(async_closure)] +#![feature(async_fn_traits)] + +use std::ops::AsyncFn; + +#[inline(never)] +fn identity(x: T) -> T { x } + +// We can't actually create a `dyn AsyncFn()`, because it's not object-safe, but we should check +// that we don't bug out when we encounter one. + +fn main() { + let f = identity(async || ()); + let _ = f.async_call(()); + let _ = f(); +} From 8cc9a912d7daf092adc53f0b6b7fac5e86fa7e0a Mon Sep 17 00:00:00 2001 From: Matthew Maurer Date: Tue, 26 Mar 2024 17:00:57 +0000 Subject: [PATCH 2/2] 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 | 113 +++++++++++++----- tests/ui/sanitizer/cfi-async-closures.rs | 3 + tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs | 23 ---- tests/ui/sanitizer/cfi-closures.rs | 83 +++++++++++++ tests/ui/sanitizer/cfi-coroutine.rs | 30 +++++ 5 files changed, 197 insertions(+), 55 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 82ceccba66a24..5963bd7c5f16c 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::{ @@ -1161,43 +1162,91 @@ pub fn typeid_for_instance<'tcx>( }; let stripped_ty = strip_receiver_auto(tcx, upcast_ty); instance.args = tcx.mk_args_trait(stripped_ty, instance.args.into_iter().skip(1)); + } else if let ty::InstanceDef::VTableShim(def_id) = instance.def + && let Some(trait_id) = tcx.trait_of_item(def_id) + { + // VTableShims may have a trait method, but a concrete Self. This is not suitable for a vtable, + // as the caller will not know the concrete Self. + let trait_ref = ty::TraitRef::new(tcx, trait_id, instance.args); + let invoke_ty = trait_object_ty(tcx, ty::Binder::dummy(trait_ref)); + instance.args = tcx.mk_args_trait(invoke_ty, trait_ref.args.into_iter().skip(1)); } - if !options.contains(EncodeTyOptions::NO_SELF_TYPE_ERASURE) - && let Some(impl_id) = tcx.impl_of_method(instance.def_id()) - && let Some(trait_ref) = tcx.impl_trait_ref(impl_id) - { - let impl_method = tcx.associated_item(instance.def_id()); - let method_id = impl_method - .trait_item_def_id - .expect("Part of a trait implementation, but not linked to the def_id?"); - let trait_method = tcx.associated_item(method_id); - let trait_id = trait_ref.skip_binder().def_id; - if traits::is_vtable_safe_method(tcx, trait_id, trait_method) - && tcx.object_safety_violations(trait_id).is_empty() + if !options.contains(EncodeTyOptions::NO_SELF_TYPE_ERASURE) { + if let Some(impl_id) = tcx.impl_of_method(instance.def_id()) + && let Some(trait_ref) = tcx.impl_trait_ref(impl_id) { - // Trait methods will have a Self polymorphic parameter, where the concreteized - // implementatation will not. We need to walk back to the more general trait method - let trait_ref = tcx.instantiate_and_normalize_erasing_regions( - instance.args, - ty::ParamEnv::reveal_all(), - trait_ref, - ); + let impl_method = tcx.associated_item(instance.def_id()); + let method_id = impl_method + .trait_item_def_id + .expect("Part of a trait implementation, but not linked to the def_id?"); + let trait_method = tcx.associated_item(method_id); + let trait_id = trait_ref.skip_binder().def_id; + if traits::is_vtable_safe_method(tcx, trait_id, trait_method) + && tcx.object_safety_violations(trait_id).is_empty() + { + // Trait methods will have a Self polymorphic parameter, where the concreteized + // implementatation will not. We need to walk back to the more general trait method + let trait_ref = tcx.instantiate_and_normalize_erasing_regions( + instance.args, + ty::ParamEnv::reveal_all(), + trait_ref, + ); + let invoke_ty = trait_object_ty(tcx, ty::Binder::dummy(trait_ref)); + + // At the call site, any call to this concrete function through a vtable will be + // `Virtual(method_id, idx)` with appropriate arguments for the method. Since we have the + // original method id, and we've recovered the trait arguments, we can make the callee + // instance we're computing the alias set for match the caller instance. + // + // Right now, our code ignores the vtable index everywhere, so we use 0 as a placeholder. + // If we ever *do* start encoding the vtable index, we will need to generate an alias set + // based on which vtables we are putting this method into, as there will be more than one + // index value when supertraits are involved. + instance.def = ty::InstanceDef::Virtual(method_id, 0); + let abstract_trait_args = + 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, inputs) = match closure_ty.kind() { + ty::Closure(..) => { + let closure_args = instance.args.as_closure(); + let trait_id = tcx.fn_trait_kind_to_def_id(closure_args.kind()).unwrap(); + let tuple_args = + tcx.instantiate_bound_regions_with_erased(closure_args.sig()).inputs()[0]; + (trait_id, tuple_args) + } + ty::Coroutine(..) => ( + tcx.require_lang_item(LangItem::Coroutine, None), + instance.args.as_coroutine().resume_ty(), + ), + ty::CoroutineClosure(..) => ( + tcx.require_lang_item(LangItem::FnOnce, None), + tcx.instantiate_bound_regions_with_erased( + instance.args.as_coroutine_closure().coroutine_closure_sig(), + ) + .tupled_inputs_ty, + ), + x => bug!("Unexpected type kind for closure-like: {x:?}"), + }; + let trait_ref = ty::TraitRef::new(tcx, trait_id, [closure_ty, inputs]); 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; - // At the call site, any call to this concrete function through a vtable will be - // `Virtual(method_id, idx)` with appropriate arguments for the method. Since we have the - // original method id, and we've recovered the trait arguments, we can make the callee - // instance we're computing the alias set for match the caller instance. - // - // Right now, our code ignores the vtable index everywhere, so we use 0 as a placeholder. - // If we ever *do* start encoding the vtable index, we will need to generate an alias set - // based on which vtables we are putting this method into, as there will be more than one - // index value when supertraits are involved. - instance.def = ty::InstanceDef::Virtual(method_id, 0); - let abstract_trait_args = - 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); + instance.def = ty::InstanceDef::Virtual(call, 0); + instance.args = abstract_args; } } diff --git a/tests/ui/sanitizer/cfi-async-closures.rs b/tests/ui/sanitizer/cfi-async-closures.rs index 50edeb852bd6e..d94f2992d8428 100644 --- a/tests/ui/sanitizer/cfi-async-closures.rs +++ b/tests/ui/sanitizer/cfi-async-closures.rs @@ -10,6 +10,7 @@ //@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0 //@ [cfi] compile-flags: -Z sanitizer=cfi //@ [kcfi] compile-flags: -Z sanitizer=kcfi +//@ [kcfi] compile-flags: -C panic=abort -Z panic-abort-tests -C prefer-dynamic=off //@ run-pass #![feature(async_closure)] @@ -27,4 +28,6 @@ fn main() { let f = identity(async || ()); let _ = f.async_call(()); let _ = f(); + let g: Box _> = Box::new(f) as _; + let _ = g(); } 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 03818544aab48..0000000000000 --- a/tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs +++ /dev/null @@ -1,23 +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. - -//@ revisions: cfi kcfi -// FIXME(#122848) Remove only-linux once OSX CFI binaries work -//@ only-linux -//@ [cfi] needs-sanitizer-cfi -//@ [kcfi] needs-sanitizer-kcfi -//@ compile-flags: -C target-feature=-crt-static -//@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0 -//@ [cfi] compile-flags: -Z sanitizer=cfi -//@ [kcfi] compile-flags: -Z sanitizer=kcfi -//@ [kcfi] compile-flags: -C panic=abort -C prefer-dynamic=off -//@ 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..54f1cc035bc5b --- /dev/null +++ b/tests/ui/sanitizer/cfi-closures.rs @@ -0,0 +1,83 @@ +// Check various forms of dynamic closure calls + +//@ revisions: cfi kcfi +// FIXME(#122848) Remove only-linux once OSX CFI binaries work +//@ only-linux +//@ [cfi] needs-sanitizer-cfi +//@ [kcfi] needs-sanitizer-kcfi +//@ compile-flags: -C target-feature=-crt-static +//@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0 +//@ [cfi] compile-flags: -Z sanitizer=cfi +//@ [kcfi] compile-flags: -Z sanitizer=kcfi +//@ [kcfi] compile-flags: -C panic=abort -Z panic-abort-tests -C prefer-dynamic=off +//@ compile-flags: --test +//@ run-pass + +#![feature(fn_traits)] +#![feature(unboxed_closures)] +#![feature(cfg_sanitize)] + +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); +} + +fn hrtb_helper(f: &dyn for<'a> Fn(&'a usize)) { + f(&10) +} + +#[test] +fn hrtb_fn() { + hrtb_helper((&|x: &usize| println!("{}", *x)) as _) +} + +#[test] +fn fnonce() { + let f: Box = Box::new(|| {}) as _; + f(); +} + +fn use_closure(call: extern "rust-call" fn(&C, ()) -> i32, f: &C) -> i32 { + call(f, ()) +} + +#[test] +// FIXME after KCFI reify support is added, remove this +// It will appear to work if you test locally, set -C opt-level=0 to see it fail. +#[cfg_attr(sanitize = "kcfi", ignore)] +fn closure_addr_taken() { + let x = 3i32; + let f = || x; + let call = Fn::<()>::call; + use_closure(call, &f); +} diff --git a/tests/ui/sanitizer/cfi-coroutine.rs b/tests/ui/sanitizer/cfi-coroutine.rs new file mode 100644 index 0000000000000..24e59cf5b4d0f --- /dev/null +++ b/tests/ui/sanitizer/cfi-coroutine.rs @@ -0,0 +1,30 @@ +// Verifies that we can call dynamic coroutines + +//@ revisions: cfi kcfi +// FIXME(#122848) Remove only-linux once OSX CFI binaries work +//@ only-linux +//@ [cfi] needs-sanitizer-cfi +//@ [kcfi] needs-sanitizer-kcfi +//@ compile-flags: -C target-feature=-crt-static +//@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0 +//@ [cfi] compile-flags: -Z sanitizer=cfi +//@ [kcfi] compile-flags: -Z sanitizer=kcfi +//@ [kcfi] compile-flags: -C panic=abort -Z panic-abort-tests -C prefer-dynamic=off +//@ compile-flags: --test +//@ 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")); +}