Repository navigation
feat(parquet/encryption): KMS support for reading PKMT1 key material - #1478
EdwardShen125 wants to merge 2 commits into
Conversation
New optional parquet/encryption package with a KmsClient interface and a KmsKeyRetriever that plugs into parquet.WithKeyRetriever, so files written by the parquet-cpp/pyarrow CryptoFactory can be read: PKMT1 key metadata is parsed, and DEKs are resolved via the KMS client (single wrapping) or by unwrapping the KEK through the KMS and decrypting the locally wrapped DEK with AES-GCM (double wrapping, AAD = decoded KEK ID bytes, matching parquet-cpp EncryptKeyLocally: nonce || ciphertext || tag). Interop tests against pyarrow CryptoFactory output cover both wrapping modes; the test KMS wraps with plain base64 so no extra python dependencies are needed and the existing cgo python CI job runs them unchanged. Also addresses the review nits from apache#1439 (renamed the requirePyarrow flag, dropped the unused footer key metadata, deduped a doc comment). Part of apache#566
zeroshade
left a comment
zeroshade
left a comment
There was a problem hiding this comment.
Thanks for this! I ran the unit tests and the interop tests locally against pyarrow 26.0.1 with PARQUET_TEST_REQUIRE_PYARROW=1, and both single and double wrapping read correctly, so the local unwrapping format matches. A few things before merging:
-
Dead fallback in
encryption_interop_pyarrow.py.CryptoFactory(kms_client_factory)has been the constructor since pyarrow added KMS support, and pyarrow has noregister_kms_client_factorymethod (I checked 26.0.1), so theexcept TypeErrorbranch could never work and the "pyarrow >= 26" comment is misleading. Please call the constructor directly. -
Example doc. "Without this package, files whose key metadata carries PKMT1 key material cannot be read" isn't quite true: a hand-written
DecryptionKeyRetrievercan read them. Maybe something like "This package handles the PKMT1 parsing and unwrapping that would otherwise need a custom retriever." -
Caching unwrapped KEKs. Within one file the reader already calls
GetKeyonce per column, but across files every open calls the KMS once per encrypted column plus the footer. parquet-cpp caches unwrapped KEKs by KEK ID (with an expiry), and that cache is what makes double wrapping worth being the default: reading many files costs one KMS call per KEK rather than one per column per file. Without it, double wrapping makes as many KMS calls as single wrapping. Since aKmsKeyRetrieveris meant to be shared across files, a small concurrency-safe cache keyed by(masterKeyID, keyEncryptionKeyID)would be good to have here.
Minor, non-blocking:
- In
doc.go, "carries no effect on code that does not import it" reads oddly; "has no effect" or just dropping the sentence would be fine. GetKeyreturning""means KMS errors surface as the footer panic or the "hidden column" error. TheDecryptionKeyRetrieverinterface can't carry an error, so pointing toUnwrapKeyMetadatain the doc is reasonable.
The retriever now caches unwrapped KEKs, so a file with double wrapping costs one KMS call per KEK instead of one per encrypted column plus the footer, matching parquet-cpp. Entries live for 10 minutes (a rotated KEK written under the same ID becomes visible after that) and the cache is bounded at 64 keys with LRU eviction; concurrent misses for the same KEK share a single KMS call. Also addresses the review feedback on apache#1478: - the interop script calls the CryptoFactory constructor directly; pyarrow has no register_kms_client_factory, so the fallback branch could never run - the example no longer claims PKMT1 files cannot be read without this package; it says the package handles the parsing and unwrapping that would otherwise need a custom retriever - the doc.go sentence about code that does not import the package is dropped, the preceding paragraph already says it The interop test now counts KMS unwrap calls: the generated files have three key slots under two master keys (pyarrow generates one KEK per master key), so double wrapping mode resolves them in two calls instead of three, single wrapping mode in three. On the previous commit the double wrapping case makes three calls and the test fails.
|
Pushed a fix. KEKs are cached per (masterKeyID, keyEncryptionKeyID) with a 10 minute expiry, LRU eviction, and one KMS call for concurrent misses. The interop script calls the CryptoFactory constructor directly now, and the example and doc.go wording are cleaned up. pyarrow generates one KEK per master key, not one per file, so the interop test counts unwrap calls now: three key slots under two master keys resolve in two calls, three on the previous commit. |
zeroshade
left a comment
zeroshade
left a comment
There was a problem hiding this comment.
Thanks, everything from my earlier review is addressed. I re-ran the unit tests with -race, and the interop tests against pyarrow 26.0.1 with PARQUET_TEST_REQUIRE_PYARROW=1; both wrapping modes pass. I also checked that the new unwrap count depends on the cache: with the previous commit's kms.go swapped back in, kms_double fails with expected 2, actual 3.
The cache looks right to me. Concurrent misses share one KMS call, the shared call re-checks the cache and stores the KEK before returning, and errors aren't cached. A crafted file that reuses a KEK ID can only get the cached KEK, and its own DEK then fails GCM authentication.
Two optional nits inline.
| } | ||
|
|
||
| // put caches kek for key, replacing any previous entry. | ||
| func (c *kekCache) put(k kekCacheKey, kek []byte) { |
There was a problem hiding this comment.
nit: this stores the slice the KmsClient returned as-is. If a client reuses or zeroes its returned buffers (some do, to wipe key material), the cached KEK gets corrupted. bytes.Clone(kek) here is cheap insurance.
| // NewKmsKeyRetriever returns a retriever that resolves keys through client. | ||
| // The retriever is safe for concurrent use and meant to be shared across | ||
| // files. | ||
| func NewKmsKeyRetriever(client KmsClient) *KmsKeyRetriever { |
There was a problem hiding this comment.
nit: the cache lasts as long as the retriever, which is right because a retriever wraps a single KmsClient identity. parquet-cpp scopes its cache per access token instead. A sentence saying to share a retriever only among callers with the same KMS access would make that explicit. Making the expiry time and size limit configurable can wait for a follow-up if anyone needs it.
Rationale for this change
#566 reported that files written by pyarrow's CryptoFactory cannot be read by arrow-go: the footer and column key metadata carry PKMT1 key material with a wrapped DEK, and arrow-go had no parser or unwrapping for it, so footer decryption failed without a hand-written
KeyRetriever. As agreed in #566, the KMS support lives in the library as an optional component.What changes are included in this PR?
A new optional
parquet/encryptionpackage:KmsClientinterface mirroring the parquet-cpp/pyarrow contract (WrapKey/UnwrapKey)KeyMaterial/ParseKeyMetadatafor PKMT1 key metadata (internal key storage)KmsKeyRetrieverimplementingparquet.DecryptionKeyRetriever: in single wrapping mode the DEK is unwrapped by theKmsClient; in double wrapping mode the KEK is unwrapped through the client and the DEK is then decrypted locally with AES-GCM (nonce || ciphertext || tag, AAD = the base64 decoded KEK ID bytes), byte-compatible with parquet-cpp'sEncryptKeyLocallyIt plugs into the existing
parquet.WithKeyRetrieveroption; code that does not import the package is unaffected. The package doc and a godoc example (ExampleKmsKeyRetriever) show the wiring and state the optional nature explicitly.The interop harness from #1439 gains a
gen-kmsmode: pyarrow's CryptoFactory writes double-wrapped (the default) and single-wrapped files using a test KMS client whose wrapping is plain base64, and Go reads both through the retriever — the footer key and a column key both take that path. No new python dependencies, so the existing cgo python CI job runs them unchanged.The review nits from #1439 are addressed (flag rename, unused footer key metadata dropped, doc comment dedupe).
The write side (Go producing PKMT1 key material via a CryptoFactory) is a natural follow-up; this PR is read-only.
Are these changes tested?
Unit tests cover PKMT1 parsing (valid and error cases), single and double unwrap round-trips, and tamper detection. Against pyarrow 26.0.0 locally, with
-raceandPARQUET_TEST_REQUIRE_PYARROW=1, both interop modes pass; without pyarrow everything skips. The interop tests also count KMS unwrap calls: three key slots under two master keys resolve in two calls in double wrapping mode, three in single wrapping mode.Are there any user-facing changes?
New optional public package
parquet/encryption. No changes to existing APIs.