feat(platform): add ML-KEM-768 post-quantum key encapsulation (DSPX-2399)#3491
feat(platform): add ML-KEM-768 post-quantum key encapsulation (DSPX-2399)#3491dmihalcik-virtru wants to merge 1 commit into
Conversation
…399)
- Add MLKEMKeyPair to lib/ocrypto using Go 1.24+ crypto/mlkem stdlib
- Wire format: wrappedKey = ciphertext[1088B] || AES-256-GCM(DEK)
- Key derivation via HKDF-SHA256 with SHA-256("TDF") salt
- Add KAS_PUBLIC_KEY_ALG_ENUM_MLKEM_768 = 13 to objects.proto; regen pb.go
- Implement decapsulation in BasicManager.Decrypt and rewrap.go
- Gate behind preview.mlkem_enabled flag; enabled in opentdf-dev.yaml
- Add algorithm mappings in sdk/granter.go, sdk/basekey.go, key_indexer.go
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Dave Mihalcik <dmihalcik@virtru.com>
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Repository UI Review profile: ASSERTIVE Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces post-quantum cryptographic support to the platform by implementing the ML-KEM-768 key encapsulation mechanism. The changes encompass the core cryptographic logic, updates to protocol definitions, and integration across the SDK and KAS services to ensure compatibility with the new algorithm. The implementation is safely gated behind a preview flag to allow for controlled rollout and testing. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Ignored Files
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Quantum threats are on the rise, ML-KEM provides the prize. Wrapped in GCM, the keys are tight, Securing data through the night. Footnotes
|
github-actions
Bot
added
comp:policy
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for post-quantum cryptography using the ML-KEM-768 algorithm. Key changes include the implementation of ML-KEM key pair generation, PEM-encoded serialization, and a new key wrapping/unwrapping mechanism utilizing HKDF-SHA256 and AES-256-GCM. The functionality is integrated across the SDK and KAS service, with the rewrap logic gated behind a new configuration flag. Review feedback highlights several improvement opportunities: strengthening cryptographic validation by checking PEM block types and AES-GCM ciphertext lengths, optimizing performance by precomputing static salts, and improving efficiency by caching key lookups during request processing.
| if len(data) < nonceSize { | ||
| return nil, errors.New("ciphertext too short for nonce") | ||
| } |
There was a problem hiding this comment.
The length check for the ciphertext is insufficient. It should ensure the data is long enough to contain both the nonce and the GCM authentication tag (typically 12 + 16 = 28 bytes).
| if len(data) < nonceSize { | |
| return nil, errors.New("ciphertext too short for nonce") | |
| } | |
| nonceSize := gcm.NonceSize() | |
| if len(data) < nonceSize+gcm.Overhead() { | |
| return nil, errors.New("ciphertext too short") | |
| } |
| block, _ := pem.Decode(privateKeyPEM) | ||
| if block == nil { | ||
| return nil, errors.New("failed to parse ML-KEM-768 PEM private key") | ||
| } |
There was a problem hiding this comment.
The PEM decoding logic should validate the Type of the PEM block to ensure it matches the expected ML-KEM-768 private key header. This prevents accidental processing of incorrect key types.
| block, _ := pem.Decode(privateKeyPEM) | |
| if block == nil { | |
| return nil, errors.New("failed to parse ML-KEM-768 PEM private key") | |
| } | |
| block, _ := pem.Decode(privateKeyPEM) | |
| if block == nil || block.Type != "ML-KEM-768 PRIVATE KEY" { | |
| return nil, errors.New("failed to parse ML-KEM-768 PEM private key") | |
| } |
| // mlkemTDFSalt returns the SHA-256("TDF") salt used for HKDF in ML-KEM key wrapping. | ||
| func mlkemTDFSalt() []byte { | ||
| h := sha256.New() | ||
| h.Write([]byte("TDF")) | ||
| return h.Sum(nil) | ||
| } |
There was a problem hiding this comment.
Computing the SHA-256 hash of the "TDF" string on every call is inefficient as the result is constant. Consider precomputing this salt as a package-level variable.
| // mlkemTDFSalt returns the SHA-256("TDF") salt used for HKDF in ML-KEM key wrapping. | |
| func mlkemTDFSalt() []byte { | |
| h := sha256.New() | |
| h.Write([]byte("TDF")) | |
| return h.Sum(nil) | |
| } | |
| var mlkemTDFSaltValue = sha256.Sum256([]byte("TDF")) | |
| // mlkemTDFSalt returns the SHA-256("TDF") salt used for HKDF in ML-KEM key wrapping. | |
| func mlkemTDFSalt() []byte { | |
| return mlkemTDFSaltValue[:] | |
| } |
| block, _ := pem.Decode(publicKeyPEM) | ||
| if block == nil { | ||
| return nil, errors.New("failed to parse ML-KEM-768 PEM public key") | ||
| } |
There was a problem hiding this comment.
The PEM decoding logic should validate the Type of the PEM block to ensure it matches the expected ML-KEM-768 public key header.
| block, _ := pem.Decode(publicKeyPEM) | |
| if block == nil { | |
| return nil, errors.New("failed to parse ML-KEM-768 PEM public key") | |
| } | |
| block, _ := pem.Decode(publicKeyPEM) | |
| if block == nil || block.Type != "ML-KEM-768 PUBLIC KEY" { | |
| return nil, errors.New("failed to parse ML-KEM-768 PEM public key") | |
| } |
| keyDetails, lookupErr := p.KeyDelegator.FindKeyByID(ctx, kid) | ||
| if lookupErr == nil && keyDetails.Algorithm() == ocrypto.MLKEM768Key { |
There was a problem hiding this comment.
Performing a FindKeyByID lookup inside the KAO loop can be inefficient if multiple Key Access Objects refer to the same key ID, as it may trigger redundant database or network calls. Consider caching the algorithm lookup results for the duration of the request processing.
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
1 participant
Summary
Implements pure ML-KEM-768 post-quantum key encapsulation mechanism support in the platform service.
MLKEMKeyPairtolib/ocryptousing Go 1.24+crypto/mlkemstandard librarywrappedKey = ciphertext[1088B] || AES-256-GCM(DEK)with HKDF-SHA256(shared_secret, SHA-256("TDF")) key derivation; noephemeralPublicKeyfieldKAS_PUBLIC_KEY_ALG_ENUM_MLKEM_768 = 13toservice/policy/objects.protoand regenerates pb.go stubsBasicManager.Decryptand gates processing inrewrap.gobehindpreview.mlkem_enabledmlkem_enabled: trueinopentdf-dev.yamlfor local development and test harnesssdk/granter.go,sdk/basekey.go,service/kas/key_indexer.go, andsdk/experimental/tdf/keysplit/attributes.goJira
https://virtru.atlassian.net/browse/DSPX-2399
Integration tests
xtest/scenarios/mechanism-mlkem.yaml
Test plan
lib/ocryptounit tests pass (go test ./...)sdkunit tests pass (go test ./...)service/kasandservice/internal/securitytests pass🤖 Generated with Claude Code