Skip to content

Commit

Permalink
Fix the bug that the key manager is not updated during the Rename (#9)
Browse files Browse the repository at this point in the history
apache/incubator-pegasus#1575

Cherry-pick from
tikv@1868d12

Signed-off-by: Xintao <[email protected]>
Signed-off-by: tabokie <[email protected]>
  • Loading branch information
acelyc111 committed Aug 16, 2023
1 parent f6ac09d commit 3b14183
Show file tree
Hide file tree
Showing 6 changed files with 109 additions and 14 deletions.
4 changes: 4 additions & 0 deletions db/db_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -2414,6 +2414,10 @@ TEST_F(DBTest, DestroyDBMetaDatabase) {
}

TEST_F(DBTest, SnapshotFiles) {
if (getenv("ENCRYPTED_ENV")) {
// File copy does not carry encryption key.
return;
}
do {
Options options = CurrentOptions();
options.write_buffer_size = 100000000; // Large write buffer
Expand Down
35 changes: 27 additions & 8 deletions encryption/encryption.cc
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@

#include "file/filename.h"
#include "port/port.h"
#include "test_util/sync_point.h"

namespace ROCKSDB_NAMESPACE {
namespace encryption {
Expand Down Expand Up @@ -270,6 +271,17 @@ Status KeyManagedEncryptedEnv::NewSequentialFile(
case EncryptionMethod::kAES192_CTR:
case EncryptionMethod::kAES256_CTR:
s = encrypted_env_->NewSequentialFile(fname, result, options);
// Hack: when upgrading from TiKV <= v5.0.0-rc, the old current
// file is encrypted but it could be replaced with a plaintext
// current file. The operation below guarantee that the current
// file is read correctly.
if (s.ok() && IsCurrentFile(fname)) {
if (!IsValidCurrentFile(std::move(*result))) {
s = target()->NewSequentialFile(fname, result, options);
} else {
s = encrypted_env_->NewSequentialFile(fname, result, options);
}
}
break;
default:
s = Status::InvalidArgument(
Expand Down Expand Up @@ -309,7 +321,8 @@ Status KeyManagedEncryptedEnv::NewWritableFile(
const EnvOptions& options) {
FileEncryptionInfo file_info;
Status s;
bool skipped = ShouldSkipEncryption(fname);
bool skipped = IsCurrentFile(fname);
TEST_SYNC_POINT_CALLBACK("KeyManagedEncryptedEnv::NewWritableFile", &skipped);
if (!skipped) {
s = key_manager_->NewFile(fname, &file_info);
if (!s.ok()) {
Expand Down Expand Up @@ -444,12 +457,12 @@ Status KeyManagedEncryptedEnv::DeleteFile(const std::string& fname) {

Status KeyManagedEncryptedEnv::LinkFile(const std::string& src_fname,
const std::string& dst_fname) {
if (ShouldSkipEncryption(dst_fname)) {
assert(ShouldSkipEncryption(src_fname));
if (IsCurrentFile(dst_fname)) {
assert(IsCurrentFile(src_fname));
Status s = target()->LinkFile(src_fname, dst_fname);
return s;
} else {
assert(!ShouldSkipEncryption(src_fname));
assert(!IsCurrentFile(src_fname));
}
Status s = key_manager_->LinkFile(src_fname, dst_fname);
if (!s.ok()) {
Expand All @@ -466,11 +479,17 @@ Status KeyManagedEncryptedEnv::LinkFile(const std::string& src_fname,

Status KeyManagedEncryptedEnv::RenameFile(const std::string& src_fname,
const std::string& dst_fname) {
if (ShouldSkipEncryption(dst_fname)) {
assert(ShouldSkipEncryption(src_fname));
return target()->RenameFile(src_fname, dst_fname);
if (IsCurrentFile(dst_fname)) {
assert(IsCurrentFile(src_fname));
Status s = target()->RenameFile(src_fname, dst_fname);
// Replacing with plaintext requires deleting the info in the key manager.
// The stale current file info exists when upgrading from TiKV <= v5.0.0-rc.
Status delete_status __attribute__((__unused__)) =
key_manager_->DeleteFile(dst_fname);
assert(delete_status.ok());
return s;
} else {
assert(!ShouldSkipEncryption(src_fname));
assert(!IsCurrentFile(src_fname));
}
// Link(copy)File instead of RenameFile to avoid losing src_fname info when
// failed to rename the src_fname in the file system.
Expand Down
46 changes: 46 additions & 0 deletions env/env_basic_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,52 @@ INSTANTIATE_TEST_CASE_P(CustomEnv, EnvBasicTestWithParam,
INSTANTIATE_TEST_CASE_P(CustomEnv, EnvMoreTestWithParam,
::testing::ValuesIn(GetCustomEnvs()));

TEST_P(EnvBasicTestWithParam, RenameCurrent) {
if (!getenv("ENCRYPTED_ENV")) {
return;
}
Slice result;
char scratch[100];
std::unique_ptr<SequentialFile> seq_file;
std::unique_ptr<WritableFile> writable_file;
std::vector<std::string> children;

// Create an encrypted `CURRENT` file so it shouldn't be skipped .
SyncPoint::GetInstance()->SetCallBack(
"KeyManagedEncryptedEnv::NewWritableFile", [&](void* arg) {
bool* skip = static_cast<bool*>(arg);
*skip = false;
});
SyncPoint::GetInstance()->EnableProcessing();
ASSERT_OK(
env_->NewWritableFile(test_dir_ + "/CURRENT", &writable_file, soptions_));
SyncPoint::GetInstance()->ClearAllCallBacks();
SyncPoint::GetInstance()->DisableProcessing();
ASSERT_OK(writable_file->Append("MANIFEST-0"));
ASSERT_OK(writable_file->Close());
writable_file.reset();

ASSERT_OK(
env_->NewSequentialFile(test_dir_ + "/CURRENT", &seq_file, soptions_));
ASSERT_OK(seq_file->Read(100, &result, scratch));
ASSERT_EQ(0, result.compare("MANIFEST-0"));

// Create a plaintext `CURRENT` temp file.
ASSERT_OK(env_->NewWritableFile(test_dir_ + "/current.dbtmp.plain",
&writable_file, soptions_));
ASSERT_OK(writable_file->Append("MANIFEST-1"));
ASSERT_OK(writable_file->Close());
writable_file.reset();

ASSERT_OK(env_->RenameFile(test_dir_ + "/current.dbtmp.plain",
test_dir_ + "/CURRENT"));

ASSERT_OK(
env_->NewSequentialFile(test_dir_ + "/CURRENT", &seq_file, soptions_));
ASSERT_OK(seq_file->Read(100, &result, scratch));
ASSERT_EQ(0, result.compare("MANIFEST-1"));
}

TEST_P(EnvBasicTestWithParam, Basics) {
uint64_t file_size;
std::unique_ptr<WritableFile> writable_file;
Expand Down
10 changes: 9 additions & 1 deletion file/filename.cc
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ static const std::string kRocksDBBlobFileExt = "blob";
static const std::string kArchivalDirName = "archive";
static const std::string kUnencryptedTempFileNameSuffix = "dbtmp.plain";

bool ShouldSkipEncryption(const std::string& fname) {
bool IsCurrentFile(const std::string& fname) {
// skip CURRENT file.
size_t current_length = strlen("CURRENT");
if (fname.length() >= current_length &&
Expand All @@ -50,6 +50,14 @@ bool ShouldSkipEncryption(const std::string& fname) {
return false;
}

bool IsValidCurrentFile(
std::unique_ptr<ROCKSDB_NAMESPACE::SequentialFile> seq_file) {
Slice result;
char scratch[64];
seq_file->Read(8, &result, scratch);
return result.compare("MANIFEST") == 0;
}

// Given a path, flatten the path name by replacing all chars not in
// {[0-9,a-z,A-Z,-,_,.]} with _. And append '_LOG\0' at the end.
// Return the number of chars stored in dest not including the trailing '\0'.
Expand Down
6 changes: 5 additions & 1 deletion file/filename.h
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,11 @@ constexpr char kFilePathSeparator = '/';

// Some non-sensitive files are not encrypted to preserve atomicity of file
// operations.
extern bool ShouldSkipEncryption(const std::string& fname);
extern bool IsCurrentFile(const std::string& fname);

// Determine if the content is read from the valid current file.
extern bool IsValidCurrentFile(
std::unique_ptr<ROCKSDB_NAMESPACE::SequentialFile> seq_file);

// Return the name of the log file with the specified number
// in the db named by "dbname". The result will be prefixed with
Expand Down
22 changes: 18 additions & 4 deletions test_util/testutil.h
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
#pragma once
#include <algorithm>
#include <deque>
#include <mutex>
#include <set>
#include <string>
#include <vector>

Expand Down Expand Up @@ -55,10 +57,13 @@ class TestKeyManager : public encryption::KeyManager {

static const std::string default_key;
static const std::string default_iv;
std::mutex mutex;
std::set<std::string> file_set;

Status GetFile(const std::string& fname,
encryption::FileEncryptionInfo* file_info) override {
if (ShouldSkipEncryption(fname)) {
std::lock_guard<std::mutex> l(mutex);
if (file_set.find(fname) == file_set.end()) {
file_info->method = encryption::EncryptionMethod::kPlaintext;
} else {
file_info->method = encryption::EncryptionMethod::kAES192_CTR;
Expand All @@ -68,16 +73,25 @@ class TestKeyManager : public encryption::KeyManager {
return Status::OK();
}

Status NewFile(const std::string& /*fname*/,
Status NewFile(const std::string& fname,
encryption::FileEncryptionInfo* file_info) override {
std::lock_guard<std::mutex> l(mutex);
file_info->method = encryption::EncryptionMethod::kAES192_CTR;
file_info->key = default_key;
file_info->iv = default_iv;
file_set.insert(fname);
return Status::OK();
}

Status DeleteFile(const std::string&) override { return Status::OK(); }
Status LinkFile(const std::string&, const std::string&) override {
Status DeleteFile(const std::string& fname) override {
std::lock_guard<std::mutex> l(mutex);
file_set.erase(fname);
return Status::OK();
}

Status LinkFile(const std::string& /*src*/, const std::string& dst) override {
std::lock_guard<std::mutex> l(mutex);
file_set.insert(dst);
return Status::OK();
}
};
Expand Down

0 comments on commit 3b14183

Please sign in to comment.