From 749cc74632ea34c4e85a954197ea4e7fd1e03701 Mon Sep 17 00:00:00 2001
From: Islam AbdelRahman <tec@fb.com>
Date: Tue, 8 Nov 2016 13:44:38 -0800
Subject: [PATCH] Fix Forward Iterator Seek()/SeekToFirst()

Summary:
In ForwardIterator::SeekInternal(), we may end up passing empty Slice representing an internal key to InternalKeyComparator::Compare.
and when we try to extract the user key from this empty Slice, we will create a slice with size = 0 - 8 ( which will overflow and cause us to read invalid memory as well )

Scenarios to reproduce these issues are in the unit tests
Closes https://github.com/facebook/rocksdb/pull/1467

Differential Revision: D4136660

Pulled By: lightmark

fbshipit-source-id: 151e128
---
 db/db_tailing_iter_test.cc | 52 ++++++++++++++++++++++++++++++++++++++
 db/forward_iterator.cc     | 12 ++++++---
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/db/db_tailing_iter_test.cc b/db/db_tailing_iter_test.cc
index bfb62926e..c8e233aba 100644
--- a/db/db_tailing_iter_test.cc
+++ b/db/db_tailing_iter_test.cc
@@ -699,6 +699,58 @@ TEST_F(DBTestTailingIterator, ForwardIteratorVersionProperty) {
   }
   ASSERT_EQ(v3, v4);
 }
+
+TEST_F(DBTestTailingIterator, SeekWithUpperBoundBug) {
+  ReadOptions read_options;
+  read_options.tailing = true;
+  const Slice upper_bound("cc", 3);
+  read_options.iterate_upper_bound = &upper_bound;
+
+
+  // 1st L0 file
+  ASSERT_OK(db_->Put(WriteOptions(), "aa", "SEEN"));
+  ASSERT_OK(Flush());
+
+  // 2nd L0 file
+  ASSERT_OK(db_->Put(WriteOptions(), "zz", "NOT-SEEN"));
+  ASSERT_OK(Flush());
+
+  std::unique_ptr<Iterator> iter(db_->NewIterator(read_options));
+
+  iter->Seek("aa");
+  ASSERT_TRUE(iter->Valid());
+  ASSERT_EQ(iter->key().ToString(), "aa");
+}
+
+TEST_F(DBTestTailingIterator, SeekToFirstWithUpperBoundBug) {
+  ReadOptions read_options;
+  read_options.tailing = true;
+  const Slice upper_bound("cc", 3);
+  read_options.iterate_upper_bound = &upper_bound;
+
+
+  // 1st L0 file
+  ASSERT_OK(db_->Put(WriteOptions(), "aa", "SEEN"));
+  ASSERT_OK(Flush());
+
+  // 2nd L0 file
+  ASSERT_OK(db_->Put(WriteOptions(), "zz", "NOT-SEEN"));
+  ASSERT_OK(Flush());
+
+  std::unique_ptr<Iterator> iter(db_->NewIterator(read_options));
+
+  iter->SeekToFirst();
+  ASSERT_TRUE(iter->Valid());
+  ASSERT_EQ(iter->key().ToString(), "aa");
+
+  iter->Next();
+  ASSERT_FALSE(iter->Valid());
+
+  iter->SeekToFirst();
+  ASSERT_TRUE(iter->Valid());
+  ASSERT_EQ(iter->key().ToString(), "aa");
+}
+
 }  // namespace rocksdb
 
 #endif  // !defined(ROCKSDB_LITE)
diff --git a/db/forward_iterator.cc b/db/forward_iterator.cc
index 2bd95e9f6..8a03d24aa 100644
--- a/db/forward_iterator.cc
+++ b/db/forward_iterator.cc
@@ -296,9 +296,15 @@ void ForwardIterator::SeekInternal(const Slice& internal_key,
   // an option to turn it off.
   if (seek_to_first || NeedToSeekImmutable(internal_key)) {
     immutable_status_ = Status::OK();
-    if ((has_iter_trimmed_for_upper_bound_) &&
-        (cfd_->internal_comparator().InternalKeyComparator::Compare(
-             prev_key_.GetKey(), internal_key) > 0)) {
+    if (has_iter_trimmed_for_upper_bound_ &&
+        (
+            // prev_ is not set yet
+            is_prev_set_ == false ||
+            // We are doing SeekToFirst() and internal_key.size() = 0
+            seek_to_first ||
+            // prev_key_ > internal_key
+            cfd_->internal_comparator().InternalKeyComparator::Compare(
+                prev_key_.GetKey(), internal_key) > 0)) {
       // Some iterators are trimmed. Need to rebuild.
       RebuildIterators(true);
       // Already seeked mutable iter, so seek again