Atomize Rename operation when encryption is enabled (#224) (#225)

Cherry-picking #224.

#222 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 <hunterlxt@live.com>

db/db_test_util.cc

@@ -14,14 +14,6 @@
 
 namespace rocksdb {
 
-#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)

db/db_test_util.h

@@ -56,38 +56,6 @@
 
 namespace rocksdb {
 
-//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

@@ -9,6 +9,7 @@
 #include <algorithm>
 #include <limits>
 
+#include "file/filename.h"
 #include "port/port.h"
 
 namespace rocksdb {
@@ -300,10 +301,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);
@@ -317,7 +325,7 @@ Status KeyManagedEncryptedEnv::NewWritableFile(
       s = Status::InvalidArgument("Unsupported encryption method: " +
                                   ToString(static_cast<int>(file_info.method)));
   }
-  if (!s.ok()) {
+  if (!s.ok() && !skipped) {
     // Ignore error
     key_manager_->DeleteFile(fname);
   }
@@ -425,6 +433,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;
@@ -440,6 +455,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

@@ -4,11 +4,12 @@
 //
 // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
 
+#include <algorithm>
 #include <memory>
 #include <string>
 #include <vector>
-#include <algorithm>
 
+#include "db/db_test_util.h"
 #include "env/mock_env.h"
 #include "rocksdb/env.h"
 #include "test_util/testharness.h"
@@ -94,8 +95,15 @@ static std::unique_ptr<Env> mem_env(NewMemEnv(Env::Default()));
 INSTANTIATE_TEST_CASE_P(MemEnv, EnvBasicTestWithParam,
                         ::testing::Values(mem_env.get()));
 
-namespace {
+#ifdef OPENSSL
+std::shared_ptr<encryption::KeyManager> key_manager(new TestKeyManager);
+static std::unique_ptr<Env> key_managed_encrypted_env(NewKeyManagedEncryptedEnv(
+    new NormalizingEnvWrapper(Env::Default()), key_manager));
+INSTANTIATE_TEST_CASE_P(KeyManagedEncryptedEnv, EnvBasicTestWithParam,
+                        ::testing::Values(key_managed_encrypted_env.get()));
+#endif  // OPENSSL
 
+namespace {
 // Returns a vector of 0 or 1 Env*, depending whether an Env is registered for
 // TEST_ENV_URI.
 //
@@ -285,7 +293,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

@@ -24,6 +24,25 @@ namespace rocksdb {
 static const std::string kRocksDbTFileExt = "sst";
 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) {
+  // 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.
@@ -159,6 +178,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) {
@@ -364,7 +387,7 @@ Status SetCurrentFile(Env* env, 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);
   Status s = WriteStringToFile(env, contents.ToString() + "\n", tmp, true);
   if (s.ok()) {
     TEST_KILL_RANDOM("SetCurrentFile:0", rocksdb_kill_odds * REDUCE_ODDS2);

file/filename.h

@@ -42,6 +42,10 @@ enum FileType {
   kBlobFile
 };
 
+// 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

@@ -19,6 +19,17 @@
 #include "util/file_reader_writer.h"
 
 namespace rocksdb {
+
+#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
+
 namespace test {
 
 const uint32_t kDefaultFormatVersion = BlockBasedTableOptions().format_version;

test_util/testutil.h

@@ -15,6 +15,8 @@
 
 #include "rocksdb/compaction_filter.h"
 #include "rocksdb/env.h"
+#include "rocksdb/db.h"
+#include "rocksdb/encryption.h"
 #include "rocksdb/iterator.h"
 #include "rocksdb/merge_operator.h"
 #include "rocksdb/options.h"
@@ -25,8 +27,48 @@
 #include "table/plain/plain_table_factory.h"
 #include "util/mutexlock.h"
 #include "util/random.h"
+#include "file/filename.h"
 
 namespace rocksdb {
+
+// 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
+
 class SequentialFile;
 class SequentialFileReader;