Make set_iterate_upper_bound method safe (#377)

src/db.rs

@@ -416,15 +416,15 @@ impl DB {
 
     pub fn iterator<'a: 'b, 'b>(&'a self, mode: IteratorMode) -> DBIterator<'b> {
         let readopts = ReadOptions::default();
-        self.iterator_opt(mode, &readopts)
+        self.iterator_opt(mode, readopts)
     }
 
     pub fn iterator_opt<'a: 'b, 'b>(
         &'a self,
         mode: IteratorMode,
-        readopts: &ReadOptions,
+        readopts: ReadOptions,
     ) -> DBIterator<'b> {
-        DBIterator::new(self, &readopts, mode)
+        DBIterator::new(self, readopts, mode)
     }
 
     /// Opens an iterator using the provided ReadOptions.
@@ -432,10 +432,10 @@ impl DB {
     pub fn iterator_cf_opt<'a: 'b, 'b>(
         &'a self,
         cf_handle: &ColumnFamily,
-        readopts: &ReadOptions,
+        readopts: ReadOptions,
         mode: IteratorMode,
     ) -> DBIterator<'b> {
-        DBIterator::new_cf(self, cf_handle, &readopts, mode)
+        DBIterator::new_cf(self, cf_handle, readopts, mode)
     }
 
     /// Opens an iterator with `set_total_order_seek` enabled.
@@ -444,7 +444,7 @@ impl DB {
     pub fn full_iterator<'a: 'b, 'b>(&'a self, mode: IteratorMode) -> DBIterator<'b> {
         let mut opts = ReadOptions::default();
         opts.set_total_order_seek(true);
-        DBIterator::new(self, &opts, mode)
+        DBIterator::new(self, opts, mode)
     }
 
     pub fn prefix_iterator<'a: 'b, 'b, P: AsRef<[u8]>>(&'a self, prefix: P) -> DBIterator<'b> {
@@ -452,7 +452,7 @@ impl DB {
         opts.set_prefix_same_as_start(true);
         DBIterator::new(
             self,
-            &opts,
+            opts,
             IteratorMode::From(prefix.as_ref(), Direction::Forward),
         )
     }
@@ -463,7 +463,7 @@ impl DB {
         mode: IteratorMode,
     ) -> DBIterator<'b> {
         let opts = ReadOptions::default();
-        DBIterator::new_cf(self, cf_handle, &opts, mode)
+        DBIterator::new_cf(self, cf_handle, opts, mode)
     }
 
     pub fn full_iterator_cf<'a: 'b, 'b>(
@@ -473,7 +473,7 @@ impl DB {
     ) -> DBIterator<'b> {
         let mut opts = ReadOptions::default();
         opts.set_total_order_seek(true);
-        DBIterator::new_cf(self, cf_handle, &opts, mode)
+        DBIterator::new_cf(self, cf_handle, opts, mode)
     }
 
     pub fn prefix_iterator_cf<'a: 'b, 'b, P: AsRef<[u8]>>(
@@ -486,7 +486,7 @@ impl DB {
         DBIterator::new_cf(
             self,
             cf_handle,
-            &opts,
+            opts,
             IteratorMode::From(prefix.as_ref(), Direction::Forward),
         )
     }
@@ -494,17 +494,17 @@ impl DB {
     /// Opens a raw iterator over the database, using the default read options
     pub fn raw_iterator<'a: 'b, 'b>(&'a self) -> DBRawIterator<'b> {
         let opts = ReadOptions::default();
-        DBRawIterator::new(self, &opts)
+        DBRawIterator::new(self, opts)
     }
 
     /// Opens a raw iterator over the given column family, using the default read options
     pub fn raw_iterator_cf<'a: 'b, 'b>(&'a self, cf_handle: &ColumnFamily) -> DBRawIterator<'b> {
         let opts = ReadOptions::default();
-        DBRawIterator::new_cf(self, cf_handle, &opts)
+        DBRawIterator::new_cf(self, cf_handle, opts)
     }
 
     /// Opens a raw iterator over the database, using the given read options
-    pub fn raw_iterator_opt<'a: 'b, 'b>(&'a self, readopts: &ReadOptions) -> DBRawIterator<'b> {
+    pub fn raw_iterator_opt<'a: 'b, 'b>(&'a self, readopts: ReadOptions) -> DBRawIterator<'b> {
         DBRawIterator::new(self, readopts)
     }
 
@@ -512,7 +512,7 @@ impl DB {
     pub fn raw_iterator_cf_opt<'a: 'b, 'b>(
         &'a self,
         cf_handle: &ColumnFamily,
-        readopts: &ReadOptions,
+        readopts: ReadOptions,
     ) -> DBRawIterator<'b> {
         DBRawIterator::new_cf(self, cf_handle, readopts)
     }
@@ -1056,6 +1056,31 @@ fn iterator_test_past_end() {
     DB::destroy(&opts, path).unwrap();
 }
 
+#[test]
+fn iterator_test_upper_bound() {
+    let path = "_rust_rocksdb_iteratortest_upper_bound";
+    {
+        let db = DB::open_default(path).unwrap();
+        db.put(b"k1", b"v1").unwrap();
+        db.put(b"k2", b"v2").unwrap();
+        db.put(b"k3", b"v3").unwrap();
+        db.put(b"k4", b"v4").unwrap();
+        db.put(b"k5", b"v5").unwrap();
+
+        let mut readopts = ReadOptions::default();
+        readopts.set_iterate_upper_bound(b"k4".to_vec());
+
+        let iter = db.iterator_opt(IteratorMode::Start, readopts);
+        let expected: Vec<_> = vec![(b"k1", b"v1"), (b"k2", b"v2"), (b"k3", b"v3")]
+            .into_iter()
+            .map(|(k, v)| (k.to_vec().into_boxed_slice(), v.to_vec().into_boxed_slice()))
+            .collect();
+        assert_eq!(expected, iter.collect::<Vec<_>>());
+    }
+    let opts = Options::default();
+    DB::destroy(&opts, path).unwrap();
+}
+
 #[test]
 fn iterator_test_tailing() {
     let path = "_rust_rocksdb_iteratortest_tailing";
@@ -1070,7 +1095,7 @@ fn iterator_test_tailing() {
         let r = db.put(k, v);
         assert!(r.is_ok());
 
-        let tail_iter = db.iterator_opt(IteratorMode::Start, &ro);
+        let tail_iter = db.iterator_opt(IteratorMode::Start, ro);
         for (k, v) in data_iter {
             let r = db.put(k, v);
             assert!(r.is_ok());

src/db_iterator.rs

@@ -67,14 +67,21 @@ use std::slice;
 /// ```
 pub struct DBRawIterator<'a> {
     inner: *mut ffi::rocksdb_iterator_t,
+
+    /// When iterate_upper_bound is set, the inner C iterator keeps a pointer to the upper bound
+    /// inside `_readopts`. Storing this makes sure the upper bound is always alive when the
+    /// iterator is being used.
+    _readopts: ReadOptions,
+
     db: PhantomData<&'a DB>,
 }
 
 impl<'a> DBRawIterator<'a> {
-    pub(crate) fn new(db: &DB, readopts: &ReadOptions) -> DBRawIterator<'a> {
+    pub(crate) fn new(db: &DB, readopts: ReadOptions) -> DBRawIterator<'a> {
         unsafe {
             DBRawIterator {
                 inner: ffi::rocksdb_create_iterator(db.inner, readopts.inner),
+                _readopts: readopts,
                 db: PhantomData,
             }
         }
@@ -83,11 +90,12 @@ impl<'a> DBRawIterator<'a> {
     pub(crate) fn new_cf(
         db: &DB,
         cf_handle: &ColumnFamily,
-        readopts: &ReadOptions,
+        readopts: ReadOptions,
     ) -> DBRawIterator<'a> {
         unsafe {
             DBRawIterator {
                 inner: ffi::rocksdb_create_iterator_cf(db.inner, readopts.inner, cf_handle.inner),
+                _readopts: readopts,
                 db: PhantomData,
             }
         }
@@ -381,7 +389,7 @@ pub enum IteratorMode<'a> {
 }
 
 impl<'a> DBIterator<'a> {
-    pub(crate) fn new(db: &DB, readopts: &ReadOptions, mode: IteratorMode) -> DBIterator<'a> {
+    pub(crate) fn new(db: &DB, readopts: ReadOptions, mode: IteratorMode) -> DBIterator<'a> {
         let mut rv = DBIterator {
             raw: DBRawIterator::new(db, readopts),
             direction: Direction::Forward, // blown away by set_mode()
@@ -394,7 +402,7 @@ impl<'a> DBIterator<'a> {
     pub(crate) fn new_cf(
         db: &DB,
         cf_handle: &ColumnFamily,
-        readopts: &ReadOptions,
+        readopts: ReadOptions,
         mode: IteratorMode,
     ) -> DBIterator<'a> {
         let mut rv = DBIterator {

src/db_options.rs

@@ -130,6 +130,7 @@ pub struct BlockBasedOptions {
 
 pub struct ReadOptions {
     pub(crate) inner: *mut ffi::rocksdb_readoptions_t,
+    iterate_upper_bound: Option<Vec<u8>>,
 }
 
 // Safety note: auto-implementing Send on most db-related types is prevented by the inner FFI
@@ -1635,21 +1636,22 @@ impl ReadOptions {
         }
     }
 
-    /// Set the upper bound for an iterator.
+    /// Sets the upper bound for an iterator.
     /// The upper bound itself is not included on the iteration result.
-    ///
-    /// # Safety
-    ///
-    /// This function will store a clone of key and will give a raw pointer of it to the
-    /// underlying C++ API, therefore, when given to any other [`DB`] method you must ensure
-    /// that this [`ReadOptions`] value does not leave the scope too early (e.g. `DB::iterator_cf_opt`).
-    pub unsafe fn set_iterate_upper_bound<K: AsRef<[u8]>>(&mut self, key: K) {
-        let key = key.as_ref();
-        ffi::rocksdb_readoptions_set_iterate_upper_bound(
-            self.inner,
-            key.as_ptr() as *const c_char,
-            key.len() as size_t,
-        );
+    pub fn set_iterate_upper_bound<K: Into<Vec<u8>>>(&mut self, key: K) {
+        self.iterate_upper_bound = Some(key.into());
+        let upper_bound = self
+            .iterate_upper_bound
+            .as_ref()
+            .expect("iterate_upper_bound must exist.");
+
+        unsafe {
+            ffi::rocksdb_readoptions_set_iterate_upper_bound(
+                self.inner,
+                upper_bound.as_ptr() as *const c_char,
+                upper_bound.len() as size_t,
+            );
+        }
     }
 
     pub fn set_prefix_same_as_start(&mut self, v: bool) {
@@ -1702,6 +1704,7 @@ impl Default for ReadOptions {
         unsafe {
             ReadOptions {
                 inner: ffi::rocksdb_readoptions_create(),
+                iterate_upper_bound: None,
             }
         }
     }

src/snapshot.rs

@@ -54,7 +54,7 @@ impl<'a> Snapshot<'a> {
 
     pub fn iterator_opt(&self, mode: IteratorMode, mut readopts: ReadOptions) -> DBIterator<'a> {
         readopts.set_snapshot(self);
-        DBIterator::new(self.db, &readopts, mode)
+        DBIterator::new(self.db, readopts, mode)
     }
 
     pub fn iterator_cf_opt(
@@ -64,7 +64,7 @@ impl<'a> Snapshot<'a> {
         mode: IteratorMode,
     ) -> DBIterator {
         readopts.set_snapshot(self);
-        DBIterator::new_cf(self.db, cf_handle, &readopts, mode)
+        DBIterator::new_cf(self.db, cf_handle, readopts, mode)
     }
 
     /// Opens a raw iterator over the data in this snapshot, using the default read options.
@@ -82,7 +82,7 @@ impl<'a> Snapshot<'a> {
     /// Opens a raw iterator over the data in this snapshot, using the given read options.
     pub fn raw_iterator_opt(&self, mut readopts: ReadOptions) -> DBRawIterator {
         readopts.set_snapshot(self);
-        DBRawIterator::new(self.db, &readopts)
+        DBRawIterator::new(self.db, readopts)
     }
 
     /// Opens a raw iterator over the data in this snapshot under the given column family, using the given read options.
@@ -92,7 +92,7 @@ impl<'a> Snapshot<'a> {
         mut readopts: ReadOptions,
     ) -> DBRawIterator {
         readopts.set_snapshot(self);
-        DBRawIterator::new_cf(self.db, cf_handle, &readopts)
+        DBRawIterator::new_cf(self.db, cf_handle, readopts)
     }
 
     pub fn get<K: AsRef<[u8]>>(&self, key: K) -> Result<Option<Vec<u8>>, Error> {