feat(encryption): introduce PegasusEnv (#1606)

#1575

This patch introduces `PegasusEnv()` to obtain the `Env` instance used by RocksDB. Then
it's possible to obtain an encrypted Env instance by `PegasusEnv(FileDataType::kSensitive)`,
the encrypted Env is used for operating on sensitive files, the writing data to the file
will be encrypted and the reading data from the file will be decrypted.

Some file operate functions and related unit tests are added as well.

src/aio/test/CMakeLists.txt

@@ -33,7 +33,7 @@ set(MY_PROJ_SRC "")
 # "GLOB" for non-recursive search
 set(MY_SRC_SEARCH_MODE "GLOB")
 
-set(MY_PROJ_LIBS gtest dsn_runtime dsn_aio)
+set(MY_PROJ_LIBS gtest dsn_runtime dsn_aio rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 

src/block_service/test/CMakeLists.txt

@@ -36,7 +36,7 @@ set(MY_PROJ_LIBS
     gtest
     gtest_main
     hdfs
-    )
+    rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 

src/client/test/CMakeLists.txt

@@ -27,7 +27,7 @@ set(MY_PROJ_LIBS
     dsn_runtime
     dsn_utils
     gtest
-)
+    rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 

src/common/test/CMakeLists.txt

@@ -30,7 +30,7 @@ set(MY_PROJ_LIBS
         dsn_replication_common
         dsn_runtime
         gtest
-        )
+        rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 

src/failure_detector/test/CMakeLists.txt

@@ -41,7 +41,7 @@ set(MY_PROJ_LIBS
     dsn.failure_detector
     gtest
     hashtable
-    )
+    rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 

src/http/test/CMakeLists.txt

@@ -26,7 +26,7 @@ set(MY_PROJ_LIBS
     dsn_runtime
     gtest
     gtest_main
-    )
+    rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 

src/meta/CMakeLists.txt

@@ -54,7 +54,7 @@ set(MY_PROJ_LIBS
     crypto
     hashtable
     hdfs
-    )
+    rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 

src/nfs/test/CMakeLists.txt

@@ -33,7 +33,7 @@ set(MY_PROJ_SRC "")
 # "GLOB" for non-recursive search
 set(MY_SRC_SEARCH_MODE "GLOB")
 
-set(MY_PROJ_LIBS dsn_nfs dsn_runtime gtest dsn_aio)
+set(MY_PROJ_LIBS dsn_nfs dsn_runtime gtest dsn_aio rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 

src/perf_counter/test/CMakeLists.txt

@@ -33,7 +33,7 @@ set(MY_PROJ_SRC "")
 # "GLOB" for non-recursive search
 set(MY_SRC_SEARCH_MODE "GLOB")
 
-set(MY_PROJ_LIBS gtest dsn_runtime)
+set(MY_PROJ_LIBS gtest dsn_runtime rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 

src/replica/backup/test/CMakeLists.txt

@@ -30,7 +30,7 @@ set(MY_PROJ_LIBS dsn_meta_server
         dsn_utils
         hashtable
         gtest
-)
+        rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 

src/replica/duplication/test/CMakeLists.txt

@@ -30,7 +30,7 @@ set(MY_PROJ_LIBS dsn_meta_server
         zookeeper
         hashtable
         gtest
-)
+        rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 

src/test_util/test_util.h

@@ -21,10 +21,21 @@
 
 #include <functional>
 
+#include "gtest/gtest.h"
+#include "utils/flags.h"
 #include "utils/test_macros.h"
 
+DSN_DECLARE_bool(encrypt_data_at_rest);
+
 namespace pegasus {
 
+// A base parameterized test class for testing enable/disable encryption at rest.
+class encrypt_data_test_base : public testing::TestWithParam<bool>
+{
+public:
+    encrypt_data_test_base() { FLAGS_encrypt_data_at_rest = GetParam(); }
+};
+
 #define ASSERT_EVENTUALLY(expr)                                                                    \
     do {                                                                                           \
         AssertEventually(expr);                                                                    \

src/utils/CMakeLists.txt

@@ -31,7 +31,7 @@ set(MY_SRC_SEARCH_MODE "GLOB")
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 
-set(MY_PROJ_LIBS dsn_http crypto)
+set(MY_PROJ_LIBS dsn_http crypto rocksdb)
 
 # Extra files that will be installed
 set(MY_BINPLACES "")

src/utils/env.cpp

@@ -0,0 +1,201 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "env.h"
+
+#include <algorithm>
+#include <memory>
+#include <string>
+
+#include <fmt/core.h>
+#include <rocksdb/convenience.h>
+#include <rocksdb/env.h>
+#include <rocksdb/env_encryption.h>
+#include <rocksdb/slice.h>
+
+#include "utils/defer.h"
+#include "utils/filesystem.h"
+#include "utils/flags.h"
+#include "utils/fmt_logging.h"
+#include "utils/utils.h"
+
+DSN_DEFINE_bool(pegasus.server,
+                encrypt_data_at_rest,
+                false,
+                "Whether the sensitive files should be encrypted on the file system.");
+
+DSN_DEFINE_string(pegasus.server,
+                  server_key_for_testing,
+                  "server_key_for_testing",
+                  "The encrypted server key to use in the filesystem. NOTE: only for testing.");
+
+DSN_DEFINE_string(pegasus.server,
+                  encryption_method,
+                  "AES128CTR",
+                  "The encryption method to use in the filesystem. Now "
+                  "supports AES128CTR, AES192CTR, AES256CTR and SM4CTR.");
+
+namespace dsn {
+namespace utils {
+
+rocksdb::Env *NewEncryptedEnv()
+{
+    // Create an encryption provider.
+    std::shared_ptr<rocksdb::EncryptionProvider> provider;
+    auto provider_id =
+        fmt::format("AES:{},{}", FLAGS_server_key_for_testing, FLAGS_encryption_method);
+    auto s = rocksdb::EncryptionProvider::CreateFromString(
+        rocksdb::ConfigOptions(), provider_id, &provider);
+    CHECK(s.ok(), "Failed to create encryption provider: {}", s.ToString());
+
+    // Create an encrypted env.
+    return NewEncryptedEnv(rocksdb::Env::Default(), provider);
+}
+
+rocksdb::Env *PegasusEnv(FileDataType type)
+{
+    // Return an encrypted env only when the file is sensitive and FLAGS_encrypt_data_at_rest
+    // is enabled at the same time.
+    if (FLAGS_encrypt_data_at_rest && type == FileDataType::kSensitive) {
+        static rocksdb::Env *env = NewEncryptedEnv();
+        return env;
+    }
+
+    // Otherwise, return a common non-encrypted env.
+    static rocksdb::Env *env = rocksdb::Env::Default();
+    return env;
+}
+
+namespace {
+rocksdb::Status do_copy_file(const std::string &src_fname,
+                             dsn::utils::FileDataType src_type,
+                             const std::string &dst_fname,
+                             dsn::utils::FileDataType dst_type,
+                             int64_t remain_size,
+                             uint64_t *total_size)
+{
+    rocksdb::EnvOptions src_env_options;
+    std::unique_ptr<rocksdb::SequentialFile> src_file;
+    auto s =
+        dsn::utils::PegasusEnv(src_type)->NewSequentialFile(src_fname, &src_file, src_env_options);
+    LOG_AND_RETURN_NOT_RDB_OK(WARNING, s, "failed to open file {} for reading", src_fname);
+
+    // Limit the size of the file to be copied.
+    int64_t src_file_size;
+    CHECK_TRUE(dsn::utils::filesystem::file_size(src_fname, src_type, src_file_size));
+    if (remain_size == -1) {
+        // Copy the whole file if 'remain_size' is -1.
+        remain_size = src_file_size;
+    } else {
+        remain_size = std::min(remain_size, src_file_size);
+    }
+
+    rocksdb::EnvOptions dst_env_options;
+    std::unique_ptr<rocksdb::WritableFile> dst_file;
+    s = dsn::utils::PegasusEnv(dst_type)->NewWritableFile(dst_fname, &dst_file, dst_env_options);
+    LOG_AND_RETURN_NOT_RDB_OK(WARNING, s, "failed to open file {} for writing", dst_fname);
+
+    // Read at most 4MB once.
+    // TODO(yingchun): make it configurable.
+    const uint64_t kCopyBlockSize = 4 << 20;
+    auto buffer = dsn::utils::make_shared_array<char>(kCopyBlockSize);
+    uint64_t offset = 0;
+    do {
+        int bytes_per_copy = std::min(remain_size, static_cast<int64_t>(kCopyBlockSize));
+        // Reach the EOF.
+        if (bytes_per_copy <= 0) {
+            break;
+        }
+
+        rocksdb::Slice result;
+        LOG_AND_RETURN_NOT_RDB_OK(WARNING,
+                                  src_file->Read(bytes_per_copy, &result, buffer.get()),
+                                  "failed to read file {}",
+                                  src_fname);
+        CHECK(!result.empty(),
+              "read file {} at offset {} with size {} failed",
+              src_fname,
+              offset,
+              bytes_per_copy);
+        LOG_AND_RETURN_NOT_RDB_OK(
+            WARNING, dst_file->Append(result), "failed to write file {}", dst_fname);
+
+        offset += result.size();
+        remain_size -= result.size();
+
+        // Reach the EOF.
+        if (result.size() < bytes_per_copy) {
+            break;
+        }
+    } while (true);
+    LOG_AND_RETURN_NOT_RDB_OK(WARNING, dst_file->Fsync(), "failed to fsync file {}", dst_fname);
+
+    if (total_size != nullptr) {
+        *total_size = offset;
+    }
+
+    LOG_INFO("copy file from {} to {}, total size {}", src_fname, dst_fname, offset);
+    return rocksdb::Status::OK();
+}
+} // anonymous namespace
+
+rocksdb::Status
+copy_file(const std::string &src_fname, const std::string &dst_fname, uint64_t *total_size)
+{
+    // TODO(yingchun): Consider to use hard link, i.e. rocksdb::Env()::LinkFile().
+    return do_copy_file(
+        src_fname, FileDataType::kSensitive, dst_fname, FileDataType::kSensitive, -1, total_size);
+}
+
+rocksdb::Status
+encrypt_file(const std::string &src_fname, const std::string &dst_fname, uint64_t *total_size)
+{
+    return do_copy_file(src_fname,
+                        FileDataType::kNonSensitive,
+                        dst_fname,
+                        FileDataType::kSensitive,
+                        -1,
+                        total_size);
+}
+
+rocksdb::Status encrypt_file(const std::string &fname, uint64_t *total_size)
+{
+    // TODO(yingchun): add timestamp to the tmp encrypted file name.
+    std::string tmp_fname = fname + ".encrypted.tmp";
+    auto cleanup = dsn::defer([tmp_fname]() { utils::filesystem::remove_path(tmp_fname); });
+    LOG_AND_RETURN_NOT_RDB_OK(
+        WARNING, encrypt_file(fname, tmp_fname, total_size), "failed to encrypt file {}", fname);
+    if (!::dsn::utils::filesystem::rename_path(tmp_fname, fname)) {
+        LOG_WARNING("rename file from {} to {} failed", tmp_fname, fname);
+        return rocksdb::Status::IOError("rename file failed");
+    }
+    return rocksdb::Status::OK();
+}
+
+rocksdb::Status
+copy_file_by_size(const std::string &src_fname, const std::string &dst_fname, int64_t limit_size)
+{
+    return do_copy_file(src_fname,
+                        FileDataType::kSensitive,
+                        dst_fname,
+                        FileDataType::kSensitive,
+                        limit_size,
+                        nullptr);
+}
+
+} // namespace utils
+} // namespace dsn

src/utils/env.h

@@ -0,0 +1,66 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <rocksdb/env.h>
+#include <rocksdb/status.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <string>
+
+namespace dsn {
+namespace utils {
+
+// Indicate whether the file is sensitive or not.
+// Only the sensitive file will be encrypted if FLAGS_encrypt_data_at_rest
+// is enabled at the same time.
+enum class FileDataType
+{
+    kSensitive = 0,
+    kNonSensitive = 1
+};
+
+static const size_t kEncryptionHeaderkSize = rocksdb::kDefaultPageSize;
+
+// Get the rocksdb::Env instance for the given file type.
+rocksdb::Env *PegasusEnv(FileDataType type);
+
+// Encrypt the original non-encrypted 'src_fname' to 'dst_fname'.
+// The 'total_size' is the total size of the file content, exclude the file encryption header
+// (typically 4KB).
+rocksdb::Status encrypt_file(const std::string &src_fname,
+                             const std::string &dst_fname,
+                             uint64_t *total_size = nullptr);
+
+// Similar to the above, but encrypt the file in the same path.
+rocksdb::Status encrypt_file(const std::string &fname, uint64_t *total_size = nullptr);
+
+// Copy the original 'src_fname' to 'dst_fname'.
+// Both 'src_fname' and 'dst_fname' are sensitive files.
+rocksdb::Status copy_file(const std::string &src_fname,
+                          const std::string &dst_fname,
+                          uint64_t *total_size = nullptr);
+
+// Similar to the above, but copy the file by a limited size.
+// Both 'src_fname' and 'dst_fname' are sensitive files, 'limit_size' is the max size of the
+// file to copy, and -1 means no limit.
+rocksdb::Status copy_file_by_size(const std::string &src_fname,
+                                  const std::string &dst_fname,
+                                  int64_t limit_size = -1);
+} // namespace utils
+} // namespace dsn