feat(encryption): add kms key management

Samunroyu · Nov 29, 2023 · 79e05fd · 79e05fd
1 parent ca1c951
commit 79e05fd
Show file tree

Hide file tree

Showing 8 changed files with 86 additions and 96 deletions.
diff --git a/src/replica/key_provider.h b/src/replica/key_provider.h
@@ -41,11 +41,6 @@ class KeyProvider
                                                std::string *iv,
                                                std::string *key_version) = 0;
 
-    // Generates an encryption key with the tenant_id.
-    virtual dsn::error_s GenerateTenantKey(const std::string &tenant_id,
-                                           std::string *encryption_key,
-                                           std::string *iv,
-                                           std::string *key_version) = 0;
 };
 } // namespace security
 } // namespace dsn
diff --git a/src/replica/kms_client.cpp b/src/replica/kms_client.cpp
@@ -35,43 +35,54 @@ dsn::error_s PegasusKMSClient::DecryptEncryptionKey(const std::string &encryptio
 {
     nlohmann::json post;
     post["name"] = cluster_key_name_;
-    std::string iv_plain = absl::BytesToHexString(iv);
+    std::string iv_plain = ::absl::HexStringToBytes(iv);
     std::string iv_b64;
-    absl::WebSafeBase64Escape(iv_plain, &iv_b64);
+    ::absl::WebSafeBase64Escape(iv_plain, &iv_b64);
     post["iv"] = iv_b64;
-    std::string eek_plain = absl::BytesToHexString(encryption_key);
+    std::string eek_plain = ::absl::HexStringToBytes(encryption_key);
     std::string eek_b64;
-    absl::WebSafeBase64Escape(eek_plain, &eek_b64);
+    ::absl::WebSafeBase64Escape(eek_plain, &eek_b64);
     post["material"] = eek_b64;
 
     http_client client;
     auto err = client.init();
-    CHECK_EQ_MSG(err, ERR_OK, "Start http client failed");
+    if(!err.is_ok()){
+        return dsn::error_s::make(ERR_CURL_FAILED, "Start http client failed");
+    }
+
     err = client.set_auth(http_auth_type::SPNEGO);
-    CHECK_EQ_MSG(err, ERR_OK, "http client set auth type failed");
+    if(!err.is_ok()){
+        return dsn::error_s::make(ERR_CURL_FAILED, "http client set auth type failed");
+    }
+
     std::vector<std::string> urls;
     urls.reserve(kms_urls_.size());
     for (const auto &url : kms_urls_) {
         urls.emplace_back(
-            absl::Substitute("$0/kms/v1/keyversion/$1/_eek?eek_op=decrypt", url, key_version));
+            ::absl::Substitute("$0/v1/keyversion/$1/_eek?eek_op=decrypt", url, key_version));
     }
 
     client.clear_header_fields();
     client.set_content_type("application/json");
     client.set_accept("*/*");
     err = client.with_post_method(post.dump());
-    CHECK_EQ_MSG(err, ERR_OK, "http client set method failed");
+    if(!err.is_ok()){
+        return dsn::error_s::make(ERR_CURL_FAILED, "http client set method failed");
+    }
 
     nlohmann::json j;
     for (const auto &url : urls) {
         err = client.set_url(url);
-        CHECK_EQ_MSG(err, ERR_OK, "http clientt set url failed");
+        if(!err.is_ok()){
+            return dsn::error_s::make(ERR_CURL_FAILED, "http clientt set url failed");
+        }
         std::string resp;
         err = client.exec_method(&resp);
-        CHECK_EQ_MSG(err, ERR_OK, "http client exec post method failed");
+        if(!err.is_ok()){
+            return dsn::error_s::make(ERR_CURL_FAILED, "http client exec post method failed");
+        }
         long http_status;
         client.get_http_status(http_status);
-        LOG_INFO("http status = ({})", http_status);
         if (http_status == 200) {
             j = nlohmann::json::parse(resp);
         }
@@ -80,15 +91,14 @@ dsn::error_s PegasusKMSClient::DecryptEncryptionKey(const std::string &encryptio
     std::string dek_b64;
     if (j.contains("material")) {
         dek_b64 = j.at("material");
-        LOG_INFO("POST material is ({})", dek_b64);
     } else {
         return dsn::error_s::make(ERR_INVALID_DATA, "Null material received");
     }
     std::string dek_plain;
-    if (!absl::WebSafeBase64Unescape(dek_b64, &dek_plain)) {
+    if (!::absl::WebSafeBase64Unescape(dek_b64, &dek_plain)) {
         return dsn::error_s::make(ERR_INVALID_DATA, "Invalid IV received");
     }
-    *decrypted_key = absl::HexStringToBytes(dek_plain);
+    *decrypted_key = ::absl::BytesToHexString(dek_plain);
     return dsn::error_s::ok();
 }
 
@@ -99,29 +109,40 @@ dsn::error_s PegasusKMSClient::GenerateEncryptionKeyFromKMS(const std::string &k
 {
     http_client client;
     auto err = client.init();
-    CHECK_EQ_MSG(err, ERR_OK, "Start http client failed");
+    if(!err.is_ok()){
+        return dsn::error_s::make(ERR_CURL_FAILED, "Start http client failed");
+    }
     err = client.set_auth(http_auth_type::SPNEGO);
-    CHECK_EQ_MSG(err, ERR_OK, "http client set auth type failed");
-
+    if(!err.is_ok()){
+        return dsn::error_s::make(ERR_CURL_FAILED, "http client set auth type failed");
+    }
     std::vector<std::string> urls;
     urls.reserve(kms_urls_.size());
     for (const auto &url : kms_urls_) {
         urls.emplace_back(
-            absl::Substitute("$0/kms/v1/key/$1/_eek?eek_op=generate&num_keys=1", url, key_name));
+            ::absl::Substitute("$0/v1/key/$1/_eek?eek_op=generate&num_keys=1", url, key_name));
     }
 
     nlohmann::json j = nlohmann::json::object();
     for (const auto &url : urls) {
         err = client.set_url(url);
-        CHECK_EQ_MSG(err, ERR_OK, "http client set url failed");
+        if(!err.is_ok()){
+            return dsn::error_s::make(ERR_CURL_FAILED, "http client set url failed");
+        }
+
         err = client.with_get_method();
-        CHECK_EQ_MSG(err, ERR_OK, "http client set get method failed");
+        if(!err.is_ok()){
+            return dsn::error_s::make(ERR_CURL_FAILED, "http client set get method failed");
+        }
+
         std::string resp;
         err = client.exec_method(&resp);
-        CHECK_EQ_MSG(err, ERR_OK, "http client exec get method failed");
+        if(!err.is_ok()){
+            return dsn::error_s::make(ERR_CURL_FAILED, "http client exec get method failed");
+        }
+
         long http_status;
         client.get_http_status(http_status);
-        LOG_INFO("http status = ({})", http_status);
         if (http_status == 200) {
             j = nlohmann::json::parse(resp);
             nlohmann::json jsonObject = j.at(0);
@@ -142,10 +163,10 @@ dsn::error_s PegasusKMSClient::GenerateEncryptionKeyFromKMS(const std::string &k
         return dsn::error_s::make(ERR_INVALID_DATA, "Null IV received");
     }
     std::string iv_plain;
-    if (!absl::WebSafeBase64Unescape(iv_b64, &iv_plain)) {
+    if (!::absl::WebSafeBase64Unescape(iv_b64, &iv_plain)) {
         return dsn::error_s::make(ERR_INVALID_DATA, "Invalid IV received");
     }
-    *iv = absl::HexStringToBytes(iv_plain);
+    *iv = ::absl::BytesToHexString(iv_plain);
     std::string key_b64;
     if (j.contains("encryptedKeyVersion") && j.at("encryptedKeyVersion").contains("material")) {
         key_b64 = j.at("encryptedKeyVersion").at("material");
@@ -154,21 +175,13 @@ dsn::error_s PegasusKMSClient::GenerateEncryptionKeyFromKMS(const std::string &k
                                   "Null encryptedKeyVersion or material received");
     }
     std::string key_plain;
-    if (!absl::WebSafeBase64Unescape(key_b64, &key_plain)) {
+    if (!::absl::WebSafeBase64Unescape(key_b64, &key_plain)) {
         return dsn::error_s::make(ERR_INVALID_DATA, "Invalid encryption key received");
     }
-    *encryption_key = absl::HexStringToBytes(key_plain);
+    *encryption_key = ::absl::BytesToHexString(key_plain);
     return dsn::error_s::ok();
 }
 
-dsn::error_s PegasusKMSClient::GenerateTenantKey(const std::string &tenant_id,
-                                                 std::string *encryption_key,
-                                                 std::string *iv,
-                                                 std::string *key_version)
-{
-    return GenerateEncryptionKeyFromKMS(tenant_id, encryption_key, iv, key_version);
-}
-
 dsn::error_s PegasusKMSClient::GenerateEncryptionKey(std::string *encryption_key,
                                                      std::string *iv,
                                                      std::string *key_version)

diff --git a/src/replica/kms_client.h b/src/replica/kms_client.h
@@ -31,7 +31,7 @@ class PegasusKMSClient
 {
 public:
     PegasusKMSClient(const std::string kms_url, std::string cluster_key_name)
-        : kms_urls_(absl::StrSplit(kms_url, ",", absl::SkipEmpty())),
+        : kms_urls_(::absl::StrSplit(kms_url, ",", ::absl::SkipEmpty())),
           cluster_key_name_(std::move(cluster_key_name))
     {
     }
@@ -44,11 +44,6 @@ class PegasusKMSClient
     dsn::error_s
     GenerateEncryptionKey(std::string *encryption_key, std::string *iv, std::string *key_version);
 
-    dsn::error_s GenerateTenantKey(const std::string &tenant_id,
-                                   std::string *encryption_key,
-                                   std::string *iv,
-                                   std::string *key_version);
-
 private:
     dsn::error_s GenerateEncryptionKeyFromKMS(const std::string &key_name,
                                               std::string *encryption_key,

diff --git a/src/replica/pegasus_kms_key_provider.cpp b/src/replica/pegasus_kms_key_provider.cpp
@@ -38,13 +38,5 @@ dsn::error_s PegasusKMSKeyProvider::GenerateEncryptionKey(std::string *encryptio
     return client_.GenerateEncryptionKey(encryption_key, iv, key_version);
 }
 
-dsn::error_s PegasusKMSKeyProvider::GenerateTenantKey(const std::string &tenant_id,
-                                                      std::string *encryption_key,
-                                                      std::string *iv,
-                                                      std::string *key_version)
-{
-    return client_.GenerateTenantKey(tenant_id, encryption_key, iv, key_version);
-}
-
 } // namespace security
 } // namespace dsn
diff --git a/src/replica/pegasus_kms_key_provider.h b/src/replica/pegasus_kms_key_provider.h
@@ -47,12 +47,6 @@ class PegasusKMSKeyProvider : public KeyProvider
                                        std::string *iv,
                                        std::string *key_version) override;
 
-    // Generates an encryption key with the tenant_id.
-    dsn::error_s GenerateTenantKey(const std::string &tenant_id,
-                                   std::string *encryption_key,
-                                   std::string *iv,
-                                   std::string *key_version) override;
-
 private:
     PegasusKMSClient client_;
 };

diff --git a/src/replica/replica_stub.cpp b/src/replica/replica_stub.cpp
@@ -92,6 +92,8 @@
 #include "remote_cmd/remote_command.h"
 #include "utils/fail_point.h"
 
+DSN_DECLARE_bool(encrypt_data_at_rest);
+DSN_DECLARE_string(server_key);
 namespace dsn {
 namespace replication {
 DSN_DEFINE_bool(replication,
@@ -189,13 +191,7 @@ DSN_DEFINE_string(pegasus.server,
                   hadoop_kms_url,
                   "",
                   "Where the server encrypted key of file system can get from.");
-DSN_DEFINE_string(pegasus.server,
-                  server_key,
-                  "server_key",
-                  "The encrypted server key to use in the filesystem.");
 
-// DSN_DECLARE_string(data_dirs);
-DSN_DECLARE_bool(encrypt_data_at_rest);
 DSN_DECLARE_bool(duplication_enabled);
 DSN_DECLARE_int32(fd_beacon_interval_seconds);
 DSN_DECLARE_int32(fd_check_interval_seconds);
@@ -595,21 +591,31 @@ void replica_stub::initialize(bool clear /* = false*/)
     _access_controller = std::make_unique<dsn::security::access_controller>();
 }
 
-dsn::error_s store_kms_key(std::string server_key){
-    // std::string file_name = FLAGS_data_dirs;
-    std::string file_name = "/home/yujingwei";
-    std::string path = "/replica_encrypted_key";
-    file_name = absl::StrCat(file_name, path);
-    std::ofstream in(file_name, std::ios::app);
+dsn::error_s store_kms_key(std::string data_dir, std::string encryption_key, std::string iv, std::string key_version){
+    std::string path = ::absl::StrCat(data_dir, "/replica_encrypted_key");
+    std::ofstream in(file_name, std::ios::binary);
     if(in.is_open()){
         in << server_key << std::endl;
+        in << iv << std::endl;
+        in << key_version << std::endl;
     } else{
         return dsn::error_s::make(ERR_FILE_OPERATION_FAILED, "Can't open replica_encrypted_key file to write");
     }
 
     return dsn::error_s::ok();
 }
 
+void get_kms_key(std::string data_dir, std::string encryption_key, std::string iv, std::string key_version){
+    std::string path = ::absl::StrCat(data_dir, "/replica_encrypted_key");
+    std::ofstream in(path, std::ios::binary);
+    if(!in.bad()){
+        getline(in, encryption_key);
+        getline(in, iv);
+        getline(in, key_version);
+    }
+
+}
+
 void replica_stub::initialize(const replication_options &opts, bool clear /* = false*/)
 {
     _primary_address = dsn_primary_address();
@@ -637,37 +643,28 @@ void replica_stub::initialize(const replication_options &opts, bool clear /* = f
         }
     }
 
+    std::string encryption_key;
+    std::string iv;
+    std::string key_version;
+    std::string server_key;
     //get and store eek from kms
     if (key_provider) {
-        CHECK(key_provider, "invalid kms url ({})", FLAGS_hadoop_kms_url);
-        std::string encryption_key;
-        std::string iv;
-        std::string key_version;
-        std::string server_key;
-        auto err = key_provider->GenerateEncryptionKey(&encryption_key, &iv, &key_version);
-        if(!err.is_ok()){
-            fmt::print(stderr, "get encryption key failed, error={} ", err.description());
-        }
-        // LOG_INFO("data dir is = ({})",FLAGS_data_dirs);
-        err = store_kms_key(encryption_key);
-        if(!err.is_ok()){
-            fmt::print(stderr, "strore decryption key failed, error={} ", err.description());
-        }
-        LOG_INFO("encryption key is ({})", encryption_key);
-        LOG_INFO("iv is ({})", iv);
-        LOG_INFO("key version is ({})", key_version);
-        // 用 flag 传
-        // 暂时看下能不能用
-        err = key_provider->DecryptEncryptionKey(encryption_key, iv, key_version, &server_key);
-        if(!err.is_ok()){
-            fmt::print(stderr, "get decryption key failed, error={} ", err.description());
+        get_kms_key(_options.data_dirs[0], encryption_key, iv, key_version);
+        if(encryption_key == nullptr){
+            CHECK(key_provider, "invalid kms url ({})", FLAGS_hadoop_kms_url);
+            CHECK(key_provider->GenerateEncryptionKey(&encryption_key, &iv, &key_version), "get encryption key failed");
         }
-        LOG_INFO("server key is {}", server_key);
+        CHECK(key_provider->DecryptEncryptionKey(encryption_key, iv, key_version, &server_key),"get decryption key failed");
+        FLAGS_server_key = server_key.c_str();
     }
 
     // Initialize the file system manager.
     _fs_manager.initialize(_options.data_dirs, _options.data_dir_tags);
 
+    if (key_provider) {
+        CHECK(store_kms_key(encryption_key), "Cant store kms key");
+    }
+
     // TODO(yingchun): remove the slog related code.
     // Create slog directory if it does not exist.
     std::string cdir;

diff --git a/src/utils/CMakeLists.txt b/src/utils/CMakeLists.txt
@@ -31,7 +31,6 @@ set(MY_SRC_SEARCH_MODE "GLOB")
 
 set(MY_BOOST_LIBS Boost::system Boost::filesystem)
 
-# set(MY_PROJ_LIBS absl::strings dsn_http curl gssapi_krb5 rocksdb)
 set(MY_PROJ_LIBS dsn_http rocksdb)
 
 # Extra files that will be installed

diff --git a/src/utils/env.cpp b/src/utils/env.cpp
@@ -49,6 +49,11 @@ DSN_DEFINE_string(pegasus.server,
                   "The encryption method to use in the filesystem. Now "
                   "supports AES128CTR, AES192CTR, AES256CTR and SM4CTR.");
 
+DSN_DEFINE_string(pegasus.server,
+                  server_key,
+                  "server_key",
+                  "The encrypted server key to use in the filesystem.");
+
 DSN_DEFINE_bool(replication,
                 enable_direct_io,
                 false,
@@ -63,7 +68,7 @@ rocksdb::Env *NewEncryptedEnv()
     // Create an encryption provider.
     std::shared_ptr<rocksdb::EncryptionProvider> provider;
     auto provider_id = fmt::format("id=AES;hex_instance_key={};method={}",
-                                   FLAGS_server_key_for_testing,
+                                   FLAGS_server_key,
                                    FLAGS_encryption_method);
     auto s = rocksdb::EncryptionProvider::CreateFromString(
         rocksdb::ConfigOptions(), provider_id, &provider);