This is an automated email from the ASF dual-hosted git repository.

laiyingchun pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-pegasus.git


The following commit(s) were added to refs/heads/master by this push:
     new 625cebcc9 feat(encryption): introduce PegasusEnv (#1606)
625cebcc9 is described below

commit 625cebcc90c100bd9ba6d2c6d55ac0d29924ae00
Author: Yingchun Lai <[email protected]>
AuthorDate: Wed Sep 20 16:25:25 2023 +0800

    feat(encryption): introduce PegasusEnv (#1606)
    
    https://github.com/apache/incubator-pegasus/issues/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                     |   2 +-
 src/block_service/test/CMakeLists.txt           |   2 +-
 src/client/test/CMakeLists.txt                  |   2 +-
 src/common/test/CMakeLists.txt                  |   2 +-
 src/failure_detector/test/CMakeLists.txt        |   2 +-
 src/http/test/CMakeLists.txt                    |   2 +-
 src/meta/CMakeLists.txt                         |   2 +-
 src/nfs/test/CMakeLists.txt                     |   2 +-
 src/perf_counter/test/CMakeLists.txt            |   2 +-
 src/replica/backup/test/CMakeLists.txt          |   2 +-
 src/replica/duplication/test/CMakeLists.txt     |   2 +-
 src/test_util/test_util.h                       |  11 +
 src/utils/CMakeLists.txt                        |   2 +-
 src/utils/env.cpp                               | 201 +++++++++++++++++
 src/utils/env.h                                 |  66 ++++++
 src/utils/filesystem.cpp                        |  33 ++-
 src/utils/filesystem.h                          |  10 +
 src/utils/fmt_logging.h                         |  11 +
 src/utils/long_adder_bench/CMakeLists.txt       |   2 +-
 src/utils/test/CMakeLists.txt                   |   3 +-
 src/utils/test/env.cpp                          | 285 +++++++++++++++++++++++-
 src/utils/test/file_system_test.cpp             |  90 +++++++-
 src/utils/test/nth_element_bench/CMakeLists.txt |   2 +-
 src/zookeeper/test/CMakeLists.txt               |   2 +-
 24 files changed, 715 insertions(+), 25 deletions(-)

diff --git a/src/aio/test/CMakeLists.txt b/src/aio/test/CMakeLists.txt
index 4228a0148..c1b0a44e4 100644
--- a/src/aio/test/CMakeLists.txt
+++ b/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)
 
diff --git a/src/block_service/test/CMakeLists.txt 
b/src/block_service/test/CMakeLists.txt
index a8e47d597..537778caa 100644
--- a/src/block_service/test/CMakeLists.txt
+++ b/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)
 
diff --git a/src/client/test/CMakeLists.txt b/src/client/test/CMakeLists.txt
index 90ce5d5e8..bcae3897a 100644
--- a/src/client/test/CMakeLists.txt
+++ b/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)
 
diff --git a/src/common/test/CMakeLists.txt b/src/common/test/CMakeLists.txt
index 78d94000c..1cdf58407 100644
--- a/src/common/test/CMakeLists.txt
+++ b/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)
 
diff --git a/src/failure_detector/test/CMakeLists.txt 
b/src/failure_detector/test/CMakeLists.txt
index d529e50ea..ed4a9703d 100644
--- a/src/failure_detector/test/CMakeLists.txt
+++ b/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)
 
diff --git a/src/http/test/CMakeLists.txt b/src/http/test/CMakeLists.txt
index 6a09e42dc..d4da7e5b3 100644
--- a/src/http/test/CMakeLists.txt
+++ b/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)
 
diff --git a/src/meta/CMakeLists.txt b/src/meta/CMakeLists.txt
index b38ced722..def16ec26 100644
--- a/src/meta/CMakeLists.txt
+++ b/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)
 
diff --git a/src/nfs/test/CMakeLists.txt b/src/nfs/test/CMakeLists.txt
index 735bb29bd..64c7967ac 100644
--- a/src/nfs/test/CMakeLists.txt
+++ b/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)
 
diff --git a/src/perf_counter/test/CMakeLists.txt 
b/src/perf_counter/test/CMakeLists.txt
index a02a6923b..434d97dc1 100644
--- a/src/perf_counter/test/CMakeLists.txt
+++ b/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)
 
diff --git a/src/replica/backup/test/CMakeLists.txt 
b/src/replica/backup/test/CMakeLists.txt
index 21ec10f50..e3dc0ec05 100644
--- a/src/replica/backup/test/CMakeLists.txt
+++ b/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)
 
diff --git a/src/replica/duplication/test/CMakeLists.txt 
b/src/replica/duplication/test/CMakeLists.txt
index 09619343e..faab20f7e 100644
--- a/src/replica/duplication/test/CMakeLists.txt
+++ b/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)
 
diff --git a/src/test_util/test_util.h b/src/test_util/test_util.h
index ec7bce6e2..d33775ff5 100644
--- a/src/test_util/test_util.h
+++ b/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);                                                
                    \
diff --git a/src/utils/CMakeLists.txt b/src/utils/CMakeLists.txt
index e17fef900..1127a868e 100644
--- a/src/utils/CMakeLists.txt
+++ b/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 "")
diff --git a/src/utils/env.cpp b/src/utils/env.cpp
new file mode 100644
index 000000000..329406117
--- /dev/null
+++ b/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
diff --git a/src/utils/env.h b/src/utils/env.h
new file mode 100644
index 000000000..839bd81ab
--- /dev/null
+++ b/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
diff --git a/src/utils/filesystem.cpp b/src/utils/filesystem.cpp
index 73aa3a50a..f2fa59ef8 100644
--- a/src/utils/filesystem.cpp
+++ b/src/utils/filesystem.cpp
@@ -33,11 +33,8 @@
  *     xxxx-xx-xx, author, fix bug about xxx
  */
 
-#include <boost/filesystem/operations.hpp>
-#include <boost/system/error_code.hpp>
 #include <errno.h>
 #include <fcntl.h>
-#include <fmt/core.h>
 #include <ftw.h>
 #include <limits.h>
 #include <openssl/md5.h>
@@ -45,11 +42,18 @@
 #include <stdlib.h>
 #include <sys/stat.h>
 // IWYU pragma: no_include <bits/struct_stat.h>
-#include <sys/stat.h> // IWYU pragma: keep
 #include <unistd.h>
-#include <fstream>
+
+#include <istream>
+
+#include <boost/filesystem/operations.hpp>
+#include <boost/system/error_code.hpp>
+#include <fmt/core.h>
+#include <rocksdb/env.h>
+#include <rocksdb/status.h>
 
 #include "utils/defer.h"
+#include "utils/env.h"
 #include "utils/fail_point.h"
 #include "utils/filesystem.h"
 #include "utils/fmt_logging.h"
@@ -388,7 +392,7 @@ bool rename_path(const std::string &path1, const 
std::string &path2)
     return ret;
 }
 
-bool file_size(const std::string &path, int64_t &sz)
+bool deprecated_file_size(const std::string &path, int64_t &sz)
 {
     struct stat_ st;
     std::string npath;
@@ -417,6 +421,23 @@ bool file_size(const std::string &path, int64_t &sz)
     return true;
 }
 
+bool file_size(const std::string &path, int64_t &sz)
+{
+    return file_size(path, dsn::utils::FileDataType::kNonSensitive, sz);
+}
+
+bool file_size(const std::string &path, FileDataType type, int64_t &sz)
+{
+    uint64_t file_size = 0;
+    auto s = dsn::utils::PegasusEnv(type)->GetFileSize(path, &file_size);
+    if (!s.ok()) {
+        LOG_ERROR("GetFileSize failed, file '{}', err = {}", path, 
s.ToString());
+        return false;
+    }
+    sz = file_size;
+    return true;
+}
+
 static int create_directory_component(const std::string &npath)
 {
     int err;
diff --git a/src/utils/filesystem.h b/src/utils/filesystem.h
index 4229f551d..9c5d3efea 100644
--- a/src/utils/filesystem.h
+++ b/src/utils/filesystem.h
@@ -61,6 +61,8 @@
 
 namespace dsn {
 namespace utils {
+enum class FileDataType;
+
 namespace filesystem {
 
 int get_normalized_path(const std::string &path, std::string &npath);
@@ -96,7 +98,15 @@ bool remove_path(const std::string &path);
 // this will always remove target path if exist
 bool rename_path(const std::string &path1, const std::string &path2);
 
+// Get the file size. The encryption header is considered as part of the file 
if it is an encrypted
+// file.
+// TODO(yingchun): refactor to use uint64_t.
 bool file_size(const std::string &path, int64_t &sz);
+// The legacy file_size(), just for testing.
+bool deprecated_file_size(const std::string &path, int64_t &sz);
+// Get the file size. The encryption header is not considered as part of the 
file if it is an
+// encrypted file and 'type' is specified as FileDataType::kSensitive.
+bool file_size(const std::string &path, FileDataType type, int64_t &sz);
 
 bool create_directory(const std::string &path);
 
diff --git a/src/utils/fmt_logging.h b/src/utils/fmt_logging.h
index f0c74bd29..35c60aadd 100644
--- a/src/utils/fmt_logging.h
+++ b/src/utils/fmt_logging.h
@@ -20,6 +20,7 @@
 #pragma once
 
 #include <fmt/ostream.h>
+#include <rocksdb/status.h>
 
 #include "utils/api_utilities.h"
 
@@ -272,6 +273,16 @@ inline const char *null_str_printer(const char *s) { 
return s == nullptr ? "(nul
         LOG_AND_RETURN_NOT_TRUE(level, _err == ::dsn::ERR_OK, _err, 
__VA_ARGS__);                  \
     } while (0)
 
+// Return the given rocksdb::Status 's' if it is not OK.
+#define LOG_AND_RETURN_NOT_RDB_OK(level, s, ...)                               
                    \
+    do {                                                                       
                    \
+        const auto &_s = (s);                                                  
                    \
+        if (dsn_unlikely(!_s.ok())) {                                          
                    \
+            LOG_##level("{}: {}", _s.ToString(), fmt::format(__VA_ARGS__));    
                    \
+            return _s;                                                         
                    \
+        }                                                                      
                    \
+    } while (0)
+
 #ifndef NDEBUG
 #define DCHECK CHECK
 #define DCHECK_NOTNULL CHECK_NOTNULL
diff --git a/src/utils/long_adder_bench/CMakeLists.txt 
b/src/utils/long_adder_bench/CMakeLists.txt
index f63efc8a9..480b04807 100644
--- a/src/utils/long_adder_bench/CMakeLists.txt
+++ b/src/utils/long_adder_bench/CMakeLists.txt
@@ -27,7 +27,7 @@ set(MY_PROJ_SRC "")
 # "GLOB" for non-recursive search
 set(MY_SRC_SEARCH_MODE "GLOB")
 
-set(MY_PROJ_LIBS dsn_runtime dsn_utils)
+set(MY_PROJ_LIBS dsn_runtime dsn_utils rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 
diff --git a/src/utils/test/CMakeLists.txt b/src/utils/test/CMakeLists.txt
index fb284b0e8..266b9beeb 100644
--- a/src/utils/test/CMakeLists.txt
+++ b/src/utils/test/CMakeLists.txt
@@ -33,7 +33,8 @@ set(MY_PROJ_LIBS dsn_http
                  dsn_runtime
                  dsn_utils
                  gtest
-                 )
+                 rocksdb
+                 test_utils)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 
diff --git a/src/utils/test/env.cpp b/src/utils/test/env.cpp
index 619e5d7eb..813465de2 100644
--- a/src/utils/test/env.cpp
+++ b/src/utils/test/env.cpp
@@ -33,18 +33,31 @@
  *     xxxx-xx-xx, author, fix bug about xxx
  */
 
+#include <fmt/core.h>
+#include <gtest/gtest-param-test.h>
 // IWYU pragma: no_include <gtest/gtest-message.h>
 // IWYU pragma: no_include <gtest/gtest-test-part.h>
 #include <gtest/gtest.h>
+#include <rocksdb/env.h>
+#include <rocksdb/slice.h>
+#include <rocksdb/status.h>
 #include <stdint.h>
+#include <algorithm>
 #include <limits>
-#include <memory>
+#include <string>
 
+#include "test_util/test_util.h"
+#include "utils/enum_helper.h"
+#include "utils/env.h"
+#include "utils/filesystem.h"
+#include "utils/flags.h"
 #include "utils/rand.h"
 
+DSN_DECLARE_bool(encrypt_data_at_rest);
+
 using namespace ::dsn;
 
-TEST(core, env)
+TEST(env_test, rand)
 {
     uint64_t xs[] = {0, std::numeric_limits<uint64_t>::max() - 1, 0xdeadbeef};
 
@@ -56,3 +69,271 @@ TEST(core, env)
         EXPECT_TRUE(r == x || r == (x + 1));
     }
 }
+
+TEST(env_test, get_env)
+{
+    FLAGS_encrypt_data_at_rest = false;
+    auto *env_no_enc1 = 
dsn::utils::PegasusEnv(dsn::utils::FileDataType::kNonSensitive);
+    auto *env_no_enc2 = 
dsn::utils::PegasusEnv(dsn::utils::FileDataType::kSensitive);
+    ASSERT_EQ(env_no_enc1, env_no_enc2);
+
+    FLAGS_encrypt_data_at_rest = true;
+    auto *env_no_enc3 = 
dsn::utils::PegasusEnv(dsn::utils::FileDataType::kNonSensitive);
+    ASSERT_EQ(env_no_enc1, env_no_enc3);
+
+    auto *env_enc1 = 
dsn::utils::PegasusEnv(dsn::utils::FileDataType::kSensitive);
+    ASSERT_NE(env_no_enc1, env_enc1);
+}
+
+class env_file_test : public pegasus::encrypt_data_test_base
+{
+public:
+    env_file_test() : pegasus::encrypt_data_test_base()
+    {
+        // The size of an actual encrypted file should plus 
kEncryptionHeaderkSize bytes if consider
+        // it as kNonSensitive.
+        if (FLAGS_encrypt_data_at_rest) {
+            _extra_size = dsn::utils::kEncryptionHeaderkSize;
+        }
+    }
+    uint64_t _extra_size = 0;
+};
+
+INSTANTIATE_TEST_CASE_P(, env_file_test, ::testing::Values(false, true));
+
+TEST_P(env_file_test, encrypt_file_2_files)
+{
+    const std::string kFileName = "encrypt_file_2_files";
+    const std::string kEncryptedFileName = kFileName + ".encrypted";
+    const uint64_t kFileContentSize = 100;
+    const std::string kFileContent(kFileContentSize, 'a');
+
+    // Prepare a non-encrypted test file.
+    auto s =
+        
rocksdb::WriteStringToFile(dsn::utils::PegasusEnv(dsn::utils::FileDataType::kNonSensitive),
+                                   rocksdb::Slice(kFileContent),
+                                   kFileName,
+                                   /* should_sync */ true);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+
+    // Check file size.
+    int64_t wfile_size;
+    ASSERT_TRUE(dsn::utils::filesystem::file_size(
+        kFileName, dsn::utils::FileDataType::kNonSensitive, wfile_size));
+    ASSERT_EQ(kFileContentSize, wfile_size);
+
+    // Check encrypt_file(src_fname, dst_fname, total_size).
+    // Loop twice to check overwrite.
+    for (int i = 0; i < 2; ++i) {
+        uint64_t encrypt_file_size;
+        s = dsn::utils::encrypt_file(kFileName, kEncryptedFileName, 
&encrypt_file_size);
+        ASSERT_TRUE(s.ok()) << s.ToString();
+        ASSERT_EQ(kFileContentSize, encrypt_file_size);
+        ASSERT_TRUE(dsn::utils::filesystem::file_size(
+            kEncryptedFileName, dsn::utils::FileDataType::kSensitive, 
wfile_size));
+        ASSERT_EQ(kFileContentSize, wfile_size);
+        ASSERT_TRUE(dsn::utils::filesystem::file_size(
+            kEncryptedFileName, dsn::utils::FileDataType::kNonSensitive, 
wfile_size));
+        ASSERT_EQ(kFileContentSize + _extra_size, wfile_size);
+        // Check file content.
+        std::string data;
+        s = 
rocksdb::ReadFileToString(dsn::utils::PegasusEnv(dsn::utils::FileDataType::kSensitive),
+                                      kEncryptedFileName,
+                                      &data);
+        ASSERT_EQ(kFileContent, data);
+    }
+}
+
+TEST_P(env_file_test, encrypt_file_1_file)
+{
+    const std::string kFileName = "encrypt_file_1_file";
+    const uint64_t kFileContentSize = 100;
+    const std::string kFileContent(kFileContentSize, 'a');
+
+    // Prepare a non-encrypted test file.
+    auto s =
+        
rocksdb::WriteStringToFile(dsn::utils::PegasusEnv(dsn::utils::FileDataType::kNonSensitive),
+                                   rocksdb::Slice(kFileContent),
+                                   kFileName,
+                                   /* should_sync */ true);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+
+    // Check file size.
+    int64_t wfile_size;
+    ASSERT_TRUE(dsn::utils::filesystem::file_size(
+        kFileName, dsn::utils::FileDataType::kNonSensitive, wfile_size));
+    ASSERT_EQ(kFileContentSize, wfile_size);
+
+    // Check encrypt_file(fname, total_size).
+    uint64_t encrypt_file_size;
+    s = dsn::utils::encrypt_file(kFileName, &encrypt_file_size);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+    ASSERT_EQ(kFileContentSize, encrypt_file_size);
+    ASSERT_TRUE(dsn::utils::filesystem::file_size(
+        kFileName, dsn::utils::FileDataType::kSensitive, wfile_size));
+    ASSERT_EQ(kFileContentSize, wfile_size);
+    ASSERT_TRUE(dsn::utils::filesystem::file_size(
+        kFileName, dsn::utils::FileDataType::kNonSensitive, wfile_size));
+    ASSERT_EQ(kFileContentSize + _extra_size, wfile_size);
+    // Check file content.
+    std::string data;
+    s = rocksdb::ReadFileToString(
+        dsn::utils::PegasusEnv(dsn::utils::FileDataType::kSensitive), 
kFileName, &data);
+    ASSERT_EQ(kFileContent, data);
+}
+
+TEST_P(env_file_test, copy_file)
+{
+    const std::string kFileName = "copy_file";
+    const std::string kCopyFileName = kFileName + ".copy";
+    const uint64_t kFileContentSize = 100;
+    const std::string kFileContent(kFileContentSize, 'a');
+
+    // Prepare an encrypted test file.
+    auto s =
+        
rocksdb::WriteStringToFile(dsn::utils::PegasusEnv(dsn::utils::FileDataType::kSensitive),
+                                   rocksdb::Slice(kFileContent),
+                                   kFileName,
+                                   /* should_sync */ true);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+
+    // Check file size.
+    int64_t wfile_size;
+    ASSERT_TRUE(dsn::utils::filesystem::file_size(
+        kFileName, dsn::utils::FileDataType::kSensitive, wfile_size));
+    ASSERT_EQ(kFileContentSize, wfile_size);
+    ASSERT_TRUE(dsn::utils::filesystem::file_size(
+        kFileName, dsn::utils::FileDataType::kNonSensitive, wfile_size));
+    ASSERT_EQ(kFileContentSize + _extra_size, wfile_size);
+
+    // Check copy_file(src_fname, dst_fname, total_size).
+    // Loop twice to check overwrite.
+    for (int i = 0; i < 2; ++i) {
+        uint64_t copy_file_size;
+        s = dsn::utils::copy_file(kFileName, kCopyFileName, &copy_file_size);
+        ASSERT_TRUE(s.ok()) << s.ToString();
+        ASSERT_EQ(kFileContentSize, copy_file_size);
+        ASSERT_TRUE(dsn::utils::filesystem::file_size(
+            kCopyFileName, dsn::utils::FileDataType::kSensitive, wfile_size));
+        ASSERT_EQ(kFileContentSize, wfile_size);
+        ASSERT_TRUE(dsn::utils::filesystem::file_size(
+            kCopyFileName, dsn::utils::FileDataType::kNonSensitive, 
wfile_size));
+        ASSERT_EQ(kFileContentSize + _extra_size, wfile_size);
+        // Check file content.
+        std::string data;
+        s = rocksdb::ReadFileToString(
+            dsn::utils::PegasusEnv(dsn::utils::FileDataType::kSensitive), 
kCopyFileName, &data);
+        ASSERT_EQ(kFileContent, data);
+    }
+}
+
+TEST_P(env_file_test, copy_file_by_size)
+{
+    const std::string kFileName = "copy_file_by_size";
+    std::string kCopyFileName = kFileName + ".copy";
+    const uint64_t kFileContentSize = 100;
+    const std::string kFileContent(kFileContentSize, 'a');
+
+    // Prepare an encrypted test file.
+    auto s =
+        
rocksdb::WriteStringToFile(dsn::utils::PegasusEnv(dsn::utils::FileDataType::kSensitive),
+                                   rocksdb::Slice(kFileContent),
+                                   kFileName,
+                                   /* should_sync */ true);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+
+    // Check file size.
+    int64_t wfile_size;
+    ASSERT_TRUE(dsn::utils::filesystem::file_size(
+        kFileName, dsn::utils::FileDataType::kSensitive, wfile_size));
+    ASSERT_EQ(kFileContentSize, wfile_size);
+    ASSERT_TRUE(dsn::utils::filesystem::file_size(
+        kFileName, dsn::utils::FileDataType::kNonSensitive, wfile_size));
+    ASSERT_EQ(kFileContentSize + _extra_size, wfile_size);
+
+    // Check copy_file_by_size(src_fname, dst_fname, limit_size).
+    struct test_case
+    {
+        int64_t limit_size;
+        int64_t expect_size;
+    } tests[] = {{-1, kFileContentSize},
+                 {0, 0},
+                 {10, 10},
+                 {kFileContentSize, kFileContentSize},
+                 {kFileContentSize + 10, kFileContentSize}};
+    for (const auto &test : tests) {
+        s = dsn::utils::copy_file_by_size(kFileName, kCopyFileName, 
test.limit_size);
+        ASSERT_TRUE(s.ok()) << s.ToString();
+
+        int64_t actual_size;
+        ASSERT_TRUE(dsn::utils::filesystem::file_size(
+            kCopyFileName, dsn::utils::FileDataType::kSensitive, actual_size));
+        ASSERT_EQ(test.expect_size, actual_size);
+        ASSERT_TRUE(dsn::utils::filesystem::file_size(
+            kCopyFileName, dsn::utils::FileDataType::kNonSensitive, 
wfile_size));
+        ASSERT_EQ(test.expect_size + _extra_size, wfile_size);
+        // Check file content.
+        std::string data;
+        s = rocksdb::ReadFileToString(
+            dsn::utils::PegasusEnv(dsn::utils::FileDataType::kSensitive), 
kCopyFileName, &data);
+        ASSERT_EQ(std::string(test.expect_size, 'a'), data);
+    }
+}
+
+TEST_P(env_file_test, copy_non_encrypt_file)
+{
+    const std::string kFileName = "copy_non_encrypt_file";
+    std::string kCopyFileName = kFileName + ".copy";
+    const uint64_t kFileContentSize = 100;
+    const std::string kFileContent(kFileContentSize, 'a');
+
+    // Prepare a non-encrypted test file.
+    auto s =
+        
rocksdb::WriteStringToFile(dsn::utils::PegasusEnv(dsn::utils::FileDataType::kNonSensitive),
+                                   rocksdb::Slice(kFileContent),
+                                   kFileName,
+                                   /* should_sync */ true);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+
+    // Check file size.
+    int64_t wfile_size;
+    ASSERT_TRUE(dsn::utils::filesystem::file_size(
+        kFileName, dsn::utils::FileDataType::kNonSensitive, wfile_size));
+    ASSERT_EQ(kFileContentSize, wfile_size);
+
+    // Check copy_file() on non-sensitive file.
+    s = dsn::utils::copy_file(kFileName, kCopyFileName);
+    if (FLAGS_encrypt_data_at_rest) {
+        // copy_file() consider the source file as encrypted, so it will fail.
+        ASSERT_TRUE(s.IsCorruption()) << s.ToString();
+        ASSERT_TRUE(s.ToString().find(
+                        fmt::format("Corruption: Invalid encryption header in 
{}", kFileName)) == 0)
+            << s.ToString();
+    } else {
+        // Although copy_file() consider the source file as non-encrypted, but 
it will succeed if
+        // FLAGS_encrypt_data_at_rest is disabled.
+        ASSERT_TRUE(s.ok()) << s.ToString();
+        int64_t copy_file_size;
+        ASSERT_TRUE(dsn::utils::filesystem::file_size(
+            kCopyFileName, dsn::utils::FileDataType::kNonSensitive, 
copy_file_size));
+        ASSERT_EQ(kFileContentSize, copy_file_size);
+    }
+
+    // Check copy_file_by_size() on non-sensitive file.
+    s = dsn::utils::copy_file_by_size(kFileName, kCopyFileName);
+    if (FLAGS_encrypt_data_at_rest) {
+        // copy_file_by_size() consider the source file as encrypted, so it 
will fail.
+        ASSERT_TRUE(s.IsCorruption()) << s.ToString();
+        ASSERT_TRUE(s.ToString().find(
+                        fmt::format("Corruption: Invalid encryption header in 
{}", kFileName)) == 0)
+            << s.ToString();
+    } else {
+        // Although copy_file_by_size() consider the source file as 
non-encrypted, but it will
+        // succeed if FLAGS_encrypt_data_at_rest is disabled.
+        ASSERT_TRUE(s.ok()) << s.ToString();
+        int64_t copy_file_size;
+        ASSERT_TRUE(dsn::utils::filesystem::file_size(
+            kCopyFileName, dsn::utils::FileDataType::kNonSensitive, 
copy_file_size));
+        ASSERT_EQ(kFileContentSize, copy_file_size);
+    }
+}
diff --git a/src/utils/test/file_system_test.cpp 
b/src/utils/test/file_system_test.cpp
index f6da3e2b3..ccfc2da2f 100644
--- a/src/utils/test/file_system_test.cpp
+++ b/src/utils/test/file_system_test.cpp
@@ -18,16 +18,104 @@
 // IWYU pragma: no_include <gtest/gtest-message.h>
 // IWYU pragma: no_include <gtest/gtest-test-part.h>
 #include <gtest/gtest.h>
+#include <rocksdb/env.h>
+#include <rocksdb/slice.h>
+#include <rocksdb/status.h>
 #include <stdint.h>
+#include <set>
 #include <string>
 
+#include "utils/env.h"
 #include "utils/filesystem.h"
+#include "utils/flags.h"
+
+DSN_DECLARE_bool(encrypt_data_at_rest);
 
 namespace dsn {
 namespace utils {
 namespace filesystem {
 
-TEST(verify_file, verify_file_test)
+TEST(filesystem_test, compare_with_legacy_file_size)
+{
+    const std::string kFileName = "file_size_test";
+    std::set<int64_t> test_file_sizes({0, 100});
+    for (const auto &test_file_size : test_file_sizes) {
+        const std::string kFileContent(test_file_size, 'a');
+
+        // Prepare a non-encrypted test file.
+        auto s = rocksdb::WriteStringToFile(
+            dsn::utils::PegasusEnv(dsn::utils::FileDataType::kNonSensitive),
+            rocksdb::Slice(kFileContent),
+            kFileName,
+            /* should_sync */ true);
+        ASSERT_TRUE(s.ok()) << s.ToString();
+
+        // The file size should be the same as the legacy file size.
+        int64_t actual_file_size;
+        ASSERT_TRUE(file_size(kFileName, actual_file_size));
+        ASSERT_EQ(test_file_size, actual_file_size);
+        ASSERT_TRUE(deprecated_file_size(kFileName, actual_file_size));
+        ASSERT_EQ(test_file_size, actual_file_size);
+    }
+}
+
+TEST(filesystem_test_p, non_encrypted_file_size)
+{
+    FLAGS_encrypt_data_at_rest = false;
+    const std::string kFileName = "non_encrypted_file_size";
+    const uint64_t kFileContentSize = 100;
+    const std::string kFileContent(kFileContentSize, 'a');
+
+    // Prepare the non-encrypted test file.
+    auto s =
+        
rocksdb::WriteStringToFile(dsn::utils::PegasusEnv(dsn::utils::FileDataType::kNonSensitive),
+                                   rocksdb::Slice(kFileContent),
+                                   kFileName,
+                                   /* should_sync */ true);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+
+    int64_t actual_file_size;
+    // Check file_size(path, sz)
+    ASSERT_TRUE(file_size(kFileName, actual_file_size));
+    ASSERT_EQ(kFileContentSize, actual_file_size);
+    // Check file_size(path, type, sz)
+    ASSERT_TRUE(file_size(kFileName, FileDataType::kNonSensitive, 
actual_file_size));
+    ASSERT_EQ(kFileContentSize, actual_file_size);
+    // Check file_size(path, type, sz) with kSensitive type.
+    // It's able to get the correct file size because 
FLAGS_encrypt_data_at_rest is disabled.
+    ASSERT_TRUE(file_size(kFileName, FileDataType::kSensitive, 
actual_file_size));
+    ASSERT_EQ(kFileContentSize, actual_file_size);
+}
+
+TEST(filesystem_test_p, encrypted_file_size)
+{
+    FLAGS_encrypt_data_at_rest = true;
+    const std::string kFileName = "encrypted_file_size";
+    const uint64_t kFileContentSize = 100;
+    const std::string kFileContent(kFileContentSize, 'a');
+
+    // Prepare the non-encrypted test file.
+    auto s =
+        
rocksdb::WriteStringToFile(dsn::utils::PegasusEnv(dsn::utils::FileDataType::kSensitive),
+                                   rocksdb::Slice(kFileContent),
+                                   kFileName,
+                                   /* should_sync */ true);
+    ASSERT_TRUE(s.ok()) << s.ToString();
+
+    int64_t actual_file_size;
+    // Check file_size(path, sz), the encryption header size is counted.
+    ASSERT_TRUE(file_size(kFileName, actual_file_size));
+    ASSERT_EQ(kFileContentSize + kEncryptionHeaderkSize, actual_file_size);
+    // Check file_size(path, type, sz) with correct type.
+    ASSERT_TRUE(file_size(kFileName, FileDataType::kSensitive, 
actual_file_size));
+    ASSERT_EQ(kFileContentSize, actual_file_size);
+    // Check file_size(path, type, sz) with kNonSensitive type, the encryption 
header size is
+    // counted.
+    ASSERT_TRUE(file_size(kFileName, FileDataType::kNonSensitive, 
actual_file_size));
+    ASSERT_EQ(kFileContentSize + kEncryptionHeaderkSize, actual_file_size);
+}
+
+TEST(filesystem_test, verify_file_test)
 {
     const std::string &fname = "test_file";
     std::string expected_md5;
diff --git a/src/utils/test/nth_element_bench/CMakeLists.txt 
b/src/utils/test/nth_element_bench/CMakeLists.txt
index 217d9c436..2bd530690 100644
--- a/src/utils/test/nth_element_bench/CMakeLists.txt
+++ b/src/utils/test/nth_element_bench/CMakeLists.txt
@@ -27,7 +27,7 @@ set(MY_PROJ_SRC "")
 # "GLOB" for non-recursive search
 set(MY_SRC_SEARCH_MODE "GLOB")
 
-set(MY_PROJ_LIBS dsn_runtime dsn_utils)
+set(MY_PROJ_LIBS dsn_runtime dsn_utils rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 
diff --git a/src/zookeeper/test/CMakeLists.txt 
b/src/zookeeper/test/CMakeLists.txt
index ed2f742ac..6179f8d1e 100644
--- a/src/zookeeper/test/CMakeLists.txt
+++ b/src/zookeeper/test/CMakeLists.txt
@@ -41,7 +41,7 @@ set(MY_PROJ_LIBS
     gtest
     ssl
     crypto
-    )
+    rocksdb)
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem Boost::regex)
 


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]


Reply via email to