#include <algorithm>
#include <iostream>
#include <map>
#include <string>
#include "proto/gen/db_operation.pb.h"
#include "rocksdb/db.h"
#include "rocksdb/file_system.h"
#include "src/libfuzzer/libfuzzer_macro.h"
#include "util.h"
protobuf_mutator::libfuzzer::PostProcessorRegistration<DBOperations> reg = {
[](DBOperations* input, unsigned int ) {
const rocksdb::Comparator* comparator = rocksdb::BytewiseComparator();
auto ops = input->mutable_operations();
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);
}
}
}
}};
DEFINE_PROTO_FUZZER(DBOperations& input) {
if (input.operations().empty()) {
return;
}
const std::string kDbPath = "/tmp/db_map_fuzzer_test";
auto fs = rocksdb::FileSystem::Default();
if (fs->FileExists(kDbPath, rocksdb::IOOptions(), nullptr).ok()) {
std::cerr << "db path " << kDbPath << " already exists" << std::endl;
abort();
}
std::map<std::string, std::string> kv;
rocksdb::DB* db = nullptr;
rocksdb::Options options;
options.create_if_missing = true;
CHECK_OK(rocksdb::DB::Open(options, kDbPath, &db));
for (const DBOperation& op : input.operations()) {
switch (op.type()) {
case OpType::PUT: {
CHECK_OK(db->Put(rocksdb::WriteOptions(), op.key(), op.value()));
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());
break;
}
case OpType::DELETE_RANGE: {
CHECK_OK(db->DeleteRange(rocksdb::WriteOptions(),
db->DefaultColumnFamily(), op.key(),
op.value()));
kv.erase(kv.lower_bound(op.key()), kv.lower_bound(op.value()));
break;
}
default: {
std::cerr << "Unsupported operation" << static_cast<int>(op.type());
return;
}
}
}
CHECK_OK(db->Close());
delete db;
db = nullptr;
CHECK_OK(rocksdb::DB::Open(options, kDbPath, &db));
auto kv_it = kv.begin();
rocksdb::Iterator* it = db->NewIterator(rocksdb::ReadOptions());
for (it->SeekToFirst(); it->Valid(); it->Next(), kv_it++) {
CHECK_TRUE(kv_it != kv.end());
CHECK_EQ(it->key().ToString(), kv_it->first);
CHECK_EQ(it->value().ToString(), kv_it->second);
}
CHECK_TRUE(kv_it == kv.end());
delete it;
CHECK_OK(db->Close());
delete db;
CHECK_OK(rocksdb::DestroyDB(kDbPath, options));
}