Fix potential data loss when using Table::get_mut()

src/tree_store/btree_base.rs

@@ -255,11 +255,12 @@ impl<V: Value + 'static> Drop for AccessGuard<'_, V> {
 }
 
 pub struct AccessGuardMut<'a, V: Value + 'static> {
-    page: PageMut,
+    value: Vec<u8>,
+    page: PageNumber,
     offset: usize,
     len: usize,
     entry_index: usize,
-    parent: Option<(PageMut, usize)>,
+    parent: Option<(PageNumber, usize)>,
     mem: Arc<TransactionalMemory>,
     allocated: Arc<Mutex<PageTrackerPolicy>>,
     root_ref: &'a mut BtreeHeader,
@@ -284,12 +285,14 @@ impl<'a, V: Value + 'static> AccessGuardMut<'a, V> {
         if let Some((ref parent_page, _)) = parent {
             assert!(mem.uncommitted(parent_page.get_page_number()));
         }
+        let value = page.memory()[offset..(offset + len)].to_vec();
         AccessGuardMut {
-            page,
+            value,
+            page: page.get_page_number(),
             offset,
             len,
             entry_index,
-            parent,
+            parent: parent.map(|(p, i)| (p.get_page_number(), i)),
             mem,
             allocated,
             root_ref,
@@ -300,33 +303,35 @@ impl<'a, V: Value + 'static> AccessGuardMut<'a, V> {
 
     /// Access the stored value
     pub fn value(&self) -> V::SelfType<'_> {
-        V::from_bytes(&self.page.memory()[self.offset..(self.offset + self.len)])
+        V::from_bytes(&self.value)
     }
 
     /// Replace the stored value
     pub fn insert<'v>(&mut self, value: impl Borrow<V::SelfType<'v>>) -> Result<()> {
         let value_bytes = V::as_bytes(value.borrow());
+        self.value = value_bytes.as_ref().to_vec();
+
+        let mut page = self.mem.get_page_mut(self.page)?;
 
         // TODO: optimize this to avoid copying the key
         let key_bytes = {
-            let accessor = LeafAccessor::new(self.page.memory(), self.key_width, V::fixed_width());
+            let accessor = LeafAccessor::new(page.memory(), self.key_width, V::fixed_width());
             accessor.key_unchecked(self.entry_index).to_vec()
         };
 
         if LeafMutator::sufficient_insert_inplace_space(
-            &self.page,
+            &page,
             self.entry_index,
             true,
             self.key_width,
             V::fixed_width(),
             key_bytes.as_slice(),
             value_bytes.as_ref(),
         ) {
-            let mut mutator =
-                LeafMutator::new(self.page.memory_mut(), self.key_width, V::fixed_width());
+            let mut mutator = LeafMutator::new(page.memory_mut(), self.key_width, V::fixed_width());
             mutator.insert(self.entry_index, true, &key_bytes, value_bytes.as_ref());
         } else {
-            let accessor = LeafAccessor::new(self.page.memory(), self.key_width, V::fixed_width());
+            let accessor = LeafAccessor::new(page.memory(), self.key_width, V::fixed_width());
             let mut builder = LeafBuilder::new(
                 &self.mem,
                 &self.allocated,
@@ -347,16 +352,18 @@ impl<'a, V: Value + 'static> AccessGuardMut<'a, V> {
             let new_page = builder.build()?;
 
             // Update parent branch page if it exists, otherwise update root
-            if let Some((ref mut parent_page, parent_entry_index)) = self.parent {
+            if let Some((ref mut parent_page_number, parent_entry_index)) = self.parent {
+                let mut parent_page = self.mem.get_page_mut(*parent_page_number)?;
                 let mut mutator = BranchMutator::new(parent_page.memory_mut());
                 mutator.write_child_page(parent_entry_index, new_page.get_page_number(), DEFERRED);
             } else {
                 self.root_ref.root = new_page.get_page_number();
                 self.root_ref.checksum = DEFERRED;
             }
 
-            let old_page_number = self.page.get_page_number();
-            self.page = new_page;
+            let old_page_number = page.get_page_number();
+            self.page = new_page.get_page_number();
+            page = new_page;
             let mut allocated = self.allocated.lock().unwrap();
             assert!(
                 self.mem
@@ -365,7 +372,7 @@ impl<'a, V: Value + 'static> AccessGuardMut<'a, V> {
         }
 
         // Update our page reference to the new page and recalculate offset/length
-        let new_accessor = LeafAccessor::new(self.page.memory(), self.key_width, V::fixed_width());
+        let new_accessor = LeafAccessor::new(page.memory(), self.key_width, V::fixed_width());
         let (new_start, new_end) = new_accessor.value_range(self.entry_index).unwrap();
 
         self.offset = new_start;

tests/integration_tests.rs

@@ -1315,6 +1315,32 @@ fn regression25() {
     }
 }
 
+#[test]
+fn regression26() {
+    let tmpfile = create_tempfile();
+    let table_def: TableDefinition<u64, (&str, &[u8])> = TableDefinition::new("issue_1117");
+
+    let db = Database::create(tmpfile.path()).unwrap();
+
+    let txn = db.begin_write().unwrap();
+    {
+        let mut table = txn.open_table(table_def).unwrap();
+        table.insert(0, ("name", &[0u8][..])).unwrap();
+    }
+    txn.commit().unwrap();
+
+    {
+        let txn = db.begin_write().unwrap();
+        let mut table = txn.open_table(table_def).unwrap();
+        let mut access = table.get_mut(&0).unwrap().unwrap();
+        let name = access.value().0.to_string();
+        let large_value = vec![1u8; 8192];
+        access.insert((&name[..], large_value.as_slice())).unwrap();
+        drop(table);
+        txn.commit().unwrap();
+    }
+}
+
 #[test]
 fn check_integrity_clean() {
     let tmpfile = create_tempfile();