Skip to content

Commit 8fd564f

Browse files
committed
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.<id> = Base64).
1 parent 49f8a87 commit 8fd564f

10 files changed

Lines changed: 555 additions & 3 deletions

‎src/iceberg/CMakeLists.txt‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ set(ICEBERG_SOURCES
3333
deletes/position_delete_index.cc
3434
deletes/position_delete_range_consumer.cc
3535
deletes/roaring_position_bitmap.cc
36+
encryption/key_management_client.cc
3637
encryption/standard_key_metadata.cc
3738
expression/aggregate.cc
3839
expression/binder.cc
@@ -327,7 +328,11 @@ if(ICEBERG_BUILD_BUNDLE)
327328
parquet/row_number_reader.cc)
328329

329330
if(ICEBERG_ENCRYPTION)
330-
list(APPEND ICEBERG_BUNDLE_SOURCES encryption/aes_gcm.cc encryption/aes_gcm_stream.cc)
331+
list(APPEND
332+
ICEBERG_BUNDLE_SOURCES
333+
encryption/aes_gcm.cc
334+
encryption/aes_gcm_stream.cc
335+
encryption/in_memory_kms.cc)
331336
endif()
332337

333338
# Libraries to link with exported libiceberg_bundle.{so,a}.

‎src/iceberg/cpp20_compatibility_internal.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
#include "iceberg/deletes/position_delete_range_consumer.h"
5050
#include "iceberg/deletes/roaring_position_bitmap.h"
5151
#include "iceberg/encryption/encrypted_key.h"
52+
#include "iceberg/encryption/key_management_client.h"
5253
#include "iceberg/encryption/standard_key_metadata.h"
5354
#include "iceberg/exception.h"
5455
#include "iceberg/expression/aggregate.h"
@@ -218,6 +219,7 @@
218219

219220
#ifdef ICEBERG_PUBLIC_HEADERS_ENCRYPTION
220221
# include "iceberg/encryption/aes_gcm_stream.h"
222+
# include "iceberg/encryption/in_memory_kms.h"
221223
#endif // ICEBERG_PUBLIC_HEADERS_ENCRYPTION
222224

223225
#ifdef ICEBERG_PUBLIC_HEADERS_HIVE

‎src/iceberg/encryption/CMakeLists.txt‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,16 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717

18-
iceberg_install_all_headers(iceberg/encryption PATTERN encrypted_key.h
18+
iceberg_install_all_headers(iceberg/encryption
19+
PATTERN
20+
encrypted_key.h
21+
key_management_client.h
1922
standard_key_metadata.h)
2023
if(ICEBERG_BUILD_BUNDLE AND ICEBERG_ENCRYPTION)
2124
iceberg_install_all_headers(iceberg/encryption
2225
PATTERN
2326
aes_gcm_stream.h
27+
in_memory_kms.h
2428
HEADER_GROUP
2529
ICEBERG_PUBLIC_HEADERS_ENCRYPTION)
2630
endif()
Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
#include "iceberg/encryption/in_memory_kms.h"
21+
22+
#include "iceberg/encryption/aes_gcm_internal.h"
23+
#include "iceberg/util/base64.h"
24+
#include "iceberg/util/macros.h"
25+
26+
namespace iceberg {
27+
28+
void InMemoryKms::Register() {
29+
KmsRegistry::Register(std::string(kName),
30+
[]() -> Result<std::unique_ptr<KeyManagementClient>> {
31+
return std::make_unique<InMemoryKms>();
32+
});
33+
}
34+
35+
Status InMemoryKms::Initialize(
36+
const std::unordered_map<std::string, std::string>& properties) {
37+
// Decode into a staging map so a bad property leaves master_keys_ untouched;
38+
// properties is unordered, so a partial update would be non-deterministic.
39+
std::unordered_map<std::string, std::vector<uint8_t>> decoded;
40+
for (const auto& [name, value] : properties) {
41+
if (!name.starts_with(kKeyPropertyPrefix)) {
42+
continue;
43+
}
44+
auto id = name.substr(kKeyPropertyPrefix.size());
45+
ICEBERG_PRECHECK(!id.empty(), "Master key property {} has an empty key id", name);
46+
ICEBERG_ASSIGN_OR_RAISE(auto key, Base64::Decode(value));
47+
ICEBERG_PRECHECK(key.size() == 16 || key.size() == 24 || key.size() == 32,
48+
"Master key {} must be 16, 24 or 32 bytes, got {}", name,
49+
key.size());
50+
decoded.insert_or_assign(std::move(id), std::vector<uint8_t>(key.begin(), key.end()));
51+
}
52+
for (auto& [id, key] : decoded) {
53+
master_keys_.insert_or_assign(id, std::move(key));
54+
}
55+
return {};
56+
}
57+
58+
Result<const std::vector<uint8_t>*> InMemoryKms::MasterKey(std::string_view id) const {
59+
auto it = master_keys_.find(std::string(id));
60+
if (it == master_keys_.end()) {
61+
return NotFound("Master key {} is not found", id);
62+
}
63+
return &it->second;
64+
}
65+
66+
Result<std::vector<uint8_t>> InMemoryKms::WrapKey(std::span<const uint8_t> key,
67+
std::string_view wrapping_key_id) {
68+
ICEBERG_ASSIGN_OR_RAISE(auto master_key, MasterKey(wrapping_key_id));
69+
ICEBERG_ASSIGN_OR_RAISE(auto cipher, encryption::AesGcmCipher::Make(*master_key));
70+
return cipher->Encrypt(key, {});
71+
}
72+
73+
Result<std::vector<uint8_t>> InMemoryKms::UnwrapKey(std::span<const uint8_t> wrapped_key,
74+
std::string_view wrapping_key_id) {
75+
ICEBERG_ASSIGN_OR_RAISE(auto master_key, MasterKey(wrapping_key_id));
76+
ICEBERG_ASSIGN_OR_RAISE(auto cipher, encryption::AesGcmCipher::Make(*master_key));
77+
return cipher->Decrypt(wrapped_key, {});
78+
}
79+
80+
} // namespace iceberg
Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
#pragma once
21+
22+
/// \file iceberg/encryption/in_memory_kms.h
23+
/// KMS client with master keys held in memory, for tests and local development.
24+
25+
#include <cstdint>
26+
#include <span>
27+
#include <string>
28+
#include <string_view>
29+
#include <unordered_map>
30+
#include <vector>
31+
32+
#include "iceberg/encryption/key_management_client.h"
33+
#include "iceberg/iceberg_bundle_export.h"
34+
35+
namespace iceberg {
36+
37+
/// \brief KMS client with in-memory master keys. Not for production use.
38+
///
39+
/// Wire compatible with Java's MemoryMockKMS (used by Java tests through UnitestKMS):
40+
/// keys are wrapped with AES-GCM under the master key, as nonce | ciphertext | tag,
41+
/// without AAD.
42+
class ICEBERG_BUNDLE_EXPORT InMemoryKms : public KeyManagementClient {
43+
public:
44+
/// \brief The name used with `encryption.kms-impl` after RegisterInMemoryKms().
45+
static constexpr std::string_view kName = "in-memory";
46+
/// \brief Catalog property prefix for master keys: `encryption.in-memory-kms.key.<id>`
47+
/// = Base64 of the key bytes.
48+
static constexpr std::string_view kKeyPropertyPrefix = "encryption.in-memory-kms.key.";
49+
50+
InMemoryKms() = default;
51+
explicit InMemoryKms(std::unordered_map<std::string, std::vector<uint8_t>> master_keys)
52+
: master_keys_(std::move(master_keys)) {}
53+
54+
/// \brief Register this client in KmsRegistry under kName.
55+
static void Register();
56+
57+
/// \brief Add master keys from `encryption.in-memory-kms.key.<id>` properties.
58+
Status Initialize(
59+
const std::unordered_map<std::string, std::string>& properties) override;
60+
61+
Result<std::vector<uint8_t>> WrapKey(std::span<const uint8_t> key,
62+
std::string_view wrapping_key_id) override;
63+
Result<std::vector<uint8_t>> UnwrapKey(std::span<const uint8_t> wrapped_key,
64+
std::string_view wrapping_key_id) override;
65+
66+
private:
67+
Result<const std::vector<uint8_t>*> MasterKey(std::string_view id) const;
68+
69+
std::unordered_map<std::string, std::vector<uint8_t>> master_keys_;
70+
};
71+
72+
} // namespace iceberg
Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
#include "iceberg/encryption/key_management_client.h"
21+
22+
#include <mutex>
23+
24+
#include "iceberg/util/macros.h"
25+
#include "iceberg/util/string_util.h"
26+
27+
namespace iceberg {
28+
29+
namespace {
30+
31+
struct RegistryState {
32+
std::mutex mutex;
33+
std::unordered_map<std::string, KmsRegistry::Factory> factories;
34+
};
35+
36+
RegistryState& State() {
37+
static RegistryState state;
38+
return state;
39+
}
40+
41+
} // namespace
42+
43+
void KmsRegistry::Register(std::string name, Factory factory) {
44+
auto& state = State();
45+
std::lock_guard lock(state.mutex);
46+
state.factories.insert_or_assign(std::move(name), std::move(factory));
47+
}
48+
49+
Result<std::shared_ptr<KeyManagementClient>> KmsRegistry::Load(
50+
std::string_view name, const Properties& properties) {
51+
Factory factory;
52+
{
53+
auto& state = State();
54+
std::lock_guard lock(state.mutex);
55+
auto it = state.factories.find(std::string(name));
56+
if (it == state.factories.end()) {
57+
return NotFound(
58+
"KMS client not found: {}. No built-in clients ship yet; register one with "
59+
"KmsRegistry::Register and select it with encryption.kms-impl",
60+
name);
61+
}
62+
factory = it->second;
63+
}
64+
ICEBERG_PRECHECK(static_cast<bool>(factory), "KMS client factory for {} is empty",
65+
name);
66+
ICEBERG_ASSIGN_OR_RAISE(auto client, factory());
67+
ICEBERG_PRECHECK(client != nullptr, "KMS client factory for {} returned null", name);
68+
ICEBERG_RETURN_UNEXPECTED(client->Initialize(properties));
69+
return std::shared_ptr<KeyManagementClient>(std::move(client));
70+
}
71+
72+
Result<std::shared_ptr<KeyManagementClient>> KmsRegistry::FromCatalogProperties(
73+
const Properties& properties) {
74+
auto type_it = properties.find(std::string(kKmsType));
75+
auto impl_it = properties.find(std::string(kKmsImpl));
76+
// Like Java, a property that is present but empty counts as set, so that a
77+
// misconfigured catalog fails instead of silently going unencrypted.
78+
bool has_type = type_it != properties.end();
79+
bool has_impl = impl_it != properties.end();
80+
81+
if (has_type && has_impl) {
82+
return InvalidArgument("Cannot set both KMS type ({}) and KMS impl ({})",
83+
type_it->second, impl_it->second);
84+
}
85+
if (has_impl) {
86+
return Load(impl_it->second, properties);
87+
}
88+
if (has_type) {
89+
auto type = StringUtils::ToLower(type_it->second);
90+
if (type != kKmsTypeAws && type != kKmsTypeGcp && type != kKmsTypeAzure) {
91+
return InvalidArgument("Unsupported KMS type: {}", type_it->second);
92+
}
93+
return Load(type, properties);
94+
}
95+
return std::shared_ptr<KeyManagementClient>{};
96+
}
97+
98+
} // namespace iceberg

0 commit comments

Comments
 (0)