Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 9802478

Browse files
hyf0hyf0
authored
fix(rust/rolldown): prevent segment fault in generate_bundle by removing unsafe code (#1905)
<!-- Thank you for contributing! --> ### Description Another fix compared to #1904 for the segment fault problem. <!-- Please insert your description here and provide especially info about the "what" this PR is solving -->
1 parent d879a8c commit 9802478

6 files changed

Lines changed: 123 additions & 41 deletions

File tree

‎crates/rolldown_binding/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,3 +24,4 @@ pub mod types;
2424
pub mod utils;
2525
mod worker_manager;
2626
pub use oxc_transform_napi;
27+
mod type_aliases;

‎crates/rolldown_binding/src/options/plugin/js_plugin.rs‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,13 @@ use crate::types::{
33
js_callback::MaybeAsyncJsCallbackExt,
44
};
55
use rolldown_plugin::{Plugin, __inner::SharedPluginable, typedmap::TypedMapKey};
6-
use std::{borrow::Cow, ops::Deref, sync::Arc};
6+
use rolldown_utils::unique_arc::UniqueArc;
7+
use std::{
8+
borrow::Cow,
9+
mem,
10+
ops::Deref,
11+
sync::{Arc, Mutex},
12+
};
713

814
use super::{
915
binding_transform_context::BindingTransformPluginContext,
@@ -295,12 +301,10 @@ impl Plugin for JsPlugin {
295301
is_write: bool,
296302
) -> rolldown_plugin::HookNoopReturn {
297303
if let Some(cb) = &self.generate_bundle {
298-
cb.await_call((
299-
Arc::clone(ctx).into(),
300-
BindingOutputs::new(unsafe { std::mem::transmute(bundle) }),
301-
is_write,
302-
))
303-
.await?;
304+
let old_bundle = UniqueArc::new(Mutex::new(mem::take(bundle)));
305+
cb.await_call((Arc::clone(ctx).into(), BindingOutputs::new(old_bundle.weak_ref()), is_write))
306+
.await?;
307+
*bundle = old_bundle.into_inner().into_inner()?;
304308
}
305309
Ok(())
306310
}
@@ -311,11 +315,9 @@ impl Plugin for JsPlugin {
311315
bundle: &mut Vec<rolldown_common::Output>,
312316
) -> rolldown_plugin::HookNoopReturn {
313317
if let Some(cb) = &self.write_bundle {
314-
cb.await_call((
315-
Arc::clone(ctx).into(),
316-
BindingOutputs::new(unsafe { std::mem::transmute(bundle) }),
317-
))
318-
.await?;
318+
let old_bundle = UniqueArc::new(Mutex::new(mem::take(bundle)));
319+
cb.await_call((Arc::clone(ctx).into(), BindingOutputs::new(old_bundle.weak_ref()))).await?;
320+
*bundle = old_bundle.into_inner().into_inner()?;
319321
}
320322
Ok(())
321323
}

‎crates/rolldown_binding/src/type_aliases.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
use std::sync::Mutex;
2+
3+
use rolldown_utils::unique_arc::{UniqueArc, WeakRef};
4+
5+
pub type UniqueArcMutex<T> = UniqueArc<Mutex<T>>;
6+
pub type WeakRefMutex<T> = WeakRef<Mutex<T>>;

‎crates/rolldown_binding/src/types/binding_outputs.rs‎

Lines changed: 44 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,32 @@
11
use napi_derive::napi;
22

3+
use crate::type_aliases::{UniqueArcMutex, WeakRefMutex};
4+
35
use super::{binding_output_asset::BindingOutputAsset, binding_output_chunk::BindingOutputChunk};
46

57
/// The `BindingOutputs` owner `Vec<Output>` the mutable reference, it avoid `Clone` at call `writeBundle/generateBundle` hook, and make it mutable.
68
#[napi]
79
pub struct BindingOutputs {
8-
inner: &'static mut Vec<rolldown_common::Output>,
10+
inner: WeakRefMutex<Vec<rolldown_common::Output>>,
911
}
1012

1113
#[napi]
1214
impl BindingOutputs {
13-
pub fn new(inner: &'static mut Vec<rolldown_common::Output>) -> Self {
15+
pub fn new(inner: WeakRefMutex<Vec<rolldown_common::Output>>) -> Self {
1416
Self { inner }
1517
}
1618

1719
#[napi(getter)]
1820
pub fn chunks(&mut self) -> Vec<BindingOutputChunk> {
1921
let mut chunks: Vec<BindingOutputChunk> = vec![];
20-
21-
self.inner.iter_mut().for_each(|o| match o {
22-
rolldown_common::Output::Chunk(chunk) => {
23-
chunks.push(BindingOutputChunk::new(unsafe { std::mem::transmute(chunk.as_mut()) }));
24-
}
25-
rolldown_common::Output::Asset(_) => {}
22+
self.inner.with_inner(|inner| {
23+
let mut inner = inner.lock().expect("PoisonError raised");
24+
inner.iter_mut().for_each(|o| match o {
25+
rolldown_common::Output::Chunk(chunk) => {
26+
chunks.push(BindingOutputChunk::new(unsafe { std::mem::transmute(chunk.as_mut()) }));
27+
}
28+
rolldown_common::Output::Asset(_) => {}
29+
});
2630
});
2731

2832
chunks
@@ -32,45 +36,53 @@ impl BindingOutputs {
3236
pub fn assets(&mut self) -> Vec<BindingOutputAsset> {
3337
let mut assets: Vec<BindingOutputAsset> = vec![];
3438

35-
self.inner.iter_mut().for_each(|o| match o {
36-
rolldown_common::Output::Asset(asset) => {
37-
assets.push(BindingOutputAsset::new(unsafe { std::mem::transmute(asset.as_mut()) }));
38-
}
39-
rolldown_common::Output::Chunk(_) => {}
39+
self.inner.with_inner(|inner| {
40+
let mut inner = inner.lock().expect("PoisonError raised");
41+
inner.iter_mut().for_each(|o| match o {
42+
rolldown_common::Output::Asset(asset) => {
43+
assets.push(BindingOutputAsset::new(unsafe { std::mem::transmute(asset.as_mut()) }));
44+
}
45+
rolldown_common::Output::Chunk(_) => {}
46+
});
4047
});
4148
assets
4249
}
4350

4451
#[napi]
4552
pub fn delete(&mut self, file_name: String) {
46-
if let Some(index) = self.inner.iter().position(|o| o.filename() == file_name) {
47-
self.inner.remove(index);
48-
}
53+
self.inner.with_inner(|inner| {
54+
let mut inner = inner.lock().expect("PoisonError raised");
55+
if let Some(index) = inner.iter().position(|o| o.filename() == file_name) {
56+
inner.remove(index);
57+
}
58+
});
4959
}
5060
}
5161

5262
/// The `FinalBindingOutputs` is used at `write()` or `generate()`, it is similar to `BindingOutputs`, if using `BindingOutputs` has unexpected behavior.
5363
/// TODO find a way to export it gracefully.
5464
#[napi]
5565
pub struct FinalBindingOutputs {
56-
inner: Vec<rolldown_common::Output>,
66+
inner: UniqueArcMutex<Vec<rolldown_common::Output>>,
5767
}
5868

5969
#[napi]
6070
impl FinalBindingOutputs {
6171
pub fn new(inner: Vec<rolldown_common::Output>) -> Self {
62-
Self { inner }
72+
Self { inner: UniqueArcMutex::new(inner.into()) }
6373
}
6474

6575
#[napi(getter)]
6676
pub fn chunks(&mut self) -> Vec<BindingOutputChunk> {
6777
let mut chunks: Vec<BindingOutputChunk> = vec![];
68-
69-
self.inner.iter_mut().for_each(|o| match o {
70-
rolldown_common::Output::Chunk(chunk) => {
71-
chunks.push(BindingOutputChunk::new(unsafe { std::mem::transmute(chunk.as_mut()) }));
72-
}
73-
rolldown_common::Output::Asset(_) => {}
78+
self.inner.weak_ref().with_inner(|inner| {
79+
let mut inner = inner.lock().expect("PoisonError raised");
80+
inner.iter_mut().for_each(|o| match o {
81+
rolldown_common::Output::Chunk(chunk) => {
82+
chunks.push(BindingOutputChunk::new(unsafe { std::mem::transmute(chunk.as_mut()) }));
83+
}
84+
rolldown_common::Output::Asset(_) => {}
85+
});
7486
});
7587

7688
chunks
@@ -80,11 +92,14 @@ impl FinalBindingOutputs {
8092
pub fn assets(&mut self) -> Vec<BindingOutputAsset> {
8193
let mut assets: Vec<BindingOutputAsset> = vec![];
8294

83-
self.inner.iter_mut().for_each(|o| match o {
84-
rolldown_common::Output::Asset(asset) => {
85-
assets.push(BindingOutputAsset::new(unsafe { std::mem::transmute(asset.as_mut()) }));
86-
}
87-
rolldown_common::Output::Chunk(_) => {}
95+
self.inner.weak_ref().with_inner(|inner| {
96+
let mut inner = inner.lock().expect("PoisonError raised");
97+
inner.iter_mut().for_each(|o| match o {
98+
rolldown_common::Output::Asset(asset) => {
99+
assets.push(BindingOutputAsset::new(unsafe { std::mem::transmute(asset.as_mut()) }));
100+
}
101+
rolldown_common::Output::Chunk(_) => {}
102+
});
88103
});
89104
assets
90105
}

‎crates/rolldown_utils/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,3 +17,4 @@ pub mod rustc_hash;
1717
pub mod sanitize_file_name;
1818
pub mod xxhash;
1919
pub use bitset::BitSet;
20+
pub mod unique_arc;

‎crates/rolldown_utils/src/unique_arc.rs‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
use std::sync::{Arc, Weak};
2+
3+
use crate::debug::pretty_type_name;
4+
5+
#[derive(Debug)]
6+
pub struct UniqueArc<T: ?Sized>(Arc<T>);
7+
8+
impl<T> UniqueArc<T> {
9+
pub fn new(value: T) -> Self {
10+
Self(Arc::new(value))
11+
}
12+
13+
/// # Panics
14+
/// - Don't call this method while calling the `WeakRef#with_inner` method.
15+
pub fn into_inner(self) -> T {
16+
Arc::into_inner(self.0).unwrap_or_else(|| {
17+
let t_name = pretty_type_name::<T>();
18+
panic!("UniqueArc<{t_name}> has multiple references")
19+
})
20+
}
21+
22+
pub fn weak_ref(&self) -> WeakRef<T> {
23+
WeakRef(Arc::downgrade(&self.0))
24+
}
25+
}
26+
27+
#[derive(Debug)]
28+
pub struct WeakRef<T: ?Sized>(Weak<T>);
29+
30+
impl<T> WeakRef<T> {
31+
/// This API pattern is intended to prevent users from getting `Arc<T>` and storing it.
32+
/// This will cause `UniqueArc#into_inner` panic if `Arc<T>` got stored.
33+
pub fn try_with_inner<F, R>(&self, f: F) -> Option<R>
34+
where
35+
F: FnOnce(&T) -> R,
36+
{
37+
self.0.upgrade().map(|arc| f(&*arc))
38+
}
39+
40+
pub fn with_inner<F, R>(&self, f: F) -> R
41+
where
42+
F: FnOnce(&T) -> R,
43+
{
44+
self.try_with_inner(f).unwrap_or_else(|| {
45+
let t_name = pretty_type_name::<T>();
46+
panic!(
47+
"UniqueArc<{t_name}> is already dropped. You can't access it by this WeakRef<{t_name}>",
48+
)
49+
})
50+
}
51+
}
52+
53+
impl<T> Clone for WeakRef<T> {
54+
fn clone(&self) -> Self {
55+
Self(Weak::clone(&self.0))
56+
}
57+
}

0 commit comments

Comments
 (0)