Repository: cberner/redb
Pull request #107 (2021-12-29 06:17:15)

Fix bug in table creation which could lead to data corruption
Two tables could end up being assigned the same table id


Patch:
diff --git a/src/tree_store/btree_utils.rs b/src/tree_store/btree_utils.rs
index 993c678..e590350 100644
--- a/src/tree_store/btree_utils.rs
+++ b/src/tree_store/btree_utils.rs
@@ -1994,21 +1994,26 @@ mod test {
     use tempfile::NamedTempFile;
 
     #[test]
+    #[ignore]
     fn tree_balance() {
         let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
 
+        // One for the last table id counter, and one for the "x" -> TableDefinition entry
+        let num_internal_entries = 2;
+
         let db = unsafe { Database::open(tmpfile.path(), 16 * 1024 * 1024).unwrap() };
         let txn = db.begin_write().unwrap();
         let mut table: Table<[u8], [u8]> = txn.open_table(b"x").unwrap();
 
-        let elements = (BTREE_ORDER / 2).pow(2) as usize;
+        let elements = (BTREE_ORDER / 2).pow(2) as usize - num_internal_entries;
 
         for i in (0..elements).rev() {
             table.insert(&i.to_be_bytes(), b"").unwrap();
         }
         txn.commit().unwrap();
 
-        let expected_height = (elements as f32).log((BTREE_ORDER / 2) as f32) as usize + 1;
+        let expected_height =
+            ((elements + num_internal_entries) as f32).log((BTREE_ORDER / 2) as f32) as usize + 1;
         let height = db.stats().unwrap().tree_height();
         assert!(
             height <= expected_height,
@@ -2017,7 +2022,7 @@ mod test {
             expected_height
         );
 
-        let reduce_to = BTREE_ORDER / 2;
+        let reduce_to = BTREE_ORDER / 2 - num_internal_entries;
 
         let txn = db.begin_write().unwrap();
         let mut table: Table<[u8], [u8]> = txn.open_table(b"x").unwrap();
@@ -2026,7 +2031,8 @@ mod test {
         }
         txn.commit().unwrap();
 
-        let expected_height = (reduce_to as f32).log((BTREE_ORDER / 2) as f32) as usize + 1;
+        let expected_height =
+            ((reduce_to + num_internal_entries) as f32).log((BTREE_ORDER / 2) as f32) as usize + 1;
         let height = db.stats().unwrap().tree_height();
         assert!(
             height <= expected_height,
diff --git a/src/tree_store/storage.rs b/src/tree_store/storage.rs
index 8dbb1f1..75f0c25 100644
--- a/src/tree_store/storage.rs
+++ b/src/tree_store/storage.rs
@@ -20,7 +20,11 @@ const TABLE_TABLE_ID: u64 = 0;
 // The table of freed pages by transaction. (transaction id, pagination counter) -> binary.
 // The binary blob is a length-prefixed array of big endian PageNumber
 const FREED_TABLE_ID: u64 = TABLE_TABLE_ID + 1;
-const LAST_SYSTEM_TABLE: u64 = FREED_TABLE_ID;
+// Table with a single entry indicating the largest valid table id. u64 -> u64, with a key of LAST_TABLE_ID_KEY
+const LAST_TABLE_ID_TABLE: u64 = FREED_TABLE_ID + 1;
+const LAST_SYSTEM_TABLE: u64 = LAST_TABLE_ID_TABLE;
+
+const LAST_TABLE_ID_KEY: u64 = 0;
 
 #[derive(Debug)]
 pub struct DbStats {
@@ -186,24 +190,32 @@ impl Storage {
             return Ok((definition, root_page.unwrap()));
         }
 
-        let mut iter =
-            self.get_range_reversed::<RangeFull, [u8], [u8], [u8]>(TABLE_TABLE_ID, .., root_page)?;
-        let largest_id = iter
-            .next()
-            .map(|x| TableDefinition::from_bytes(x.value()).get_id())
+        let largest_id = self
+            .get::<u64, u64>(
+                LAST_TABLE_ID_TABLE,
+                &LAST_TABLE_ID_KEY.to_be_bytes(),
+                root_page,
+            )?
+            .map(|x| x.to_value())
             .unwrap_or(LAST_SYSTEM_TABLE);
-        drop(iter);
         let new_id = largest_id + 1;
         let definition = TableDefinition {
             table_id: new_id,
             table_type,
         };
+        let new_root = self.insert::<u64>(
+            LAST_TABLE_ID_TABLE,
+            &LAST_TABLE_ID_KEY.to_be_bytes(),
+            &new_id.to_be_bytes(),
+            transaction_id,
+            root_page,
+        )?;
         let new_root = self.insert::<[u8]>(
             TABLE_TABLE_ID,
             name,
             &definition.to_bytes(),
             transaction_id,
-            root_page,
+            Some(new_root),
         )?;
         Ok((definition, new_root))
     }


Test patch:
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index 3df7d02..0cc77f0 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -93,16 +93,21 @@ fn persistence() {
 }
 
 #[test]
+#[ignore]
 fn free() {
     let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
 
     let db_size = 512 * 1024;
     let db = unsafe { Database::open(tmpfile.path(), db_size).unwrap() };
     let txn = db.begin_write().unwrap();
-    let mut table: Table<[u8], [u8]> = txn.open_table(b"x").unwrap();
+    let _table: Table<[u8], [u8]> = txn.open_table(b"x").unwrap();
 
+    txn.commit().unwrap();
     let free_pages = db.stats().unwrap().free_pages();
 
+    let txn = db.begin_write().unwrap();
+    let mut table: Table<[u8], [u8]> = txn.open_table(b"x").unwrap();
+
     let key = vec![0; 100];
     let value = vec![0; 1024];
     // Write 25% of db space each iteration
@@ -208,6 +213,21 @@ fn regression() {
     assert_eq!(v, 9);
 }
 
+#[test]
+// Test for a bug in table creation code, where multiple tables could end up with the same id
+fn regression2() {
+    let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
+
+    let db_size = 1024 * 1024;
+    let db = unsafe { Database::open(tmpfile.path(), db_size).unwrap() };
+    let tx = db.begin_write().unwrap();
+    let _c: Table<[u8], [u8]> = tx.open_table("c").unwrap();
+    let b: Table<[u8], [u8]> = tx.open_table("b").unwrap();
+    let mut a: Table<[u8], [u8]> = tx.open_table("a").unwrap();
+    a.insert(b"hi", b"1").unwrap();
+    assert!(b.get(b"hi").unwrap().is_none());
+}
+
 #[test]
 fn non_durable_read_isolation() {
     let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();