[DRAFT] PQC xwing (gemini generated) - #3212
Conversation
|
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 system by implementing ML-KEM and Hybrid X-Wing key wrapping. These changes involve updating the core cryptographic library, extending the KAS (Key Access Server) protocols, and modifying the SDK to handle these new algorithm types during key generation, encryption, and decryption processes. The update ensures the system is prepared for quantum-resistant key exchange while maintaining compatibility with existing workflows. Highlights
🧠 New Feature in Public Preview: You can now enable Memory 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 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 counter productive. 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. The quantum threat is drawing near, But lattice math will calm our fear. With X-Wing flight and KEM in hand, We guard the keys across the land. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request implements support for Post-Quantum Cryptography (PQC) algorithms, including ML-KEM (768 and 1024) and Hybrid X-Wing (X25519 + ML-KEM-768), across the crypto library, SDK, and KAS service. The feedback identifies critical bugs where multiple "%w" verbs are used in "fmt.Errorf" calls, which is invalid in Go. Furthermore, several suggestions are provided to reduce extensive code duplication in the new encryption, decryption, and key formatting logic by extracting shared helper functions and consolidating identical case blocks.
| decap1024, err1024 := mlkem.NewDecapsulationKey1024(block.Bytes) | ||
| if err1024 != nil { | ||
| return nil, fmt.Errorf("mlkem.NewDecapsulationKey1024 failed after mlkem.NewDecapsulationKey768 failed: %w / %w", err, err1024) | ||
| } |
There was a problem hiding this comment.
The use of multiple %w format specifiers in fmt.Errorf is not valid. In Go 1.20+ this is a compile-time error, and in older versions it can lead to a runtime panic. To combine multiple errors while preserving their types for inspection, you should use errors.Join.
| decap1024, err1024 := mlkem.NewDecapsulationKey1024(block.Bytes) | |
| if err1024 != nil { | |
| return nil, fmt.Errorf("mlkem.NewDecapsulationKey1024 failed after mlkem.NewDecapsulationKey768 failed: %w / %w", err, err1024) | |
| } | |
| decap1024, err1024 := mlkem.NewDecapsulationKey1024(block.Bytes) | |
| if err1024 != nil { | |
| return nil, fmt.Errorf("mlkem.NewDecapsulationKey failed for both 768 and 1024: %w", errors.Join(err, err1024)) | |
| } |
| encap1024, err1024 := mlkem.NewEncapsulationKey1024(block.Bytes) | ||
| if err1024 != nil { | ||
| return nil, fmt.Errorf("mlkem.NewEncapsulationKey1024 failed after mlkem.NewEncapsulationKey768 failed: %w / %w", err, err1024) | ||
| } |
There was a problem hiding this comment.
The use of multiple %w format specifiers in fmt.Errorf is not valid. In Go 1.20+ this is a compile-time error, and in older versions it can lead to a runtime panic. To combine multiple errors while preserving their types for inspection, you should use errors.Join.
| encap1024, err1024 := mlkem.NewEncapsulationKey1024(block.Bytes) | |
| if err1024 != nil { | |
| return nil, fmt.Errorf("mlkem.NewEncapsulationKey1024 failed after mlkem.NewEncapsulationKey768 failed: %w / %w", err, err1024) | |
| } | |
| encap1024, err1024 := mlkem.NewEncapsulationKey1024(block.Bytes) | |
| if err1024 != nil { | |
| return nil, fmt.Errorf("mlkem.NewEncapsulationKey failed for both 768 and 1024: %w", errors.Join(err, err1024)) | |
| } |
| func (d MLKEMDecryptor768) DecryptWithEphemeralKey(data, ephemeral []byte) ([]byte, error) { | ||
| if d.decap == nil { | ||
| return nil, errors.New("mlkem decapsulation key is nil") | ||
| } | ||
| if len(ephemeral) == 0 { | ||
| return nil, errors.New("ciphertext encapsulation is required for ML-KEM decryption") | ||
| } | ||
|
|
||
| sharedSecret, err := d.decap.Decapsulate(ephemeral) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("mlkem.Decapsulate failed: %w", err) | ||
| } | ||
|
|
||
| block, err := aes.NewCipher(sharedSecret) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("aes.NewCipher failure: %w", err) | ||
| } | ||
|
|
||
| gcm, err := cipher.NewGCM(block) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("cipher.NewGCM failure: %w", err) | ||
| } | ||
|
|
||
| nonceSize := gcm.NonceSize() | ||
| if len(data) < nonceSize { | ||
| return nil, errors.New("ciphertext too short") | ||
| } | ||
|
|
||
| nonce, ciphertext := data[:nonceSize], data[nonceSize:] | ||
| plaintext, err := gcm.Open(nil, nonce, ciphertext, nil) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("gcm.Open failure: %w", err) | ||
| } | ||
|
|
||
| return plaintext, nil | ||
| } |
There was a problem hiding this comment.
There is significant code duplication between MLKEMDecryptor768.DecryptWithEphemeralKey, MLKEMDecryptor1024.DecryptWithEphemeralKey, and HybridXWingDecryptorWrapper.DecryptWithEphemeralKey. The core AES-GCM decryption logic is identical in all of them.
To improve maintainability and reduce redundancy, I suggest extracting this common logic into a private helper function.
For example, you could create a function decryptAESGCM:
func decryptAESGCM(data, sharedSecret []byte) ([]byte, error) {
block, err := aes.NewCipher(sharedSecret)
if err != nil {
return nil, fmt.Errorf("aes.NewCipher failure: %w", err)
}
gcm, err := cipher.NewGCM(block)
if err != nil {
return nil, fmt.Errorf("cipher.NewGCM failure: %w", err)
}
nonceSize := gcm.NonceSize()
if len(data) < nonceSize {
return nil, errors.New("ciphertext too short")
}
nonce, ciphertext := data[:nonceSize], data[nonceSize:]
plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
if err != nil {
return nil, fmt.Errorf("gcm.Open failure: %w", err)
}
return plaintext, nil
}Then, this function and the other similar ones can be simplified to:
func (d MLKEMDecryptor768) DecryptWithEphemeralKey(data, ephemeral []byte) ([]byte, error) {
if d.decap == nil {
return nil, errors.New("mlkem decapsulation key is nil")
}
if len(ephemeral) == 0 {
return nil, errors.New("ciphertext encapsulation is required for ML-KEM decryption")
}
sharedSecret, err := d.decap.Decapsulate(ephemeral)
if err != nil {
return nil, fmt.Errorf("mlkem.Decapsulate failed: %w", err)
}
return decryptAESGCM(data, sharedSecret)
}This change would apply to MLKEMDecryptor1024.DecryptWithEphemeralKey and HybridXWingDecryptorWrapper.DecryptWithEphemeralKey as well.
| func (e MLKEMEncryptor768) Encrypt(data []byte) ([]byte, error) { | ||
| block, err := aes.NewCipher(e.sharedSecret) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("aes.NewCipher failed: %w", err) | ||
| } | ||
|
|
||
| gcm, err := cipher.NewGCM(block) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("cipher.NewGCM failed: %w", err) | ||
| } | ||
|
|
||
| nonce := make([]byte, gcm.NonceSize()) | ||
| if _, err := io.ReadFull(rand.Reader, nonce); err != nil { | ||
| return nil, fmt.Errorf("nonce generation failed: %w", err) | ||
| } | ||
|
|
||
| return gcm.Seal(nonce, nonce, data, nil), nil | ||
| } |
There was a problem hiding this comment.
There is significant code duplication between the Encrypt methods for MLKEMEncryptor768, MLKEMEncryptor1024, and HybridXWingEncryptorWrapper. The core AES-GCM encryption logic is identical.
To improve maintainability, I suggest extracting this common logic into a private helper function.
For example, you could create a function encryptAESGCM:
func encryptAESGCM(data, sharedSecret []byte) ([]byte, error) {
block, err := aes.NewCipher(sharedSecret)
if err != nil {
return nil, fmt.Errorf("aes.NewCipher failed: %w", err)
}
gcm, err := cipher.NewGCM(block)
if err != nil {
return nil, fmt.Errorf("cipher.NewGCM failed: %w", err)
}
nonce := make([]byte, gcm.NonceSize())
if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
return nil, fmt.Errorf("nonce generation failed: %w", err)
}
return gcm.Seal(nonce, nonce, data, nil), nil
}Then, this Encrypt method can be simplified to:
func (e MLKEMEncryptor768) Encrypt(data []byte) ([]byte, error) {
return encryptAESGCM(data, e.sharedSecret)
}This change would apply to MLKEMEncryptor1024.Encrypt and HybridXWingEncryptorWrapper.Encrypt as well.
| func (keyPair MLKEMKeyPair) PrivateKeyInPemFormat() (string, error) { | ||
| if keyPair.PrivateKey == nil { | ||
| return "", errors.New("failed to generate PEM formatted private key") | ||
| } | ||
|
|
||
| privateKeyPEM := pem.EncodeToMemory( | ||
| &pem.Block{ | ||
| Type: "MLKEM DECAPSULATION KEY", | ||
| Bytes: keyPair.PrivateKey.Bytes(), | ||
| }, | ||
| ) | ||
| return string(privateKeyPEM), nil | ||
| } | ||
|
|
||
| func (keyPair MLKEMKeyPair) PublicKeyInPemFormat() (string, error) { | ||
| if keyPair.PrivateKey == nil { | ||
| return "", errors.New("failed to generate PEM formatted public key") | ||
| } | ||
|
|
||
| publicKeyPEM := pem.EncodeToMemory( | ||
| &pem.Block{ | ||
| Type: "MLKEM ENCAPSULATOR", | ||
| Bytes: keyPair.PrivateKey.EncapsulationKey().Bytes(), | ||
| }, | ||
| ) | ||
| return string(publicKeyPEM), nil | ||
| } |
There was a problem hiding this comment.
The PrivateKeyInPemFormat and PublicKeyInPemFormat methods for MLKEMKeyPair and MLKEM1024KeyPair are nearly identical, leading to code duplication.
To improve maintainability, you could use helper functions that operate on an interface. Both *mlkem.DecapsulationKey768 and *mlkem.DecapsulationKey1024 have Bytes() and EncapsulationKey().Bytes() methods, so they can satisfy a common interface.
Example refactoring:
type mlkemPrivateKey interface {
Bytes() []byte
EncapsulationKey() interface {
Bytes() []byte
}
}
func mlkemPrivateKeyInPemFormat(key mlkemPrivateKey) (string, error) {
if key == nil {
return "", errors.New("failed to generate PEM formatted private key")
}
privateKeyPEM := pem.EncodeToMemory(&pem.Block{
Type: "MLKEM DECAPSULATION KEY",
Bytes: key.Bytes(),
})
return string(privateKeyPEM), nil
}
func mlkemPublicKeyInPemFormat(key mlkemPrivateKey) (string, error) {
if key == nil {
return "", errors.New("failed to generate PEM formatted public key")
}
publicKeyPEM := pem.EncodeToMemory(&pem.Block{
Type: "MLKEM ENCAPSULATOR",
Bytes: key.EncapsulationKey().Bytes(),
})
return string(publicKeyPEM), nil
}
// Then the methods become one-liners:
func (keyPair MLKEMKeyPair) PrivateKeyInPemFormat() (string, error) {
return mlkemPrivateKeyInPemFormat(keyPair.PrivateKey)
}
func (keyPair MLKEMKeyPair) PublicKeyInPemFormat() (string, error) {
return mlkemPublicKeyInPemFormat(keyPair.PrivateKey)
}This would make the code more DRY and easier to maintain.
| func generateWrapKeyWithMLKEM(publicKey string, symKey []byte) (string, string, error) { | ||
| publicKeyEncryptor, err := ocrypto.FromPublicPEM(publicKey) | ||
| if err != nil { | ||
| return "", "", fmt.Errorf("generateWrapKeyWithMLKEM: ocrypto.FromPublicPEM failed:%w", err) | ||
| } | ||
|
|
||
| wrappedKey, err := publicKeyEncryptor.Encrypt(symKey) | ||
| if err != nil { | ||
| return "", "", fmt.Errorf("generateWrapKeyWithMLKEM: encrypt failed:%w", err) | ||
| } | ||
|
|
||
| encapsulatedKey := publicKeyEncryptor.EphemeralKey() | ||
| if len(encapsulatedKey) == 0 { | ||
| return "", "", errors.New("generateWrapKeyWithMLKEM: encapsulated key missing") | ||
| } | ||
|
|
||
| return string(ocrypto.Base64Encode(wrappedKey)), string(ocrypto.Base64Encode(encapsulatedKey)), nil | ||
| } | ||
|
|
||
| func generateWrapKeyWithHybrid(publicKey string, symKey []byte) (string, string, error) { | ||
| publicKeyEncryptor, err := ocrypto.FromPublicPEM(publicKey) | ||
| if err != nil { | ||
| return "", "", fmt.Errorf("generateWrapKeyWithHybrid: ocrypto.FromPublicPEM failed:%w", err) | ||
| } | ||
|
|
||
| wrappedKey, err := publicKeyEncryptor.Encrypt(symKey) | ||
| if err != nil { | ||
| return "", "", fmt.Errorf("generateWrapKeyWithHybrid: encrypt failed:%w", err) | ||
| } | ||
|
|
||
| encapsulatedKey := publicKeyEncryptor.EphemeralKey() | ||
| if len(encapsulatedKey) == 0 { | ||
| return "", "", errors.New("generateWrapKeyWithHybrid: encapsulated key missing") | ||
| } | ||
|
|
||
| return string(ocrypto.Base64Encode(wrappedKey)), string(ocrypto.Base64Encode(encapsulatedKey)), nil | ||
| } |
There was a problem hiding this comment.
The functions generateWrapKeyWithMLKEM and generateWrapKeyWithHybrid are almost identical. The only difference is the name used in the error message. This code duplication can be avoided by merging them into a single function that takes the algorithm name as a parameter for error reporting.
For example, a single function generateWrapKeyWithKEM could be created:
func generateWrapKeyWithKEM(publicKey string, symKey []byte, name string) (string, string, error) {
publicKeyEncryptor, err := ocrypto.FromPublicPEM(publicKey)
if err != nil {
return "", "", fmt.Errorf("generateWrapKeyWithKEM for %s: ocrypto.FromPublicPEM failed:%w", name, err)
}
wrappedKey, err := publicKeyEncryptor.Encrypt(symKey)
if err != nil {
return "", "", fmt.Errorf("generateWrapKeyWithKEM for %s: encrypt failed:%w", name, err)
}
encapsulatedKey := publicKeyEncryptor.EphemeralKey()
if len(encapsulatedKey) == 0 {
return "", "", fmt.Errorf("generateWrapKeyWithKEM for %s: encapsulated key missing", name)
}
return string(ocrypto.Base64Encode(wrappedKey)), string(ocrypto.Base64Encode(encapsulatedKey)), nil
}You can then call this from createKeyAccess for both ML-KEM and Hybrid cases, passing the appropriate name.
| case ocrypto.MLKEM768Key, ocrypto.MLKEM1024Key: | ||
| plaintext, err := decrypter.DecryptWithEphemeralKey(ciphertext, ephemeralPublicKey) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to decrypt with ML-KEM: %w", err) | ||
| } | ||
| protectedKey, err := ocrypto.NewAESProtectedKey(plaintext) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to create protected key: %w", err) | ||
| } | ||
| return protectedKey, nil | ||
| case ocrypto.HybridXWingKey: | ||
| plaintext, err := decrypter.DecryptWithEphemeralKey(ciphertext, ephemeralPublicKey) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to decrypt with Hybrid X-Wing: %w", err) | ||
| } | ||
| protectedKey, err := ocrypto.NewAESProtectedKey(plaintext) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to create protected key: %w", err) | ||
| } | ||
| return protectedKey, nil |
There was a problem hiding this comment.
The logic inside the case blocks for MLKEM768Key, MLKEM1024Key and HybridXWingKey is identical. You can combine these cases to reduce code duplication and improve readability.
| case ocrypto.MLKEM768Key, ocrypto.MLKEM1024Key: | |
| plaintext, err := decrypter.DecryptWithEphemeralKey(ciphertext, ephemeralPublicKey) | |
| if err != nil { | |
| return nil, fmt.Errorf("failed to decrypt with ML-KEM: %w", err) | |
| } | |
| protectedKey, err := ocrypto.NewAESProtectedKey(plaintext) | |
| if err != nil { | |
| return nil, fmt.Errorf("failed to create protected key: %w", err) | |
| } | |
| return protectedKey, nil | |
| case ocrypto.HybridXWingKey: | |
| plaintext, err := decrypter.DecryptWithEphemeralKey(ciphertext, ephemeralPublicKey) | |
| if err != nil { | |
| return nil, fmt.Errorf("failed to decrypt with Hybrid X-Wing: %w", err) | |
| } | |
| protectedKey, err := ocrypto.NewAESProtectedKey(plaintext) | |
| if err != nil { | |
| return nil, fmt.Errorf("failed to create protected key: %w", err) | |
| } | |
| return protectedKey, nil | |
| case ocrypto.MLKEM768Key, ocrypto.MLKEM1024Key, ocrypto.HybridXWingKey: | |
| plaintext, err := decrypter.DecryptWithEphemeralKey(ciphertext, ephemeralPublicKey) | |
| if err != nil { | |
| return nil, fmt.Errorf("failed to decrypt with KEM: %w", err) | |
| } | |
| protectedKey, err := ocrypto.NewAESProtectedKey(plaintext) | |
| if err != nil { | |
| return nil, fmt.Errorf("failed to create protected key: %w", err) | |
| } | |
| return protectedKey, nil |
| case string(ocrypto.MLKEM): | ||
| ephemeralCiphertext := kao.GetKeyAccessObject().GetEphemeralPublicKey() | ||
| if ephemeralCiphertext == "" { | ||
| p.Logger.WarnContext(ctx, "missing encapsulated key for ml-kem rewrap") | ||
| failedKAORewrap(results, kao, err400("bad request")) | ||
| continue | ||
| } | ||
|
|
||
| encapsulatedKey, err := ocrypto.Base64Decode([]byte(ephemeralCiphertext)) | ||
| if err != nil { | ||
| p.Logger.WarnContext(ctx, "failed to decode encapsulated key for ml-kem rewrap", slog.Any("error", err)) | ||
| failedKAORewrap(results, kao, err400("bad request")) | ||
| continue | ||
| } | ||
|
|
||
| kid := trust.KeyIdentifier(kao.GetKeyAccessObject().GetKid()) | ||
| dek, err = p.KeyDelegator.Decrypt(ctx, kid, kao.GetKeyAccessObject().GetWrappedKey(), encapsulatedKey) | ||
| if err != nil { | ||
| p.Logger.WarnContext(ctx, "failed to decrypt ML-KEM key", slog.Any("error", err)) | ||
| failedKAORewrap(results, kao, err400("bad request")) | ||
| continue | ||
| } | ||
| case string(ocrypto.Hybrid): | ||
| ephemeralCiphertext := kao.GetKeyAccessObject().GetEphemeralPublicKey() | ||
| if ephemeralCiphertext == "" { | ||
| p.Logger.WarnContext(ctx, "missing encapsulated key for hybrid rewrap") | ||
| failedKAORewrap(results, kao, err400("bad request")) | ||
| continue | ||
| } | ||
|
|
||
| encapsulatedKey, err := ocrypto.Base64Decode([]byte(ephemeralCiphertext)) | ||
| if err != nil { | ||
| p.Logger.WarnContext(ctx, "failed to decode encapsulated key for hybrid rewrap", slog.Any("error", err)) | ||
| failedKAORewrap(results, kao, err400("bad request")) | ||
| continue | ||
| } | ||
|
|
||
| kid := trust.KeyIdentifier(kao.GetKeyAccessObject().GetKid()) | ||
| dek, err = p.KeyDelegator.Decrypt(ctx, kid, kao.GetKeyAccessObject().GetWrappedKey(), encapsulatedKey) | ||
| if err != nil { | ||
| p.Logger.WarnContext(ctx, "failed to decrypt hybrid key", slog.Any("error", err)) | ||
| failedKAORewrap(results, kao, err400("bad request")) | ||
| continue | ||
| } |
There was a problem hiding this comment.
The logic for handling ocrypto.MLKEM and ocrypto.Hybrid key types is identical. These two case blocks can be combined to reduce code duplication and improve maintainability.
case string(ocrypto.MLKEM), string(ocrypto.Hybrid):
ephemeralCiphertext := kao.GetKeyAccessObject().GetEphemeralPublicKey()
if ephemeralCiphertext == "" {
p.Logger.WarnContext(ctx, "missing encapsulated key for KEM rewrap")
failedKAORewrap(results, kao, err400("bad request"))
continue
}
encapsulatedKey, err := ocrypto.Base64Decode([]byte(ephemeralCiphertext))
if err != nil {
p.Logger.WarnContext(ctx, "failed to decode encapsulated key for KEM rewrap", slog.Any("error", err))
failedKAORewrap(results, kao, err400("bad request"))
continue
}
kid := trust.KeyIdentifier(kao.GetKeyAccessObject().GetKid())
dek, err = p.KeyDelegator.Decrypt(ctx, kid, kao.GetKeyAccessObject().GetWrappedKey(), encapsulatedKey)
if err != nil {
p.Logger.WarnContext(ctx, "failed to decrypt KEM key", slog.Any("error", err))
failedKAORewrap(results, kao, err400("bad request"))
continue
}Signed-off-by: David Mihalcik <dmihalcik@virtru.com>
Signed-off-by: David Mihalcik <dmihalcik@virtru.com>
This commit introduces support for the Hybrid X-Wing KEM scheme as defined in draft-connolly-cfrg-xwing-kem-10. - Added HybridXWingKey (hpqt:xwing) to lib/ocrypto. - Implemented HybridXWingEncryptor/Decryptor using ASN.1 for composite keys and ciphertexts. - Integrated hybrid scheme support into KAS rewrap and Policy service. - Added 'hybrid' scheme type to SDK and manifest schema. - Added comprehensive unit tests for Hybrid X-Wing round-trips. Signed-off-by: Dave Mihalcik <dmihalcik@virtru.com>
5802096 to
493092f
Compare
|
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
|
Merged in #3276 |
Proposed Changes
Checklist
Testing Instructions