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 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) // 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 // 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 // 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). // https://github.com/rust-lang/rust/issues/99188 (and after other `evaluate_obligations`-related changes).
@ -8,41 +9,27 @@
#[macro_export] #[macro_export]
macro_rules! render_resource_wrapper { macro_rules! render_resource_wrapper {
($wrapper_type:ident, $wgpu_type:ty) => { ($wrapper_type:ident, $wgpu_type:ty) => {
#[derive(Clone, Debug)] #[derive(Debug)]
pub struct $wrapper_type(Option<std::sync::Arc<Box<()>>>); // 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 { impl $wrapper_type {
pub fn new(value: $wgpu_type) -> Self { pub fn new(value: $wgpu_type) -> Self {
unsafe { let arc = std::sync::Arc::new(value);
Self(Some(std::sync::Arc::new(std::mem::transmute(Box::new( let value_ptr = std::sync::Arc::into_raw(arc);
value, let unit_ptr = value_ptr.cast::<()>();
))))) Self(unit_ptr)
}
} }
pub fn try_unwrap(mut self) -> Option<$wgpu_type> { pub fn try_unwrap(self) -> Option<$wgpu_type> {
let inner = self.0.take(); let value_ptr = self.0.cast::<$wgpu_type>();
if let Some(inner) = inner { // SAFETY: pointer refers to a valid Arc, and was created from Arc::into_raw.
match std::sync::Arc::try_unwrap(inner) { let arc = unsafe { std::sync::Arc::from_raw(value_ptr) };
Ok(untyped_box) => {
let typed_box = unsafe { // we forget ourselves here since the reconstructed arc will be dropped/decremented within this scope
std::mem::transmute::<Box<()>, Box<$wgpu_type>>(untyped_box) std::mem::forget(self);
};
Some(*typed_box) std::sync::Arc::try_unwrap(arc).ok()
}
Err(inner) => {
let _ = unsafe {
std::mem::transmute::<
std::sync::Arc<Box<()>>,
std::sync::Arc<Box<$wgpu_type>>,
>(inner)
};
None
}
}
} else {
None
}
} }
} }
@ -50,41 +37,46 @@ macro_rules! render_resource_wrapper {
type Target = $wgpu_type; type Target = $wgpu_type;
fn deref(&self) -> &Self::Target { fn deref(&self) -> &Self::Target {
let untyped_box = self let value_ptr = self.0.cast::<$wgpu_type>();
.0 // SAFETY: the arc lives for 'self, so the ref lives for 'self
.as_ref() let value_ref = unsafe { value_ptr.as_ref() };
.expect("render_resource_wrapper inner value has already been taken (via drop or try_unwrap") value_ref.unwrap()
.as_ref();
let typed_box =
unsafe { std::mem::transmute::<&Box<()>, &Box<$wgpu_type>>(untyped_box) };
typed_box.as_ref()
} }
} }
impl Drop for $wrapper_type { impl Drop for $wrapper_type {
fn drop(&mut self) { fn drop(&mut self) {
let inner = self.0.take(); let value_ptr = self.0.cast::<$wgpu_type>();
if let Some(inner) = inner { // SAFETY: pointer refers to a valid Arc, and was created from Arc::into_raw.
let _ = unsafe { // this reconstructed arc is dropped/decremented within this scope.
std::mem::transmute::< unsafe { std::sync::Arc::from_raw(value_ptr) };
std::sync::Arc<Box<()>>,
std::sync::Arc<Box<$wgpu_type>>,
>(inner)
};
}
} }
} }
// 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. // 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 // 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 // we can implement a macro variant that omits these manual Send + Sync impls
// `impl !Sync for $wrapper_type {}` unsafe impl Send for $wrapper_type {}
unsafe impl Sync for $wrapper_type {}
const _: () = { const _: () = {
trait AssertSendSyncBound: Send + Sync {} trait AssertSendSyncBound: Send + Sync {}
impl AssertSendSyncBound for $wgpu_type {} 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)
}
}
}; };
} }