`CombinatorSystem` is unsound

[SkiFire13]

Bevy version

0.14.1

What you did

use bevy_ecs::system::{CombinatorSystem, Combine, IntoSystem, Resource, System};
use bevy_ecs::world::World;
use scoped_tls_hkt::scoped_thread_local; // scoped_tls_hkt = "0.1.4" in Cargo.toml

#[derive(Clone, Copy)]
struct BadFn<'a>(&'a dyn Fn());

scoped_thread_local!(static BAD_FN: for<'a> BadFn<'a>);

struct BadCombine;

impl<A: System<In = (), Out = ()>, B: System<In = (), Out = ()>> Combine<A, B> for BadCombine {
    type In = ();
    type Out = ();

    fn combine(
        _input: Self::In,
        a: impl FnOnce(A::In) -> A::Out,
        b: impl FnOnce(B::In) -> B::Out,
    ) -> Self::Out {
        let a = std::cell::Cell::new(Some(a));
        let bad_fn = BadFn(Box::leak(Box::new(move || (a.take().unwrap())(()))));
        BAD_FN.set(bad_fn, || b(()));
    }
}

#[derive(Resource)]
struct Foo(String);

fn a(world: &mut World) {
    world.remove_resource::<Foo>();
}

fn b(world: &mut World) {
    world.insert_resource(Foo("foo".to_string()));
    let foo = world.resource::<Foo>();
    BAD_FN.with(|bad_fn| (bad_fn.0)());
    println!("{}", foo.0);
}

fn main() {
    let mut world = World::new();

    let a_system = IntoSystem::into_system(a);
    let b_system = IntoSystem::into_system(b);
    let mut combinator_system =
        CombinatorSystem::<BadCombine, _, _>::new(a_system, b_system, "".into());

    combinator_system.initialize(&mut world);
    combinator_system.run((), &mut world);
}

What went wrong

This should either fail to compile or throw an error at runtime. Instead, this produced UB due to b holding a reference to a resource that has been removed. MIRI also flags this as UB:

``` error: Undefined Behavior: not granting access to tag <24319> because that would remove [Unique for <19222>] which is strongly protected because it is an argument of call 4818 --> E:\Programmi\Rust\.cargo\registry\src\index.crates.io-6f17d22bba15001f\bevy_ptr-0.14.1\src/lib.rs:574:18 | 574 | unsafe { &mut *self.get() } | ^^^^^^^^^^^^^^^^ not granting access to tag <24319> because that would remove [Unique for <19222>] which is strongly protected because it is an argument of call 4818 | = help: this indicates a potential bug in the program: it performed an invalid operation, but the Stacked Borrows rules it violated are still experimental = help: see https://github.com/rust-lang/unsafe-code-guidelines/blob/master/wip/stacked-borrows.md for further information help: <24319> was created by a SharedReadWrite retag at offsets [0x0..0x458] --> src/main.rs:22:55 | 22 | let bad_fn = BadFn(Box::leak(Box::new(move || (a.take().unwrap())(())))); | ^^^^^^^^^^^^^^^^^^^^^^^ help: <19222> is this argument --> src/main.rs:34:6 | 34 | fn b(world: &mut World) { | ^^^^^ = note: BACKTRACE (of the first span): = note: inside `<&std::cell::UnsafeCell as bevy_ecs::bevy_ptr::UnsafeCellDeref<'_, bevy_ecs::world::World>>::deref_mut` at E:\Programmi\Rust\.cargo\registry\src\index.crates.io-6f17d22bba15001f\bevy_ptr-0.14.1\src/lib.rs:574:18: 574:34 = note: inside closure at E:\Programmi\Rust\.cargo\registry\src\index.crates.io-6f17d22bba15001f\bevy_ecs-0.14.1\src\system\combinator.rs:194:48: 194:65 note: inside closure --> src/main.rs:22:55 | 22 | let bad_fn = BadFn(Box::leak(Box::new(move || (a.take().unwrap())(())))); | ^^^^^^^^^^^^^^^^^^^^^^^ note: inside closure --> src/main.rs:37:26 | 37 | BAD_FN.with(|bad_fn| (bad_fn.0)()); | ^^^^^^^^^^^^ note: inside `BAD_FN::>::with::<{closure@src/main.rs:37:17: 37:25}, ()>` --> src/main.rs:8:1 | 8 | scoped_thread_local!(static BAD_FN: for<'a> BadFn<'a>); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ note: inside `b` --> src/main.rs:37:5 | 37 | BAD_FN.with(|bad_fn| (bad_fn.0)()); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ = note: inside ` fn(&'a mut bevy_ecs::world::World) {b} as std::ops::FnMut<(&mut bevy_ecs::world::World,)>>::call_mut - shim(for<'a> fn(&'a mut bevy_ecs::world::World) {b})` at E:\Programmi\Rust\.rustup\toolchains\nightly-x86_64-pc-windows-msvc\lib\rustlib\src\rust\library\core\src\ops\function.rs:166:5: 166:75 = note: inside `std::ops::function::impls:: for &mut for<'a> fn(&'a mut bevy_ecs::world::World) {b}>::call_mut` at E:\Programmi\Rust\.rustup\toolchains\nightly-x86_64-pc-windows-msvc\lib\rustlib\src\rust\library\core\src\ops\function.rs:294:13: 294:35 = note: inside ` Out>>::run::call_inner::<(), &mut for<'a> fn(&'a mut bevy_ecs::world::World) {b}>` at E:\Programmi\Rust\.cargo\registry\src\index.crates.io-6f17d22bba15001f\bevy_ecs-0.14.1\src\system\exclusive_function_system.rs:218:21: 218:42 = note: inside ` fn(&'a mut bevy_ecs::world::World) {b} as bevy_ecs::system::ExclusiveSystemParamFunction>::run` at E:\Programmi\Rust\.cargo\registry\src\index.crates.io-6f17d22bba15001f\bevy_ecs-0.14.1\src\system\exclusive_function_system.rs:221:17: 221:53 = note: inside closure at E:\Programmi\Rust\.cargo\registry\src\index.crates.io-6f17d22bba15001f\bevy_ecs-0.14.1\src\system\exclusive_function_system.rs:111:23: 111:58 = note: inside `bevy_ecs::world::World::last_change_tick_scope::<(), {closure@ fn(&'a mut bevy_ecs::world::World) {b}> as bevy_ecs::system::System>::run::{closure#0}}>` at E:\Programmi\Rust\.cargo\registry\src\index.crates.io-6f17d22bba15001f\bevy_ecs-0.14.1\src\world\mod.rs:2215:9: 2215:23 = note: inside ` fn(&'a mut bevy_ecs::world::World) {b}> as bevy_ecs::system::System>::run` at E:\Programmi\Rust\.cargo\registry\src\index.crates.io-6f17d22bba15001f\bevy_ecs-0.14.1\src\system\exclusive_function_system.rs:103:9: 119:11 = note: inside closure at E:\Programmi\Rust\.cargo\registry\src\index.crates.io-6f17d22bba15001f\bevy_ecs-0.14.1\src\system\combinator.rs:196:21: 196:68 note: inside closure --> src/main.rs:23:31 | 23 | BAD_FN.set(bad_fn, || b(())); | ^^^^^ note: inside `BAD_FN::>::set::<{closure@src/main.rs:23:28: 23:30}, ()>` --> src/main.rs:8:1 | 8 | scoped_thread_local!(static BAD_FN: for<'a> BadFn<'a>); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ note: inside ` fn(&'a mut bevy_ecs::world::World) {a}>, bevy_ecs::system::ExclusiveFunctionSystem fn(&'a mut bevy_ecs::world::World) {b}>>>::combine::<{closure@ fn(&'a mut bevy_ecs::world::World) {a}>, bevy_ecs::system::ExclusiveFunctionSystem fn(&'a mut bevy_ecs::world::World) {b}>> as bevy_ecs::system::System>::run::{closure#0}}, {closure@ fn(&'a mut bevy_ecs::world::World) {a}>, bevy_ecs::system::ExclusiveFunctionSystem fn(&'a mut bevy_ecs::world::World) {b}>> as bevy_ecs::system::System>::run::{closure#1}}>` --> src/main.rs:23:9 | 23 | BAD_FN.set(bad_fn, || b(())); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ = note: inside ` fn(&'a mut bevy_ecs::world::World) {a}>, bevy_ecs::system::ExclusiveFunctionSystem fn(&'a mut bevy_ecs::world::World) {b}>> as bevy_ecs::system::System>::run` at E:\Programmi\Rust\.cargo\registry\src\index.crates.io-6f17d22bba15001f\bevy_ecs-0.14.1\src\system\combinator.rs:189:9: 197:10 note: inside `main` --> src/main.rs:50:5 | 50 | combinator_system.run((), &mut world); | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ = note: this error originates in the macro `scoped_thread_local` (in Nightly builds, run with -Z macro-backtrace for more info) ```

Additional information

The issue is that CombinatorSystem gives to the Func::combine implementation two closures, each of which holds an UnsafeWorldCell that assumes is valid for calling that system:

https://github.com/bevyengine/bevy/blob/91fa4bb64905121a18b40df0062dbd85714aa3ce/crates/bevy_ecs/src/system/combinator.rs#L185-L198

The safety comment assumes that since the functions are !Send + !Sync + !'static then they cannot possibly be called in a parallel way. However this fails to see that they can still be called in a reentrant way, that is one function calling the other while it's running. This effectively also creates overlapping references to the world or its contents and thus leads to UB.

Actually exploiting this issue is not that easy though, as can be seen in the code above, since one closure cannot be smuggled directly to the other, thus having to go through some global state. This might seem that it requires 'static (and in fact using only the stdlib is does), but there are libraries that allow doing this like scoped-tls-hkt, which allows to store non-'static data in a thread local, ensuring it will be overwritten before it goes out of scope. With this one closure can actually be smuggled into the other and lead to the overlapping access to the world and thus UB.

[hymm]

The easy fix here might be to make Combine an unsafe trait.

[SkiFire13]

A more complex alternative which would however keep Combine safe would be to change Combine's parameters to a single one that allows running either A or B, and returns a handle to call the other one.

It would look something like this (name totally up for discussions):

struct RunBoth<'a, A, B> {
    world: UnsafeWorldCell<'a>,
    a: &'a mut A,
    b: &'a mut A,
}

impl<'a, A; System, B: System> RunBoth<'a, A, B> {
    fn run_a(self, input: A::In) -> (A::Out, RunOne<'a, B>) {
        let output = unsafe { self.a.run_unsafe(input, self.world) };
        (output, RunBoth { world: self.world, s: self.b })
    }
    fn run_b(self, input: B::In) -> (B::Out, RunOne<'a, A>) {
        let output = unsafe { self.b.run_unsafe(input, self.world) };
        (output, RunBoth { world: self.world, s: self.a })
    }
}

struct RunOne<'a, S> {
    world: UnsafeWorldCell<'a>,
    s: &'a mut S
}

impl<'a, S: System> RunOne<'a, S> {
    fn run(self, input: S::In) -> S::Out {
        unsafe { self.s.run_unsafe(input, self.world) }
    }
}

pub trait Combine<A: System, B: System> {
    type In;
    type Out;

    fn combine(
        input: Self::In,
        systems: RunBoth<'_, A, B>,
    ) -> Self::Out;
}