From 668275358f81df3e64b7233bc16f08b7853767f7 Mon Sep 17 00:00:00 2001 From: huanghsiang_cheng Date: Tue, 12 May 2026 22:45:16 -0700 Subject: [PATCH 1/2] AWS KMS implementation of `KeyManagementClient` --- Cargo.lock | 385 +++++++++++++++++-- Cargo.toml | 2 + crates/iceberg/Cargo.toml | 4 + crates/iceberg/src/encryption/kms/aws_kms.rs | 223 +++++++++++ crates/iceberg/src/encryption/kms/mod.rs | 2 + 5 files changed, 574 insertions(+), 42 deletions(-) create mode 100644 crates/iceberg/src/encryption/kms/aws_kms.rs diff --git a/Cargo.lock b/Cargo.lock index 19ac987caa..93736d5b15 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -664,6 +664,30 @@ dependencies = [ "tracing", ] +[[package]] +name = "aws-sdk-kms" +version = "1.104.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c41ae6a33da941457e89075ef8ca5b4870c8009fe4dceeba82fce2f30f313ac6" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.0", + "regex-lite", + "tracing", +] + [[package]] name = "aws-sdk-s3tables" version = "1.54.0" @@ -822,19 +846,31 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a2f165a7feee6f263028b899d0a181987f4fa7179a6411a32a439fba7c5f769" dependencies = [ "aws-smithy-async", + "aws-smithy-protocol-test", "aws-smithy-runtime-api", "aws-smithy-types", - "h2", + "bytes", + "h2 0.3.27", + "h2 0.4.13", + "http 0.2.12", "http 1.4.0", - "hyper", - "hyper-rustls", + "http-body 0.4.6", + "http-body 1.0.1", + "hyper 0.14.32", + "hyper 1.8.1", + "hyper-rustls 0.24.2", + "hyper-rustls 0.27.7", "hyper-util", + "indexmap 2.13.0", "pin-project-lite", - "rustls", + "rustls 0.21.12", + "rustls 0.23.37", "rustls-native-certs", "rustls-pki-types", + "serde", + "serde_json", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tower", "tracing", ] @@ -848,6 +884,18 @@ dependencies = [ "aws-smithy-types", ] +[[package]] +name = "aws-smithy-mocks" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d9a9490933f8faa0aeb1eb9fbb8bf4f1c214b3149c61b3a4074b68030b21bd5" +dependencies = [ + "aws-smithy-http-client", + "aws-smithy-runtime-api", + "aws-smithy-types", + "http 1.4.0", +] + [[package]] name = "aws-smithy-observability" version = "0.2.6" @@ -857,6 +905,25 @@ dependencies = [ "aws-smithy-runtime-api", ] +[[package]] +name = "aws-smithy-protocol-test" +version = "0.63.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b227aa94af99a8e5ee52551cc7e3ee30a217019ef99207b6f0b7a1527685941" +dependencies = [ + "assert-json-diff", + "aws-smithy-runtime-api", + "base64-simd", + "cbor-diag", + "ciborium", + "http 0.2.12", + "pretty_assertions", + "regex-lite", + "roxmltree", + "serde_json", + "thiserror 2.0.18", +] + [[package]] name = "aws-smithy-query" version = "0.60.15" @@ -890,6 +957,7 @@ dependencies = [ "pin-utils", "tokio", "tracing", + "tracing-subscriber", ] [[package]] @@ -1087,7 +1155,7 @@ version = "3.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "519bd3116aeeb42d5372c29d982d16d0170d3d4a5ed85fc7dd91642ffff3c67c" dependencies = [ - "darling 0.20.11", + "darling 0.23.0", "ident_case", "prettyplease", "proc-macro2", @@ -1117,6 +1185,15 @@ dependencies = [ "alloc-stdlib", ] +[[package]] +name = "bs58" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf88ba1141d185c399bee5288d850d63b8369520c1eafc32a0430b5b6c287bf4" +dependencies = [ + "tinyvec", +] + [[package]] name = "bumpalo" version = "3.20.2" @@ -1186,6 +1263,25 @@ dependencies = [ "cipher", ] +[[package]] +name = "cbor-diag" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc245b6ecd09b23901a4fbad1ad975701fd5061ceaef6afa93a2d70605a64429" +dependencies = [ + "bs58", + "chrono", + "data-encoding", + "half", + "nom", + "num-bigint", + "num-rational", + "num-traits", + "separator", + "url", + "uuid", +] + [[package]] name = "cc" version = "1.2.57" @@ -1245,6 +1341,33 @@ dependencies = [ "phf", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + [[package]] name = "cipher" version = "0.4.4" @@ -1325,7 +1448,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -1668,6 +1791,12 @@ dependencies = [ "parking_lot_core", ] +[[package]] +name = "data-encoding" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" + [[package]] name = "datafusion" version = "53.1.0" @@ -2545,7 +2674,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2713,7 +2842,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3087,6 +3216,25 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "h2" +version = "0.3.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0beca50380b1fc32983fc1cb4587bfa4bb9e78fc259aad4a0032d2080309222d" +dependencies = [ + "bytes", + "fnv", + "futures-core", + "futures-sink", + "futures-util", + "http 0.2.12", + "indexmap 2.13.0", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "h2" version = "0.4.13" @@ -3292,6 +3440,30 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" +[[package]] +name = "hyper" +version = "0.14.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41dfc780fdec9373c01bae43289ea34c972e40ee3c9f6b3c8801a35f35586ce7" +dependencies = [ + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "h2 0.3.27", + "http 0.2.12", + "http-body 0.4.6", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "socket2 0.5.10", + "tokio", + "tower-service", + "tracing", + "want", +] + [[package]] name = "hyper" version = "1.8.1" @@ -3302,7 +3474,7 @@ dependencies = [ "bytes", "futures-channel", "futures-core", - "h2", + "h2 0.4.13", "http 1.4.0", "http-body 1.0.1", "httparse", @@ -3315,6 +3487,21 @@ dependencies = [ "want", ] +[[package]] +name = "hyper-rustls" +version = "0.24.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec3efd23720e2049821a693cbc7e65ea87c72f1c58ff2f9522ff332b1491e590" +dependencies = [ + "futures-util", + "http 0.2.12", + "hyper 0.14.32", + "log", + "rustls 0.21.12", + "tokio", + "tokio-rustls 0.24.1", +] + [[package]] name = "hyper-rustls" version = "0.27.7" @@ -3322,13 +3509,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58" dependencies = [ "http 1.4.0", - "hyper", + "hyper 1.8.1", "hyper-util", - "rustls", + "rustls 0.23.37", "rustls-native-certs", "rustls-pki-types", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tower-service", ] @@ -3344,12 +3531,12 @@ dependencies = [ "futures-util", "http 1.4.0", "http-body 1.0.1", - "hyper", + "hyper 1.8.1", "ipnet", "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.3", "tokio", "tower-service", "tracing", @@ -3397,6 +3584,9 @@ dependencies = [ "arrow-string", "as-any", "async-trait", + "aws-config", + "aws-sdk-kms", + "aws-smithy-mocks", "backon", "base64", "bimap", @@ -3923,7 +4113,7 @@ dependencies = [ "portable-atomic-util", "serde_core", "wasm-bindgen", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4267,6 +4457,15 @@ dependencies = [ "serde", ] +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + [[package]] name = "md-5" version = "0.10.6" @@ -4339,6 +4538,12 @@ dependencies = [ "serde", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -4399,7 +4604,7 @@ dependencies = [ "http 1.4.0", "http-body 1.0.1", "http-body-util", - "hyper", + "hyper 1.8.1", "hyper-util", "log", "pin-project-lite", @@ -4526,13 +4731,23 @@ dependencies = [ "libc", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4597,6 +4812,17 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -4655,7 +4881,7 @@ dependencies = [ "http 1.4.0", "http-body-util", "humantime", - "hyper", + "hyper 1.8.1", "itertools 0.14.0", "md-5", "parking_lot", @@ -5453,8 +5679,8 @@ dependencies = [ "quinn-proto", "quinn-udp", "rustc-hash", - "rustls", - "socket2 0.5.10", + "rustls 0.23.37", + "socket2 0.6.3", "thiserror 2.0.18", "tokio", "tracing", @@ -5474,7 +5700,7 @@ dependencies = [ "rand 0.9.4", "ring", "rustc-hash", - "rustls", + "rustls 0.23.37", "rustls-pki-types", "slab", "thiserror 2.0.18", @@ -5492,9 +5718,9 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.3", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -5864,19 +6090,19 @@ dependencies = [ "bytes", "futures-core", "futures-util", - "h2", + "h2 0.4.13", "http 1.4.0", "http-body 1.0.1", "http-body-util", - "hyper", - "hyper-rustls", + "hyper 1.8.1", + "hyper-rustls 0.27.7", "hyper-util", "js-sys", "log", "percent-encoding", "pin-project-lite", "quinn", - "rustls", + "rustls 0.23.37", "rustls-native-certs", "rustls-pki-types", "serde", @@ -5884,7 +6110,7 @@ dependencies = [ "serde_urlencoded", "sync_wrapper", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tokio-util", "tower", "tower-http", @@ -5909,20 +6135,20 @@ dependencies = [ "http 1.4.0", "http-body 1.0.1", "http-body-util", - "hyper", - "hyper-rustls", + "hyper 1.8.1", + "hyper-rustls 0.27.7", "hyper-util", "js-sys", "log", "percent-encoding", "pin-project-lite", "quinn", - "rustls", + "rustls 0.23.37", "rustls-pki-types", "rustls-platform-verifier", "sync_wrapper", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tokio-util", "tower", "tower-http", @@ -5987,6 +6213,15 @@ dependencies = [ "byteorder", ] +[[package]] +name = "roxmltree" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "921904a62e410e37e215c40381b7117f830d9d89ba60ab5236170541dd25646b" +dependencies = [ + "xmlparser", +] + [[package]] name = "rsa" version = "0.9.10" @@ -6075,7 +6310,19 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.59.0", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls" +version = "0.21.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f56a14d1f48b391359b22f731fd4bd7e43c97f3c50eee276f3aa09c94784d3e" +dependencies = [ + "log", + "ring", + "rustls-webpki 0.101.7", + "sct", ] [[package]] @@ -6088,7 +6335,7 @@ dependencies = [ "once_cell", "ring", "rustls-pki-types", - "rustls-webpki", + "rustls-webpki 0.103.13", "subtle", "zeroize", ] @@ -6126,14 +6373,14 @@ dependencies = [ "jni", "log", "once_cell", - "rustls", + "rustls 0.23.37", "rustls-native-certs", "rustls-platform-verifier-android", - "rustls-webpki", + "rustls-webpki 0.103.13", "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -6142,6 +6389,16 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" +[[package]] +name = "rustls-webpki" +version = "0.101.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b6275d1ee7a1cd780b64aca7726599a1dbc893b1e64144529e55c3c2f745765" +dependencies = [ + "ring", + "untrusted 0.9.0", +] + [[package]] name = "rustls-webpki" version = "0.103.13" @@ -6280,6 +6537,16 @@ dependencies = [ "sha2", ] +[[package]] +name = "sct" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da046153aa2352493d6cb7da4b6e5c0c057d8a1d0a9aa8560baffdd945acd414" +dependencies = [ + "ring", + "untrusted 0.9.0", +] + [[package]] name = "security-framework" version = "3.7.0" @@ -6313,6 +6580,12 @@ dependencies = [ "serde_core", ] +[[package]] +name = "separator" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f97841a747eef040fcd2e7b3b9a220a7205926e60488e673d9e4926d27772ce5" + [[package]] name = "seq-macro" version = "0.3.6" @@ -6400,6 +6673,7 @@ version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ + "indexmap 2.13.0", "itoa", "memchr", "serde", @@ -6779,7 +7053,7 @@ dependencies = [ "memchr", "once_cell", "percent-encoding", - "rustls", + "rustls 0.23.37", "serde", "serde_json", "sha2", @@ -7075,7 +7349,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -7237,13 +7511,23 @@ dependencies = [ "syn", ] +[[package]] +name = "tokio-rustls" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" +dependencies = [ + "rustls 0.21.12", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" dependencies = [ - "rustls", + "rustls 0.23.37", "tokio", ] @@ -7432,18 +7716,35 @@ dependencies = [ "tracing-core", ] +[[package]] +name = "tracing-serde" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704b1aeb7be0d0a84fc9828cae51dab5970fee5088f83d1dd7ee6f6246fc6ff1" +dependencies = [ + "serde", + "tracing-core", +] + [[package]] name = "tracing-subscriber" version = "0.3.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" dependencies = [ + "matchers", "nu-ansi-term", + "once_cell", + "regex-automata", + "serde", + "serde_json", "sharded-slab", "smallvec", "thread_local", + "tracing", "tracing-core", "tracing-log", + "tracing-serde", ] [[package]] @@ -7999,7 +8300,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index d23e013f1c..846954eadd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -58,6 +58,8 @@ async-trait = "0.1.89" aws-config = "1.8.7" aws-sdk-glue = { version = "1.85", default-features = false, features = ["default-https-client", "rt-tokio"] } aws-sdk-s3tables = { version = "1.28", default-features = false, features = ["default-https-client", "rt-tokio"] } +aws-sdk-kms = "1.97.0" +aws-smithy-mocks = "0.2.6" backon = "1.5.1" base64 = "0.22.1" bimap = "0.6" diff --git a/crates/iceberg/Cargo.toml b/crates/iceberg/Cargo.toml index b3f082e0d1..86fe0fc794 100644 --- a/crates/iceberg/Cargo.toml +++ b/crates/iceberg/Cargo.toml @@ -47,6 +47,8 @@ arrow-select = { workspace = true } arrow-string = { workspace = true } as-any = { workspace = true } async-trait = { workspace = true } +aws-config = { workspace = true } +aws-sdk-kms = { workspace = true } backon = { workspace = true } base64 = { workspace = true } bimap = { workspace = true } @@ -92,6 +94,8 @@ regex = { workspace = true } tempfile = { workspace = true } minijinja = { workspace = true } serde_arrow = { version = "0.14", features = ["arrow-58"] } +aws-sdk-kms = { workspace = true, features = ["test-util"]} +aws-smithy-mocks = { workspace = true } [package.metadata.cargo-machete] # These dependencies are added to ensure minimal dependency version diff --git a/crates/iceberg/src/encryption/kms/aws_kms.rs b/crates/iceberg/src/encryption/kms/aws_kms.rs new file mode 100644 index 0000000000..405b60f932 --- /dev/null +++ b/crates/iceberg/src/encryption/kms/aws_kms.rs @@ -0,0 +1,223 @@ +// 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. + +use std::fmt; + +use async_trait::async_trait; +use aws_config::BehaviorVersion; +use aws_config::meta::region::RegionProviderChain; +use aws_sdk_kms::Client; +use aws_sdk_kms::primitives::Blob; +use aws_sdk_kms::types::DataKeySpec; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64; + +use crate::encryption::{GeneratedKey, KeyManagementClient, SensitiveBytes}; +use crate::{Error, ErrorKind}; + +/// AWS KMS implementation +/// +/// ``` +/// use aws_sdk_kms::types::DataKeySpec; +/// use iceberg::encryption::KeyManagementClient; +/// use iceberg::encryption::kms::AwsKeyManagementClient; +/// +/// async fn example() -> iceberg::Result<()> { +/// let kms = AwsKeyManagementClient::new(DataKeySpec::Aes128).await; +/// let dek = vec![0u8; 32]; +/// let wrapped = kms.wrap_key(&dek, "arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab").await?; +/// let unwrapped = kms.unwrap_key(&wrapped, "arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab").await?; +/// assert_eq!(dek.as_slice(), unwrapped.as_bytes()); +/// Ok(()) +/// } +/// ``` +#[derive(Clone)] +pub struct AwsKeyManagementClient { + kms_client: Client, + data_key_spec: DataKeySpec, +} + +impl fmt::Debug for AwsKeyManagementClient { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AwsKeyManagementClient").finish() + } +} + +impl AwsKeyManagementClient { + /// Creates an AwsKeyManagementClient with a default AWS KMS client and DataKeySpec + pub async fn new(data_key_spec: DataKeySpec) -> Self { + let region_provider = RegionProviderChain::default_provider().or_else("us-west-2"); + + // This line uses the default credential provider chain + let config = aws_config::defaults(BehaviorVersion::v2026_01_12()) + // region can also be loaded from AWS_DEFAULT_REGION, just remove this line. + .region(region_provider) + .load() + .await; + + Self { + kms_client: Client::new(&config), + data_key_spec, + } + } + + /// Creates an AwsKeyManagementClient with a provided AWS KMS client and DataKeySpec + pub fn new_with_client(kms_client: Client, data_key_spec: DataKeySpec) -> Self { + Self { + kms_client, + data_key_spec, + } + } +} + +#[async_trait] +impl KeyManagementClient for AwsKeyManagementClient { + async fn wrap_key(&self, key: &[u8], wrapping_key_id: &str) -> crate::Result> { + let blob = Blob::new(key); + + let resp = self + .kms_client + .encrypt() + .key_id(wrapping_key_id) + .plaintext(blob) + .send() + .await; + let blob = resp + .map(|r| r.ciphertext_blob.expect("Could not get encrypted text")) + .map_err(|e| { + Error::new( + ErrorKind::Unexpected, + format!("Error encrypting wrapped key: {e}"), + ) + })?; + + Ok(blob.into_inner()) + } + + async fn unwrap_key( + &self, + wrapped_key: &[u8], + wrapping_key_id: &str, + ) -> crate::Result { + let wrapped_key = BASE64.decode(wrapped_key).map(Blob::new).map_err(|e| { + Error::new( + ErrorKind::Unexpected, + format!("Error base64 decoding wrapped key: {e}"), + ) + })?; + + let resp = self + .kms_client + .decrypt() + .key_id(wrapping_key_id) + .ciphertext_blob(wrapped_key) + .send() + .await; + + let inner = resp.map(|r| r.plaintext.unwrap()).map_err(|e| { + Error::new( + ErrorKind::Unexpected, + format!("Error decrypting wrapped key: {e}"), + ) + })?; + + Ok(SensitiveBytes::new(inner.into_inner())) + } + + fn supports_key_generation(&self) -> bool { + true + } + + async fn generate_key(&self, wrapping_key_id: &str) -> crate::Result { + let resp = self + .kms_client + .generate_data_key() + .key_id(wrapping_key_id) + .key_spec(self.data_key_spec.clone()) + .send() + .await; + + let key = resp + .map(|r| { + let wrapped_key = r + .ciphertext_blob + .expect("Could not get encrypted text") + .into_inner(); + let plaintext = r.plaintext.expect("Could not get plaintext"); + GeneratedKey::new(SensitiveBytes::new(plaintext.into_inner()), wrapped_key) + }) + .map_err(|e| { + Error::new( + ErrorKind::Unexpected, + format!("Error generating wrapped key: {e}"), + ) + })?; + + Ok(key) + } +} + +#[cfg(test)] +mod tests { + use aws_sdk_kms::operation::decrypt::DecryptOutput; + use aws_sdk_kms::operation::encrypt::EncryptOutput; + use aws_smithy_mocks::{mock, mock_client}; + + use super::*; + + #[tokio::test] + async fn test_wrap_key() { + let encrypt_rule = mock!(Client::encrypt) + .match_requests(|req| req.key_id == Some(String::from("master-key-id"))) + .then_output(|| { + EncryptOutput::builder() + .ciphertext_blob(Blob::from(b"hello world".to_vec())) + .build() + }); + + let client = mock_client!(aws_sdk_kms, [&encrypt_rule]); + + let kms = AwsKeyManagementClient::new_with_client(client, DataKeySpec::Aes256); + + assert_eq!( + kms.wrap_key(b"123", "master-key-id").await.unwrap(), + b"hello world".to_vec() + ); + } + + #[tokio::test] + async fn test_unwrap_key() { + let decrypt_rule = mock!(Client::decrypt) + .match_requests(|req| req.key_id == Some(String::from("master-key-id"))) + .then_output(|| { + DecryptOutput::builder() + .plaintext(Blob::from(b"hello world".to_vec())) + .build() + }); + + let client = mock_client!(aws_sdk_kms, [&decrypt_rule]); + + let kms = AwsKeyManagementClient::new_with_client(client, DataKeySpec::Aes128); + + assert_eq!( + kms.unwrap_key(BASE64.encode(b"123").as_bytes(), "master-key-id") + .await + .unwrap(), + SensitiveBytes::new(b"hello world".to_vec()) + ); + } +} diff --git a/crates/iceberg/src/encryption/kms/mod.rs b/crates/iceberg/src/encryption/kms/mod.rs index 160e692550..6d11632c03 100644 --- a/crates/iceberg/src/encryption/kms/mod.rs +++ b/crates/iceberg/src/encryption/kms/mod.rs @@ -20,8 +20,10 @@ //! This module provides the [`KeyManagementClient`] trait for pluggable KMS //! integration and implementations for different key management systems. +mod aws_kms; mod client; mod memory; +pub use aws_kms::AwsKeyManagementClient; pub use client::{GeneratedKey, KeyManagementClient}; pub use memory::MemoryKeyManagementClient; From 861e9df7d813c9dc4daf9113572be927efdc2baf Mon Sep 17 00:00:00 2001 From: huanghsiang_cheng Date: Thu, 14 May 2026 10:54:19 -0700 Subject: [PATCH 2/2] Refactor generate_key API --- crates/iceberg/src/encryption/kms/aws_kms.rs | 157 ++++++++++--------- crates/iceberg/src/encryption/kms/client.rs | 16 +- crates/iceberg/src/encryption/kms/memory.rs | 8 +- crates/iceberg/src/encryption/manager.rs | 5 +- 4 files changed, 109 insertions(+), 77 deletions(-) diff --git a/crates/iceberg/src/encryption/kms/aws_kms.rs b/crates/iceberg/src/encryption/kms/aws_kms.rs index 405b60f932..e3095aa653 100644 --- a/crates/iceberg/src/encryption/kms/aws_kms.rs +++ b/crates/iceberg/src/encryption/kms/aws_kms.rs @@ -15,72 +15,60 @@ // specific language governing permissions and limitations // under the License. -use std::fmt; - use async_trait::async_trait; use aws_config::BehaviorVersion; use aws_config::meta::region::RegionProviderChain; use aws_sdk_kms::Client; use aws_sdk_kms::primitives::Blob; use aws_sdk_kms::types::DataKeySpec; -use base64::Engine; -use base64::engine::general_purpose::STANDARD as BASE64; -use crate::encryption::{GeneratedKey, KeyManagementClient, SensitiveBytes}; +use crate::encryption::{AesKeySize, GeneratedKey, KeyManagementClient, SensitiveBytes}; use crate::{Error, ErrorKind}; -/// AWS KMS implementation +/// AWS KMS implementation of [`KeyManagementClient`]. /// -/// ``` -/// use aws_sdk_kms::types::DataKeySpec; +/// ```no_run /// use iceberg::encryption::KeyManagementClient; /// use iceberg::encryption::kms::AwsKeyManagementClient; /// /// async fn example() -> iceberg::Result<()> { -/// let kms = AwsKeyManagementClient::new(DataKeySpec::Aes128).await; +/// let kms = AwsKeyManagementClient::new().await; /// let dek = vec![0u8; 32]; -/// let wrapped = kms.wrap_key(&dek, "arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab").await?; -/// let unwrapped = kms.unwrap_key(&wrapped, "arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab").await?; +/// const KMS_KEY_ID: &str = +/// "arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab"; +/// let wrapped = kms.wrap_key(&dek, KMS_KEY_ID).await?; +/// let unwrapped = kms.unwrap_key(&wrapped, KMS_KEY_ID).await?; /// assert_eq!(dek.as_slice(), unwrapped.as_bytes()); /// Ok(()) /// } /// ``` -#[derive(Clone)] +#[derive(Clone, Debug)] pub struct AwsKeyManagementClient { kms_client: Client, - data_key_spec: DataKeySpec, -} - -impl fmt::Debug for AwsKeyManagementClient { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("AwsKeyManagementClient").finish() - } } impl AwsKeyManagementClient { - /// Creates an AwsKeyManagementClient with a default AWS KMS client and DataKeySpec - pub async fn new(data_key_spec: DataKeySpec) -> Self { - let region_provider = RegionProviderChain::default_provider().or_else("us-west-2"); + /// Creates an `AwsKeyManagementClient` using AWS SDK default credential + /// and region resolution. + /// + /// AWS SDK config resolution (env vars, profiles, IMDS) runs once on construction; + /// callers should reuse a single instance rather than building per request. + pub async fn new() -> Self { + let region_provider = RegionProviderChain::default_provider(); - // This line uses the default credential provider chain let config = aws_config::defaults(BehaviorVersion::v2026_01_12()) - // region can also be loaded from AWS_DEFAULT_REGION, just remove this line. .region(region_provider) .load() .await; Self { kms_client: Client::new(&config), - data_key_spec, } } - /// Creates an AwsKeyManagementClient with a provided AWS KMS client and DataKeySpec - pub fn new_with_client(kms_client: Client, data_key_spec: DataKeySpec) -> Self { - Self { - kms_client, - data_key_spec, - } + /// Creates an `AwsKeyManagementClient` from a pre-configured AWS KMS client. + pub fn new_with_client(kms_client: Client) -> Self { + Self { kms_client } } } @@ -95,17 +83,20 @@ impl KeyManagementClient for AwsKeyManagementClient { .key_id(wrapping_key_id) .plaintext(blob) .send() - .await; - let blob = resp - .map(|r| r.ciphertext_blob.expect("Could not get encrypted text")) + .await .map_err(|e| { - Error::new( - ErrorKind::Unexpected, - format!("Error encrypting wrapped key: {e}"), - ) + Error::new(ErrorKind::Unexpected, "Failed to encrypt key via AWS KMS") + .with_source(e) })?; - Ok(blob.into_inner()) + let ciphertext = resp.ciphertext_blob.ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + "AWS KMS encrypt response missing ciphertext_blob", + ) + })?; + + Ok(ciphertext.into_inner()) } async fn unwrap_key( @@ -113,61 +104,87 @@ impl KeyManagementClient for AwsKeyManagementClient { wrapped_key: &[u8], wrapping_key_id: &str, ) -> crate::Result { - let wrapped_key = BASE64.decode(wrapped_key).map(Blob::new).map_err(|e| { - Error::new( - ErrorKind::Unexpected, - format!("Error base64 decoding wrapped key: {e}"), - ) - })?; + let blob = Blob::new(wrapped_key); let resp = self .kms_client .decrypt() .key_id(wrapping_key_id) - .ciphertext_blob(wrapped_key) + .ciphertext_blob(blob) .send() - .await; + .await + .map_err(|e| { + Error::new(ErrorKind::Unexpected, "Failed to decrypt key via AWS KMS") + .with_source(e) + })?; - let inner = resp.map(|r| r.plaintext.unwrap()).map_err(|e| { + let plaintext = resp.plaintext.ok_or_else(|| { Error::new( ErrorKind::Unexpected, - format!("Error decrypting wrapped key: {e}"), + "AWS KMS decrypt response missing plaintext", ) })?; - Ok(SensitiveBytes::new(inner.into_inner())) + Ok(SensitiveBytes::new(plaintext.into_inner())) } fn supports_key_generation(&self) -> bool { true } - async fn generate_key(&self, wrapping_key_id: &str) -> crate::Result { + async fn generate_key( + &self, + wrapping_key_id: &str, + aes_key_size: AesKeySize, + ) -> crate::Result { + let data_key_spec = match aes_key_size { + AesKeySize::Bits128 => Ok(DataKeySpec::Aes128), + AesKeySize::Bits192 => Err(Error::new( + ErrorKind::FeatureUnsupported, + "AWS KMS DataKeySpec doesn't support AES-192", + )), + AesKeySize::Bits256 => Ok(DataKeySpec::Aes256), + }?; + let resp = self .kms_client .generate_data_key() .key_id(wrapping_key_id) - .key_spec(self.data_key_spec.clone()) + .key_spec(data_key_spec) .send() - .await; - - let key = resp - .map(|r| { - let wrapped_key = r - .ciphertext_blob - .expect("Could not get encrypted text") - .into_inner(); - let plaintext = r.plaintext.expect("Could not get plaintext"); - GeneratedKey::new(SensitiveBytes::new(plaintext.into_inner()), wrapped_key) - }) + .await .map_err(|e| { Error::new( ErrorKind::Unexpected, - format!("Error generating wrapped key: {e}"), + "Failed to generate data key via AWS KMS", ) + .with_source(e) })?; - Ok(key) + let wrapped_key = resp + .ciphertext_blob + .ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + "AWS KMS generate_data_key response missing ciphertext_blob", + ) + })? + .into_inner(); + + let plaintext = resp + .plaintext + .ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + "AWS KMS generate_data_key response missing plaintext", + ) + })? + .into_inner(); + + Ok(GeneratedKey::new( + SensitiveBytes::new(plaintext), + wrapped_key, + )) } } @@ -182,7 +199,7 @@ mod tests { #[tokio::test] async fn test_wrap_key() { let encrypt_rule = mock!(Client::encrypt) - .match_requests(|req| req.key_id == Some(String::from("master-key-id"))) + .match_requests(|req| req.key_id.as_deref() == Some("master-key-id")) .then_output(|| { EncryptOutput::builder() .ciphertext_blob(Blob::from(b"hello world".to_vec())) @@ -191,7 +208,7 @@ mod tests { let client = mock_client!(aws_sdk_kms, [&encrypt_rule]); - let kms = AwsKeyManagementClient::new_with_client(client, DataKeySpec::Aes256); + let kms = AwsKeyManagementClient::new_with_client(client); assert_eq!( kms.wrap_key(b"123", "master-key-id").await.unwrap(), @@ -202,7 +219,7 @@ mod tests { #[tokio::test] async fn test_unwrap_key() { let decrypt_rule = mock!(Client::decrypt) - .match_requests(|req| req.key_id == Some(String::from("master-key-id"))) + .match_requests(|req| req.key_id.as_deref() == Some("master-key-id")) .then_output(|| { DecryptOutput::builder() .plaintext(Blob::from(b"hello world".to_vec())) @@ -211,10 +228,10 @@ mod tests { let client = mock_client!(aws_sdk_kms, [&decrypt_rule]); - let kms = AwsKeyManagementClient::new_with_client(client, DataKeySpec::Aes128); + let kms = AwsKeyManagementClient::new_with_client(client); assert_eq!( - kms.unwrap_key(BASE64.encode(b"123").as_bytes(), "master-key-id") + kms.unwrap_key(b"encrypted-blob", "master-key-id") .await .unwrap(), SensitiveBytes::new(b"hello world".to_vec()) diff --git a/crates/iceberg/src/encryption/kms/client.rs b/crates/iceberg/src/encryption/kms/client.rs index 85cd511758..c0c1ea40ad 100644 --- a/crates/iceberg/src/encryption/kms/client.rs +++ b/crates/iceberg/src/encryption/kms/client.rs @@ -22,7 +22,7 @@ use async_trait::async_trait; use crate::Result; -use crate::encryption::SensitiveBytes; +use crate::encryption::{AesKeySize, SensitiveBytes}; /// Result of a server-side key generation operation. /// @@ -71,7 +71,11 @@ pub trait KeyManagementClient: Send + Sync + std::fmt::Debug { /// /// This is only supported when [`supports_key_generation`](Self::supports_key_generation) /// returns `true`. - async fn generate_key(&self, wrapping_key_id: &str) -> Result; + async fn generate_key( + &self, + wrapping_key_id: &str, + key_size: AesKeySize, + ) -> Result; } #[async_trait] @@ -92,7 +96,11 @@ impl + Send + Sync + std::fmt::Debug> KeyManag self.as_ref().supports_key_generation() } - async fn generate_key(&self, wrapping_key_id: &str) -> Result { - self.as_ref().generate_key(wrapping_key_id).await + async fn generate_key( + &self, + wrapping_key_id: &str, + key_size: AesKeySize, + ) -> Result { + self.as_ref().generate_key(wrapping_key_id, key_size).await } } diff --git a/crates/iceberg/src/encryption/kms/memory.rs b/crates/iceberg/src/encryption/kms/memory.rs index 65319831dd..e3b7b166c8 100644 --- a/crates/iceberg/src/encryption/kms/memory.rs +++ b/crates/iceberg/src/encryption/kms/memory.rs @@ -168,7 +168,11 @@ impl KeyManagementClient for MemoryKeyManagementClient { false } - async fn generate_key(&self, _wrapping_key_id: &str) -> Result { + async fn generate_key( + &self, + _wrapping_key_id: &str, + _aes_key_size: AesKeySize, + ) -> Result { Err(Error::new( ErrorKind::FeatureUnsupported, "MemoryKeyManagementClient does not support server-side key generation", @@ -218,7 +222,7 @@ mod tests { let kms = MemoryKeyManagementClient::new(); assert!(!kms.supports_key_generation()); - let result = kms.generate_key("master-1").await; + let result = kms.generate_key("master-1", AesKeySize::Bits128).await; assert!(result.is_err()); } diff --git a/crates/iceberg/src/encryption/manager.rs b/crates/iceberg/src/encryption/manager.rs index a4c5b9c645..f003a49a80 100644 --- a/crates/iceberg/src/encryption/manager.rs +++ b/crates/iceberg/src/encryption/manager.rs @@ -214,7 +214,10 @@ impl EncryptionManager { /// the manager's `encryption_keys` map. async fn create_kek(&self) -> Result { let (plaintext_kek, wrapped_kek) = if self.kms_client.supports_key_generation() { - let result = self.kms_client.generate_key(&self.table_key_id).await?; + let result = self + .kms_client + .generate_key(&self.table_key_id, self.key_size) + .await?; (result.key().clone(), result.wrapped_key().to_vec()) } else { let plaintext_key = SecureKey::generate(self.key_size);