[#3455] Modify ldb Tool to be Encryption Aware

Summary: Add `--key_file` option to ldb to pass in a `<key_id>:</path/to/key/file>` and an option `--only_verify_checksums` flag to `scan` to not print out entries and only do checksum validation. Also provide a yb-admin `write_universe_key_to_file` command to retrieve the universe key and write it to a local file.

Test Plan:
Jenkins: build platform: mac
Jenkins: compile only

Reviewers: bogdan, mikhail, timur

Reviewed By: timur

Subscribers: ybase

Differential Revision: https://phabricator.dev.yugabyte.com/D7876

ent/src/yb/tools/yb-admin_cli_ent.cc

@@ -237,6 +237,17 @@ void ClusterAdminCli::RegisterCommandHandlers(ClusterAdminClientClass* client) {
         return Status::OK();
       });
 
+  Register(
+      "write_universe_key_to_file", " <key_id> <file_name>",
+      [client](const CLIArguments& args) -> Status {
+        if (args.size() != 4) {
+          return ClusterAdminCli::kInvalidArguments;
+        }
+        RETURN_NOT_OK_PREPEND(client->WriteUniverseKeyToFile(args[2], args[3]),
+                              "Unable to write key to file");
+        return Status::OK();
+      });
+
   Register(
       "create_cdc_stream", " <table_id>",
       [client](const CLIArguments& args) -> Status {

ent/src/yb/tools/yb-admin_client.h

@@ -65,6 +65,8 @@ class ClusterAdminClient : public yb::tools::ClusterAdminClient {
 
   CHECKED_STATUS DisableEncryptionInMemory();
 
+  CHECKED_STATUS WriteUniverseKeyToFile(const std::string& key_id, const std::string& file_name);
+
   CHECKED_STATUS CreateCDCStream(const TableId& table_id);
 
   CHECKED_STATUS SetupUniverseReplication(const std::string& producer_uuid,

ent/src/yb/tools/yb-admin_client_ent.cc

@@ -632,6 +632,32 @@ Status ClusterAdminClient::DisableEncryptionInMemory() {
   return Status::OK();
 }
 
+Status ClusterAdminClient::WriteUniverseKeyToFile(
+    const std::string& key_id, const std::string& file_name) {
+  RETURN_NOT_OK_PREPEND(WaitUntilMasterLeaderReady(), "Wait for master leader failed!");
+  rpc::RpcController rpc;
+  rpc.set_timeout(timeout_);
+
+  master::GetUniverseKeyRegistryRequestPB req;
+  master::GetUniverseKeyRegistryResponsePB resp;
+  RETURN_NOT_OK_PREPEND(master_proxy_->GetUniverseKeyRegistry(req, &resp, &rpc),
+                        "MasterServiceImpl::ChangeEncryptionInfo call fails.");
+  if (resp.has_error()) {
+    return StatusFromPB(resp.error().status());
+  }
+
+  auto universe_keys = resp.universe_keys();
+  const auto& it = universe_keys.map().find(key_id);
+  if (it == universe_keys.map().end()) {
+    return STATUS_FORMAT(NotFound, "Could not find key with id $0", key_id);
+  }
+
+  RETURN_NOT_OK(WriteStringToFile(Env::Default(), Slice(it->second), file_name));
+
+  std::cout << "Finished writing to file\n";
+  return Status::OK();
+}
+
 Status ClusterAdminClient::CreateCDCStream(const TableId& table_id) {
   master::CreateCDCStreamRequestPB req;
   master::CreateCDCStreamResponsePB resp;

ent/src/yb/tserver/CMakeLists-include.txt

@@ -20,7 +20,6 @@ string(REPLACE ${CMAKE_SOURCE_DIR} ${CMAKE_SOURCE_DIR}/ent
 set(TSERVER_SRCS_EXTENSIONS
   ${YB_ENT_CURRENT_SOURCE_DIR}/backup_service.cc
   ${YB_ENT_CURRENT_SOURCE_DIR}/tablet_server_ent.cc
-  ${YB_ENT_CURRENT_SOURCE_DIR}/header_manager_impl.cc
   ${YB_ENT_CURRENT_SOURCE_DIR}/cdc_consumer.cc
   ${YB_ENT_CURRENT_SOURCE_DIR}/twodc_output_client.cc 
   ${YB_ENT_CURRENT_SOURCE_DIR}/cdc_poller.cc
@@ -38,6 +37,5 @@ set(TSERVER_EXTENSIONS_TESTS
   backup_service-test
   remote_bootstrap_rocksdb_session-test_ent
   remote_bootstrap_rocksdb_client-test_ent
-  header_manager_impl-test
   PARENT_SCOPE)
 

ent/src/yb/tserver/tablet_server_ent.cc

@@ -20,7 +20,6 @@
 #include "yb/rpc/rpc.h"
 #include "yb/tserver/backup_service.h"
 #include "yb/tserver/cdc_consumer.h"
-#include "yb/tserver/header_manager_impl.h"
 #include "yb/tserver/tablet_server.h"
 #include "yb/tserver/ts_tablet_manager.h"
 
@@ -29,6 +28,7 @@
 #include "yb/util/ntp_clock.h"
 #include "yb/util/encrypted_file_factory.h"
 #include "yb/util/universe_key_manager.h"
+#include "yb/util/header_manager_impl.h"
 
 #include "yb/rocksutil/rocksdb_encrypted_file_factory.h"
 

src/yb/rocksdb/CMakeLists.txt

@@ -224,7 +224,7 @@ add_library(rocksdb_tools
   tools/ldb_cmd.cc
   tools/ldb_tool.cc
   tools/sst_dump_tool.cc)
-target_link_libraries(rocksdb_tools rocksdb)
+target_link_libraries(rocksdb_tools rocksdb yb_rocksutil)
 
 macro(ADD_YB_ROCKSDB_TOOL tool)
   add_executable(${tool} tools/${tool}.cc)
@@ -243,7 +243,7 @@ if (NOT "${CMAKE_BUILD_TYPE}" STREQUAL "RELEASE")
 endif()
 
 add_library(rocksdb_test_util ${ROCKSDB_TEST_UTIL_SRCS})
-target_link_libraries(rocksdb_test_util gmock rocksdb)
+target_link_libraries(rocksdb_test_util gmock rocksdb yb_rocksutil)
 
 ADD_YB_ROCKSDB_TOOL(sst_dump)
 add_executable(db_bench tools/db_bench.cc tools/db_bench_tool.cc)
@@ -259,7 +259,7 @@ target_link_libraries(rocksdb_dump rocksdb_tools)
 add_executable(rocksdb_undump tools/dump/rocksdb_undump.cc)
 target_link_libraries(rocksdb_undump rocksdb_tools)
 
-set(YB_TEST_LINK_LIBS rocksdb_test_util rocksdb ${YB_MIN_TEST_LIBS})
+set(YB_TEST_LINK_LIBS rocksdb_test_util rocksdb ${YB_MIN_TEST_LIBS} rocksdb_tools)
 
 ADD_YB_TEST(db/compaction_job_test)
 ADD_YB_TEST(db/compaction_picker_test)
@@ -290,7 +290,6 @@ ADD_YB_TEST(table/full_filter_block_test)
 ADD_YB_TEST(table/fixed_size_filter_block_test)
 ADD_YB_TEST(table/merger_test)
 ADD_YB_TEST(table/table_test)
-ADD_YB_TEST(tools/ldb_cmd_test)
 ADD_YB_TEST(tools/sst_dump_test)
 ADD_YB_TEST(util/arena_test)
 ADD_YB_TEST(util/autovector_test)
@@ -357,6 +356,7 @@ ADD_YB_TEST(db/version_edit_test)
 ADD_YB_TEST(db/write_callback_test)
 ADD_YB_TEST(table/cuckoo_table_builder_test)
 ADD_YB_TEST(table/cuckoo_table_reader_test)
+ADD_YB_TEST(tools/ldb_cmd_test)
 ADD_YB_TEST(tools/reduce_levels_test)
 ADD_YB_TEST(util/delete_scheduler_test)
 ADD_YB_TEST(util/thread_local_test)

src/yb/rocksdb/db/db_test_util.cc

@@ -23,6 +23,12 @@
 
 #include "yb/rocksdb/db/db_test_util.h"
 
+#include "yb/util/encryption_util.h"
+#include "yb/util/random_util.h"
+#include "yb/util/header_manager_impl.h"
+#include "yb/util/universe_key_manager.h"
+#include "yb/rocksutil/rocksdb_encrypted_file_factory.h"
+
 namespace rocksdb {
 
 // Special Env used to delay background operations
@@ -53,10 +59,16 @@ SpecialEnv::SpecialEnv(Env* base)
   table_write_callback_ = nullptr;
 }
 
-DBTestBase::DBTestBase(const std::string path)
+const string DBTestBase::kKeyId = "key_id";
+const string DBTestBase::kKeyFile = "universe_key_file";
+
+DBTestBase::DBTestBase(const std::string path, bool encryption_enabled)
     : option_config_(kDefault),
       mem_env_(!getenv("MEM_ENV") ? nullptr : new MockEnv(Env::Default())),
       env_(new SpecialEnv(mem_env_ ? mem_env_ : Env::Default())) {
+  if (encryption_enabled) {
+    CreateEncryptedEnv();
+  }
   env_->SetBackgroundThreads(1, Env::LOW);
   env_->SetBackgroundThreads(1, Env::HIGH);
   dbname_ = test::TmpDir(env_) + path;
@@ -74,6 +86,7 @@ DBTestBase::DBTestBase(const std::string path)
 }
 
 DBTestBase::~DBTestBase() {
+  Env::Default()->DeleteFile(kKeyFile);
   rocksdb::SyncPoint::GetInstance()->DisableProcessing();
   rocksdb::SyncPoint::GetInstance()->LoadDependency({});
   rocksdb::SyncPoint::GetInstance()->ClearAllCallBacks();
@@ -87,6 +100,25 @@ DBTestBase::~DBTestBase() {
   delete env_;
 }
 
+void DBTestBase::CreateEncryptedEnv() {
+  auto bytes = yb::RandomBytes(32);
+  yb::Slice key(bytes.data(), bytes.size());
+  auto status = yb::WriteStringToFile(yb::Env::Default(), key, kKeyFile);
+  if (!status.ok()) {
+    LOG(FATAL) << "Could not write slice to file:" << status.ToString();
+  }
+
+  auto res = yb::enterprise::UniverseKeyManager::FromKey(kKeyId, key);
+  if (!res.ok()) {
+    LOG(FATAL) << "Could not get key from bytes:" << res.status().ToString();
+  }
+  universe_key_manager_ = std::move(*res);
+  encrypted_env_ = yb::enterprise::NewRocksDBEncryptedEnv(
+      yb::enterprise::DefaultHeaderManager(universe_key_manager_.get()));
+  delete env_;
+  env_ = new rocksdb::SpecialEnv(encrypted_env_.get());
+}
+
 bool DBTestBase::ShouldSkipOptions(int option_config, int skip_mask) {
 #ifdef ROCKSDB_LITE
     // These options are not supported in ROCKSDB_LITE
@@ -400,6 +432,7 @@ Options DBTestBase::CurrentOptions(
     options.table_factory.reset(NewBlockBasedTableFactory(table_options));
   }
   options.env = env_;
+  options.checkpoint_env = env_->IsPlainText() ? env_ : Env::Default();
   options.create_if_missing = true;
   options.fail_if_options_file_error = true;
   return options;

src/yb/rocksdb/db/db_test_util.h

@@ -75,6 +75,14 @@
 #include "yb/rocksdb/util/xfunc.h"
 #include "yb/rocksdb/utilities/merge_operators.h"
 
+namespace yb {
+namespace enterprise {
+
+class UniverseKeyManager;
+
+} // namespace enterprise
+} // namespace yb
+
 namespace rocksdb {
 
 uint64_t TestGetTickerCount(const Options& options, Tickers ticker_type) {
@@ -591,6 +599,13 @@ class DBTestBase : public testing::Test {
 
   Options last_options_;
 
+  // For encryption
+  std::unique_ptr<yb::enterprise::UniverseKeyManager> universe_key_manager_;
+  std::unique_ptr<rocksdb::Env> encrypted_env_;
+
+  static const std::string kKeyId;
+  static const std::string kKeyFile;
+
   // Skip some options, as they may not be applicable to a specific test.
   // To add more skip constants, use values 4, 8, 16, etc.
   enum OptionSkip {
@@ -606,10 +621,12 @@ class DBTestBase : public testing::Test {
     kSkipMmapReads = 256,
   };
 
-  explicit DBTestBase(const std::string path);
+  explicit DBTestBase(const std::string path, bool encryption_enabled = false);
 
   ~DBTestBase();
 
+  void CreateEncryptedEnv();
+
   static std::string Key(int i) {
     char buf[100];
     snprintf(buf, sizeof(buf), "key%06d", i);

src/yb/rocksdb/tools/ldb_cmd.cc

@@ -76,6 +76,8 @@ const string LDBCommand::ARG_WRITE_BUFFER_SIZE = "write_buffer_size";
 const string LDBCommand::ARG_FILE_SIZE = "file_size";
 const string LDBCommand::ARG_CREATE_IF_MISSING = "create_if_missing";
 const string LDBCommand::ARG_NO_VALUE = "no_value";
+const string LDBCommand::ARG_UNIVERSE_KEY_FILE = "key_file";
+const string LDBCommand::ARG_ONLY_VERIFY_CHECKSUMS = "only_verify_checksums";
 
 const char* LDBCommand::DELIM = " ==> ";
 
@@ -126,12 +128,12 @@ LDBCommand* LDBCommand::InitFromCmdLineArgs(
 
   for (const auto& arg : args) {
     if (arg[0] == '-' && arg[1] == '-') {
-      vector<string> splits = StringSplit(arg, '=');
-      if (splits.size() == 2) {
-        string optionKey = splits[0].substr(OPTION_PREFIX.size());
-        option_map[optionKey] = splits[1];
+      auto pos = arg.find('=', 0);
+      if (pos != string::npos) {
+        string optionKey = arg.substr(OPTION_PREFIX.size(), pos - OPTION_PREFIX.size());
+        option_map[optionKey] = arg.substr(pos + 1);
       } else {
-        string optionKey = splits[0].substr(OPTION_PREFIX.size());
+        string optionKey = arg.substr(OPTION_PREFIX.size());
         flags.push_back(optionKey);
       }
     } else {
@@ -257,7 +259,6 @@ bool LDBCommand::ParseStringOption(const map<string, string>& options,
 }
 
 Options LDBCommand::PrepareOptionsForOpenDB() {
-
   Options opt = options_;
   opt.create_if_missing = false;
 
@@ -368,6 +369,9 @@ Options LDBCommand::PrepareOptionsForOpenDB() {
     }
   }
 
+  opt.env = env_ ? env_.get() : Env::Default();
+  opt.checkpoint_env = Env::Default();
+
   return opt;
 }
 
@@ -1808,7 +1812,7 @@ ScanCommand::ScanCommand(const vector<string>& params,
                  BuildCmdLineOptions(
                      {ARG_TTL,      ARG_NO_VALUE,  ARG_HEX,    ARG_KEY_HEX,
                       ARG_TO,       ARG_VALUE_HEX, ARG_FROM,   ARG_TIMESTAMP,
-                      ARG_MAX_KEYS, ARG_TTL_START, ARG_TTL_END})),
+                      ARG_MAX_KEYS, ARG_TTL_START, ARG_TTL_END, ARG_ONLY_VERIFY_CHECKSUMS})),
       start_key_specified_(false),
       end_key_specified_(false),
       max_keys_scanned_(-1),
@@ -1837,6 +1841,12 @@ ScanCommand::ScanCommand(const vector<string>& params,
     no_value_ = true;
   }
 
+  vitr =  std::find(flags.begin(), flags.end(), ARG_ONLY_VERIFY_CHECKSUMS);
+  if (vitr != flags.end()) {
+    LOG(INFO) << "Only verify checksums, don't print entries.";
+    only_verify_checksums_ = true;
+  }
+
   itr = options.find(ARG_MAX_KEYS);
   if (itr != options.end()) {
     try {
@@ -1865,6 +1875,7 @@ void ScanCommand::Help(string& ret) {
   ret.append(" [--" + ARG_TTL_START + "=<N>:- is inclusive]");
   ret.append(" [--" + ARG_TTL_END + "=<N>:- is exclusive]");
   ret.append(" [--" + ARG_NO_VALUE + "]");
+  ret.append(" [--" + ARG_ONLY_VERIFY_CHECKSUMS + "]");
   ret.append("\n");
 }
 
@@ -1908,7 +1919,7 @@ void ScanCommand::DoCommand() {
       if (rawtime < ttl_start || rawtime >= ttl_end) {
         continue;
       }
-      if (timestamp_) {
+      if (timestamp_ && !only_verify_checksums_) {
         fprintf(stdout, "%s ", ReadableTime(rawtime).c_str());
       }
     }
@@ -1924,7 +1935,7 @@ void ScanCommand::DoCommand() {
       key_slice = formatted_key;
     }
 
-    if (no_value_) {
+    if (no_value_ && !only_verify_checksums_) {
       fprintf(stdout, "%.*s\n", static_cast<int>(key_slice.size()),
               key_slice.data());
     } else {
@@ -1934,9 +1945,11 @@ void ScanCommand::DoCommand() {
         formatted_value = "0x" + val_slice.ToString(true /* hex */);
         val_slice = formatted_value;
       }
-      fprintf(stdout, "%.*s : %.*s\n", static_cast<int>(key_slice.size()),
-              key_slice.data(), static_cast<int>(val_slice.size()),
-              val_slice.data());
+      if (!only_verify_checksums_) {
+        fprintf(stdout, "%.*s : %.*s\n", static_cast<int>(key_slice.size()),
+                key_slice.data(), static_cast<int>(val_slice.size()),
+                val_slice.data());
+      }
     }
 
     num_keys_scanned++;