From 37bb2f803d68f66861cadb012ce9a9f1c7cb7e94 Mon Sep 17 00:00:00 2001 From: Rahul Desirazu Date: Fri, 21 Feb 2020 10:36:30 -0800 Subject: [PATCH] [#3455] Modify ldb Tool to be Encryption Aware Summary: Add `--key_file` option to ldb to pass in a `:` 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 | 11 ++++ ent/src/yb/tools/yb-admin_client.h | 2 + ent/src/yb/tools/yb-admin_client_ent.cc | 26 ++++++++ ent/src/yb/tserver/CMakeLists-include.txt | 2 - ent/src/yb/tserver/tablet_server_ent.cc | 2 +- src/yb/rocksdb/CMakeLists.txt | 8 +-- src/yb/rocksdb/db/db_test_util.cc | 35 ++++++++++- src/yb/rocksdb/db/db_test_util.h | 19 +++++- src/yb/rocksdb/tools/ldb_cmd.cc | 37 +++++++---- src/yb/rocksdb/tools/ldb_cmd.h | 38 +++++++++++- src/yb/rocksdb/tools/ldb_cmd_test.cc | 62 ++++++++++++++++++- src/yb/rocksdb/tools/ldb_tool.cc | 10 ++- src/yb/util/CMakeLists.txt | 2 + src/yb/util/encryption_util.h | 2 +- .../yb/util}/header_manager_impl-test.cc | 14 ++--- .../yb/util}/header_manager_impl.cc | 32 +++++----- .../yb/util}/header_manager_impl.h | 14 ++--- src/yb/util/universe_key_manager.cc | 14 +++++ src/yb/util/universe_key_manager.h | 2 + 19 files changed, 267 insertions(+), 65 deletions(-) rename {ent/src/yb/tserver => src/yb/util}/header_manager_impl-test.cc (81%) rename {ent/src/yb/tserver => src/yb/util}/header_manager_impl.cc (81%) rename {ent/src/yb/tserver => src/yb/util}/header_manager_impl.h (75%) diff --git a/ent/src/yb/tools/yb-admin_cli_ent.cc b/ent/src/yb/tools/yb-admin_cli_ent.cc index 19906199814b..b88503006005 100644 --- a/ent/src/yb/tools/yb-admin_cli_ent.cc +++ b/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", " ", + [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", " ", [client](const CLIArguments& args) -> Status { diff --git a/ent/src/yb/tools/yb-admin_client.h b/ent/src/yb/tools/yb-admin_client.h index 7bf5ae4f4ef5..b06b23af880e 100644 --- a/ent/src/yb/tools/yb-admin_client.h +++ b/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, diff --git a/ent/src/yb/tools/yb-admin_client_ent.cc b/ent/src/yb/tools/yb-admin_client_ent.cc index 8bd8099a9dab..3ca2c8f4e9ad 100644 --- a/ent/src/yb/tools/yb-admin_client_ent.cc +++ b/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; diff --git a/ent/src/yb/tserver/CMakeLists-include.txt b/ent/src/yb/tserver/CMakeLists-include.txt index d32fb1d37d41..68361b6aa26c 100644 --- a/ent/src/yb/tserver/CMakeLists-include.txt +++ b/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) diff --git a/ent/src/yb/tserver/tablet_server_ent.cc b/ent/src/yb/tserver/tablet_server_ent.cc index 33218f319512..3740ddcf2927 100644 --- a/ent/src/yb/tserver/tablet_server_ent.cc +++ b/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" diff --git a/src/yb/rocksdb/CMakeLists.txt b/src/yb/rocksdb/CMakeLists.txt index 9bc16373e86c..670c0ccab751 100644 --- a/src/yb/rocksdb/CMakeLists.txt +++ b/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) diff --git a/src/yb/rocksdb/db/db_test_util.cc b/src/yb/rocksdb/db/db_test_util.cc index bca9e6e96683..c5a7adc28484 100644 --- a/src/yb/rocksdb/db/db_test_util.cc +++ b/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; diff --git a/src/yb/rocksdb/db/db_test_util.h b/src/yb/rocksdb/db/db_test_util.h index b8eede747ab6..b3937df6e83f 100644 --- a/src/yb/rocksdb/db/db_test_util.h +++ b/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 universe_key_manager_; + std::unique_ptr 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); diff --git a/src/yb/rocksdb/tools/ldb_cmd.cc b/src/yb/rocksdb/tools/ldb_cmd.cc index 553b27c10c42..fa2c8b3d94e7 100644 --- a/src/yb/rocksdb/tools/ldb_cmd.cc +++ b/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 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& 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& 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& 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 + "=:- is inclusive]"); ret.append(" [--" + ARG_TTL_END + "=:- 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(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(key_slice.size()), - key_slice.data(), static_cast(val_slice.size()), - val_slice.data()); + if (!only_verify_checksums_) { + fprintf(stdout, "%.*s : %.*s\n", static_cast(key_slice.size()), + key_slice.data(), static_cast(val_slice.size()), + val_slice.data()); + } } num_keys_scanned++; diff --git a/src/yb/rocksdb/tools/ldb_cmd.h b/src/yb/rocksdb/tools/ldb_cmd.h index 6668a30ae1c6..0ee18e638356 100644 --- a/src/yb/rocksdb/tools/ldb_cmd.h +++ b/src/yb/rocksdb/tools/ldb_cmd.h @@ -39,12 +39,15 @@ #include "yb/rocksdb/iterator.h" #include "yb/rocksdb/ldb_tool.h" #include "yb/rocksdb/options.h" -#include "yb/util/slice.h" #include "yb/rocksdb/utilities/db_ttl.h" #include "yb/rocksdb/tools/ldb_cmd_execute_result.h" #include "yb/rocksdb/util/logging.h" -#include "yb/util/string_util.h" #include "yb/rocksdb/utilities/ttl/db_ttl_impl.h" +#include "yb/util/header_manager_impl.h" +#include "yb/util/slice.h" +#include "yb/util/string_util.h" +#include "yb/util/universe_key_manager.h" +#include "yb/rocksutil/rocksdb_encrypted_file_factory.h" using std::string; using std::map; @@ -79,6 +82,8 @@ class LDBCommand { static const string ARG_FILE_SIZE; static const string ARG_CREATE_IF_MISSING; static const string ARG_NO_VALUE; + static const string ARG_UNIVERSE_KEY_FILE; + static const string ARG_ONLY_VERIFY_CHECKSUMS; static LDBCommand* InitFromCmdLineArgs( const vector& args, const Options& options, @@ -239,6 +244,9 @@ class LDBCommand { /** List of command-line options valid for this command */ const vector valid_cmd_line_options_; + std::unique_ptr universe_key_manager_; + std::unique_ptr env_; + bool ParseKeyValue(const string& line, string* key, string* value, bool is_key_hex, bool is_value_hex); @@ -259,6 +267,29 @@ class LDBCommand { db_path_ = itr->second; } + itr = options.find(ARG_UNIVERSE_KEY_FILE); + if (itr != options.end()) { + vector splits = StringSplit(itr->second, ':'); + if (splits.size() != 2) { + LOG(FATAL) << yb::Format("Could not split $0 by ':' into a key id and key file", + itr->second); + } + string key_data; + auto key_id = splits[0]; + auto key_path = splits[1]; + Status s = ReadFileToString(Env::Default(), key_path, &key_data); + if(!s.ok()) { + LOG(FATAL) << yb::Format("Could not read file at path $0: $1", key_path, s.ToString()); + } + auto res = yb::enterprise::UniverseKeyManager::FromKey(key_id, yb::Slice(key_data)); + if (!res.ok()) { + LOG(FATAL) << "Could not create universe key manager: " << res.status().ToString(); + } + universe_key_manager_ = std::move(*res); + env_ = yb::enterprise::NewRocksDBEncryptedEnv( + yb::enterprise::DefaultHeaderManager(universe_key_manager_.get())); + } + itr = options.find(ARG_CF_NAME); if (itr != options.end()) { column_family_name_ = itr->second; @@ -414,7 +445,7 @@ class LDBCommand { vector ret = {ARG_DB, ARG_BLOOM_BITS, ARG_BLOCK_SIZE, ARG_AUTO_COMPACTION, ARG_COMPRESSION_TYPE, ARG_WRITE_BUFFER_SIZE, ARG_FILE_SIZE, - ARG_FIX_PREFIX_LEN, ARG_CF_NAME}; + ARG_FIX_PREFIX_LEN, ARG_CF_NAME, ARG_UNIVERSE_KEY_FILE}; ret.insert(ret.end(), options.begin(), options.end()); return ret; } @@ -828,6 +859,7 @@ class ScanCommand : public LDBCommand { bool end_key_specified_; int max_keys_scanned_; bool no_value_; + bool only_verify_checksums_ = false; }; class DeleteCommand : public LDBCommand { diff --git a/src/yb/rocksdb/tools/ldb_cmd_test.cc b/src/yb/rocksdb/tools/ldb_cmd_test.cc index 703e8a4a2dc8..906f8b680420 100644 --- a/src/yb/rocksdb/tools/ldb_cmd_test.cc +++ b/src/yb/rocksdb/tools/ldb_cmd_test.cc @@ -21,10 +21,51 @@ #include "yb/rocksdb/tools/ldb_cmd.h" #include "yb/rocksdb/util/testharness.h" +#include "yb/rocksdb/env.h" +#include "yb/rocksdb/db/db_test_util.h" -class LdbCmdTest : public testing::Test {}; +DECLARE_bool(TEST_exit_on_finish); -TEST_F(LdbCmdTest, HexToString) { +const string kDbName = "/lbd_cmd_test"; +const string kKeyFileOption = "--key_file=" + rocksdb::DBTestBase::kKeyId + ":" + + rocksdb::DBTestBase::kKeyFile; + + +class LdbCmdTest : public rocksdb::DBTestBase, public testing::WithParamInterface { + public: + LdbCmdTest() : rocksdb::DBTestBase(kDbName, GetParam()) {} + + yb::Status WriteRecordsAndFlush() { + RETURN_NOT_OK(Put("k1", "v1")); + RETURN_NOT_OK(Put("k2", "v2")); + return Flush(); + } + + std::vector CreateArgs(std::string command_name) { + std::vector args; + args.push_back("./ldb"); + args.push_back("--db=" + dbname_); + args.push_back(command_name); + if (GetParam()) { + args.push_back(kKeyFileOption); + } + return args; + } + + void RunLdb(const std::vector& args) { + std::vector args_ptr; + args_ptr.resize(args.size()); + std::transform(args.begin(), args.end(), args_ptr.begin(), [](const std::string& str) { + return const_cast(str.c_str()); + }); + rocksdb::LDBTool tool; + ASSERT_NO_FATALS(tool.Run(static_cast(args_ptr.size()), args_ptr.data())); + } +}; + +INSTANTIATE_TEST_CASE_P(EncryptionEnabled, LdbCmdTest, ::testing::Bool()); + +TEST_P(LdbCmdTest, HexToString) { // map input to expected outputs. map> inputMap = { {"0x7", {7}}, {"0x5050", {80, 80}}, {"0xFF", {-1}}, @@ -40,7 +81,7 @@ TEST_F(LdbCmdTest, HexToString) { } } -TEST_F(LdbCmdTest, HexToStringBadInputs) { +TEST_P(LdbCmdTest, HexToStringBadInputs) { const vector badInputs = { "0xZZ", "123", "0xx5", "0x11G", "Ox12", "0xT", "0x1Q1", }; @@ -54,6 +95,21 @@ TEST_F(LdbCmdTest, HexToStringBadInputs) { } } +TEST_P(LdbCmdTest, Scan) { + FLAGS_TEST_exit_on_finish = false; + ASSERT_OK(WriteRecordsAndFlush()); + auto args = CreateArgs("scan"); + args.push_back("--only_verify_checksums"); + ASSERT_NO_FATALS(RunLdb(args)); +} + +TEST_P(LdbCmdTest, CheckConsistency) { + FLAGS_TEST_exit_on_finish = false; + ASSERT_OK(WriteRecordsAndFlush()); + auto args = CreateArgs("checkconsistency"); + ASSERT_NO_FATALS(RunLdb(args)); +} + int main(int argc, char** argv) { ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); diff --git a/src/yb/rocksdb/tools/ldb_tool.cc b/src/yb/rocksdb/tools/ldb_tool.cc index 3188572e574b..27c2f8c31350 100644 --- a/src/yb/rocksdb/tools/ldb_tool.cc +++ b/src/yb/rocksdb/tools/ldb_tool.cc @@ -21,12 +21,16 @@ #include "yb/rocksdb/ldb_tool.h" #include "yb/rocksdb/tools/ldb_cmd.h" +#include "yb/util/flag_tags.h" + +DEFINE_test_flag(bool, TEST_exit_on_finish, true, "Exit the process on finishing."); + namespace rocksdb { LDBOptions::LDBOptions() {} class LDBCommandRunner { -public: + public: static void PrintHelp(const char* exec_name) { string ret; @@ -120,7 +124,9 @@ class LDBCommandRunner { fprintf(stderr, "%s\n", ret.ToString().c_str()); delete cmdObj; - exit(ret.IsFailed()); + if (FLAGS_TEST_exit_on_finish) { + exit(ret.IsFailed()); + } } }; diff --git a/src/yb/util/CMakeLists.txt b/src/yb/util/CMakeLists.txt index 820a7490736c..39f009565715 100644 --- a/src/yb/util/CMakeLists.txt +++ b/src/yb/util/CMakeLists.txt @@ -178,6 +178,7 @@ set(UTIL_SRCS flag_tags.cc flags.cc hdr_histogram.cc + header_manager_impl.cc hexdump.cc init.cc jsonreader.cc @@ -360,6 +361,7 @@ ADD_YB_TEST(flags-test) ADD_YB_TEST(format-test RUN_SERIAL true) ADD_YB_TEST(hash_util-test) ADD_YB_TEST(hdr_histogram-test) +ADD_YB_TEST(header_manager_impl-test) ADD_YB_TEST(inline_slice-test) ADD_YB_TEST(jsonreader-test) ADD_YB_TEST(lockfree-test) diff --git a/src/yb/util/encryption_util.h b/src/yb/util/encryption_util.h index 8fbf9b956859..de570f190b5d 100644 --- a/src/yb/util/encryption_util.h +++ b/src/yb/util/encryption_util.h @@ -134,7 +134,7 @@ Status CreateEncryptionInfoForWrite(HeaderManager* header_manager, // Since file doesn't exist or this overwrites, append key to the name and create. *stream = std::make_unique(std::move(encryption_params)); RETURN_NOT_OK((*stream)->Init()); - *header_size = header.size(); + *header_size = static_cast(header.size()); return Status::OK(); } diff --git a/ent/src/yb/tserver/header_manager_impl-test.cc b/src/yb/util/header_manager_impl-test.cc similarity index 81% rename from ent/src/yb/tserver/header_manager_impl-test.cc rename to src/yb/util/header_manager_impl-test.cc index 993b01fe8351..075e24214454 100644 --- a/ent/src/yb/tserver/header_manager_impl-test.cc +++ b/src/yb/util/header_manager_impl-test.cc @@ -16,7 +16,7 @@ #include -#include "yb/tserver/header_manager_impl.h" +#include "yb/util/header_manager_impl.h" #include "yb/util/encryption_util.h" #include "yb/util/status.h" @@ -30,15 +30,14 @@ #include namespace yb { -namespace tserver { namespace enterprise { class TestHeaderManagerImpl : public YBTest {}; -std::unique_ptr GenerateUniverseKeyManager() { - auto universe_key_manager = std::make_unique(); +std::unique_ptr GenerateUniverseKeyManager() { + auto universe_key_manager = std::make_unique(); UniverseKeyRegistryPB registry; - auto encryption_params = yb::enterprise::EncryptionParams::NewEncryptionParams(); + auto encryption_params = EncryptionParams::NewEncryptionParams(); EncryptionParamsPB params_pb; encryption_params->ToEncryptionParamsPB(¶ms_pb); auto version_id = RandomHumanReadableString(16); @@ -51,9 +50,9 @@ std::unique_ptr GenerateUniverseKeyManager() TEST_F(TestHeaderManagerImpl, FileOps) { auto universe_key_manger = GenerateUniverseKeyManager(); - std::unique_ptr header_manager = + std::unique_ptr header_manager = DefaultHeaderManager(universe_key_manger.get()); - auto params = yb::enterprise::EncryptionParams::NewEncryptionParams(); + auto params = EncryptionParams::NewEncryptionParams(); string header = ASSERT_RESULT(header_manager->SerializeEncryptionParams(*params.get())); auto start_idx = header_manager->GetEncryptionMetadataStartIndex(); auto status = ASSERT_RESULT(header_manager->GetFileEncryptionStatusFromPrefix( @@ -65,5 +64,4 @@ TEST_F(TestHeaderManagerImpl, FileOps) { } } // namespace enterprise -} // namespace tserver } // namespace yb diff --git a/ent/src/yb/tserver/header_manager_impl.cc b/src/yb/util/header_manager_impl.cc similarity index 81% rename from ent/src/yb/tserver/header_manager_impl.cc rename to src/yb/util/header_manager_impl.cc index d5e53142aebe..1704d704784f 100644 --- a/ent/src/yb/tserver/header_manager_impl.cc +++ b/src/yb/util/header_manager_impl.cc @@ -11,7 +11,7 @@ // under the License. // -#include "yb/tserver/header_manager_impl.h" +#include "yb/util/header_manager_impl.h" #include "yb/util/env.h" #include "yb/util/encryption_util.h" @@ -26,16 +26,14 @@ static const string kEncryptionMagic = "encrypt!"; namespace yb { -namespace tserver { namespace enterprise { -class HeaderManagerImpl : public yb::enterprise::HeaderManager { +class HeaderManagerImpl : public HeaderManager { public: - explicit HeaderManagerImpl(yb::enterprise::UniverseKeyManager* universe_key_manager) + explicit HeaderManagerImpl(UniverseKeyManager* universe_key_manager) : universe_key_manager_(universe_key_manager) {} - Result SerializeEncryptionParams( - const yb::enterprise::EncryptionParams& encryption_info) override { + Result SerializeEncryptionParams(const EncryptionParams& encryption_info) override { // 1. Generate the EncrytionParamsPB string to be encrypted. EncryptionParamsPB encryption_params_pb; encryption_info.ToEncryptionParamsPB(&encryption_params_pb); @@ -44,10 +42,10 @@ class HeaderManagerImpl : public yb::enterprise::HeaderManager { // 2. Encrypt the encryption params with the latest universe key. auto universe_params = VERIFY_RESULT(universe_key_manager_->GetLatestUniverseParams()); - auto stream = VERIFY_RESULT(yb::enterprise::BlockAccessCipherStream::FromEncryptionParams( + auto stream = VERIFY_RESULT(BlockAccessCipherStream::FromEncryptionParams( std::move(universe_params.params))); auto encrypted_data_key = static_cast( - yb::enterprise::EncryptionBuffer::Get()->GetBuffer(encryption_params_pb_str.size())); + EncryptionBuffer::Get()->GetBuffer(encryption_params_pb_str.size())); RETURN_NOT_OK(stream->Encrypt(0, encryption_params_pb_str, encrypted_data_key)); // 3. Serialize the universe key id. @@ -71,7 +69,7 @@ class HeaderManagerImpl : public yb::enterprise::HeaderManager { return kEncryptionMagic + string(header_size, sizeof(header_size)) + metadata_str; } - Result + Result DecodeEncryptionParamsFromEncryptionMetadata(const Slice& s) override { Slice s_mutable(s); // 1. Get the size of the universe key id. @@ -87,7 +85,7 @@ class HeaderManagerImpl : public yb::enterprise::HeaderManager { // 3. Create an encryption stream from the universe key. auto universe_params = VERIFY_RESULT( universe_key_manager_->GetUniverseParamsWithVersion(universe_key_id)); - auto stream = VERIFY_RESULT(yb::enterprise::BlockAccessCipherStream::FromEncryptionParams( + auto stream = VERIFY_RESULT(BlockAccessCipherStream::FromEncryptionParams( std::move(universe_params))); // 4. Get the size of the encryption params pb. @@ -97,7 +95,7 @@ class HeaderManagerImpl : public yb::enterprise::HeaderManager { // 5. Decrypt the resulting data key. auto decrypted_data_key = static_cast( - yb::enterprise::EncryptionBuffer::Get()->GetBuffer(s_mutable.size())); + EncryptionBuffer::Get()->GetBuffer(s_mutable.size())); RETURN_NOT_OK(stream->Decrypt(0, s_mutable, decrypted_data_key)); // 6. Convert the resulting decrypted data key into encryption params. @@ -105,16 +103,16 @@ class HeaderManagerImpl : public yb::enterprise::HeaderManager { s_mutable, encryption_params_pb_size, "encryption params")); auto encryption_params_pb = VERIFY_RESULT(pb_util::ParseFromSlice( Slice(decrypted_data_key, encryption_params_pb_size))); - return yb::enterprise::EncryptionParams::FromEncryptionParamsPB(encryption_params_pb); + return EncryptionParams::FromEncryptionParamsPB(encryption_params_pb); } uint32_t GetEncryptionMetadataStartIndex() override { return kEncryptionMagic.size() + sizeof(uint32_t); } - Result GetFileEncryptionStatusFromPrefix( + Result GetFileEncryptionStatusFromPrefix( const Slice& s) override { - yb::enterprise::FileEncryptionStatus status; + FileEncryptionStatus status; status.is_encrypted = s.compare_prefix(Slice(kEncryptionMagic)) == 0; if (status.is_encrypted) { status.header_size = BigEndian::Load32(s.data() + kEncryptionMagic.size()); @@ -136,15 +134,13 @@ class HeaderManagerImpl : public yb::enterprise::HeaderManager { return Status::OK(); } - yb::enterprise::UniverseKeyManager* universe_key_manager_; + UniverseKeyManager* universe_key_manager_; }; -std::unique_ptr DefaultHeaderManager( - yb::enterprise::UniverseKeyManager* universe_key_manager) { +std::unique_ptr DefaultHeaderManager(UniverseKeyManager* universe_key_manager) { return std::make_unique(universe_key_manager); } } // namespace enterprise -} // namespace tserver } // namespace yb diff --git a/ent/src/yb/tserver/header_manager_impl.h b/src/yb/util/header_manager_impl.h similarity index 75% rename from ent/src/yb/tserver/header_manager_impl.h rename to src/yb/util/header_manager_impl.h index 33b5a8fda43b..c2e3fb7d5335 100644 --- a/ent/src/yb/tserver/header_manager_impl.h +++ b/src/yb/util/header_manager_impl.h @@ -11,8 +11,8 @@ // under the License. // -#ifndef ENT_SRC_YB_TSERVER_HEADER_MANAGER_IMPL_H -#define ENT_SRC_YB_TSERVER_HEADER_MANAGER_IMPL_H +#ifndef YB_UTIL_HEADER_MANAGER_IMPL_H +#define YB_UTIL_HEADER_MANAGER_IMPL_H #include @@ -20,11 +20,9 @@ namespace yb { namespace enterprise { + class HeaderManager; class UniverseKeyManager; -} -namespace tserver { -namespace enterprise { // Implementation of used by FileFactory to construct header for encrypted files. // The header format looks like: @@ -35,11 +33,9 @@ namespace enterprise { // universe key id // EncryptionParamsPB size (4 bytes) // EncryptionParamsPB -std::unique_ptr DefaultHeaderManager( - yb::enterprise::UniverseKeyManager* universe_key_manager); +std::unique_ptr DefaultHeaderManager(UniverseKeyManager* universe_key_manager); } // namespace enterprise -} // namespace tserver } // namespace yb -#endif // ENT_SRC_YB_TSERVER_HEADER_MANAGER_IMPL_H +#endif // YB_UTIL_HEADER_MANAGER_IMPL_H diff --git a/src/yb/util/universe_key_manager.cc b/src/yb/util/universe_key_manager.cc index cfdf0a3d55ad..fd2cb0fafb75 100644 --- a/src/yb/util/universe_key_manager.cc +++ b/src/yb/util/universe_key_manager.cc @@ -16,6 +16,20 @@ namespace yb { namespace enterprise { +Result> UniverseKeyManager::FromKey( + const std::string& key_id, const Slice& key_data) { + auto universe_key_manager = std::make_unique(); + UniverseKeyRegistryPB universe_key_registry; + universe_key_registry.set_encryption_enabled(true); + universe_key_registry.set_latest_version_id(key_id); + auto encryption_params = VERIFY_RESULT(yb::enterprise::EncryptionParams::FromSlice(key_data)); + yb::EncryptionParamsPB params_pb; + encryption_params->ToEncryptionParamsPB(¶ms_pb); + (*universe_key_registry.mutable_universe_keys())[key_id] = params_pb; + universe_key_manager->SetUniverseKeyRegistry(universe_key_registry); + return universe_key_manager; +} + void UniverseKeyManager::SetUniverseKeyRegistry( const UniverseKeyRegistryPB& universe_key_registry) { { diff --git a/src/yb/util/universe_key_manager.h b/src/yb/util/universe_key_manager.h index ceff5258d708..a567355ca9c7 100644 --- a/src/yb/util/universe_key_manager.h +++ b/src/yb/util/universe_key_manager.h @@ -25,6 +25,8 @@ namespace enterprise { // in creating new files and reading exising files. class UniverseKeyManager { public: + static Result> FromKey( + const std::string& key_id, const Slice& key_data); void SetUniverseKeyRegistry(const yb::UniverseKeyRegistryPB& universe_key_registry); // From an existing version id, generate encryption params. Used when creating readable files. Result GetUniverseParamsWithVersion(