remove potential ub in render_resource_wrapper (#7279)

# Objective

[as noted](https://github.com/bevyengine/bevy/pull/5950#discussion_r1080762807) by james, transmuting arcs may be UB.
 
we now store a `*const ()` pointer internally, and only rely on `ptr.cast::<()>().cast::<T>() == ptr`.

as a happy side effect this removes the need for boxing the value, so todo: potentially use this for release mode as well
This commit is contained in:
robtfm 2023-02-06 17:14:00 +00:00
parent 4fd092fbec
commit 8f81be9845

View File

@ -1,6 +1,7 @@
// structs containing wgpu types take a long time to compile. this is particularly bad for generic
// structs containing wgpu structs. we avoid that in debug builds (and for cargo check and rust analyzer)
// by boxing and type-erasing with the `render_resource_wrapper` macro.
// by type-erasing with the `render_resource_wrapper` macro. The resulting type behaves like Arc<$wgpu_type>,
// but avoids explicitly storing an Arc<$wgpu_type> member.
// analysis from https://github.com/bevyengine/bevy/pull/5950#issuecomment-1243473071 indicates this is
// due to `evaluate_obligations`. we should check if this can be removed after a fix lands for
// https://github.com/rust-lang/rust/issues/99188 (and after other `evaluate_obligations`-related changes).
@ -8,41 +9,27 @@
#[macro_export]
macro_rules! render_resource_wrapper {
($wrapper_type:ident, $wgpu_type:ty) => {
#[derive(Clone, Debug)]
pub struct $wrapper_type(Option<std::sync::Arc<Box<()>>>);
#[derive(Debug)]
// SAFETY: while self is live, self.0 comes from `into_raw` of an Arc<$wgpu_type> with a strong ref.
pub struct $wrapper_type(*const ());
impl $wrapper_type {
pub fn new(value: $wgpu_type) -> Self {
unsafe {
Self(Some(std::sync::Arc::new(std::mem::transmute(Box::new(
value,
)))))
}
let arc = std::sync::Arc::new(value);
let value_ptr = std::sync::Arc::into_raw(arc);
let unit_ptr = value_ptr.cast::<()>();
Self(unit_ptr)
}
pub fn try_unwrap(mut self) -> Option<$wgpu_type> {
let inner = self.0.take();
if let Some(inner) = inner {
match std::sync::Arc::try_unwrap(inner) {
Ok(untyped_box) => {
let typed_box = unsafe {
std::mem::transmute::<Box<()>, Box<$wgpu_type>>(untyped_box)
};
Some(*typed_box)
}
Err(inner) => {
let _ = unsafe {
std::mem::transmute::<
std::sync::Arc<Box<()>>,
std::sync::Arc<Box<$wgpu_type>>,
>(inner)
};
None
}
}
} else {
None
}
pub fn try_unwrap(self) -> Option<$wgpu_type> {
let value_ptr = self.0.cast::<$wgpu_type>();
// SAFETY: pointer refers to a valid Arc, and was created from Arc::into_raw.
let arc = unsafe { std::sync::Arc::from_raw(value_ptr) };
// we forget ourselves here since the reconstructed arc will be dropped/decremented within this scope
std::mem::forget(self);
std::sync::Arc::try_unwrap(arc).ok()
}
}
@ -50,41 +37,46 @@ macro_rules! render_resource_wrapper {
type Target = $wgpu_type;
fn deref(&self) -> &Self::Target {
let untyped_box = self
.0
.as_ref()
.expect("render_resource_wrapper inner value has already been taken (via drop or try_unwrap")
.as_ref();
let typed_box =
unsafe { std::mem::transmute::<&Box<()>, &Box<$wgpu_type>>(untyped_box) };
typed_box.as_ref()
let value_ptr = self.0.cast::<$wgpu_type>();
// SAFETY: the arc lives for 'self, so the ref lives for 'self
let value_ref = unsafe { value_ptr.as_ref() };
value_ref.unwrap()
}
}
impl Drop for $wrapper_type {
fn drop(&mut self) {
let inner = self.0.take();
if let Some(inner) = inner {
let _ = unsafe {
std::mem::transmute::<
std::sync::Arc<Box<()>>,
std::sync::Arc<Box<$wgpu_type>>,
>(inner)
};
}
let value_ptr = self.0.cast::<$wgpu_type>();
// SAFETY: pointer refers to a valid Arc, and was created from Arc::into_raw.
// this reconstructed arc is dropped/decremented within this scope.
unsafe { std::sync::Arc::from_raw(value_ptr) };
}
}
// Arc<Box<()>> and Arc<()> will be Sync and Send even when $wgpu_type is not Sync or Send.
// SAFETY: We manually implement Send and Sync, which is valid for Arc<T> when T: Send + Sync.
// We ensure correctness by checking that $wgpu_type does implement Send and Sync.
// If in future there is a case where a wrapper is required for a non-send/sync type
// we can implement a macro variant that also does `impl !Send for $wrapper_type {}` and
// `impl !Sync for $wrapper_type {}`
// we can implement a macro variant that omits these manual Send + Sync impls
unsafe impl Send for $wrapper_type {}
unsafe impl Sync for $wrapper_type {}
const _: () = {
trait AssertSendSyncBound: Send + Sync {}
impl AssertSendSyncBound for $wgpu_type {}
};
impl Clone for $wrapper_type {
fn clone(&self) -> Self {
let value_ptr = self.0.cast::<$wgpu_type>();
// SAFETY: pointer refers to a valid Arc, and was created from Arc::into_raw.
let arc = unsafe { std::sync::Arc::from_raw(value_ptr.cast::<$wgpu_type>()) };
let cloned = std::sync::Arc::clone(&arc);
// we forget the reconstructed Arc to avoid decrementing the ref counter, as self is still live.
std::mem::forget(arc);
let cloned_value_ptr = std::sync::Arc::into_raw(cloned);
let cloned_unit_ptr = cloned_value_ptr.cast::<()>();
Self(cloned_unit_ptr)
}
}
};
}