From 8fd564ff2b2b463b31d45ddba925d182d9e8caf1 Mon Sep 17 00:00:00 2001 From: alice-shang Date: Thu, 8 Oct 2026 21:03:38 +0900 Subject: [PATCH] feat: add KeyManagementClient, KmsRegistry and InMemoryKms - KeyManagementClient (core): wraps and unwraps keys with master keys held in a KMS, mirroring Java's interface, including optional key generation. - KmsRegistry (core): creates the catalog's KMS client from the same catalog properties as Java, `encryption.kms-type` (aws, gcp, azure) or `encryption.kms-impl` (a registered name). As in Java, a property that is present but empty is a misconfiguration rather than "no KMS". Clients are initialized with the catalog properties. - InMemoryKms (bundle): master keys in memory, for tests and local use. Wire compatible with Java's MemoryMockKMS, so it unwraps keys written by Java tests. Master keys can be given as catalog properties (encryption.in-memory-kms.key. = Base64). --- src/iceberg/CMakeLists.txt | 7 +- src/iceberg/cpp20_compatibility_internal.h | 2 + src/iceberg/encryption/CMakeLists.txt | 6 +- src/iceberg/encryption/in_memory_kms.cc | 80 ++++++++++++ src/iceberg/encryption/in_memory_kms.h | 72 +++++++++++ .../encryption/key_management_client.cc | 98 ++++++++++++++ .../encryption/key_management_client.h | 122 ++++++++++++++++++ src/iceberg/test/CMakeLists.txt | 2 +- src/iceberg/test/aes_gcm_test.cc | 57 ++++++++ src/iceberg/test/kms_test.cc | 112 ++++++++++++++++ 10 files changed, 555 insertions(+), 3 deletions(-) create mode 100644 src/iceberg/encryption/in_memory_kms.cc create mode 100644 src/iceberg/encryption/in_memory_kms.h create mode 100644 src/iceberg/encryption/key_management_client.cc create mode 100644 src/iceberg/encryption/key_management_client.h create mode 100644 src/iceberg/test/kms_test.cc diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt index c751cd7a6..f7cc0eec2 100644 --- a/src/iceberg/CMakeLists.txt +++ b/src/iceberg/CMakeLists.txt @@ -33,6 +33,7 @@ set(ICEBERG_SOURCES deletes/position_delete_index.cc deletes/position_delete_range_consumer.cc deletes/roaring_position_bitmap.cc + encryption/key_management_client.cc encryption/standard_key_metadata.cc expression/aggregate.cc expression/binder.cc @@ -327,7 +328,11 @@ if(ICEBERG_BUILD_BUNDLE) parquet/row_number_reader.cc) if(ICEBERG_ENCRYPTION) - list(APPEND ICEBERG_BUNDLE_SOURCES encryption/aes_gcm.cc encryption/aes_gcm_stream.cc) + list(APPEND + ICEBERG_BUNDLE_SOURCES + encryption/aes_gcm.cc + encryption/aes_gcm_stream.cc + encryption/in_memory_kms.cc) endif() # Libraries to link with exported libiceberg_bundle.{so,a}. diff --git a/src/iceberg/cpp20_compatibility_internal.h b/src/iceberg/cpp20_compatibility_internal.h index bd9a9475d..9eac44f9e 100644 --- a/src/iceberg/cpp20_compatibility_internal.h +++ b/src/iceberg/cpp20_compatibility_internal.h @@ -49,6 +49,7 @@ #include "iceberg/deletes/position_delete_range_consumer.h" #include "iceberg/deletes/roaring_position_bitmap.h" #include "iceberg/encryption/encrypted_key.h" +#include "iceberg/encryption/key_management_client.h" #include "iceberg/encryption/standard_key_metadata.h" #include "iceberg/exception.h" #include "iceberg/expression/aggregate.h" @@ -218,6 +219,7 @@ #ifdef ICEBERG_PUBLIC_HEADERS_ENCRYPTION # include "iceberg/encryption/aes_gcm_stream.h" +# include "iceberg/encryption/in_memory_kms.h" #endif // ICEBERG_PUBLIC_HEADERS_ENCRYPTION #ifdef ICEBERG_PUBLIC_HEADERS_HIVE diff --git a/src/iceberg/encryption/CMakeLists.txt b/src/iceberg/encryption/CMakeLists.txt index 73072dbc1..c490aa79c 100644 --- a/src/iceberg/encryption/CMakeLists.txt +++ b/src/iceberg/encryption/CMakeLists.txt @@ -15,12 +15,16 @@ # specific language governing permissions and limitations # under the License. -iceberg_install_all_headers(iceberg/encryption PATTERN encrypted_key.h +iceberg_install_all_headers(iceberg/encryption + PATTERN + encrypted_key.h + key_management_client.h standard_key_metadata.h) if(ICEBERG_BUILD_BUNDLE AND ICEBERG_ENCRYPTION) iceberg_install_all_headers(iceberg/encryption PATTERN aes_gcm_stream.h + in_memory_kms.h HEADER_GROUP ICEBERG_PUBLIC_HEADERS_ENCRYPTION) endif() diff --git a/src/iceberg/encryption/in_memory_kms.cc b/src/iceberg/encryption/in_memory_kms.cc new file mode 100644 index 000000000..31e1eedb7 --- /dev/null +++ b/src/iceberg/encryption/in_memory_kms.cc @@ -0,0 +1,80 @@ +/* + * 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 "iceberg/encryption/in_memory_kms.h" + +#include "iceberg/encryption/aes_gcm_internal.h" +#include "iceberg/util/base64.h" +#include "iceberg/util/macros.h" + +namespace iceberg { + +void InMemoryKms::Register() { + KmsRegistry::Register(std::string(kName), + []() -> Result> { + return std::make_unique(); + }); +} + +Status InMemoryKms::Initialize( + const std::unordered_map& properties) { + // Decode into a staging map so a bad property leaves master_keys_ untouched; + // properties is unordered, so a partial update would be non-deterministic. + std::unordered_map> decoded; + for (const auto& [name, value] : properties) { + if (!name.starts_with(kKeyPropertyPrefix)) { + continue; + } + auto id = name.substr(kKeyPropertyPrefix.size()); + ICEBERG_PRECHECK(!id.empty(), "Master key property {} has an empty key id", name); + ICEBERG_ASSIGN_OR_RAISE(auto key, Base64::Decode(value)); + ICEBERG_PRECHECK(key.size() == 16 || key.size() == 24 || key.size() == 32, + "Master key {} must be 16, 24 or 32 bytes, got {}", name, + key.size()); + decoded.insert_or_assign(std::move(id), std::vector(key.begin(), key.end())); + } + for (auto& [id, key] : decoded) { + master_keys_.insert_or_assign(id, std::move(key)); + } + return {}; +} + +Result*> InMemoryKms::MasterKey(std::string_view id) const { + auto it = master_keys_.find(std::string(id)); + if (it == master_keys_.end()) { + return NotFound("Master key {} is not found", id); + } + return &it->second; +} + +Result> InMemoryKms::WrapKey(std::span key, + std::string_view wrapping_key_id) { + ICEBERG_ASSIGN_OR_RAISE(auto master_key, MasterKey(wrapping_key_id)); + ICEBERG_ASSIGN_OR_RAISE(auto cipher, encryption::AesGcmCipher::Make(*master_key)); + return cipher->Encrypt(key, {}); +} + +Result> InMemoryKms::UnwrapKey(std::span wrapped_key, + std::string_view wrapping_key_id) { + ICEBERG_ASSIGN_OR_RAISE(auto master_key, MasterKey(wrapping_key_id)); + ICEBERG_ASSIGN_OR_RAISE(auto cipher, encryption::AesGcmCipher::Make(*master_key)); + return cipher->Decrypt(wrapped_key, {}); +} + +} // namespace iceberg diff --git a/src/iceberg/encryption/in_memory_kms.h b/src/iceberg/encryption/in_memory_kms.h new file mode 100644 index 000000000..b97025cac --- /dev/null +++ b/src/iceberg/encryption/in_memory_kms.h @@ -0,0 +1,72 @@ +/* + * 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 + +/// \file iceberg/encryption/in_memory_kms.h +/// KMS client with master keys held in memory, for tests and local development. + +#include +#include +#include +#include +#include +#include + +#include "iceberg/encryption/key_management_client.h" +#include "iceberg/iceberg_bundle_export.h" + +namespace iceberg { + +/// \brief KMS client with in-memory master keys. Not for production use. +/// +/// Wire compatible with Java's MemoryMockKMS (used by Java tests through UnitestKMS): +/// keys are wrapped with AES-GCM under the master key, as nonce | ciphertext | tag, +/// without AAD. +class ICEBERG_BUNDLE_EXPORT InMemoryKms : public KeyManagementClient { + public: + /// \brief The name used with `encryption.kms-impl` after RegisterInMemoryKms(). + static constexpr std::string_view kName = "in-memory"; + /// \brief Catalog property prefix for master keys: `encryption.in-memory-kms.key.` + /// = Base64 of the key bytes. + static constexpr std::string_view kKeyPropertyPrefix = "encryption.in-memory-kms.key."; + + InMemoryKms() = default; + explicit InMemoryKms(std::unordered_map> master_keys) + : master_keys_(std::move(master_keys)) {} + + /// \brief Register this client in KmsRegistry under kName. + static void Register(); + + /// \brief Add master keys from `encryption.in-memory-kms.key.` properties. + Status Initialize( + const std::unordered_map& properties) override; + + Result> WrapKey(std::span key, + std::string_view wrapping_key_id) override; + Result> UnwrapKey(std::span wrapped_key, + std::string_view wrapping_key_id) override; + + private: + Result*> MasterKey(std::string_view id) const; + + std::unordered_map> master_keys_; +}; + +} // namespace iceberg diff --git a/src/iceberg/encryption/key_management_client.cc b/src/iceberg/encryption/key_management_client.cc new file mode 100644 index 000000000..85f681f44 --- /dev/null +++ b/src/iceberg/encryption/key_management_client.cc @@ -0,0 +1,98 @@ +/* + * 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 "iceberg/encryption/key_management_client.h" + +#include + +#include "iceberg/util/macros.h" +#include "iceberg/util/string_util.h" + +namespace iceberg { + +namespace { + +struct RegistryState { + std::mutex mutex; + std::unordered_map factories; +}; + +RegistryState& State() { + static RegistryState state; + return state; +} + +} // namespace + +void KmsRegistry::Register(std::string name, Factory factory) { + auto& state = State(); + std::lock_guard lock(state.mutex); + state.factories.insert_or_assign(std::move(name), std::move(factory)); +} + +Result> KmsRegistry::Load( + std::string_view name, const Properties& properties) { + Factory factory; + { + auto& state = State(); + std::lock_guard lock(state.mutex); + auto it = state.factories.find(std::string(name)); + if (it == state.factories.end()) { + return NotFound( + "KMS client not found: {}. No built-in clients ship yet; register one with " + "KmsRegistry::Register and select it with encryption.kms-impl", + name); + } + factory = it->second; + } + ICEBERG_PRECHECK(static_cast(factory), "KMS client factory for {} is empty", + name); + ICEBERG_ASSIGN_OR_RAISE(auto client, factory()); + ICEBERG_PRECHECK(client != nullptr, "KMS client factory for {} returned null", name); + ICEBERG_RETURN_UNEXPECTED(client->Initialize(properties)); + return std::shared_ptr(std::move(client)); +} + +Result> KmsRegistry::FromCatalogProperties( + const Properties& properties) { + auto type_it = properties.find(std::string(kKmsType)); + auto impl_it = properties.find(std::string(kKmsImpl)); + // Like Java, a property that is present but empty counts as set, so that a + // misconfigured catalog fails instead of silently going unencrypted. + bool has_type = type_it != properties.end(); + bool has_impl = impl_it != properties.end(); + + if (has_type && has_impl) { + return InvalidArgument("Cannot set both KMS type ({}) and KMS impl ({})", + type_it->second, impl_it->second); + } + if (has_impl) { + return Load(impl_it->second, properties); + } + if (has_type) { + auto type = StringUtils::ToLower(type_it->second); + if (type != kKmsTypeAws && type != kKmsTypeGcp && type != kKmsTypeAzure) { + return InvalidArgument("Unsupported KMS type: {}", type_it->second); + } + return Load(type, properties); + } + return std::shared_ptr{}; +} + +} // namespace iceberg diff --git a/src/iceberg/encryption/key_management_client.h b/src/iceberg/encryption/key_management_client.h new file mode 100644 index 000000000..6d8c087b4 --- /dev/null +++ b/src/iceberg/encryption/key_management_client.h @@ -0,0 +1,122 @@ +/* + * 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 + +/// \file iceberg/encryption/key_management_client.h +/// Client of a key management service (KMS), and the registry that creates one from +/// catalog properties. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "iceberg/iceberg_export.h" +#include "iceberg/result.h" + +namespace iceberg { + +/// \brief Wraps and unwraps secret keys with master keys kept in a KMS. +/// +/// Mirrors Java's org.apache.iceberg.encryption.KeyManagementClient. The table master +/// key (table property `encryption.key-id`) never leaves the KMS; it only wraps and +/// unwraps the table's key encryption keys. +class ICEBERG_EXPORT KeyManagementClient { + public: + /// \brief A generated key and its wrapped form. + struct GeneratedKey { + std::vector key; + std::vector wrapped_key; + }; + + virtual ~KeyManagementClient() = default; + + /// \brief Initialize the client with the catalog properties. + virtual Status Initialize( + const std::unordered_map& properties) { + return {}; + } + + /// \brief Wrap (encrypt) a secret key with the master key `wrapping_key_id`. + /// + /// The result may include KMS-specific metadata needed to unwrap it. + virtual Result> WrapKey(std::span key, + std::string_view wrapping_key_id) = 0; + + /// \brief Unwrap a key returned by WrapKey. + virtual Result> UnwrapKey(std::span wrapped_key, + std::string_view wrapping_key_id) = 0; + + /// \brief Whether the KMS can generate keys (GenerateKey). + virtual bool SupportsKeyGeneration() const { return false; } + + /// \brief Generate a new key in the KMS, returned in plain and wrapped form. + virtual Result GenerateKey(std::string_view wrapping_key_id) { + return NotSupported("Key generation is not supported by this KMS client"); + } +}; + +/// \brief Creates KMS clients by name. +/// +/// Catalogs create their KMS client from the catalog properties, with the same +/// property names as Java: +/// - `encryption.kms-type`: a built-in client, `aws`, `gcp` or `azure`; +/// - `encryption.kms-impl`: the name a custom client was registered under. +/// +/// Java has no built-in clients yet and rejects every `encryption.kms-type` +/// (EncryptionUtil.createKmsClient); here the type names resolve through the same +/// registry, so a client registered under one of them is used. +/// +/// Built-in clients register themselves when their library is linked and initialized; +/// applications register custom clients before creating catalogs. +class ICEBERG_EXPORT KmsRegistry { + public: + static constexpr std::string_view kKmsType = "encryption.kms-type"; + static constexpr std::string_view kKmsImpl = "encryption.kms-impl"; + static constexpr std::string_view kKmsTypeAws = "aws"; + static constexpr std::string_view kKmsTypeGcp = "gcp"; + static constexpr std::string_view kKmsTypeAzure = "azure"; + + using Properties = std::unordered_map; + using Factory = std::function>()>; + + /// \brief Register a client factory under `name`, replacing any previous one. + /// + /// An empty `factory` is accepted here but makes the corresponding Load fail + /// with an invalid-argument error rather than being invoked. + static void Register(std::string name, Factory factory); + + /// \brief Create and initialize the client registered under `name`. + static Result> Load(std::string_view name, + const Properties& properties); + + /// \brief Create the KMS client configured by catalog properties. + /// + /// \return the initialized client, or nullptr if neither `encryption.kms-type` nor + /// `encryption.kms-impl` is set. + static Result> FromCatalogProperties( + const Properties& properties); +}; + +} // namespace iceberg diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 8622a1a67..d5bca7eac 100644 --- a/src/iceberg/test/CMakeLists.txt +++ b/src/iceberg/test/CMakeLists.txt @@ -116,7 +116,7 @@ add_iceberg_test(expression_test residual_evaluator_test.cc strict_metrics_evaluator_test.cc) -add_iceberg_test(encryption_test SOURCES standard_key_metadata_test.cc) +add_iceberg_test(encryption_test SOURCES kms_test.cc standard_key_metadata_test.cc) add_iceberg_test(json_serde_test SOURCES diff --git a/src/iceberg/test/aes_gcm_test.cc b/src/iceberg/test/aes_gcm_test.cc index 9b21a2c49..e3d1248e7 100644 --- a/src/iceberg/test/aes_gcm_test.cc +++ b/src/iceberg/test/aes_gcm_test.cc @@ -31,6 +31,8 @@ #include "iceberg/arrow/arrow_io_util.h" #include "iceberg/encryption/aes_gcm_internal.h" #include "iceberg/encryption/aes_gcm_stream.h" +#include "iceberg/encryption/in_memory_kms.h" +#include "iceberg/encryption/key_management_client.h" #include "iceberg/encryption/standard_key_metadata.h" #include "iceberg/test/matchers.h" #include "iceberg/test/temp_file_test_base.h" @@ -373,4 +375,59 @@ TEST_F(AesGcmStreamTest, RejectsInvalidEncryptedLength) { EXPECT_THAT(file->Open(), IsError(ErrorKind::kInvalid)); } +// The key encryption key of the Java fixture table, wrapped by Java's UnitestKMS under +// its keyA master key. +const std::vector kJavaKekWrappedByKeyA = { + 0x04, 0x0d, 0xdf, 0x90, 0xe8, 0x41, 0xa3, 0xff, 0xf7, 0xec, 0x77, + 0xf1, 0x04, 0x0d, 0x65, 0x50, 0x46, 0xd2, 0xae, 0x0a, 0xef, 0xde, + 0x4d, 0x3a, 0xa2, 0xf5, 0x06, 0xca, 0x2e, 0xf4, 0x00, 0x83, 0x4b, + 0xbe, 0xf0, 0x94, 0xf6, 0x20, 0x94, 0x76, 0x1c, 0x11, 0xf6, 0x53}; + +// The plaintext of the key above: the 12-byte nonce, ciphertext and 16-byte tag of +// that blob decrypted under keyA with AES-GCM, verified against an independent +// AES-GCM implementation rather than this one. +const std::vector kJavaKek = {0xdb, 0xba, 0x8c, 0xde, 0xaa, 0xb0, 0x33, 0x61, + 0xfb, 0x64, 0x0a, 0x1b, 0xdd, 0x2b, 0x53, 0x67}; + +// Java's UnitestKMS.MASTER_KEY1, the master key it registers under "keyA". +constexpr std::string_view kJavaMasterKeyA = "0123456789012345"; + +TEST(InMemoryKmsTest, UnwrapsJavaWrappedKey) { + auto master_key = AsBytes(kJavaMasterKeyA); + InMemoryKms kms({{"keyA", std::vector(master_key.begin(), master_key.end())}}); + ICEBERG_UNWRAP_OR_FAIL(auto kek, kms.UnwrapKey(kJavaKekWrappedByKeyA, "keyA")); + EXPECT_EQ(kek, kJavaKek); + + ICEBERG_UNWRAP_OR_FAIL(auto wrapped, kms.WrapKey(kek, "keyA")); + ICEBERG_UNWRAP_OR_FAIL(auto unwrapped, kms.UnwrapKey(wrapped, "keyA")); + EXPECT_EQ(unwrapped, kJavaKek); + EXPECT_THAT(kms.UnwrapKey(wrapped, "keyB"), IsError(ErrorKind::kNotFound)); +} + +TEST(InMemoryKmsTest, RegistryAndProperties) { + InMemoryKms::Register(); + ICEBERG_UNWRAP_OR_FAIL( + auto kms, KmsRegistry::FromCatalogProperties( + {{"encryption.kms-impl", "in-memory"}, + // Base64 of "0123456789012345", UnitestKMS's keyA + {"encryption.in-memory-kms.key.keyA", "MDEyMzQ1Njc4OTAxMjM0NQ=="}})); + ASSERT_NE(kms, nullptr); + ICEBERG_UNWRAP_OR_FAIL(auto kek, kms->UnwrapKey(kJavaKekWrappedByKeyA, "keyA")); + EXPECT_EQ(kek, kJavaKek); +} + +TEST(InMemoryKmsTest, RejectsInvalidMasterKeyProperties) { + InMemoryKms::Register(); + auto load = [](std::string id, std::string value) { + return KmsRegistry::FromCatalogProperties( + {{"encryption.kms-impl", "in-memory"}, + {"encryption.in-memory-kms.key." + id, std::move(value)}}); + }; + + EXPECT_THAT(load("keyA", "not base64!"), IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT(load("", "MDEyMzQ1Njc4OTAxMjM0NQ=="), IsError(ErrorKind::kInvalidArgument)); + // Base64 of "0123456", seven bytes, not a valid AES key length. + EXPECT_THAT(load("keyA", "MDEyMzQ1Ng=="), IsError(ErrorKind::kInvalidArgument)); +} + } // namespace iceberg diff --git a/src/iceberg/test/kms_test.cc b/src/iceberg/test/kms_test.cc new file mode 100644 index 000000000..9561de097 --- /dev/null +++ b/src/iceberg/test/kms_test.cc @@ -0,0 +1,112 @@ +/* + * 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 + +#include "iceberg/encryption/key_management_client.h" +#include "iceberg/test/matchers.h" + +namespace iceberg { + +namespace { + +/// Records the properties it was initialized with; "wraps" by reversing the bytes. +class FakeKms : public KeyManagementClient { + public: + Status Initialize(const std::unordered_map& props) override { + properties = props; + return {}; + } + Result> WrapKey(std::span key, + std::string_view) override { + return std::vector(key.rbegin(), key.rend()); + } + Result> UnwrapKey(std::span wrapped, + std::string_view) override { + return std::vector(wrapped.rbegin(), wrapped.rend()); + } + + std::unordered_map properties; +}; + +void RegisterFake(const std::string& name) { + KmsRegistry::Register(name, []() -> Result> { + return std::make_unique(); + }); +} + +} // namespace + +TEST(KmsRegistryTest, NotConfigured) { + ICEBERG_UNWRAP_OR_FAIL(auto client, KmsRegistry::FromCatalogProperties({})); + EXPECT_EQ(client, nullptr); +} + +TEST(KmsRegistryTest, LoadsCustomImplementation) { + RegisterFake("fake-kms"); + KmsRegistry::Properties properties{{"encryption.kms-impl", "fake-kms"}, + {"fake.endpoint", "localhost"}}; + ICEBERG_UNWRAP_OR_FAIL(auto client, KmsRegistry::FromCatalogProperties(properties)); + ASSERT_NE(client, nullptr); + auto* fake = dynamic_cast(client.get()); + ASSERT_NE(fake, nullptr); + // Initialized with the catalog properties + EXPECT_EQ(fake->properties.at("fake.endpoint"), "localhost"); + + std::vector key{1, 2, 3}; + ICEBERG_UNWRAP_OR_FAIL(auto wrapped, client->WrapKey(key, "master")); + ICEBERG_UNWRAP_OR_FAIL(auto unwrapped, client->UnwrapKey(wrapped, "master")); + EXPECT_EQ(unwrapped, key); + EXPECT_FALSE(client->SupportsKeyGeneration()); + EXPECT_THAT(client->GenerateKey("master"), IsError(ErrorKind::kNotSupported)); +} + +TEST(KmsRegistryTest, LoadsBuiltInType) { + // Built-in types resolve to the client registered under the type name. The registry + // is process-wide, so this shadows a real "aws" client for the rest of the binary. + RegisterFake("aws"); + ICEBERG_UNWRAP_OR_FAIL( + auto client, KmsRegistry::FromCatalogProperties({{"encryption.kms-type", "AWS"}})); + EXPECT_NE(dynamic_cast(client.get()), nullptr); +} + +TEST(KmsRegistryTest, EmptyFactoryIsRejected) { + // An empty std::function must surface as an error, not a std::bad_function_call + // escaping through the Result-returning Load. + KmsRegistry::Register("empty-kms", KmsRegistry::Factory{}); + EXPECT_THAT(KmsRegistry::FromCatalogProperties({{"encryption.kms-impl", "empty-kms"}}), + IsError(ErrorKind::kInvalidArgument)); +} + +TEST(KmsRegistryTest, InvalidConfiguration) { + EXPECT_THAT(KmsRegistry::FromCatalogProperties( + {{"encryption.kms-type", "aws"}, {"encryption.kms-impl", "fake-kms"}}), + IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT(KmsRegistry::FromCatalogProperties({{"encryption.kms-type", "vault"}}), + IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT(KmsRegistry::FromCatalogProperties({{"encryption.kms-impl", "missing"}}), + IsError(ErrorKind::kNotFound)); + // Set but empty is a misconfiguration, not "no KMS", as in Java + EXPECT_THAT(KmsRegistry::FromCatalogProperties({{"encryption.kms-type", ""}}), + IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT(KmsRegistry::FromCatalogProperties({{"encryption.kms-impl", ""}}), + IsError(ErrorKind::kNotFound)); +} + +} // namespace iceberg