feat: add KeyManagementClient, KmsRegistry and InMemoryKms

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}.

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

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()

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<std::unique_ptr<KeyManagementClient>> {
+                          return std::make_unique<InMemoryKms>();
+                        });
+}
+
+Status InMemoryKms::Initialize(
+    const std::unordered_map<std::string, std::string>& 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<std::string, std::vector<uint8_t>> 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<uint8_t>(key.begin(), key.end()));
+  }
+  for (auto& [id, key] : decoded) {
+    master_keys_.insert_or_assign(id, std::move(key));
+  }
+  return {};
+}
+
+Result<const std::vector<uint8_t>*> 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<std::vector<uint8_t>> InMemoryKms::WrapKey(std::span<const uint8_t> 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<std::vector<uint8_t>> InMemoryKms::UnwrapKey(std::span<const uint8_t> 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

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 <cstdint>
+#include <span>
+#include <string>
+#include <string_view>
+#include <unordered_map>
+#include <vector>
+
+#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.<id>`
+  /// = Base64 of the key bytes.
+  static constexpr std::string_view kKeyPropertyPrefix = "encryption.in-memory-kms.key.";
+
+  InMemoryKms() = default;
+  explicit InMemoryKms(std::unordered_map<std::string, std::vector<uint8_t>> 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.<id>` properties.
+  Status Initialize(
+      const std::unordered_map<std::string, std::string>& properties) override;
+
+  Result<std::vector<uint8_t>> WrapKey(std::span<const uint8_t> key,
+                                       std::string_view wrapping_key_id) override;
+  Result<std::vector<uint8_t>> UnwrapKey(std::span<const uint8_t> wrapped_key,
+                                         std::string_view wrapping_key_id) override;
+
+ private:
+  Result<const std::vector<uint8_t>*> MasterKey(std::string_view id) const;
+
+  std::unordered_map<std::string, std::vector<uint8_t>> master_keys_;
+};
+
+}  // namespace iceberg

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 <mutex>
+
+#include "iceberg/util/macros.h"
+#include "iceberg/util/string_util.h"
+
+namespace iceberg {
+
+namespace {
+
+struct RegistryState {
+  std::mutex mutex;
+  std::unordered_map<std::string, KmsRegistry::Factory> 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<std::shared_ptr<KeyManagementClient>> 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<bool>(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<KeyManagementClient>(std::move(client));
+}
+
+Result<std::shared_ptr<KeyManagementClient>> 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<KeyManagementClient>{};
+}
+
+}  // namespace iceberg