From bcf7899efde5db118ae24a040d8f1db80644e7ba Mon Sep 17 00:00:00 2001
From: Dmitrii Novikov <novikov.dm.al@gmail.com>
Date: Fri, 11 Oct 2024 12:38:50 +0400
Subject: [PATCH] fix(ethexe): Fix critical bug of prefix iterating RocksDB
 (#4285)

---
 ethexe/db/src/lib.rs   | 30 +++++++++++++++++++++--------
 ethexe/db/src/rocks.rs | 43 ++++++++++++++++++++++++++++++++++--------
 2 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/ethexe/db/src/lib.rs b/ethexe/db/src/lib.rs
index 7a7d93bc44a..9a4ae67156e 100644
--- a/ethexe/db/src/lib.rs
+++ b/ethexe/db/src/lib.rs
@@ -109,18 +109,32 @@ mod tests {
     }
 
     pub fn kv_iter_prefix<DB: KVDatabase>(db: DB) {
+        let testcase = |prefix: &str, expectations: &[(&str, &str)]| {
+            let actual: BTreeSet<_> = db.iter_prefix(prefix.as_bytes()).collect();
+            let expected: BTreeSet<_> = expectations
+                .iter()
+                .map(|(k, v)| (k.as_bytes().to_vec(), v.as_bytes().to_vec()))
+                .collect();
+            assert_eq!(actual, expected);
+        };
+
         db.put(b"prefix_foo", b"hello".to_vec());
         db.put(b"prefix_bar", b"world".to_vec());
 
-        let values: BTreeSet<(Vec<u8>, Vec<u8>)> = db.iter_prefix(b"prefix_").collect();
-        assert_eq!(
-            values,
-            [
-                (b"prefix_foo".to_vec(), b"hello".to_vec()),
-                (b"prefix_bar".to_vec(), b"world".to_vec()),
-            ]
-            .into(),
+        testcase(
+            "prefix_",
+            &[("prefix_foo", "hello"), ("prefix_bar", "world")],
         );
+
+        testcase("", &[("prefix_foo", "hello"), ("prefix_bar", "world")]);
+
+        testcase("0", &[]);
+
+        testcase("prefix_foobar", &[]);
+
+        testcase("prefix_foo", &[("prefix_foo", "hello")]);
+
+        testcase("prefix_bar", &[("prefix_bar", "world")]);
     }
 
     pub fn cas_multi_thread<DB: CASDatabase>(db: DB) {
diff --git a/ethexe/db/src/rocks.rs b/ethexe/db/src/rocks.rs
index 7d2f48a8b89..c9b77880801 100644
--- a/ethexe/db/src/rocks.rs
+++ b/ethexe/db/src/rocks.rs
@@ -19,7 +19,7 @@
 use crate::{CASDatabase, KVDatabase};
 use anyhow::Result;
 use gprimitives::H256;
-use rocksdb::{Options, DB};
+use rocksdb::{DBIteratorWithThreadMode, Options, DB};
 use std::{path::PathBuf, sync::Arc};
 
 /// Database for storing states and codes in memory.
@@ -85,13 +85,40 @@ impl KVDatabase for RocksDatabase {
             .expect("Failed to write data, database is not in valid state");
     }
 
-    fn iter_prefix(&self, prefix: &[u8]) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + '_> {
-        Box::new(
-            self.inner
-                .prefix_iterator(prefix)
-                .map(|kv| kv.expect("unexpected error during iteration"))
-                .map(|(k, v)| (<[u8]>::into_vec(k), <[u8]>::into_vec(v))),
-        )
+    fn iter_prefix<'a>(
+        &'a self,
+        prefix: &'a [u8],
+    ) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + '_> {
+        Box::new(PrefixIterator {
+            prefix,
+            prefix_iter: self.inner.prefix_iterator(prefix),
+            done: false,
+        })
+    }
+}
+
+pub struct PrefixIterator<'a> {
+    prefix: &'a [u8],
+    prefix_iter: DBIteratorWithThreadMode<'a, DB>,
+    done: bool,
+}
+
+impl<'a> Iterator for PrefixIterator<'a> {
+    type Item = (Vec<u8>, Vec<u8>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.done {
+            return None;
+        }
+
+        match self.prefix_iter.next() {
+            Some(Ok((k, v))) if k.starts_with(self.prefix) => Some((k.to_vec(), v.to_vec())),
+            Some(Err(e)) => panic!("Failed to read data, database is not in valid state: {e:?}"),
+            _ => {
+                self.done = true;
+                None
+            }
+        }
     }
 }