From 8f81be984595db14ca3c8c20afbb14a9c503472e Mon Sep 17 00:00:00 2001 From: robtfm <50659922+robtfm@users.noreply.github.com> Date: Mon, 6 Feb 2023 17:14:00 +0000 Subject: [PATCH] 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::() == ptr`. as a happy side effect this removes the need for boxing the value, so todo: potentially use this for release mode as well --- .../src/render_resource/resource_macros.rs | 96 +++++++++---------- 1 file changed, 44 insertions(+), 52 deletions(-) diff --git a/crates/bevy_render/src/render_resource/resource_macros.rs b/crates/bevy_render/src/render_resource/resource_macros.rs index 26aaceb4fc92f..a827fcc7ba3fb 100644 --- a/crates/bevy_render/src/render_resource/resource_macros.rs +++ b/crates/bevy_render/src/render_resource/resource_macros.rs @@ -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>>); + #[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<$wgpu_type>>(untyped_box) - }; - Some(*typed_box) - } - Err(inner) => { - let _ = unsafe { - std::mem::transmute::< - std::sync::Arc>, - std::sync::Arc>, - >(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>, - std::sync::Arc>, - >(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> 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 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) + } + } }; }