From ff4b3fb5b4a698b752855ab72ff443814b2cbac1 Mon Sep 17 00:00:00 2001
From: Islam AbdelRahman <tec@fb.com>
Date: Tue, 3 May 2016 16:50:01 -0700
Subject: [PATCH] Fix Iterator::Prev memory pinning bug

Summary: We should not use IterKey::SetKey with copy = false except if we are pinning the iterator thru it's life time, otherwise we may release the temporarily pinned blocks and in this case the IterKey will be pointing to freed memory

Test Plan: added a new test

Reviewers: sdong, andrewkr

Reviewed By: andrewkr

Subscribers: andrewkr, dhruba

Differential Revision: https://reviews.facebook.net/D57561
---
 db/db_iter.cc          | 17 +++++++------
 db/db_iterator_test.cc | 57 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 7 deletions(-)

diff --git a/db/db_iter.cc b/db/db_iter.cc
index 08aaef333..27762232e 100644
--- a/db/db_iter.cc
+++ b/db/db_iter.cc
@@ -369,21 +369,24 @@ void DBIter::FindNextUserEntryInternal(bool skipping, bool prefix_check) {
             case kTypeSingleDeletion:
               // Arrange to skip all upcoming entries for this key since
               // they are hidden by this deletion.
-              saved_key_.SetKey(ikey.user_key,
-                                !iter_->IsKeyPinned() /* copy */);
+              saved_key_.SetKey(
+                  ikey.user_key,
+                  !iter_->IsKeyPinned() || !pin_thru_lifetime_ /* copy */);
               skipping = true;
               num_skipped = 0;
               PERF_COUNTER_ADD(internal_delete_skipped_count, 1);
               break;
             case kTypeValue:
               valid_ = true;
-              saved_key_.SetKey(ikey.user_key,
-                                !iter_->IsKeyPinned() /* copy */);
+              saved_key_.SetKey(
+                  ikey.user_key,
+                  !iter_->IsKeyPinned() || !pin_thru_lifetime_ /* copy */);
               return;
             case kTypeMerge:
               // By now, we are sure the current ikey is going to yield a value
-              saved_key_.SetKey(ikey.user_key,
-                                !iter_->IsKeyPinned() /* copy */);
+              saved_key_.SetKey(
+                  ikey.user_key,
+                  !iter_->IsKeyPinned() || !pin_thru_lifetime_ /* copy */);
               current_entry_is_merged_ = true;
               valid_ = true;
               MergeValuesNewToOld();  // Go to a different state machine
@@ -547,7 +550,7 @@ void DBIter::PrevInternal() {
 
   while (iter_->Valid()) {
     saved_key_.SetKey(ExtractUserKey(iter_->key()),
-                      !iter_->IsKeyPinned() /* copy */);
+                      !iter_->IsKeyPinned() || !pin_thru_lifetime_ /* copy */);
     if (FindValueForCurrentKey()) {
       valid_ = true;
       if (!iter_->Valid()) {
diff --git a/db/db_iterator_test.cc b/db/db_iterator_test.cc
index 3a00e5db6..48f7cf9e1 100644
--- a/db/db_iterator_test.cc
+++ b/db/db_iterator_test.cc
@@ -1384,6 +1384,63 @@ TEST_F(DBIteratorTest, IterPrevKeyCrossingBlocksRandomized) {
 
     delete iter;
   }
+
+  {
+    ReadOptions ro;
+    ro.fill_cache = false;
+    Iterator* iter = db_->NewIterator(ro);
+    auto data_iter = true_data.rbegin();
+
+    int entries_right = 0;
+    std::string seek_key;
+    for (iter->SeekToLast(); iter->Valid(); iter->Prev()) {
+      // Verify key/value of current position
+      ASSERT_EQ(iter->key().ToString(), data_iter->first);
+      ASSERT_EQ(iter->value().ToString(), data_iter->second);
+
+      bool restore_position_with_seek = rnd.Uniform(2);
+      if (restore_position_with_seek) {
+        seek_key = iter->key().ToString();
+      }
+
+      // Do some Next() operations the restore the iterator to orignal position
+      int next_count =
+          entries_right > 0 ? rnd.Uniform(std::min(entries_right, 10)) : 0;
+      for (int i = 0; i < next_count; i++) {
+        iter->Next();
+        data_iter--;
+
+        ASSERT_EQ(iter->key().ToString(), data_iter->first);
+        ASSERT_EQ(iter->value().ToString(), data_iter->second);
+      }
+
+      if (restore_position_with_seek) {
+        // Restore orignal position using Seek()
+        iter->Seek(seek_key);
+        for (int i = 0; i < next_count; i++) {
+          data_iter++;
+        }
+
+        ASSERT_EQ(iter->key().ToString(), data_iter->first);
+        ASSERT_EQ(iter->value().ToString(), data_iter->second);
+      } else {
+        // Restore original position using Prev()
+        for (int i = 0; i < next_count; i++) {
+          iter->Prev();
+          data_iter++;
+
+          ASSERT_EQ(iter->key().ToString(), data_iter->first);
+          ASSERT_EQ(iter->value().ToString(), data_iter->second);
+        }
+      }
+
+      entries_right++;
+      data_iter++;
+    }
+    ASSERT_EQ(data_iter, true_data.rend());
+
+    delete iter;
+  }
 }
 
 TEST_F(DBIteratorTest, IteratorWithLocalStatistics) {
-- 
GitLab