Fix the bug that the key manager is not updated during the Rename (#227)

Signed-off-by: Xintao <hunterlxt@live.com>

encryption/encryption.cc

@@ -11,6 +11,7 @@
 
 #include "file/filename.h"
 #include "port/port.h"
+#include "test_util/sync_point.h"
 
 namespace rocksdb {
 namespace encryption {
@@ -265,6 +266,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("Unsupported encryption method: " +
@@ -302,7 +314,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()) {
@@ -433,12 +446,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()) {
@@ -455,11 +468,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.

env/env_basic_test.cc

@@ -129,6 +129,52 @@ INSTANTIATE_TEST_CASE_P(CustomEnv, EnvMoreTestWithParam,
 
 #endif // ROCKSDB_LITE
 
+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;

file/filename.cc

@@ -26,7 +26,7 @@ static const std::string kLevelDbTFileExt = "ldb";
 static const std::string kRocksDBBlobFileExt = "blob";
 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 &&
@@ -44,6 +44,13 @@ bool ShouldSkipEncryption(const std::string& fname) {
  return false;
 }
 
+bool IsValidCurrentFile(std::unique_ptr<rocksdb::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'.

file/filename.h

@@ -44,7 +44,11 @@ enum FileType {
 
 // 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::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

test_util/testutil.h

@@ -12,6 +12,7 @@
 #include <deque>
 #include <string>
 #include <vector>
+#include <set>
 
 #include "rocksdb/compaction_filter.h"
 #include "rocksdb/env.h"
@@ -40,10 +41,11 @@ class TestKeyManager : public encryption::KeyManager {
 
  static const std::string default_key;
  static const std::string default_iv;
+ std::set<std::string> file_set;
 
  Status GetFile(const std::string& fname,
  encryption::FileEncryptionInfo* file_info) override {
- if (ShouldSkipEncryption(fname)) {
+ if (file_set.find(fname) == file_set.end()) {
  file_info->method = encryption::EncryptionMethod::kPlaintext;
  } else {
  file_info->method = encryption::EncryptionMethod::kAES192_CTR;
@@ -53,16 +55,22 @@ 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 {
  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 {
+ file_set.erase(fname);
+ return Status::OK();
+ }
+
+ Status LinkFile(const std::string&, const std::string& dst) override {
+ file_set.insert(dst);
  return Status::OK();
  }
 };