Atomize Rename operation when encryption is enabled (facebook#8)

apache/incubator-pegasus#1575 Cherry-pick from tikv@bbd27cf used LinkFile instead of RenameFile api of key manager. But LinkFile needs check the dst file information, in RenameFile logic, we don't care about that. So just skip encryption for current file. Signed-off-by: Xintao [[email protected]](mailto:[email protected])
acelyc111 · Sep 15, 2023 · 60d0a7b · 60d0a7b
1 parent 2a2181e
commit 60d0a7b
Show file tree

Hide file tree

Showing 8 changed files with 122 additions and 47 deletions.
diff --git a/db/db_test_util.cc b/db/db_test_util.cc
@@ -31,14 +31,6 @@ int64_t MaybeCurrentTime(Env* env) {
 }
 }  // anonymous namespace
 
-#ifdef OPENSSL
-const std::string TestKeyManager::default_key =
-    "\x12\x34\x56\x78\x12\x34\x56\x78\x12\x34\x56\x78\x12\x34\x56\x78\x12\x34"
-    "\x56\x78\x12\x34\x56\x78";
-const std::string TestKeyManager::default_iv =
-    "\xaa\xbb\xcc\xdd\xaa\xbb\xcc\xdd\xaa\xbb\xcc\xdd\xaa\xbb\xcc\xdd";
-#endif
-
 // Special Env used to delay background operations
 
 SpecialEnv::SpecialEnv(Env* base, bool time_elapse_only_sleep)
@@ -80,7 +72,8 @@ DBTestBase::DBTestBase(const std::string path, bool env_do_fsync)
   }
   if (getenv("ENCRYPTED_ENV")) {
 #ifdef OPENSSL
-    std::shared_ptr<encryption::KeyManager> key_manager(new TestKeyManager);
+    std::shared_ptr<encryption::KeyManager> key_manager(
+        new test::TestKeyManager);
     encrypted_env_ = NewKeyManagedEncryptedEnv(Env::Default(), key_manager);
 #else
     fprintf(stderr, "EncryptedEnv is not available without OpenSSL.");

diff --git a/db/db_test_util.h b/db/db_test_util.h
@@ -56,38 +56,6 @@
 namespace ROCKSDB_NAMESPACE {
 class MockEnv;
 
-// TODO(yiwu): Use InMemoryKeyManager instead for tests.
-#ifdef OPENSSL
-class TestKeyManager : public encryption::KeyManager {
- public:
-  virtual ~TestKeyManager() = default;
-
-  static const std::string default_key;
-  static const std::string default_iv;
-
-  Status GetFile(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;
-    return Status::OK();
-  }
-
-  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;
-    return Status::OK();
-  }
-
-  Status DeleteFile(const std::string&) override { return Status::OK(); }
-  Status LinkFile(const std::string&, const std::string&) override {
-    return Status::OK();
-  }
-};
-#endif
-
 namespace anon {
 class AtomicCounter {
  public:

diff --git a/encryption/encryption.cc b/encryption/encryption.cc
@@ -10,6 +10,7 @@
 #include <algorithm>
 #include <limits>
 
+#include "file/filename.h"
 #include "port/port.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -307,10 +308,17 @@ Status KeyManagedEncryptedEnv::NewWritableFile(
     const std::string& fname, std::unique_ptr<WritableFile>* result,
     const EnvOptions& options) {
   FileEncryptionInfo file_info;
-  Status s = key_manager_->NewFile(fname, &file_info);
-  if (!s.ok()) {
-    return s;
+  Status s;
+  bool skipped = ShouldSkipEncryption(fname);
+  if (!skipped) {
+    s = key_manager_->NewFile(fname, &file_info);
+    if (!s.ok()) {
+      return s;
+    }
+  } else {
+    file_info.method = EncryptionMethod::kPlaintext;
   }
+
   switch (file_info.method) {
     case EncryptionMethod::kPlaintext:
       s = target()->NewWritableFile(fname, result, options);
@@ -325,7 +333,7 @@ Status KeyManagedEncryptedEnv::NewWritableFile(
           "Unsupported encryption method: " +
           std::to_string(static_cast<int>(file_info.method)));
   }
-  if (!s.ok()) {
+  if (!s.ok() && !skipped) {
     // Ignore error
     key_manager_->DeleteFile(fname);
   }
@@ -436,6 +444,13 @@ 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));
+    Status s = target()->LinkFile(src_fname, dst_fname);
+    return s;
+  } else {
+    assert(!ShouldSkipEncryption(src_fname));
+  }
   Status s = key_manager_->LinkFile(src_fname, dst_fname);
   if (!s.ok()) {
     return s;
@@ -451,6 +466,12 @@ 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);
+  } else {
+    assert(!ShouldSkipEncryption(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.
   Status s = key_manager_->LinkFile(src_fname, dst_fname);

diff --git a/env/env_basic_test.cc b/env/env_basic_test.cc
@@ -9,6 +9,7 @@
 #include <string>
 #include <vector>
 
+#include "db/db_test_util.h"
 #include "env/mock_env.h"
 #include "file/file_util.h"
 #include "rocksdb/convenience.h"
@@ -81,6 +82,16 @@ static Env* GetTestFS() {
   return fs_env;
 }
 
+#ifdef OPENSSL
+static Env* GetKeyManagedEncryptedEnv() {
+  static std::shared_ptr<encryption::KeyManager> key_manager(
+      new test::TestKeyManager);
+  static std::unique_ptr<Env> key_managed_encrypted_env(
+      NewKeyManagedEncryptedEnv(Env::Default(), key_manager));
+  return key_managed_encrypted_env.get();
+}
+#endif  // OPENSSL
+
 }  // namespace
 class EnvBasicTestWithParam
     : public testing::Test,
@@ -118,8 +129,12 @@ INSTANTIATE_TEST_CASE_P(EncryptedEnv, EnvMoreTestWithParam,
 INSTANTIATE_TEST_CASE_P(MemEnv, EnvBasicTestWithParam,
                         ::testing::Values(&GetMemoryEnv));
 
-namespace {
+#ifdef OPENSSL
+INSTANTIATE_TEST_CASE_P(KeyManagedEncryptedEnv, EnvBasicTestWithParam,
+                        ::testing::Values(&GetKeyManagedEncryptedEnv));
+#endif  // OPENSSL
 
+namespace {
 // Returns a vector of 0 or 1 Env*, depending whether an Env is registered for
 // TEST_ENV_URI.
 //

diff --git a/file/filename.cc b/file/filename.cc
@@ -30,6 +30,25 @@ static const std::string kRocksDbTFileExt = "sst";
 static const std::string kLevelDbTFileExt = "ldb";
 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) {
+  // skip CURRENT file.
+  size_t current_length = strlen("CURRENT");
+  if (fname.length() >= current_length &&
+      !fname.compare(fname.length() - current_length, current_length,
+                     "CURRENT")) {
+    return true;
+  }
+  // skip temporary file for CURRENT file.
+  size_t temp_length = kUnencryptedTempFileNameSuffix.length();
+  if (fname.length() >= temp_length &&
+      !fname.compare(fname.length() - temp_length, temp_length,
+                     kUnencryptedTempFileNameSuffix)) {
+    return true;
+  }
+  return false;
+}
 
 // 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.
@@ -183,6 +202,10 @@ std::string TempFileName(const std::string& dbname, uint64_t number) {
   return MakeFileName(dbname, number, kTempFileNameSuffix.c_str());
 }
 
+std::string TempPlainFileName(const std::string& dbname, uint64_t number) {
+  return MakeFileName(dbname, number, kUnencryptedTempFileNameSuffix.c_str());
+}
+
 InfoLogPrefix::InfoLogPrefix(bool has_log_dir,
                              const std::string& db_absolute_path) {
   if (!has_log_dir) {
@@ -393,7 +416,7 @@ IOStatus SetCurrentFile(FileSystem* fs, const std::string& dbname,
   Slice contents = manifest;
   assert(contents.starts_with(dbname + "/"));
   contents.remove_prefix(dbname.size() + 1);
-  std::string tmp = TempFileName(dbname, descriptor_number);
+  std::string tmp = TempPlainFileName(dbname, descriptor_number);
   IOStatus s = WriteStringToFile(fs, contents.ToString() + "\n", tmp, true);
   TEST_SYNC_POINT_CALLBACK("SetCurrentFile:BeforeRename", &s);
   if (s.ok()) {

diff --git a/file/filename.h b/file/filename.h
@@ -37,6 +37,10 @@ constexpr char kFilePathSeparator = '\\';
 constexpr char kFilePathSeparator = '/';
 #endif
 
+// Some non-sensitive files are not encrypted to preserve atomicity of file
+// operations.
+extern bool ShouldSkipEncryption(const std::string& fname);
+
 // Return the name of the log file with the specified number
 // in the db named by "dbname".  The result will be prefixed with
 // "dbname".

diff --git a/test_util/testutil.cc b/test_util/testutil.cc
@@ -37,6 +37,16 @@ void RegisterCustomObjects(int /*argc*/, char** /*argv*/) {}
 namespace ROCKSDB_NAMESPACE {
 namespace test {
 
+#ifdef OPENSSL
+#ifndef ROCKSDB_LITE
+const std::string TestKeyManager::default_key =
+    "\x12\x34\x56\x78\x12\x34\x56\x78\x12\x34\x56\x78\x12\x34\x56\x78\x12\x34"
+    "\x56\x78\x12\x34\x56\x78";
+const std::string TestKeyManager::default_iv =
+    "\xaa\xbb\xcc\xdd\xaa\xbb\xcc\xdd\xaa\xbb\xcc\xdd\xaa\xbb\xcc\xdd";
+#endif
+#endif
+
 const uint32_t kDefaultFormatVersion = BlockBasedTableOptions().format_version;
 const std::set<uint32_t> kFooterFormatVersionsToTest{
     5U,

diff --git a/test_util/testutil.h b/test_util/testutil.h
@@ -14,8 +14,11 @@
 #include <vector>
 
 #include "env/composite_env_wrapper.h"
+#include "file/filename.h"
 #include "file/writable_file_writer.h"
 #include "rocksdb/compaction_filter.h"
+#include "rocksdb/db.h"
+#include "rocksdb/encryption.h"
 #include "rocksdb/env.h"
 #include "rocksdb/iterator.h"
 #include "rocksdb/merge_operator.h"
@@ -43,6 +46,44 @@ class SequentialFileReader;
 
 namespace test {
 
+// TODO(yiwu): Use InMemoryKeyManager instead for tests.
+#ifdef OPENSSL
+#ifndef ROCKSDB_LITE
+class TestKeyManager : public encryption::KeyManager {
+ public:
+  virtual ~TestKeyManager() = default;
+
+  static const std::string default_key;
+  static const std::string default_iv;
+
+  Status GetFile(const std::string& fname,
+                 encryption::FileEncryptionInfo* file_info) override {
+    if (ShouldSkipEncryption(fname)) {
+      file_info->method = encryption::EncryptionMethod::kPlaintext;
+    } else {
+      file_info->method = encryption::EncryptionMethod::kAES192_CTR;
+    }
+    file_info->key = default_key;
+    file_info->iv = default_iv;
+    return Status::OK();
+  }
+
+  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;
+    return Status::OK();
+  }
+
+  Status DeleteFile(const std::string&) override { return Status::OK(); }
+  Status LinkFile(const std::string&, const std::string&) override {
+    return Status::OK();
+  }
+};
+#endif
+#endif
+
 extern const uint32_t kDefaultFormatVersion;
 extern const std::set<uint32_t> kFooterFormatVersionsToTest;