Atomize Rename operation when encryption is enabled (#224)

Signed-off-by: Xintao <hunterlxt@live.com>
Signed-off-by: tabokie <xy.tao@outlook.com>

db/db_test_util.cc

@@ -28,14 +28,6 @@ int64_t MaybeCurrentTime(Env* env) {
 }
 }  // 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)
@@ -78,7 +70,8 @@ DBTestBase::DBTestBase(const std::string path, bool env_do_fsync)
 #ifndef ROCKSDB_LITE
   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.");

db/db_test_util.h

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

encryption/encryption.cc

@@ -10,6 +10,7 @@
 #include <algorithm>
 #include <limits>
 
+#include "file/filename.h"
 #include "port/port.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -304,10 +305,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);
@@ -322,7 +330,7 @@ Status KeyManagedEncryptedEnv::NewWritableFile(
           "Unsupported encryption method: " +
           ROCKSDB_NAMESPACE::ToString(static_cast<int>(file_info.method)));
   }
-  if (!s.ok()) {
+  if (!s.ok() && !skipped) {
     // Ignore error
     key_manager_->DeleteFile(fname);
   }
@@ -433,6 +441,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;
@@ -448,6 +463,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);

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"
@@ -117,6 +118,17 @@ static Env* GetInspectedEnv() {
       Env::Default(), std::make_shared<DummyFileSystemInspector>(1)));
   return inspected_env.get();
 }
+
+#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
+
 #endif  // ROCKSDB_LITE
 
 }  // namespace
@@ -160,8 +172,12 @@ INSTANTIATE_TEST_CASE_P(MemEnv, EnvBasicTestWithParam,
 INSTANTIATE_TEST_CASE_P(InspectedEnv, EnvBasicTestWithParam,
                         ::testing::Values(&GetInspectedEnv));
 
-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.
 //
@@ -344,7 +360,7 @@ TEST_P(EnvBasicTestWithParam, LargeWrite) {
     read += result.size();
   }
   ASSERT_TRUE(write_data == read_data);
-  delete [] scratch;
+  delete[] scratch;
 }
 
 TEST_P(EnvMoreTestWithParam, GetModTime) {

file/filename.cc

@@ -28,6 +28,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.
@@ -184,6 +203,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) {
@@ -394,7 +417,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()) {

file/filename.h

@@ -36,6 +36,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".

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,

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;