Repository: cberner/redb
Pull request #144 (2022-01-24 05:00:45)

Fix memory safety issues


Patch:
diff --git a/src/db.rs b/src/db.rs
index 80423b0..2af7467 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -9,6 +9,7 @@ use memmap2::MmapRaw;
 use std::cell::{Cell, RefCell};
 use std::collections::HashMap;
 use std::fs::{File, OpenOptions};
+use std::panic;
 use std::path::Path;
 use std::rc::Rc;
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -132,6 +133,7 @@ pub struct DatabaseTransaction<'a> {
     root_page: Cell<Option<PageNumber>>,
     pending_table_root_changes:
         RefCell<HashMap<(Vec<u8>, TableType), Rc<Cell<Option<PageNumber>>>>>,
+    open_tables: RefCell<HashMap<(Vec<u8>, TableType), &'static panic::Location<'static>>>,
     completed: AtomicBool,
 }
 
@@ -144,6 +146,7 @@ impl<'a> DatabaseTransaction<'a> {
             transaction_id,
             root_page: Cell::new(root_page),
             pending_table_root_changes: RefCell::new(Default::default()),
+            open_tables: RefCell::new(Default::default()),
             completed: Default::default(),
         })
     }
@@ -157,6 +160,14 @@ impl<'a> DatabaseTransaction<'a> {
     ) -> Result<Table<'a, K, V>, Error> {
         assert!(!name.as_ref().is_empty());
         assert_ne!(name.as_ref(), FREED_TABLE);
+        let key = (name.as_ref().to_vec(), TableType::Normal);
+        if let Some(location) = self.open_tables.borrow().get(&key) {
+            return Err(Error::TableAlreadyOpen(location));
+        }
+        self.open_tables
+            .borrow_mut()
+            .insert(key.clone(), panic::Location::caller());
+
         let (definition, root) = self.storage.get_or_create_table(
             name.as_ref(),
             TableType::Normal,
@@ -165,7 +176,6 @@ impl<'a> DatabaseTransaction<'a> {
         )?;
         self.root_page.set(Some(root));
 
-        let key = (name.as_ref().to_vec(), TableType::Normal);
         let root = self
             .pending_table_root_changes
             .borrow_mut()
@@ -185,6 +195,14 @@ impl<'a> DatabaseTransaction<'a> {
     ) -> Result<MultimapTable<'a, K, V>, Error> {
         assert!(!name.as_ref().is_empty());
         assert_ne!(name.as_ref(), FREED_TABLE);
+        let key = (name.as_ref().to_vec(), TableType::Multimap);
+        if let Some(location) = self.open_tables.borrow().get(&key) {
+            return Err(Error::TableAlreadyOpen(location));
+        }
+        self.open_tables
+            .borrow_mut()
+            .insert(key.clone(), panic::Location::caller());
+
         let (definition, root) = self.storage.get_or_create_table(
             name.as_ref(),
             TableType::Multimap,
@@ -193,7 +211,6 @@ impl<'a> DatabaseTransaction<'a> {
         )?;
         self.root_page.set(Some(root));
 
-        let key = (name.as_ref().to_vec(), TableType::Multimap);
         let root = self
             .pending_table_root_changes
             .borrow_mut()
diff --git a/src/error.rs b/src/error.rs
index 0139d66..6dd060a 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -13,6 +13,9 @@ pub enum Error {
     },
     DoesNotExist(String),
     LeakedWriteTransaction(&'static panic::Location<'static>),
+    // Tables cannot be opened for writing multiple times, since they could retrieve immutable &
+    // mutable references to the same dirty pages, or multiple mutable references via insert_reserve()
+    TableAlreadyOpen(&'static panic::Location<'static>),
     OutOfSpace,
     Io(io::Error),
     LockPoisoned(&'static panic::Location<'static>),
@@ -56,6 +59,9 @@ impl Display for Error {
             Error::LeakedWriteTransaction(location) => {
                 write!(f, "Leaked write transaction: {}", location)
             }
+            Error::TableAlreadyOpen(location) => {
+                write!(f, "Table already opened at: {}", location)
+            }
             Error::OutOfSpace => {
                 write!(f, "Database is out of space")
             }
diff --git a/src/lib.rs b/src/lib.rs
index 6fa911e..9e3a716 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,3 +1,5 @@
+extern crate core;
+
 pub use db::{Database, DatabaseBuilder, DatabaseTransaction, ReadOnlyDatabaseTransaction};
 pub use error::Error;
 pub use multimap_table::{
diff --git a/src/multimap_table.rs b/src/multimap_table.rs
index 613c439..18d1937 100644
--- a/src/multimap_table.rs
+++ b/src/multimap_table.rs
@@ -369,7 +369,7 @@ impl<'s, K: RedbKey + ?Sized, V: RedbKey + ?Sized> MultimapTable<'s, K, V> {
     }
 }
 
-impl<'s, K: RedbKey + ?Sized, V: RedbKey + ?Sized> ReadableMultimapTable<'s, K, V>
+impl<'s, K: RedbKey + ?Sized, V: RedbKey + ?Sized> ReadableMultimapTable<K, V>
     for MultimapTable<'s, K, V>
 {
     fn get<'a>(&'a self, key: &'a K) -> Result<MultimapGetIterType<'a, K, V>, Error> {
@@ -426,7 +426,7 @@ type MultimapGetIterType<'a, K, V> =
 type MultimapGetRangeIterType<'a, K, V> =
     MultimapRangeIter<'a, (Bound<MultimapKVPair<K, V>>, Bound<MultimapKVPair<K, V>>), K, V>;
 
-pub trait ReadableMultimapTable<'s, K: RedbKey + ?Sized, V: RedbKey + ?Sized> {
+pub trait ReadableMultimapTable<K: RedbKey + ?Sized, V: RedbKey + ?Sized> {
     fn get<'a>(&'a self, key: &'a K) -> Result<MultimapGetIterType<'a, K, V>, Error>;
 
     fn get_range<'a, T: RangeBounds<&'a K> + 'a>(
@@ -465,7 +465,7 @@ impl<'s, K: RedbKey + ?Sized, V: RedbKey + ?Sized> ReadOnlyMultimapTable<'s, K,
     }
 }
 
-impl<'s, K: RedbKey + ?Sized, V: RedbKey + ?Sized> ReadableMultimapTable<'s, K, V>
+impl<'s, K: RedbKey + ?Sized, V: RedbKey + ?Sized> ReadableMultimapTable<K, V>
     for ReadOnlyMultimapTable<'s, K, V>
 {
     fn get<'a>(&'a self, key: &'a K) -> Result<MultimapGetIterType<'a, K, V>, Error> {
diff --git a/src/table.rs b/src/table.rs
index 8cbde29..8b7dc68 100644
--- a/src/table.rs
+++ b/src/table.rs
@@ -76,8 +76,8 @@ impl<'s, K: RedbKey + ?Sized, V: RedbValue + ?Sized> Table<'s, K, V> {
     }
 }
 
-impl<'s, K: RedbKey + ?Sized, V: RedbValue + ?Sized> ReadableTable<'s, K, V> for Table<'s, K, V> {
-    fn get(&self, key: &K) -> Result<Option<AccessGuard<'s, V>>, Error> {
+impl<'s, K: RedbKey + ?Sized, V: RedbValue + ?Sized> ReadableTable<K, V> for Table<'s, K, V> {
+    fn get(&self, key: &K) -> Result<Option<AccessGuard<V>>, Error> {
         self.storage
             .get::<K, V>(key.as_bytes().as_ref(), self.table_root.get())
     }
@@ -109,8 +109,8 @@ impl<'s, K: RedbKey + ?Sized, V: RedbValue + ?Sized> ReadableTable<'s, K, V> for
     }
 }
 
-pub trait ReadableTable<'s, K: RedbKey + ?Sized, V: RedbValue + ?Sized> {
-    fn get(&self, key: &K) -> Result<Option<AccessGuard<'s, V>>, Error>;
+pub trait ReadableTable<K: RedbKey + ?Sized, V: RedbValue + ?Sized> {
+    fn get(&self, key: &K) -> Result<Option<AccessGuard<V>>, Error>;
 
     fn get_range<'a, T: RangeBounds<KR> + 'a, KR: AsRef<K>>(
         &'a self,
@@ -148,10 +148,10 @@ impl<'s, K: RedbKey + ?Sized, V: RedbValue + ?Sized> ReadOnlyTable<'s, K, V> {
     }
 }
 
-impl<'s, K: RedbKey + ?Sized, V: RedbValue + ?Sized> ReadableTable<'s, K, V>
+impl<'s, K: RedbKey + ?Sized, V: RedbValue + ?Sized> ReadableTable<K, V>
     for ReadOnlyTable<'s, K, V>
 {
-    fn get(&self, key: &K) -> Result<Option<AccessGuard<'s, V>>, Error> {
+    fn get(&self, key: &K) -> Result<Option<AccessGuard<V>>, Error> {
         self.storage
             .get::<K, V>(key.as_bytes().as_ref(), self.table_root)
     }
diff --git a/src/tree_store/btree_utils.rs b/src/tree_store/btree_utils.rs
index e181533..7d75e28 100644
--- a/src/tree_store/btree_utils.rs
+++ b/src/tree_store/btree_utils.rs
@@ -484,9 +484,13 @@ fn tree_delete_helper<'a, K: RedbKey + ?Sized>(
                     let entry = accessor.entry(i).unwrap();
                     partial.push((entry.key().to_vec(), entry.value().to_vec()));
                 }
-                // TODO: This call seems like it may be unsound. Another Table instance in the
-                // same transaction could have a reference to this page. We probably need to
-                // ensure this is only a single instance of each table
+                // Safety: No other references to this table can exist.
+                // Tables can only be opened mutably in one location (see Error::TableAlreadyOpen),
+                // and all mutation method (e.g. insert() & remove()) borrow the Table object &mut self.
+                // open_table() & delete_table() which modify the master table, have no readers they
+                // can conflict with
+                // TODO: this method should be unsafe, with the caller expected to uphold the above
+                // guarantees
                 unsafe {
                     if !manager.free_if_uncommitted(page.get_page_number())? {
                         freed.push(page.get_page_number());
@@ -512,9 +516,6 @@ fn tree_delete_helper<'a, K: RedbKey + ?Sized>(
                     builder.append(entry.key(), entry.value());
                 }
                 drop(builder);
-                // TODO: This call seems like it may be unsound. Another Table instance in the
-                // same transaction could have a reference to this page. We probably need to
-                // ensure this is only a single instance of each table
                 unsafe {
                     if !manager.free_if_uncommitted(page.get_page_number())? {
                         freed.push(page.get_page_number());
@@ -567,9 +568,6 @@ fn tree_delete_helper<'a, K: RedbKey + ?Sized>(
             }
             assert!(found);
             assert!(children.len() > 1);
-            // TODO: This call seems like it may be unsound. Another Table instance in the
-            // same transaction could have a reference to this page. We probably need to
-            // ensure this is only a single instance of each table
             unsafe {
                 if !manager.free_if_uncommitted(original_page_number)? {
                     freed.push(original_page_number);
@@ -726,9 +724,13 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
 
                 let page_number = page.get_page_number();
                 drop(page);
-                // TODO: This call seems like it may be unsound. Another Table instance in the
-                // same transaction could have a reference to this page. We probably need to
-                // ensure this is only a single instance of each table
+                // Safety: No other references to this table can exist.
+                // Tables can only be opened mutably in one location (see Error::TableAlreadyOpen),
+                // and all mutation method (e.g. insert() & remove()) borrow the Table object &mut self.
+                // open_table() & delete_table() which modify the master table, have no readers they
+                // can conflict with
+                // TODO: this method should be unsafe, with the caller expected to uphold the above
+                // guarantees
                 unsafe {
                     if !manager.free_if_uncommitted(page_number)? {
                         freed.push(page_number);
@@ -796,9 +798,6 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
 
                     let page_number = page.get_page_number();
                     drop(page);
-                    // TODO: This call seems like it may be unsound. Another Table instance in the
-                    // same transaction could have a reference to this page. We probably need to
-                    // ensure this is only a single instance of each table
                     unsafe {
                         if !manager.free_if_uncommitted(page_number)? {
                             freed.push(page_number);
@@ -846,9 +845,6 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
 
                     let page_number = page.get_page_number();
                     drop(page);
-                    // TODO: This call seems like it may be unsound. Another Table instance in the
-                    // same transaction could have a reference to this page. We probably need to
-                    // ensure this is only a single instance of each table
                     unsafe {
                         if !manager.free_if_uncommitted(page_number)? {
                             freed.push(page_number);


Test patch:
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index 54d3442..f89bba3 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -467,6 +467,18 @@ fn range_query_reversed() {
     assert!(iter.next().is_none());
 }
 
+#[test]
+fn alias_table() {
+    let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
+    let db = unsafe { Database::open(tmpfile.path(), 1024 * 1024).unwrap() };
+
+    let write_txn = db.begin_write().unwrap();
+    let table: Table<[u8], [u8]> = write_txn.open_table(b"x").unwrap();
+    let result: Result<Table<[u8], [u8]>, Error> = write_txn.open_table(b"x");
+    assert!(matches!(result.err().unwrap(), Error::TableAlreadyOpen(_)));
+    drop(table);
+}
+
 #[test]
 fn delete_table() {
     let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();