Update SstFileWriter fuzzer to iterate and check all key-value pairs (#7761)

Summary:
as title

Pull Request resolved: https://github.com/facebook/rocksdb/pull/7761

Test Plan: cd fuzz && make sst_file_writer_fuzzer && ./sst_file_writer_fuzzer

Reviewed By: pdillinger

Differential Revision: D25430802

Pulled By: cheng-chang

fbshipit-source-id: 01436307df6f4c434bb608f44e1c8e4a1119f94f

6 files changed, 175 insertions, 81 deletions

diff --git a/fuzz/Makefile b/fuzz/Makefile
index 990fd2a54..fa45b9e78 100644
--- a/fuzz/Makefile
+++ b/fuzz/Makefile
@@ -3,7 +3,9 @@
 # COPYING file in the root directory) and Apache 2.0 License
 # (found in the LICENSE.Apache file in the root directory).
 
-CWD = $(shell pwd)
+ROOT_DIR = $(abspath $(shell pwd)/../)
+
+include $(ROOT_DIR)/make_config.mk
 
 PROTOBUF_CFLAGS = `pkg-config --cflags protobuf`
 PROTOBUF_LDFLAGS = `pkg-config --libs protobuf`
@@ -11,17 +13,17 @@ PROTOBUF_LDFLAGS = `pkg-config --libs protobuf`
 PROTOBUF_MUTATOR_CFLAGS = `pkg-config --cflags libprotobuf-mutator`
 PROTOBUF_MUTATOR_LDFLAGS = `pkg-config --libs libprotobuf-mutator`
 
-ROCKSDB_INCLUDE_DIR = $(CWD)/../include
-ROCKSDB_LIB_DIR = $(CWD)/..
+ROCKSDB_INCLUDE_DIR = $(ROOT_DIR)/include
+ROCKSDB_LIB_DIR = $(ROOT_DIR)
 
-PROTO_IN = $(CWD)/proto
-PROTO_OUT = $(CWD)/proto/gen
+PROTO_IN = $(ROOT_DIR)/fuzz/proto
+PROTO_OUT = $(ROOT_DIR)/fuzz/proto/gen
 
 ifneq ($(FUZZ_ENV), ossfuzz)
 CC = clang++
-CCFLAGS += -std=c++14 -Wall -fsanitize=address,fuzzer
-CFLAGS += $(PROTOBUF_CFLAGS) $(PROTOBUF_MUTATOR_CFLAGS) -I$(PROTO_OUT) -I$(ROCKSDB_INCLUDE_DIR)
-LDFLAGS += $(PROTOBUF_LDFLAGS) $(PROTOBUF_MUTATOR_LDFLAGS) -L$(ROCKSDB_LIB_DIR) -lrocksdb -lz -lbz2
+CCFLAGS += -Wall -fsanitize=address,fuzzer
+CFLAGS += $(PLATFORM_CXXFLAGS) $(PROTOBUF_CFLAGS) $(PROTOBUF_MUTATOR_CFLAGS) -I$(PROTO_OUT) -I$(ROCKSDB_INCLUDE_DIR) -I$(ROCKSDB_LIB_DIR)
+LDFLAGS += $(PLATFORM_LDFLAGS) $(PROTOBUF_LDFLAGS) $(PROTOBUF_MUTATOR_LDFLAGS) -L$(ROCKSDB_LIB_DIR) -lrocksdb
 else
 # OSS-Fuzz sets various environment flags that are used for compilation.
 # These environment flags depend on which type of sanitizer build is being
@@ -36,14 +38,11 @@ else
 # LIB_FUZZING_ENGINE="-fsanitize=fuzzer"
 CC = $(CXX)
 CCFLAGS = $(CXXFLAGS)
-CFLAGS += $(PROTOBUF_CFLAGS) $(PROTOBUF_MUTATOR_CFLAGS) -I$(PROTO_OUT) -I$(ROCKSDB_INCLUDE_DIR)
-LDFLAGS += $(LIB_FUZZING_ENGINE) $(PROTOBUF_MUTATOR_LDFLAGS) $(PROTOBUF_LDFLAGS) -L$(ROCKSDB_LIB_DIR) -lrocksdb -lz -lbz2
+CFLAGS += $(PROTOBUF_CFLAGS) $(PROTOBUF_MUTATOR_CFLAGS) -I$(PROTO_OUT) -I$(ROCKSDB_INCLUDE_DIR) -I$(ROCKSDB_LIB_DIR)
+LDFLAGS += $(PLATFORM_LDFLAGS) $(LIB_FUZZING_ENGINE) $(PROTOBUF_MUTATOR_LDFLAGS) $(PROTOBUF_LDFLAGS) -L$(ROCKSDB_LIB_DIR) -lrocksdb
 endif
 
-.PHONY: librocksdb gen_proto
-
-librocksdb:
-	cd $(CWD)/.. && make static_lib
+.PHONY: gen_proto
 
 gen_proto:
 	mkdir -p $(PROTO_OUT)
@@ -52,11 +51,11 @@ gen_proto:
 		--cpp_out=$(PROTO_OUT) \
 		$(PROTO_IN)/*.proto
 
-db_fuzzer: librocksdb db_fuzzer.cc
+db_fuzzer: db_fuzzer.cc
 	$(CC) $(CCFLAGS) -o db_fuzzer db_fuzzer.cc $(CFLAGS) $(LDFLAGS)
 
-db_map_fuzzer: librocksdb gen_proto db_map_fuzzer.cc proto/gen/db_operation.pb.cc
+db_map_fuzzer: gen_proto db_map_fuzzer.cc proto/gen/db_operation.pb.cc
 	$(CC) $(CCFLAGS) -o db_map_fuzzer db_map_fuzzer.cc proto/gen/db_operation.pb.cc $(CFLAGS) $(LDFLAGS)
 
-sst_file_writer_fuzzer: librocksdb gen_proto sst_file_writer_fuzzer.cc proto/gen/db_operation.pb.cc
+sst_file_writer_fuzzer: gen_proto sst_file_writer_fuzzer.cc proto/gen/db_operation.pb.cc
 	$(CC) $(CCFLAGS) -o sst_file_writer_fuzzer sst_file_writer_fuzzer.cc proto/gen/db_operation.pb.cc $(CFLAGS) $(LDFLAGS)
diff --git a/fuzz/README.md b/fuzz/README.md
index 48158e039..79b89bbc3 100644
--- a/fuzz/README.md
+++ b/fuzz/README.md
@@ -56,6 +56,8 @@ protobuf_mutator::libfuzzer::PostProcessorRegistration<DBOperations>
 
 ### Compile and link the fuzzer
 
+In the rocksdb root directory, compile rocksdb library by `make static_lib`.
+
 Go to the `fuzz` directory,
 run `make sst_file_writer_fuzzer` to generate the fuzzer,
 it will compile rocksdb static library, generate protobuf, then compile and link `sst_file_writer_fuzzer`.
diff --git a/fuzz/db_map_fuzzer.cc b/fuzz/db_map_fuzzer.cc
index b9672018a..4d9637ad9 100644
--- a/fuzz/db_map_fuzzer.cc
+++ b/fuzz/db_map_fuzzer.cc
@@ -12,28 +12,7 @@
 #include "rocksdb/db.h"
 #include "rocksdb/file_system.h"
 #include "src/libfuzzer/libfuzzer_macro.h"
-
-#define CHECK_OK(expression)                       \
-  do {                                             \
-    auto status = (expression);                    \
-    if (!status.ok()) {                            \
-      std::cerr << status.ToString() << std::endl; \
-      abort();                                     \
-    }                                              \
-  } while (0)
-
-#define CHECK_EQ(a, b)                                                      \
-  if (a != b) {                                                             \
-    std::cerr << "(" << #a << "=" << a << ") != (" << #b << "=" << b << ")" \
-              << std::endl;                                                 \
-    abort();                                                                \
-  }
-
-#define CHECK_TRUE(cond)                                      \
-  if (!(cond)) {                                              \
-    std::cerr << "\"" << #cond << "\" is false" << std::endl; \
-    abort();                                                  \
-  }
+#include "util.h"
 
 protobuf_mutator::libfuzzer::PostProcessorRegistration<DBOperations> reg = {
     [](DBOperations* input, unsigned int /* seed */) {
@@ -81,6 +60,9 @@ DEFINE_PROTO_FUZZER(DBOperations& input) {
         kv[op.key()] = op.value();
         break;
       }
+      case OpType::MERGE: {
+        break;
+      }
       case OpType::DELETE: {
         CHECK_OK(db->Delete(rocksdb::WriteOptions(), op.key()));
         kv.erase(op.key());
diff --git a/fuzz/proto/db_operation.proto b/fuzz/proto/db_operation.proto
index d6921593b..20a55eaa5 100644
--- a/fuzz/proto/db_operation.proto
+++ b/fuzz/proto/db_operation.proto
@@ -10,6 +10,7 @@ syntax = "proto2";
 
 enum OpType {
   PUT = 0;
+  MERGE = 1;
   DELETE = 2;
   DELETE_RANGE = 3;
 }
diff --git a/fuzz/sst_file_writer_fuzzer.cc b/fuzz/sst_file_writer_fuzzer.cc
index 61721b815..7bc128fc0 100644
--- a/fuzz/sst_file_writer_fuzzer.cc
+++ b/fuzz/sst_file_writer_fuzzer.cc
@@ -5,20 +5,39 @@
 
 #include <algorithm>
 #include <iostream>
+#include <memory>
+#include <string>
 
 #include "proto/gen/db_operation.pb.h"
 #include "rocksdb/file_system.h"
-#include "rocksdb/sst_file_reader.h"
 #include "rocksdb/sst_file_writer.h"
 #include "src/libfuzzer/libfuzzer_macro.h"
+#include "table/table_reader.h"
+#include "util.h"
+
+using namespace ROCKSDB_NAMESPACE;
 
 // Keys in SST file writer operations must be unique and in ascending order.
 // For each DBOperation generated by the fuzzer, this function is called on
 // it to deduplicate and sort the keys in the DBOperations.
 protobuf_mutator::libfuzzer::PostProcessorRegistration<DBOperations> reg = {
     [](DBOperations* input, unsigned int /* seed */) {
-      const rocksdb::Comparator* comparator = rocksdb::BytewiseComparator();
+      const Comparator* comparator = BytewiseComparator();
       auto ops = input->mutable_operations();
+
+      // Make sure begin <= end for DELETE_RANGE.
+      for (DBOperation& op : *ops) {
+        if (op.type() == OpType::DELETE_RANGE) {
+          auto begin = op.key();
+          auto end = op.value();
+          if (comparator->Compare(begin, end) > 0) {
+            std::swap(begin, end);
+            op.set_key(begin);
+            op.set_value(end);
+          }
+        }
+      }
+
       std::sort(ops->begin(), ops->end(),
                 [&comparator](const DBOperation& a, const DBOperation& b) {
                   return comparator->Compare(a.key(), b.key()) < 0;
@@ -32,15 +51,60 @@ protobuf_mutator::libfuzzer::PostProcessorRegistration<DBOperations> reg = {
       ops->erase(last, ops->end());
     }};
 
-#define CHECK_OK(status)                         \
-  if (!status.ok()) {                            \
-    std::cerr << status.ToString() << std::endl; \
-    abort();                                     \
+TableReader* NewTableReader(const std::string& sst_file_path,
+                            const Options& options,
+                            const EnvOptions& env_options,
+                            const ImmutableCFOptions& cf_ioptions) {
+  // This code block is similar to SstFileReader::Open.
+
+  uint64_t file_size = 0;
+  std::unique_ptr<RandomAccessFile> file;
+  std::unique_ptr<RandomAccessFileReader> file_reader;
+  std::unique_ptr<TableReader> table_reader;
+  Status s = options.env->GetFileSize(sst_file_path, &file_size);
+  if (s.ok()) {
+    s = options.env->NewRandomAccessFile(sst_file_path, &file, env_options);
+  }
+  if (s.ok()) {
+    file_reader.reset(new RandomAccessFileReader(
+        NewLegacyRandomAccessFileWrapper(file), sst_file_path));
   }
+  if (s.ok()) {
+    TableReaderOptions t_opt(cf_ioptions, /*prefix_extractor=*/nullptr,
+                             env_options, cf_ioptions.internal_comparator);
+    t_opt.largest_seqno = kMaxSequenceNumber;
+    s = options.table_factory->NewTableReader(t_opt, std::move(file_reader),
+                                              file_size, &table_reader,
+                                              /*prefetch=*/false);
+  }
+  if (!s.ok()) {
+    std::cerr << "Failed to create TableReader for " << sst_file_path << ": "
+              << s.ToString() << std::endl;
+    abort();
+  }
+  return table_reader.release();
+}
+
+ValueType ToValueType(OpType op_type) {
+  switch (op_type) {
+    case OpType::PUT:
+      return ValueType::kTypeValue;
+    case OpType::MERGE:
+      return ValueType::kTypeMerge;
+    case OpType::DELETE:
+      return ValueType::kTypeDeletion;
+    case OpType::DELETE_RANGE:
+      return ValueType::kTypeRangeDeletion;
+    default:
+      std::cerr << "Unknown operation type " << static_cast<int>(op_type)
+                << std::endl;
+      abort();
+  }
+}
 
 // Fuzzes DB operations as input, let SstFileWriter generate a SST file
-// according to the operations, then let SstFileReader read the generated SST
-// file to check its checksum.
+// according to the operations, then let TableReader read and check all the
+// key-value pairs from the generated SST file.
 DEFINE_PROTO_FUZZER(DBOperations& input) {
   if (input.operations().empty()) {
     return;
@@ -48,53 +112,76 @@ DEFINE_PROTO_FUZZER(DBOperations& input) {
 
   std::string sstfile;
   {
-    auto fs = rocksdb::FileSystem::Default();
+    auto fs = FileSystem::Default();
     std::string dir;
-    rocksdb::IOOptions opt;
-    rocksdb::IOStatus s = fs->GetTestDirectory(opt, &dir, nullptr);
-    CHECK_OK(s);
+    IOOptions opt;
+    CHECK_OK(fs->GetTestDirectory(opt, &dir, nullptr));
     sstfile = dir + "/SstFileWriterFuzzer.sst";
   }
 
+  Options options;
+  EnvOptions env_options(options);
+  ImmutableCFOptions cf_ioptions(options);
+
   // Generate sst file.
-  rocksdb::Options options;
-  rocksdb::EnvOptions env_options;
-  rocksdb::SstFileWriter writer(env_options, options);
-  rocksdb::Status s = writer.Open(sstfile);
-  CHECK_OK(s);
+  SstFileWriter writer(env_options, options);
+  CHECK_OK(writer.Open(sstfile));
   for (const DBOperation& op : input.operations()) {
     switch (op.type()) {
-      case OpType::PUT:
-        s = writer.Put(op.key(), op.value());
-        CHECK_OK(s);
+      case OpType::PUT: {
+        CHECK_OK(writer.Put(op.key(), op.value()));
         break;
-      case OpType::MERGE:
-        s = writer.Merge(op.key(), op.value());
-        CHECK_OK(s);
+      }
+      case OpType::MERGE: {
+        CHECK_OK(writer.Merge(op.key(), op.value()));
         break;
-      case OpType::DELETE:
-        s = writer.Delete(op.key());
-        CHECK_OK(s);
+      }
+      case OpType::DELETE: {
+        CHECK_OK(writer.Delete(op.key()));
         break;
-      case OpType::DELETE_RANGE:
-        s = writer.DeleteRange(op.key(), op.value());
-        CHECK_OK(s);
+      }
+      case OpType::DELETE_RANGE: {
+        CHECK_OK(writer.DeleteRange(op.key(), op.value()));
         break;
-      default:
-        std::cerr << "Unsupported operation" << static_cast<int>(op.type());
-        return;
+      }
+      default: {
+        std::cerr << "Unsupported operation" << static_cast<int>(op.type())
+                  << std::endl;
+        abort();
+      }
+    }
+  }
+  ExternalSstFileInfo info;
+  CHECK_OK(writer.Finish(&info));
+
+  // Iterate and verify key-value pairs.
+  std::unique_ptr<TableReader> table_reader(
+      NewTableReader(sstfile, options, env_options, cf_ioptions));
+  ReadOptions roptions;
+  CHECK_OK(table_reader->VerifyChecksum(roptions,
+                                        TableReaderCaller::kUncategorized));
+  std::unique_ptr<InternalIterator> it(
+      table_reader->NewIterator(roptions, /*prefix_extractor=*/nullptr,
+                                /*arena=*/nullptr, /*skip_filters=*/true,
+                                TableReaderCaller::kUncategorized));
+  it->SeekToFirst();
+  for (const DBOperation& op : input.operations()) {
+    if (op.type() == OpType::DELETE_RANGE) {
+      // InternalIterator cannot iterate over DELETE_RANGE entries.
+      continue;
+    }
+    CHECK_TRUE(it->Valid());
+    ParsedInternalKey ikey;
+    CHECK_OK(ParseInternalKey(it->key(), &ikey, /*log_err_key=*/true));
+    CHECK_EQ(ikey.user_key.ToString(), op.key());
+    CHECK_EQ(ikey.sequence, 0);
+    CHECK_EQ(ikey.type, ToValueType(op.type()));
+    if (op.type() != OpType::DELETE) {
+      CHECK_EQ(op.value(), it->value().ToString());
     }
+    it->Next();
   }
-  rocksdb::ExternalSstFileInfo info;
-  s = writer.Finish(&info);
-  CHECK_OK(s);
-
-  // Verify checksum.
-  rocksdb::SstFileReader reader(options);
-  s = reader.Open(sstfile);
-  CHECK_OK(s);
-  s = reader.VerifyChecksum();
-  CHECK_OK(s);
+  CHECK_TRUE(!it->Valid());
 
   // Delete sst file.
   remove(sstfile.c_str());
diff --git a/fuzz/util.h b/fuzz/util.h
new file mode 100644
index 000000000..44ffaf536
--- /dev/null
+++ b/fuzz/util.h
@@ -0,0 +1,23 @@
+#pragma once
+
+#define CHECK_OK(expression)                       \
+  do {                                             \
+    auto status = (expression);                    \
+    if (!status.ok()) {                            \
+      std::cerr << status.ToString() << std::endl; \
+      abort();                                     \
+    }                                              \
+  } while (0)
+
+#define CHECK_EQ(a, b)                                                      \
+  if (a != b) {                                                             \
+    std::cerr << "(" << #a << "=" << a << ") != (" << #b << "=" << b << ")" \
+              << std::endl;                                                 \
+    abort();                                                                \
+  }
+
+#define CHECK_TRUE(cond)                                      \
+  if (!(cond)) {                                              \
+    std::cerr << "\"" << #cond << "\" is false" << std::endl; \
+    abort();                                                  \
+  }