fix: soundness issue in Somr

y-octo/src/doc/common/somr.rs

@@ -1,8 +1,10 @@
 use std::{
+    cell::UnsafeCell,
     fmt::{self, Write},
     hash::{Hash, Hasher},
     marker::PhantomData,
     mem,
+    ops::{Deref, DerefMut},
     ptr::NonNull,
 };
 
@@ -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> {}
 
@@ -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,
         });
@@ -57,55 +78,76 @@ 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> {
+    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())
+    }
+
+    pub 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()
     }
 
     #[allow(clippy::mut_from_ref)]
-    pub unsafe fn get_mut_unchecked(&self) -> &mut T {
+    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")
@@ -238,7 +280,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()
     }
 }
 
@@ -366,7 +408,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));
@@ -375,7 +417,7 @@ mod tests {
 
             // only owner can mut ref
             assert!(five_ref.get().is_some());
-            assert!(five_ref.get_mut().is_none());
+            assert!(five_ref.get_mut_from_ref().is_none());
 
             drop(five);
         });
@@ -442,4 +484,16 @@ mod tests {
             assert!(!five.is_owned());
         });
     }
+
+    // This is UB if we didn't use `UnsafeCell` in `Somr`
+    #[test]
+    fn test_inner_mut() {
+        let five = Somr::new(5);
+        fn add(a: &Somr<i32>, b: &Somr<i32>) {
+            a.get_mut_from_ref().map(|mut x| *x += *b.get().unwrap()).unwrap();
+        }
+
+        add(&five, &five);
+        assert_eq!(five.get().copied().unwrap(), 10);
+    }
 }

y-octo/src/doc/document.rs

@@ -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)?;
             }

y-octo/src/doc/store.rs

@@ -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())
@@ -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;
@@ -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));
@@ -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))