Skip to content

Commit

Permalink
fix: soundness issue in Somr
Browse files Browse the repository at this point in the history
  • Loading branch information
zxch3n committed Aug 26, 2023
1 parent 39aaace commit cb78ffe
Show file tree
Hide file tree
Showing 3 changed files with 82 additions and 23 deletions.
89 changes: 74 additions & 15 deletions y-octo/src/doc/common/somr.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
use std::{
cell::UnsafeCell,
fmt::{self, Write},
hash::{Hash, Hasher},
marker::PhantomData,
mem,
ops::{Deref, DerefMut},
ptr::NonNull,
};

Expand All @@ -25,13 +27,32 @@ pub(crate) struct Owned<T>(NonNull<SomrInner<T>>);
pub(crate) struct Ref<T>(NonNull<SomrInner<T>>);

pub(crate) struct SomrInner<T> {
data: Option<T>,
data: Option<UnsafeCell<T>>,
/// increase the size when we really meet the the secenerio with refs more
/// then u16::MAX(65535) times
refs: AtomicU32,
_marker: PhantomData<Option<T>>,
}

pub(crate) struct InnerRefMut<'a, T> {
inner: NonNull<T>,
_marker: PhantomData<&'a mut T>,
}

impl<'a, T> Deref for InnerRefMut<'a, T> {
type Target = T;

fn deref(&self) -> &Self::Target {
unsafe { &*self.inner.as_ptr() }
}
}

impl<'a, T> DerefMut for InnerRefMut<'a, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { &mut *self.inner.as_ptr() }
}
}

unsafe impl<T: Send> Send for Somr<T> {}
unsafe impl<T: Sync> Sync for Somr<T> {}

Expand All @@ -44,7 +65,7 @@ impl<T> Default for Somr<T> {
impl<T> Somr<T> {
pub fn new(data: T) -> Self {
let inner = Box::new(SomrInner {
data: Some(data),
data: Some(UnsafeCell::new(data)),
refs: AtomicU32::new(1),
_marker: PhantomData,
});
Expand All @@ -57,55 +78,77 @@ impl<T> Somr<T> {
}
}

impl<T> SomrInner<T> {
fn data_ref(&self) -> Option<&T> {
self.data.as_ref().map(|x| unsafe { &*x.get() })
}

fn data_mut(&self) -> Option<InnerRefMut<T>> {
self.data.as_ref().map(|x| InnerRefMut {
inner: unsafe { NonNull::new_unchecked(x.get()) },
_marker: PhantomData,
})
}
}

impl<T> Somr<T> {
#[inline]
pub fn is_none(&self) -> bool {
self.dangling() || self.inner().data.is_none()
self.dangling() || self.inner().data_ref().is_none()
}

#[inline]
pub fn is_some(&self) -> bool {
!self.dangling() && self.inner().data.is_some()
!self.dangling() && self.inner().data_ref().is_some()
}

pub fn get(&self) -> Option<&T> {
if self.dangling() {
return None;
}

self.inner().data.as_ref()
self.inner().data_ref()
}

pub unsafe fn get_unchecked(&self) -> &T {
if self.dangling() {
panic!("Try to visit Somr data that has already been dropped.")
}

match &self.inner().data {
match &self.inner().data_ref() {
Some(data) => data,
None => {
panic!("Try to unwrap on None")
}
}
}

#[allow(dead_code)]
pub fn get_mut(&self) -> Option<&mut T> {
#[allow(unused)]
pub fn get_mut(&mut self) -> Option<&mut T> {
if !self.is_owned() || self.dangling() {
return None;
}

let inner = self.inner_mut();
inner.data.as_mut()
inner.data.as_mut().map(|x| x.get_mut())
}

#[allow(clippy::mut_from_ref)]
pub unsafe fn get_mut_unchecked(&self) -> &mut T {
#[allow(unused)]
pub unsafe fn get_mut_from_ref(&self) -> Option<InnerRefMut<T>> {
if !self.is_owned() || self.dangling() {
return None;
}

let inner = self.inner_mut();
inner.data_mut()
}

pub unsafe fn get_mut_unchecked(&self) -> InnerRefMut<'_, T> {
if self.dangling() {
panic!("Try to visit Somr data that has already been dropped.")
}

match &mut self.inner_mut().data {
match self.inner_mut().data_mut() {
Some(data) => data,
None => {
panic!("Try to unwrap on None")
Expand Down Expand Up @@ -238,7 +281,7 @@ impl<T: PartialEq> PartialEq for Somr<T> {

impl<T: PartialEq> PartialEq for SomrInner<T> {
fn eq(&self, other: &Self) -> bool {
self.data == other.data
self.data_ref() == other.data_ref()
}
}

Expand Down Expand Up @@ -366,7 +409,7 @@ mod tests {
#[test]
fn acquire_mut_ref() {
loom_model!({
let five = Somr::new(5);
let mut five = Somr::new(5);

*five.get_mut().unwrap() += 1;
assert_eq!(five.get(), Some(&6));
Expand All @@ -375,7 +418,7 @@ mod tests {

// only owner can mut ref
assert!(five_ref.get().is_some());
assert!(five_ref.get_mut().is_none());
assert!(unsafe { five_ref.get_mut_from_ref() }.is_none());

drop(five);
});
Expand Down Expand Up @@ -442,4 +485,20 @@ mod tests {
assert!(!five.is_owned());
});
}

// This is UB if we didn't use `UnsafeCell` in `Somr`
#[test]
fn test_inner_mut() {
loom_model!({
let five = Somr::new(5);
fn add(a: &Somr<i32>, b: &Somr<i32>) {
unsafe { a.get_mut_from_ref() }
.map(|mut x| *x += *b.get().unwrap())
.unwrap();
}

add(&five, &five);
assert_eq!(five.get().copied().unwrap(), 10);
});
}
}
4 changes: 2 additions & 2 deletions y-octo/src/doc/document.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,8 +101,8 @@ impl Doc {
for (mut s, offset) in update.iter(store.get_state_vector()) {
if let Node::Item(item) = &mut s {
debug_assert!(item.is_owned());
let item = unsafe { item.get_mut_unchecked() };
store.repair(item, self.store.clone())?;
let mut item = unsafe { item.get_mut_unchecked() };
store.repair(&mut item, self.store.clone())?;
}
store.integrate(s, offset, None)?;
}
Expand Down
12 changes: 6 additions & 6 deletions y-octo/src/doc/store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -186,15 +186,15 @@ impl DocStore {
// we make sure store is the only entry of mutating an item,
// and we already hold mutable reference of store, so it's safe to do so
unsafe {
let left_item = raw_left_item.get_mut_unchecked();
let right_item = right_item.get_mut_unchecked();
let mut left_item = raw_left_item.get_mut_unchecked();
let mut right_item = right_item.get_mut_unchecked();
left_item.content = new_left_item.get_unchecked().content.clone();

// we had the correct left/right content
// now build the references
let right_right_ref = ItemRef::from(&left_item.right);
right_item.left = if right_right_ref.is_some() {
let right_right = right_right_ref.get_mut_unchecked();
let mut right_right = right_right_ref.get_mut_unchecked();
right_right.left.replace(right_node.clone())
} else {
Some(node.clone())
Expand Down Expand Up @@ -387,7 +387,7 @@ impl DocStore {
// before we integrate struct into store,
// the struct => Arc<Item> is owned reference actually,
// no one else refer to such item yet, we can safely mutable refer to it now.
let this = unsafe { item.get_mut_unchecked() };
let this = &mut *unsafe { item.get_mut_unchecked() };

if offset > 0 {
this.id.clock += offset;
Expand Down Expand Up @@ -485,7 +485,7 @@ impl DocStore {
if left.is_some() {
unsafe {
// SAFETY: we get store write lock, no way the left get dropped by owner
let left = left.get_mut_unchecked();
let mut left = left.get_mut_unchecked();
right = left.right.replace(Node::Item(item.clone())).into();
}
this.left = Some(Node::Item(left));
Expand All @@ -508,7 +508,7 @@ impl DocStore {
if right.is_some() {
unsafe {
// SAFETY: we get store write lock, no way the left get dropped by owner
let right = right.get_mut_unchecked();
let mut right = right.get_mut_unchecked();
right.left = Some(Node::Item(item.clone()));
}
this.right = Some(Node::Item(right))
Expand Down

0 comments on commit cb78ffe

Please sign in to comment.