fix(commitment-tree): reject negative sqlite positions

Author: thepastaclaw

grovedb-commitment-tree/src/client/sqlite_store/sql_helpers.rs

@@ -18,10 +18,18 @@ use super::{
     SqliteShardStoreError, SHARD_HEIGHT,
 };
 
+fn non_negative_i64_to_u64(value_name: &str, value: i64) -> Result<u64, SqliteShardStoreError> {
+    u64::try_from(value).map_err(|_| {
+        SqliteShardStoreError::Serialization(format!(
+            "invalid negative {value_name} in sqlite row: {value}"
+        ))
+    })
+}
+
 pub(crate) fn create_tables(conn: &Connection) -> Result<(), SqliteShardStoreError> {
     conn.execute_batch(
         "CREATE TABLE IF NOT EXISTS commitment_tree_shards (
-            shard_index INTEGER PRIMARY KEY,
+            shard_index INTEGER PRIMARY KEY CHECK (shard_index >= 0),
             shard_data  BLOB NOT NULL
         );
         CREATE TABLE IF NOT EXISTS commitment_tree_cap (
@@ -30,11 +38,11 @@ pub(crate) fn create_tables(conn: &Connection) -> Result<(), SqliteShardStoreErr
         );
         CREATE TABLE IF NOT EXISTS commitment_tree_checkpoints (
             checkpoint_id INTEGER PRIMARY KEY,
-            position      INTEGER
+            position      INTEGER CHECK (position IS NULL OR position >= 0)
         );
         CREATE TABLE IF NOT EXISTS commitment_tree_checkpoint_marks_removed (
             checkpoint_id INTEGER NOT NULL,
-            position      INTEGER NOT NULL,
+            position      INTEGER NOT NULL CHECK (position >= 0),
             PRIMARY KEY (checkpoint_id, position),
             FOREIGN KEY (checkpoint_id) REFERENCES commitment_tree_checkpoints(checkpoint_id)
         );",
@@ -85,7 +93,8 @@ pub(crate) fn sql_last_shard(
     match row {
         None => Ok(None),
         Some((index, data)) => {
-            let addr = Address::from_parts(Level::from(SHARD_HEIGHT), index as u64);
+            let index = non_negative_i64_to_u64("shard_index", index)?;
+            let addr = Address::from_parts(Level::from(SHARD_HEIGHT), index);
             let mut pos = 0;
             let tree = deserialize_tree(&data, &mut pos)?;
             let located = LocatedTree::from_parts(addr, tree).map_err(|addr| {
@@ -116,13 +125,24 @@ pub(crate) fn sql_get_shard_roots(
 ) -> Result<Vec<Address>, SqliteShardStoreError> {
     let mut stmt =
         conn.prepare("SELECT shard_index FROM commitment_tree_shards ORDER BY shard_index")?;
-    let rows = stmt.query_map([], |row| {
-        let index: i64 = row.get(0)?;
-        Ok(Address::from_parts(Level::from(SHARD_HEIGHT), index as u64))
-    })?;
+    let rows = stmt
+        .query_map([], |row| row.get::<_, i64>(0))
+        .map_err(|err| {
+            SqliteShardStoreError::Serialization(format!(
+                "failed to query shard roots from sqlite: {err}"
+            ))
+        })?;
     let mut result = Vec::new();
-    for addr in rows {
-        result.push(addr?);
+    for index in rows {
+        let index = index.map_err(|err| {
+            SqliteShardStoreError::Serialization(format!(
+                "failed to read shard root row from sqlite: {err}"
+            ))
+        })?;
+        result.push(Address::from_parts(
+            Level::from(SHARD_HEIGHT),
+            non_negative_i64_to_u64("shard_index", index)?,
+        ));
     }
     Ok(result)
 }
@@ -375,18 +395,32 @@ fn sql_load_checkpoint(
 ) -> Result<Checkpoint, SqliteShardStoreError> {
     let tree_state = match position {
         None => TreeState::Empty,
-        Some(p) => TreeState::AtPosition(Position::from(p as u64)),
+        Some(p) => TreeState::AtPosition(Position::from(non_negative_i64_to_u64("position", p)?)),
     };
 
     let mut stmt = conn.prepare(
         "SELECT position FROM commitment_tree_checkpoint_marks_removed WHERE checkpoint_id = ?1",
     )?;
-    let marks: BTreeSet<Position> = stmt
-        .query_map(params![checkpoint_id], |row| {
-            let p: i64 = row.get(0)?;
-            Ok(Position::from(p as u64))
-        })?
-        .collect::<Result<BTreeSet<_>, _>>()?;
+    let rows = stmt
+        .query_map(params![checkpoint_id], |row| row.get::<_, i64>(0))
+        .map_err(|err| {
+            SqliteShardStoreError::Serialization(format!(
+                "failed to query checkpoint removed marks from sqlite for checkpoint \
+                 {checkpoint_id}: {err}"
+            ))
+        })?;
+    let mut marks = BTreeSet::new();
+    for position in rows {
+        let position = position.map_err(|err| {
+            SqliteShardStoreError::Serialization(format!(
+                "failed to read checkpoint removed mark row from sqlite for checkpoint \
+                 {checkpoint_id}: {err}"
+            ))
+        })?;
+        marks.insert(Position::from(non_negative_i64_to_u64(
+            "position", position,
+        )?));
+    }
 
     Ok(Checkpoint::from_parts(tree_state, marks))
 }

grovedb-commitment-tree/src/client/sqlite_store_tests.rs

@@ -23,6 +23,32 @@ mod tests {
         SqliteShardStore::new(conn).expect("create store")
     }
 
+    fn legacy_schema_conn() -> Connection {
+        let conn = Connection::open_in_memory().expect("open legacy in-memory sqlite");
+        conn.execute_batch(
+            "CREATE TABLE commitment_tree_shards (
+                shard_index INTEGER PRIMARY KEY,
+                shard_data  BLOB NOT NULL
+            );
+            CREATE TABLE commitment_tree_cap (
+                id       INTEGER PRIMARY KEY CHECK (id = 0),
+                cap_data BLOB NOT NULL
+            );
+            CREATE TABLE commitment_tree_checkpoints (
+                checkpoint_id INTEGER PRIMARY KEY,
+                position      INTEGER
+            );
+            CREATE TABLE commitment_tree_checkpoint_marks_removed (
+                checkpoint_id INTEGER NOT NULL,
+                position      INTEGER NOT NULL,
+                PRIMARY KEY (checkpoint_id, position),
+                FOREIGN KEY (checkpoint_id) REFERENCES commitment_tree_checkpoints(checkpoint_id)
+            );",
+        )
+        .expect("create legacy schema");
+        conn
+    }
+
     fn test_hash(i: u8) -> MerkleHashOrchard {
         let empty = MerkleHashOrchard::empty_leaf();
         MerkleHashOrchard::combine(Level::from(i % 31 + 1), &empty, &empty)
@@ -80,6 +106,68 @@ mod tests {
         assert_eq!(count, 1);
     }
 
+    #[test]
+    fn test_schema_rejects_negative_shard_index() {
+        let store = test_store();
+        let data = serialize_tree(&Tree::leaf((test_hash(1), RetentionFlags::MARKED)));
+
+        let err = store
+            .with_conn(|conn| {
+                conn.execute(
+                    "INSERT INTO commitment_tree_shards (shard_index, shard_data) VALUES (?1, ?2)",
+                    rusqlite::params![-1i64, data],
+                )
+            })
+            .expect_err("negative shard index should be rejected");
+
+        assert!(
+            err.to_string().contains("CHECK constraint failed"),
+            "unexpected sqlite error: {err}"
+        );
+    }
+
+    #[test]
+    fn test_schema_rejects_negative_checkpoint_positions() {
+        let store = test_store();
+
+        let err = store
+            .with_conn(|conn| {
+                conn.execute(
+                    "INSERT INTO commitment_tree_checkpoints (checkpoint_id, position) VALUES (?1, ?2)",
+                    rusqlite::params![1u32, -1i64],
+                )
+            })
+            .expect_err("negative checkpoint position should be rejected");
+
+        assert!(
+            err.to_string().contains("CHECK constraint failed"),
+            "unexpected sqlite error: {err}"
+        );
+
+        store
+            .with_conn(|conn| {
+                conn.execute(
+                    "INSERT INTO commitment_tree_checkpoints (checkpoint_id, position) VALUES (?1, NULL)",
+                    rusqlite::params![2u32],
+                )
+            })
+            .expect("insert empty checkpoint");
+
+        let err = store
+            .with_conn(|conn| {
+                conn.execute(
+                    "INSERT INTO commitment_tree_checkpoint_marks_removed (checkpoint_id, position) VALUES (?1, ?2)",
+                    rusqlite::params![2u32, -1i64],
+                )
+            })
+            .expect_err("negative mark position should be rejected");
+
+        assert!(
+            err.to_string().contains("CHECK constraint failed"),
+            "unexpected sqlite error: {err}"
+        );
+    }
+
     // -- Serialization round-trip tests --
 
     #[test]
@@ -264,6 +352,48 @@ mod tests {
         assert_eq!(roots.last().expect("last root").index(), 2);
     }
 
+    #[test]
+    fn test_last_shard_rejects_negative_shard_index_from_legacy_schema() {
+        let conn = legacy_schema_conn();
+        let data = serialize_tree(&Tree::leaf((test_hash(7), RetentionFlags::MARKED)));
+        conn.execute(
+            "INSERT INTO commitment_tree_shards (shard_index, shard_data) VALUES (?1, ?2)",
+            rusqlite::params![-1i64, data],
+        )
+        .expect("insert corrupt shard");
+
+        let store = SqliteShardStore::new(conn).expect("open legacy store");
+        let err = store
+            .last_shard()
+            .expect_err("negative shard index should fail");
+        assert!(
+            err.to_string()
+                .contains("invalid negative shard_index in sqlite row: -1"),
+            "unexpected store error: {err}"
+        );
+    }
+
+    #[test]
+    fn test_get_shard_roots_rejects_negative_shard_index_from_legacy_schema() {
+        let conn = legacy_schema_conn();
+        let data = serialize_tree(&Tree::leaf((test_hash(8), RetentionFlags::MARKED)));
+        conn.execute(
+            "INSERT INTO commitment_tree_shards (shard_index, shard_data) VALUES (?1, ?2)",
+            rusqlite::params![-1i64, data],
+        )
+        .expect("insert corrupt shard");
+
+        let store = SqliteShardStore::new(conn).expect("open legacy store");
+        let err = store
+            .get_shard_roots()
+            .expect_err("negative shard index should fail");
+        assert!(
+            err.to_string()
+                .contains("invalid negative shard_index in sqlite row: -1"),
+            "unexpected store error: {err}"
+        );
+    }
+
     // -- Cap tests --
 
     #[test]
@@ -370,6 +500,51 @@ mod tests {
         assert_eq!(retrieved.tree_state(), TreeState::Empty);
     }
 
+    #[test]
+    fn test_get_checkpoint_rejects_negative_position_from_legacy_schema() {
+        let conn = legacy_schema_conn();
+        conn.execute(
+            "INSERT INTO commitment_tree_checkpoints (checkpoint_id, position) VALUES (?1, ?2)",
+            rusqlite::params![1u32, -1i64],
+        )
+        .expect("insert corrupt checkpoint");
+
+        let store = SqliteShardStore::new(conn).expect("open legacy store");
+        let err = store
+            .get_checkpoint(&1)
+            .expect_err("negative checkpoint position should fail");
+        assert!(
+            err.to_string()
+                .contains("invalid negative position in sqlite row: -1"),
+            "unexpected store error: {err}"
+        );
+    }
+
+    #[test]
+    fn test_get_checkpoint_rejects_negative_mark_position_from_legacy_schema() {
+        let conn = legacy_schema_conn();
+        conn.execute(
+            "INSERT INTO commitment_tree_checkpoints (checkpoint_id, position) VALUES (?1, ?2)",
+            rusqlite::params![1u32, 10i64],
+        )
+        .expect("insert checkpoint");
+        conn.execute(
+            "INSERT INTO commitment_tree_checkpoint_marks_removed (checkpoint_id, position) VALUES (?1, ?2)",
+            rusqlite::params![1u32, -1i64],
+        )
+        .expect("insert corrupt mark");
+
+        let store = SqliteShardStore::new(conn).expect("open legacy store");
+        let err = store
+            .get_checkpoint(&1)
+            .expect_err("negative mark position should fail");
+        assert!(
+            err.to_string()
+                .contains("invalid negative position in sqlite row: -1"),
+            "unexpected store error: {err}"
+        );
+    }
+
     #[test]
     fn test_remove_checkpoint() {
         let mut store = test_store();