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>

Author: hunterlxt

db/db_test.cc

@@ -2292,6 +2292,10 @@ TEST_F(DBTest, DestroyDBMetaDatabase) {
 
 #ifndef ROCKSDB_LITE
 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

encryption/encryption.cc

@@ -12,6 +12,7 @@
 
 #include "file/filename.h"
 #include "port/port.h"
+#include "test_util/sync_point.h"
 
 namespace ROCKSDB_NAMESPACE {
 namespace encryption {
@@ -267,6 +268,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(
@@ -306,7 +318,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()) {
@@ -441,12 +454,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()) {
@@ -463,11 +476,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

@@ -205,6 +205,52 @@ INSTANTIATE_TEST_CASE_P(CustomEnv, EnvMoreTestWithParam,
                         ::testing::ValuesIn(GetCustomEnvs()));
 #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

@@ -30,7 +30,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 &&
@@ -48,6 +48,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

@@ -38,7 +38,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::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

@@ -10,6 +10,8 @@
 #pragma once
 #include <algorithm>
 #include <deque>
+#include <mutex>
+#include <set>
 #include <string>
 #include <vector>
 
@@ -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;
@@ -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();
   }
 };