Skip to content

Commit

Permalink
Update may_enter flag handling in components (#4354)
Browse files Browse the repository at this point in the history 
This commit updates the management of the `may_enter` flag in line with
WebAssembly/component-model#57. Namely the `may_enter` flag is now
exclusively managed in the `canon lift` function (which is
`TypedFunc::call`) and is only unset after post-return completes
successfully. This implements semantics where if any trap happens for
any reason (lifting, lowering, execution, imports, etc) then the
instance is considered permanently poisoned and can no longer be
entered.

Tests needed many updates to create new instances where previously the
same instance was reused after it had an erroneous state.
  • Loading branch information
@alexcrichton
alexcrichton authored Jun 29, 2022
1 parent 816e7f7 commit e179e73
Show file tree
Hide file tree
Showing 5 changed files with 248 additions and 123 deletions.
33 changes: 5 additions & 28 deletions crates/wasmtime/src/component/func/host.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -154,10 +154,6 @@ where
bail!("cannot leave component instance");
}

// While we're lifting and lowering this instance cannot be reentered, so
// unset the flag here. This is also reset back to `true` on exit.
let _reset_may_enter = unset_and_reset_on_drop(flags, VMComponentFlags::set_may_enter);

// There's a 2x2 matrix of whether parameters and results are stored on the
// stack or on the heap. Each of the 4 branches here have a different
// representation of the storage of arguments/returns which is represented
Expand All @@ -168,21 +164,20 @@ where
// trivially DCE'd by LLVM. Perhaps one day with enough const programming in
// Rust we can make monomorphizations of this function codegen only one
// branch, but today is not that day.
let reset_may_leave;
if Params::flatten_count() <= MAX_STACK_PARAMS {
if Return::flatten_count() <= MAX_STACK_RESULTS {
let storage = cast_storage::<ReturnStack<Params::Lower, Return::Lower>>(storage);
let params = Params::lift(cx.0, &options, &storage.assume_init_ref().args)?;
let ret = closure(cx.as_context_mut(), params)?;
reset_may_leave = unset_and_reset_on_drop(flags, VMComponentFlags::set_may_leave);
(*flags).set_may_leave(false);
ret.lower(&mut cx, &options, map_maybe_uninit!(storage.ret))?;
} else {
let storage = cast_storage::<ReturnPointer<Params::Lower>>(storage).assume_init_ref();
let params = Params::lift(cx.0, &options, &storage.args)?;
let ret = closure(cx.as_context_mut(), params)?;
let mut memory = MemoryMut::new(cx.as_context_mut(), &options);
let ptr = validate_inbounds::<Return>(memory.as_slice_mut(), &storage.retptr)?;
reset_may_leave = unset_and_reset_on_drop(flags, VMComponentFlags::set_may_leave);
(*flags).set_may_leave(false);
ret.store(&mut memory, ptr)?;
}
} else {
Expand All @@ -193,7 +188,7 @@ where
validate_inbounds::<Params>(memory.as_slice(), &storage.assume_init_ref().args)?;
let params = Params::load(&memory, &memory.as_slice()[ptr..][..Params::SIZE32])?;
let ret = closure(cx.as_context_mut(), params)?;
reset_may_leave = unset_and_reset_on_drop(flags, VMComponentFlags::set_may_leave);
(*flags).set_may_leave(false);
ret.lower(&mut cx, &options, map_maybe_uninit!(storage.ret))?;
} else {
let storage = cast_storage::<ReturnPointer<ValRaw>>(storage).assume_init_ref();
Expand All @@ -202,32 +197,14 @@ where
let ret = closure(cx.as_context_mut(), params)?;
let mut memory = MemoryMut::new(cx.as_context_mut(), &options);
let ptr = validate_inbounds::<Return>(memory.as_slice_mut(), &storage.retptr)?;
reset_may_leave = unset_and_reset_on_drop(flags, VMComponentFlags::set_may_leave);
(*flags).set_may_leave(false);
ret.store(&mut memory, ptr)?;
}
}

drop(reset_may_leave);
(*flags).set_may_leave(true);

return Ok(());

unsafe fn unset_and_reset_on_drop(
slot: *mut VMComponentFlags,
set: fn(&mut VMComponentFlags, bool),
) -> impl Drop {
set(&mut *slot, false);
return Reset(slot, set);

struct Reset(*mut VMComponentFlags, fn(&mut VMComponentFlags, bool));

impl Drop for Reset {
fn drop(&mut self) {
unsafe {
(self.1)(&mut *self.0, true);
}
}
}
}
}

fn validate_inbounds<T: ComponentType>(memory: &[u8], ptr: &ValRaw) -> Result<usize> {
Expand Down
93 changes: 44 additions & 49 deletions crates/wasmtime/src/component/func/typed.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -333,11 +333,19 @@ where
let flags = instance.flags(component_instance);

unsafe {
// Test the "may enter" flag which is a "lock" on this instance.
// This is immediately set to `false` afterwards and note that
// there's no on-cleanup setting this flag back to true. That's an
// intentional design aspect where if anything goes wrong internally
// from this point on the instance is considered "poisoned" and can
// never be entered again. The only time this flag is set to `true`
// again is after post-return logic has completed successfully.
if !(*flags).may_enter() {
bail!("cannot reenter component instance");
}
debug_assert!((*flags).may_leave());
(*flags).set_may_enter(false);

debug_assert!((*flags).may_leave());
(*flags).set_may_leave(false);
let result = lower(store, &options, params, map_maybe_uninit!(space.params));
(*flags).set_may_leave(true);
Expand Down Expand Up @@ -370,41 +378,21 @@ where
// Lift the result into the host while managing post-return state
// here as well.
//
// Initially the `may_enter` flag is set to `false` for this
// component instance and additionally we set a flag indicating that
// a post-return is required. This isn't specified by the component
// model itself but is used for our implementation of the API of
// `post_return` as a separate function call.
//
// FIXME(WebAssembly/component-model#55) it's not really clear what
// the semantics should be in the face of a lift error/trap. For now
// the flags are reset so the instance can continue to be reused in
// tests but that probably isn't what's desired.
//
// Otherwise though after a successful lift the return value of the
// function, which is currently required to be 0 or 1 values
// according to the canonical ABI, is saved within the `Store`'s
// `FuncData`. This'll later get used in post-return.
(*flags).set_may_enter(false);
// After a successful lift the return value of the function, which
// is currently required to be 0 or 1 values according to the
// canonical ABI, is saved within the `Store`'s `FuncData`. This'll
// later get used in post-return.
(*flags).set_needs_post_return(true);
match lift(store.0, &options, ret) {
Ok(val) => {
let ret_slice = cast_storage(ret);
let data = &mut store.0[self.func.0];
assert!(data.post_return_arg.is_none());
match ret_slice.len() {
0 => data.post_return_arg = Some(ValRaw::i32(0)),
1 => data.post_return_arg = Some(ret_slice[0]),
_ => unreachable!(),
}
return Ok(val);
}
Err(err) => {
(*flags).set_may_enter(true);
(*flags).set_needs_post_return(false);
return Err(err);
}
let val = lift(store.0, &options, ret)?;
let ret_slice = cast_storage(ret);
let data = &mut store.0[self.func.0];
assert!(data.post_return_arg.is_none());
match ret_slice.len() {
0 => data.post_return_arg = Some(ValRaw::i32(0)),
1 => data.post_return_arg = Some(ret_slice[0]),
_ => unreachable!(),
}
return Ok(val);
}

unsafe fn cast_storage<T>(storage: &T) -> &[ValRaw] {
Expand Down Expand Up @@ -480,23 +468,30 @@ where
// This is a sanity-check assert which shouldn't ever trip.
assert!(!(*flags).may_enter());

// With the state of the world validated these flags are updated to
// their component-model-defined states.
(*flags).set_may_enter(true);
// Unset the "needs post return" flag now that post-return is being
// processed. This will cause future invocations of this method to
// panic, even if the function call below traps.
(*flags).set_needs_post_return(false);

// And finally if the function actually had a `post-return`
// configured in its canonical options that's executed here.
let (func, trampoline) = match post_return {
Some(pair) => pair,
None => return Ok(()),
};
crate::Func::call_unchecked_raw(
&mut store,
func.anyfunc,
trampoline,
&post_return_arg as *const ValRaw as *mut ValRaw,
)?;
// If the function actually had a `post-return` configured in its
// canonical options that's executed here.
//
// Note that if this traps (returns an error) this function
// intentionally leaves the instance in a "poisoned" state where it
// can no longer be entered because `may_enter` is `false`.
if let Some((func, trampoline)) = post_return {
crate::Func::call_unchecked_raw(
&mut store,
func.anyfunc,
trampoline,
&post_return_arg as *const ValRaw as *mut ValRaw,
)?;
}

// And finally if everything completed successfully then the "may
// enter" flag is set to `true` again here which enables further use
// of the component.
(*flags).set_may_enter(true);
}
Ok(())
}
Expand Down
Loading

0 comments on commit e179e73

Please sign in to comment.