Skip to content

Commit

Permalink
CFI: Rewrite closure and coroutine instances to their trait method
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
maurer committed Mar 26, 2024
1 parent 536606b commit ad0b8be
Show file tree
Hide file tree
Showing 4 changed files with 100 additions and 18 deletions.
Expand Up @@ -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::{
Expand Down Expand Up @@ -1148,6 +1149,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
Expand Down
18 changes: 0 additions & 18 deletions tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs

This file was deleted.

45 changes: 45 additions & 0 deletions 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<dyn Fn(&'a T) -> &'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<F: FnMut()>(mut f: F) {
f()
}

#[test]
fn fn_to_fnmut() {
let f: &(dyn Fn()) = &(|| {}) as _;
use_fnmut(f);
}
23 changes: 23 additions & 0 deletions 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<i32,Yield=i32,Return=&'static str>> = pin!(coro);
assert_eq!(abstract_coro.as_mut().resume(2), CoroutineState::Yielded(2));
assert_eq!(abstract_coro.as_mut().resume(0), CoroutineState::Complete("done"));
}

0 comments on commit ad0b8be

Please sign in to comment.