1

acelyc111 · Aug 14, 2023 · adc0458 · adc0458
1 parent d8723d5
commit adc0458
Show file tree

Hide file tree

Showing 9 changed files with 59 additions and 20 deletions.
diff --git a/.github/workflows/jobs-linux-run-tests.yml b/.github/workflows/jobs-linux-run-tests.yml
@@ -85,7 +85,7 @@ jobs:
     steps:
     - uses: actions/[email protected]
     - uses: "./.github/actions/pre-steps"
-    - run: mkdir build && cd build && cmake -DWITH_OPENSSL=0 -DROCKSDB_DISABLE_SNAPPY=1 -DROCKSDB_DISABLE_ZLIB=1 -DROCKSDB_DISABLE_BZIP=1 -DROCKSDB_DISABLE_LZ4=1 -DROCKSDB_DISABLE_ZSTD=1 .. && make V=1 -j5 && ctest -j5 -V
+    - run: mkdir build && cd build && cmake -DWITH_OPENSSL=0 -DWITH_SNAPPY=0 -DWITH_ZLIB=0 -DWITH_BZ2=0 -DWITH_LZ4=0 -DWITH_ZSTD=0 .. && make V=1 -j5 && ctest -j5 -V
     - run: "cd build/tools && ./sst_dump --help | grep -E -q 'Supported compression types: kNoCompression'"
     - uses: "./.github/actions/post-steps"
   build-linux-encrypted_env-openssl:

diff --git a/encryption/encryption.cc b/encryption/encryption.cc
@@ -15,13 +15,27 @@
 #include "file/filename.h"
 #include "port/likely.h"
 #include "port/port.h"
+#include "rocksdb/encryption.h"
 #include "rocksdb/utilities/options_type.h"
 #include "test_util/sync_point.h"
 #include "util/aligned_buffer.h"
 
 namespace ROCKSDB_NAMESPACE {
 namespace encryption {
 
+EncryptionMethod GetEncryptionMethod(const std::string& method) {
+  if (!strcasecmp(method.c_str(), "AES128CTR")) {
+    return ROCKSDB_NAMESPACE::encryption::EncryptionMethod::kAES128_CTR;
+  } else if (!strcasecmp(method.c_str(), "AES192CTR")) {
+    return ROCKSDB_NAMESPACE::encryption::EncryptionMethod::kAES192_CTR;
+  } else if (!strcasecmp(method.c_str(), "AES256CTR")) {
+    return ROCKSDB_NAMESPACE::encryption::EncryptionMethod::kAES256_CTR;
+  } else if (!strcasecmp(method.c_str(), "SM4CTR")) {
+    return ROCKSDB_NAMESPACE::encryption::EncryptionMethod::kSM4_CTR;
+  }
+  return ROCKSDB_NAMESPACE::encryption::EncryptionMethod::kUnknown;
+}
+
 namespace {
 const char* const kEncryptionHeaderMagic = "pegsenc";
 

diff --git a/env/env_encryption.cc b/env/env_encryption.cc
@@ -1204,6 +1204,7 @@ Status BlockCipher::CreateFromString(const ConfigOptions& config_options,
   return LoadSharedObject<BlockCipher>(config_options, value, result);
 }
 
+// TODO(yingchun): when will this function be called?
 Status EncryptionProvider::CreateFromString(
     const ConfigOptions& config_options, const std::string& value,
     std::shared_ptr<EncryptionProvider>* result) {

diff --git a/env/file_system.cc b/env/file_system.cc
@@ -5,6 +5,7 @@
 //
 #include "rocksdb/file_system.h"
 
+#include "encryption/encryption.h"
 #include "env/composite_env_wrapper.h"
 #include "env/env_chroot.h"
 #include "env/env_encryption_ctr.h"
@@ -13,6 +14,7 @@
 #include "logging/env_logger.h"
 #include "options/db_options.h"
 #include "rocksdb/convenience.h"
+#include "rocksdb/encryption.h"
 #include "rocksdb/utilities/customizable_util.h"
 #include "rocksdb/utilities/object_registry.h"
 #include "rocksdb/utilities/options_type.h"
@@ -44,9 +46,19 @@ static int RegisterBuiltinFileSystems(ObjectLibrary& library,
       });
   library.AddFactory<FileSystem>(
       EncryptedFileSystem::kClassName(),
-      [](const std::string& /*uri*/, std::unique_ptr<FileSystem>* guard,
+      [](const std::string& uri, std::unique_ptr<FileSystem>* guard,
          std::string* errmsg) {
-        Status s = NewEncryptedFileSystemImpl(nullptr, nullptr, guard);
+        // "EncryptedFileSystem:server_key,server_key_iv,method"
+        auto args1 = StringSplit(uri, ':');
+        assert(args1.size() == 2);
+        auto args = StringSplit(args1[1], ',');
+        assert(args.size() == 3);
+        encryption::EncryptionMethod method =
+            encryption::GetEncryptionMethod(args[2]);
+        assert(method != encryption::EncryptionMethod::kUnknown);
+        auto provider = std::make_shared<encryption::AESEncryptionProvider>(
+            args[0], args[1], method);
+        Status s = NewEncryptedFileSystemImpl(nullptr, provider, guard);
         if (!s.ok()) {
           *errmsg = s.ToString();
         }

diff --git a/include/rocksdb/encryption.h b/include/rocksdb/encryption.h
@@ -44,6 +44,8 @@ struct FileEncryptionInfo {
   std::string iv;  // TODO(yingchun): not used yet
 };
 
+EncryptionMethod GetEncryptionMethod(const std::string& method);
+
 }  // namespace encryption
 }  // namespace ROCKSDB_NAMESPACE
 

diff --git a/include/rocksdb/utilities/ldb_cmd.h b/include/rocksdb/utilities/ldb_cmd.h
@@ -72,6 +72,9 @@ class LDBCommand {
   static const std::string ARG_PREPOPULATE_BLOB_CACHE;
   static const std::string ARG_DECODE_BLOB_INDEX;
   static const std::string ARG_DUMP_UNCOMPRESSED_BLOBS;
+  static const std::string ARG_ENCRYPTION_METHOD;
+  static const std::string ARG_ENCRYPTION_SERVER_KEY;
+  static const std::string ARG_ENCRYPTION_SERVER_KEY_IV;
 
   struct ParsedParams {
     std::string cmd;

diff --git a/options/options_helper.cc b/options/options_helper.cc
@@ -1369,6 +1369,7 @@ Status OptionTypeInfo::Prepare(const ConfigOptions& config_options,
       if (config != nullptr) {
         return config->PrepareOptions(config_options);
       } else if (!CanBeNull()) {
+        assert(0);
         return Status::NotFound("Missing configurable object", name);
       }
     }
@@ -1389,6 +1390,7 @@ Status OptionTypeInfo::Validate(const DBOptions& db_opts,
       if (config != nullptr) {
         return config->ValidateOptions(db_opts, cf_opts);
       } else if (!CanBeNull()) {
+        assert(0);
         return Status::NotFound("Missing configurable object", name);
       }
     }

diff --git a/tools/db_bench_tool.cc b/tools/db_bench_tool.cc
@@ -1742,6 +1742,16 @@ DEFINE_string(
     encryption_method, "",
     "If non-empty, enable encryption with the specific encryption method.");
 
+DEFINE_string(
+    encryption_server_key, "",
+    "Server key in plain-text to create the encrypted Env. Only used in "
+    "tests.");
+
+DEFINE_string(
+    encryption_server_key_iv, "",
+    "Server key IV in plain-text create the encrypted Env. Only used in "
+    "tests.");
+
 namespace ROCKSDB_NAMESPACE {
 namespace {
 static Status CreateMemTableRepFactory(
@@ -8544,26 +8554,15 @@ int db_bench_tool(int argc, char** argv) {
 
 #ifdef OPENSSL
   if (!FLAGS_encryption_method.empty()) {
-    ROCKSDB_NAMESPACE::encryption::EncryptionMethod method =
-        ROCKSDB_NAMESPACE::encryption::EncryptionMethod::kUnknown;
-    if (!strcasecmp(FLAGS_encryption_method.c_str(), "AES128CTR")) {
-      method = ROCKSDB_NAMESPACE::encryption::EncryptionMethod::kAES128_CTR;
-    } else if (!strcasecmp(FLAGS_encryption_method.c_str(), "AES192CTR")) {
-      method = ROCKSDB_NAMESPACE::encryption::EncryptionMethod::kAES192_CTR;
-    } else if (!strcasecmp(FLAGS_encryption_method.c_str(), "AES256CTR")) {
-      method = ROCKSDB_NAMESPACE::encryption::EncryptionMethod::kAES256_CTR;
-    } else if (!strcasecmp(FLAGS_encryption_method.c_str(), "SM4CTR")) {
-      method = ROCKSDB_NAMESPACE::encryption::EncryptionMethod::kSM4_CTR;
-    }
-    if (method == ROCKSDB_NAMESPACE::encryption::EncryptionMethod::kUnknown) {
+    encryption::EncryptionMethod method =
+        encryption::GetEncryptionMethod(FLAGS_encryption_method);
+    if (method == encryption::EncryptionMethod::kUnknown) {
       fprintf(stderr, "Unknown encryption method %s\n",
               FLAGS_encryption_method.c_str());
       exit(1);
     }
-    std::shared_ptr<encryption::AESEncryptionProvider> provider(
-        new encryption::AESEncryptionProvider(
-            // TODO(yingchun): use google flags.
-            "test_serverkey", "test_serverkey_iv", method));
+    auto provider = std::make_shared<encryption::AESEncryptionProvider>(
+        FLAGS_encryption_server_key, FLAGS_encryption_server_key_iv, method);
     FLAGS_env = NewEncryptedEnv(FLAGS_env, provider);
   }
 #endif  // OPENSSL

diff --git a/tools/ldb_cmd.cc b/tools/ldb_cmd.cc
@@ -106,6 +106,9 @@ const std::string LDBCommand::ARG_PREPOPULATE_BLOB_CACHE =
 const std::string LDBCommand::ARG_DECODE_BLOB_INDEX = "decode_blob_index";
 const std::string LDBCommand::ARG_DUMP_UNCOMPRESSED_BLOBS =
     "dump_uncompressed_blobs";
+const std::string LDBCommand::ARG_ENCRYPTION_METHOD = "encryption_method";
+const std::string LDBCommand::ARG_ENCRYPTION_SERVER_KEY = "encryption_server_key";
+const std::string LDBCommand::ARG_ENCRYPTION_SERVER_KEY_IV = "encryption_server_key_iv";
 
 const char* LDBCommand::DELIM = " ==> ";
 
@@ -559,7 +562,10 @@ std::vector<std::string> LDBCommand::BuildCmdLineOptions(
                                   ARG_BLOB_FILE_STARTING_LEVEL,
                                   ARG_PREPOPULATE_BLOB_CACHE,
                                   ARG_IGNORE_UNKNOWN_OPTIONS,
-                                  ARG_CF_NAME};
+                                  ARG_CF_NAME,
+                                  ARG_ENCRYPTION_METHOD,
+                                  ARG_ENCRYPTION_SERVER_KEY,
+                                  ARG_ENCRYPTION_SERVER_KEY_IV};
   ret.insert(ret.end(), options.begin(), options.end());
   return ret;
 }