From 31815af86b69e5ffd58228e3d94d5f2e50f6fb96 Mon Sep 17 00:00:00 2001 From: Yingchun Lai Date: Wed, 13 Sep 2023 23:26:53 +0800 Subject: [PATCH] feat(encryption): introduce PegasusEnv --- .github/workflows/lint_and_test_cpp.yaml | 2 +- src/test_util/CMakeLists.txt | 2 +- src/test_util/test_util.h | 11 ++ src/utils/CMakeLists.txt | 2 +- src/utils/env.cpp | 186 +++++++++++++++++++ src/utils/env.h | 61 ++++++ src/utils/filesystem.cpp | 36 ++-- src/utils/filesystem.h | 3 + src/utils/fmt_logging.h | 11 ++ src/utils/fmt_utils.h | 2 +- src/utils/test/CMakeLists.txt | 3 +- src/utils/test/env.cpp | 224 ++++++++++++++++++++++- 12 files changed, 514 insertions(+), 29 deletions(-) create mode 100644 src/utils/env.cpp create mode 100644 src/utils/env.h diff --git a/.github/workflows/lint_and_test_cpp.yaml b/.github/workflows/lint_and_test_cpp.yaml index 6c125727c2..d49b95fcc3 100644 --- a/.github/workflows/lint_and_test_cpp.yaml +++ b/.github/workflows/lint_and_test_cpp.yaml @@ -676,7 +676,7 @@ jobs: run: | export JAVA_HOME="${JAVA_HOME_8_X64}" mkdir -p build - cmake -DCMAKE_BUILD_TYPE=Release -B build/ -DMACOS_OPENSSL_ROOT_DIR=${OPENSSL_ROOT_DIR} -DROCKSDB_PORTABLE=1 + cmake -DCMAKE_BUILD_TYPE=Release -B build/ -DMACOS_OPENSSL_ROOT_DIR=${OPENSSL_ROOT_DIR} -DROCKSDB_PORTABLE=native cmake --build build/ -j $(sysctl -n hw.physicalcpu) - name: Compilation run: | diff --git a/src/test_util/CMakeLists.txt b/src/test_util/CMakeLists.txt index 267c825e40..b1e7ac2dff 100644 --- a/src/test_util/CMakeLists.txt +++ b/src/test_util/CMakeLists.txt @@ -22,6 +22,6 @@ set(MY_PROJ_NAME test_utils) # "GLOB" for non-recursive search set(MY_SRC_SEARCH_MODE "GLOB") -set(MY_PROJ_LIBS gtest) +set(MY_PROJ_LIBS gtest rocksdb) dsn_add_static_library() diff --git a/src/test_util/test_util.h b/src/test_util/test_util.h index ec7bce6e2e..18e2d38947 100644 --- a/src/test_util/test_util.h +++ b/src/test_util/test_util.h @@ -20,11 +20,22 @@ #pragma once #include +#include +#include "gtest/gtest.h" +#include "utils/flags.h" #include "utils/test_macros.h" +DSN_DECLARE_bool(encrypt_data_at_rest); + namespace pegasus { +class encrypt_data_test_base : public testing::TestWithParam +{ +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 e17fef9001..1127a868ef 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 0000000000..120c08e045 --- /dev/null +++ b/src/utils/env.cpp @@ -0,0 +1,186 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include + +#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, + true, + "Whether 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() +{ + std::string provider_id = + fmt::format("AES:{},{}", FLAGS_server_key_for_testing, FLAGS_encryption_method); + std::shared_ptr provider; + auto s = rocksdb::EncryptionProvider::CreateFromString( + rocksdb::ConfigOptions(), provider_id, &provider); + CHECK(s.ok(), "Failed to create encryption provider: {}", s.ToString()); + return NewEncryptedEnv(rocksdb::Env::Default(), provider); +} + +rocksdb::Env *PegasusEnv(FileDataType type) +{ + if (FLAGS_encrypt_data_at_rest && type == FileDataType::kSensitive) { + static rocksdb::Env *env = NewEncryptedEnv(); + return env; + } + + static rocksdb::Env *env = rocksdb::Env::Default(); + return env; +} + +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 rd_env_options; + std::unique_ptr sfile; + auto s = dsn::utils::PegasusEnv(src_type)->NewSequentialFile(src_fname, &sfile, rd_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(dsn::utils::filesystem::file_size(src_fname, src_type, src_file_size), ""); + if (remain_size == -1) { + remain_size = src_file_size; + } + remain_size = std::min(remain_size, src_file_size); + + rocksdb::EnvOptions wt_env_options; + std::unique_ptr wfile; + s = dsn::utils::PegasusEnv(dst_type)->NewWritableFile(dst_fname, &wfile, wt_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 kBlockSize = 4 << 20; + auto buffer = dsn::utils::make_shared_array(kBlockSize); + uint64_t offset = 0; + do { + int bytes_per_copy = std::min(remain_size, static_cast(kBlockSize)); + if (bytes_per_copy <= 0) { + break; + } + + rocksdb::Slice result; + LOG_AND_RETURN_NOT_RDB_OK(WARNING, + sfile->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, wfile->Append(result), "failed to write file {}", dst_fname); + + offset += result.size(); + remain_size -= result.size(); + + // Reach the end of the file. + if (result.size() < bytes_per_copy) { + break; + } + } while (true); + LOG_AND_RETURN_NOT_RDB_OK(WARNING, wfile->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(); +} + +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, 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"; + 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 0000000000..9c1cb50331 --- /dev/null +++ b/src/utils/env.h @@ -0,0 +1,61 @@ +// 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 +#include +#include +#include + +namespace rocksdb { +class Env; +} // namespace rocksdb + +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 +}; + +rocksdb::Env *PegasusEnv(FileDataType type); + +// The 'total_size' is the total size of the file content, exclude the file encryption header. +// 'src_fname' is not encrypted and 'dst_fname' is encrypted. +rocksdb::Status encrypt_file(const std::string &src_fname, + const std::string &dst_fname, + uint64_t *total_size = nullptr); +// Encrypt the original non-encrypted 'fname'. +rocksdb::Status encrypt_file(const std::string &fname, uint64_t *total_size = nullptr); +// Both 'src_fname' and 'dst_fname' are encrypted files. +rocksdb::Status copy_file(const std::string &src_fname, + const std::string &dst_fname, + uint64_t *total_size = nullptr); +// Both 'src_fname' and 'dst_fname' are encrypted files. 'limit_size' is the max size of the +// size 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 73aa3a50aa..8c3a783f66 100644 --- a/src/utils/filesystem.cpp +++ b/src/utils/filesystem.cpp @@ -47,9 +47,12 @@ // IWYU pragma: no_include #include // IWYU pragma: keep #include -#include +#include +#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,32 +391,21 @@ bool rename_path(const std::string &path1, const std::string &path2) return ret; } +// TODO(yingchun): refactor to use uint64_t. bool file_size(const std::string &path, int64_t &sz) { - struct stat_ st; - std::string npath; - int err; - - if (path.empty()) { - return false; - } - - err = get_normalized_path(path, npath); - if (err != 0) { - return false; - } - - err = dsn::utils::filesystem::get_stat_internal(npath, st); - if (err != 0) { - return false; - } + return file_size(path, dsn::utils::FileDataType::kNonSensitive, sz); +} - if (!S_ISREG(st.st_mode)) { +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 = st.st_size; - + sz = file_size; return true; } diff --git a/src/utils/filesystem.h b/src/utils/filesystem.h index 4229f551d2..a49d6982e8 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); @@ -97,6 +99,7 @@ bool remove_path(const std::string &path); bool rename_path(const std::string &path1, const std::string &path2); bool file_size(const std::string &path, int64_t &sz); +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 f0c74bd293..35c60aaddf 100644 --- a/src/utils/fmt_logging.h +++ b/src/utils/fmt_logging.h @@ -20,6 +20,7 @@ #pragma once #include +#include #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/fmt_utils.h b/src/utils/fmt_utils.h index 9624bb881a..07ba9242fd 100644 --- a/src/utils/fmt_utils.h +++ b/src/utils/fmt_utils.h @@ -21,7 +21,7 @@ #define USER_DEFINED_STRUCTURE_FORMATTER(type) \ template <> \ - struct fmt::formatter : ostream_formatter \ + struct fmt::formatter : fmt::ostream_formatter \ { \ } diff --git a/src/utils/test/CMakeLists.txt b/src/utils/test/CMakeLists.txt index fb284b0e86..266b9beeb6 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 619e5d7eba..927f0f4d9b 100644 --- a/src/utils/test/env.cpp +++ b/src/utils/test/env.cpp @@ -33,18 +33,29 @@ * xxxx-xx-xx, author, fix bug about xxx */ +#include // IWYU pragma: no_include // IWYU pragma: no_include #include +#include +#include +#include #include #include -#include +#include +#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::max() - 1, 0xdeadbeef}; @@ -56,3 +67,212 @@ 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); + + FLAGS_encrypt_data_at_rest = true; + auto *env_no_enc3 = dsn::utils::PegasusEnv(dsn::utils::FileDataType::kNonSensitive); + auto *env_enc1 = dsn::utils::PegasusEnv(dsn::utils::FileDataType::kSensitive); + + ASSERT_EQ(env_no_enc1, env_no_enc2); + ASSERT_EQ(env_no_enc1, env_no_enc3); + 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 file size should plus 4096 if consider it as kNonSensitive when the if is actually + // encrypted. + if (FLAGS_encrypt_data_at_rest) { + extra_size = 4096; + } + } + uint64_t extra_size = 0; +}; + +INSTANTIATE_TEST_CASE_P(, env_file_test, ::testing::Values(false)); + +TEST_P(env_file_test, encrypt_file_2_files) +{ + const std::string kFileName = "encrypt_file_2_files"; + 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, kFileName + ".encrypted", &encrypt_file_size); + ASSERT_TRUE(s.ok()) << s.ToString(); + ASSERT_EQ(kFileContentSize, encrypt_file_size); + ASSERT_TRUE(dsn::utils::filesystem::file_size( + kFileName + ".encrypted", dsn::utils::FileDataType::kSensitive, wfile_size)); + ASSERT_EQ(kFileContentSize, wfile_size); + ASSERT_TRUE(dsn::utils::filesystem::file_size( + kFileName + ".encrypted", 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 + ".encrypted", + &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 uint64_t kFileContentSize = 100; + const std::string kFileContent(kFileContentSize, 'a'); + + // Prepare a 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, kFileName + ".copy", ©_file_size); + ASSERT_TRUE(s.ok()) << s.ToString(); + ASSERT_EQ(kFileContentSize, copy_file_size); + ASSERT_TRUE(dsn::utils::filesystem::file_size( + kFileName + ".copy", dsn::utils::FileDataType::kSensitive, wfile_size)); + ASSERT_EQ(kFileContentSize, wfile_size); + ASSERT_TRUE(dsn::utils::filesystem::file_size( + kFileName + ".copy", 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 + ".copy", + &data); + ASSERT_EQ(kFileContent, data); + } +} + +TEST_P(env_file_test, copy_file_by_size) +{ + const std::string kFileName = "copy_file_by_size"; + 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::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) { + std::string copy_file_name = kFileName + ".copy"; + s = dsn::utils::copy_file_by_size(kFileName, copy_file_name, test.limit_size); + ASSERT_TRUE(s.ok()) << s.ToString(); + + int64_t actual_size; + ASSERT_TRUE(dsn::utils::filesystem::file_size( + copy_file_name, dsn::utils::FileDataType::kSensitive, actual_size)); + ASSERT_EQ(test.expect_size, actual_size); + ASSERT_TRUE(dsn::utils::filesystem::file_size( + copy_file_name, 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), copy_file_name, &data); + ASSERT_EQ(std::string(test.expect_size, 'a'), data); + } +}