Skip to content

Commit

Permalink
remove potential ub in render_resource_wrapper (#7279)
Browse files Browse the repository at this point in the history
# Objective

[as noted](#5950 (comment)) 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
  • Loading branch information
robtfm committed Feb 6, 2023
1 parent 4fd092f commit 8f81be9
Showing 1 changed file with 44 additions and 52 deletions.
96 changes: 44 additions & 52 deletions crates/bevy_render/src/render_resource/resource_macros.rs
Original file line number Diff line number Diff line change
@@ -1,90 +1,82 @@
// 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).
#[cfg(debug_assertions)]
#[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()
}
}

impl std::ops::Deref for $wrapper_type {
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)
}
}
};
}

Expand Down

0 comments on commit 8f81be9

Please sign in to comment.