Skip to content

Commit 1868d12

Browse files
hunterlxttabokie
authored andcommitted
Fix the bug that the key manager is not updated during the Rename (#227)
Signed-off-by: Xintao <hunterlxt@live.com> Signed-off-by: tabokie <xy.tao@outlook.com>
1 parent bbd27cf commit 1868d12

File tree

6 files changed

+108
-14
lines changed

6 files changed

+108
-14
lines changed

db/db_test.cc

+4
Original file line numberDiff line numberDiff line change
@@ -2292,6 +2292,10 @@ TEST_F(DBTest, DestroyDBMetaDatabase) {
22922292

22932293
#ifndef ROCKSDB_LITE
22942294
TEST_F(DBTest, SnapshotFiles) {
2295+
if (getenv("ENCRYPTED_ENV")) {
2296+
// File copy does not carry encryption key.
2297+
return;
2298+
}
22952299
do {
22962300
Options options = CurrentOptions();
22972301
options.write_buffer_size = 100000000; // Large write buffer

encryption/encryption.cc

+27-8
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212

1313
#include "file/filename.h"
1414
#include "port/port.h"
15+
#include "test_util/sync_point.h"
1516

1617
namespace ROCKSDB_NAMESPACE {
1718
namespace encryption {
@@ -267,6 +268,17 @@ Status KeyManagedEncryptedEnv::NewSequentialFile(
267268
case EncryptionMethod::kAES192_CTR:
268269
case EncryptionMethod::kAES256_CTR:
269270
s = encrypted_env_->NewSequentialFile(fname, result, options);
271+
// Hack: when upgrading from TiKV <= v5.0.0-rc, the old current
272+
// file is encrypted but it could be replaced with a plaintext
273+
// current file. The operation below guarantee that the current
274+
// file is read correctly.
275+
if (s.ok() && IsCurrentFile(fname)) {
276+
if (!IsValidCurrentFile(std::move(*result))) {
277+
s = target()->NewSequentialFile(fname, result, options);
278+
} else {
279+
s = encrypted_env_->NewSequentialFile(fname, result, options);
280+
}
281+
}
270282
break;
271283
default:
272284
s = Status::InvalidArgument(
@@ -306,7 +318,8 @@ Status KeyManagedEncryptedEnv::NewWritableFile(
306318
const EnvOptions& options) {
307319
FileEncryptionInfo file_info;
308320
Status s;
309-
bool skipped = ShouldSkipEncryption(fname);
321+
bool skipped = IsCurrentFile(fname);
322+
TEST_SYNC_POINT_CALLBACK("KeyManagedEncryptedEnv::NewWritableFile", &skipped);
310323
if (!skipped) {
311324
s = key_manager_->NewFile(fname, &file_info);
312325
if (!s.ok()) {
@@ -441,12 +454,12 @@ Status KeyManagedEncryptedEnv::DeleteFile(const std::string& fname) {
441454

442455
Status KeyManagedEncryptedEnv::LinkFile(const std::string& src_fname,
443456
const std::string& dst_fname) {
444-
if (ShouldSkipEncryption(dst_fname)) {
445-
assert(ShouldSkipEncryption(src_fname));
457+
if (IsCurrentFile(dst_fname)) {
458+
assert(IsCurrentFile(src_fname));
446459
Status s = target()->LinkFile(src_fname, dst_fname);
447460
return s;
448461
} else {
449-
assert(!ShouldSkipEncryption(src_fname));
462+
assert(!IsCurrentFile(src_fname));
450463
}
451464
Status s = key_manager_->LinkFile(src_fname, dst_fname);
452465
if (!s.ok()) {
@@ -463,11 +476,17 @@ Status KeyManagedEncryptedEnv::LinkFile(const std::string& src_fname,
463476

464477
Status KeyManagedEncryptedEnv::RenameFile(const std::string& src_fname,
465478
const std::string& dst_fname) {
466-
if (ShouldSkipEncryption(dst_fname)) {
467-
assert(ShouldSkipEncryption(src_fname));
468-
return target()->RenameFile(src_fname, dst_fname);
479+
if (IsCurrentFile(dst_fname)) {
480+
assert(IsCurrentFile(src_fname));
481+
Status s = target()->RenameFile(src_fname, dst_fname);
482+
// Replacing with plaintext requires deleting the info in the key manager.
483+
// The stale current file info exists when upgrading from TiKV <= v5.0.0-rc.
484+
Status delete_status __attribute__((__unused__)) =
485+
key_manager_->DeleteFile(dst_fname);
486+
assert(delete_status.ok());
487+
return s;
469488
} else {
470-
assert(!ShouldSkipEncryption(src_fname));
489+
assert(!IsCurrentFile(src_fname));
471490
}
472491
// Link(copy)File instead of RenameFile to avoid losing src_fname info when
473492
// failed to rename the src_fname in the file system.

env/env_basic_test.cc

+46
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,52 @@ INSTANTIATE_TEST_CASE_P(CustomEnv, EnvMoreTestWithParam,
205205
::testing::ValuesIn(GetCustomEnvs()));
206206
#endif // ROCKSDB_LITE
207207

208+
TEST_P(EnvBasicTestWithParam, RenameCurrent) {
209+
if (!getenv("ENCRYPTED_ENV")) {
210+
return;
211+
}
212+
Slice result;
213+
char scratch[100];
214+
std::unique_ptr<SequentialFile> seq_file;
215+
std::unique_ptr<WritableFile> writable_file;
216+
std::vector<std::string> children;
217+
218+
// Create an encrypted `CURRENT` file so it shouldn't be skipped .
219+
SyncPoint::GetInstance()->SetCallBack(
220+
"KeyManagedEncryptedEnv::NewWritableFile", [&](void* arg) {
221+
bool* skip = static_cast<bool*>(arg);
222+
*skip = false;
223+
});
224+
SyncPoint::GetInstance()->EnableProcessing();
225+
ASSERT_OK(
226+
env_->NewWritableFile(test_dir_ + "/CURRENT", &writable_file, soptions_));
227+
SyncPoint::GetInstance()->ClearAllCallBacks();
228+
SyncPoint::GetInstance()->DisableProcessing();
229+
ASSERT_OK(writable_file->Append("MANIFEST-0"));
230+
ASSERT_OK(writable_file->Close());
231+
writable_file.reset();
232+
233+
ASSERT_OK(
234+
env_->NewSequentialFile(test_dir_ + "/CURRENT", &seq_file, soptions_));
235+
ASSERT_OK(seq_file->Read(100, &result, scratch));
236+
ASSERT_EQ(0, result.compare("MANIFEST-0"));
237+
238+
// Create a plaintext `CURRENT` temp file.
239+
ASSERT_OK(env_->NewWritableFile(test_dir_ + "/current.dbtmp.plain",
240+
&writable_file, soptions_));
241+
ASSERT_OK(writable_file->Append("MANIFEST-1"));
242+
ASSERT_OK(writable_file->Close());
243+
writable_file.reset();
244+
245+
ASSERT_OK(env_->RenameFile(test_dir_ + "/current.dbtmp.plain",
246+
test_dir_ + "/CURRENT"));
247+
248+
ASSERT_OK(
249+
env_->NewSequentialFile(test_dir_ + "/CURRENT", &seq_file, soptions_));
250+
ASSERT_OK(seq_file->Read(100, &result, scratch));
251+
ASSERT_EQ(0, result.compare("MANIFEST-1"));
252+
}
253+
208254
TEST_P(EnvBasicTestWithParam, Basics) {
209255
uint64_t file_size;
210256
std::unique_ptr<WritableFile> writable_file;

file/filename.cc

+8-1
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ static const std::string kRocksDBBlobFileExt = "blob";
3030
static const std::string kArchivalDirName = "archive";
3131
static const std::string kUnencryptedTempFileNameSuffix = "dbtmp.plain";
3232

33-
bool ShouldSkipEncryption(const std::string& fname) {
33+
bool IsCurrentFile(const std::string& fname) {
3434
// skip CURRENT file.
3535
size_t current_length = strlen("CURRENT");
3636
if (fname.length() >= current_length &&
@@ -48,6 +48,13 @@ bool ShouldSkipEncryption(const std::string& fname) {
4848
return false;
4949
}
5050

51+
bool IsValidCurrentFile(std::unique_ptr<rocksdb::SequentialFile> seq_file) {
52+
Slice result;
53+
char scratch[64];
54+
seq_file->Read(8, &result, scratch);
55+
return result.compare("MANIFEST") == 0;
56+
}
57+
5158
// Given a path, flatten the path name by replacing all chars not in
5259
// {[0-9,a-z,A-Z,-,_,.]} with _. And append '_LOG\0' at the end.
5360
// Return the number of chars stored in dest not including the trailing '\0'.

file/filename.h

+5-1
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,11 @@ constexpr char kFilePathSeparator = '/';
3838

3939
// Some non-sensitive files are not encrypted to preserve atomicity of file
4040
// operations.
41-
extern bool ShouldSkipEncryption(const std::string& fname);
41+
extern bool IsCurrentFile(const std::string& fname);
42+
43+
// Determine if the content is read from the valid current file.
44+
extern bool IsValidCurrentFile(
45+
std::unique_ptr<rocksdb::SequentialFile> seq_file);
4246

4347
// Return the name of the log file with the specified number
4448
// in the db named by "dbname". The result will be prefixed with

test_util/testutil.h

+18-4
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
#pragma once
1111
#include <algorithm>
1212
#include <deque>
13+
#include <mutex>
14+
#include <set>
1315
#include <string>
1416
#include <vector>
1517

@@ -55,10 +57,13 @@ class TestKeyManager : public encryption::KeyManager {
5557

5658
static const std::string default_key;
5759
static const std::string default_iv;
60+
std::mutex mutex;
61+
std::set<std::string> file_set;
5862

5963
Status GetFile(const std::string& fname,
6064
encryption::FileEncryptionInfo* file_info) override {
61-
if (ShouldSkipEncryption(fname)) {
65+
std::lock_guard<std::mutex> l(mutex);
66+
if (file_set.find(fname) == file_set.end()) {
6267
file_info->method = encryption::EncryptionMethod::kPlaintext;
6368
} else {
6469
file_info->method = encryption::EncryptionMethod::kAES192_CTR;
@@ -68,16 +73,25 @@ class TestKeyManager : public encryption::KeyManager {
6873
return Status::OK();
6974
}
7075

71-
Status NewFile(const std::string& /*fname*/,
76+
Status NewFile(const std::string& fname,
7277
encryption::FileEncryptionInfo* file_info) override {
78+
std::lock_guard<std::mutex> l(mutex);
7379
file_info->method = encryption::EncryptionMethod::kAES192_CTR;
7480
file_info->key = default_key;
7581
file_info->iv = default_iv;
82+
file_set.insert(fname);
7683
return Status::OK();
7784
}
7885

79-
Status DeleteFile(const std::string&) override { return Status::OK(); }
80-
Status LinkFile(const std::string&, const std::string&) override {
86+
Status DeleteFile(const std::string& fname) override {
87+
std::lock_guard<std::mutex> l(mutex);
88+
file_set.erase(fname);
89+
return Status::OK();
90+
}
91+
92+
Status LinkFile(const std::string& /*src*/, const std::string& dst) override {
93+
std::lock_guard<std::mutex> l(mutex);
94+
file_set.insert(dst);
8195
return Status::OK();
8296
}
8397
};

0 commit comments

Comments
 (0)