feat(core): Deprecate ECKeyPair in favor of asym_*#3293
feat(core): Deprecate ECKeyPair in favor of asym_*#3293dmihalcik-virtru wants to merge 6 commits intomainfrom
Conversation
📝 WalkthroughWalkthroughThis pull request refactors the cryptographic key handling APIs across the ocrypto package and dependent SDK modules. It replaces the Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsTimed out fetching pipeline failures after 30000ms 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 refactors the cryptographic key management system by deprecating the 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. 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. Old key pairs now must fade, New interfaces take their place, Decryption's path is made. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request introduces the PrivateKeyDecryptor interface to unify RSA and EC private key operations, such as decryption and shared secret derivation, across the ocrypto package and SDK. It refactors RsaKeyPair, ECKeyPair, and KASClient to implement this interface while deprecating legacy functions. Feedback suggests that several nil checks in the KeyType methods are redundant and should be removed to simplify the code and avoid hiding initialization issues.
| if asymDecryption.PrivateKey == nil { | ||
| return KeyType("rsa:[unknown]") | ||
| } |
There was a problem hiding this comment.
The check for asymDecryption.PrivateKey == nil is redundant here because PrivateKey is a field of the struct and the method is called on an instance. If the instance is initialized, this field should be checked at the point of creation or usage, not inside every getter method. This adds unnecessary complexity.
There was a problem hiding this comment.
yeah, that would make sense.. unfortunately FromRSA and the corresponding SDK WithSession...RSA options don't return error so it would be hard to wire that in at the moment
| if e.sk == nil { | ||
| return KeyType("ec:[unknown]") | ||
| } |
| if keyPair.PrivateKey == nil { | ||
| return KeyType("ec:[unknown]") | ||
| } |
| if keyPair.privateKey == nil { | ||
| return KeyType("rsa:[unknown]") | ||
| } |
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
- Remove DeriveSharedKey from PrivateKeyDecryptor interface, eliminating false mutual-exclusion between RSA and EC implementations - Replace SDK encrypt-side DeriveSharedKey with FromPublicPEMWithSalt+Encrypt in generateWrapKeyWithEC (tdf.go) - Replace SDK decrypt-side DeriveSharedKey with DecryptWithEphemeralKey in handleECKeyResponse/processECResponse (kas_client.go) - Service-side and experimental XOR callers type-assert to ECDecryptor and call DeriveSharedKey on the concrete type directly - Fix silent failure chain in newECIES, EphemeralKey, and Metadata - Fix convCurve to return error instead of nil (prevents nil-panic) - Extract parseEphemeralPublicKey helper to fix nestif lint violation - Fix ECKeyPair.Decrypt/Public/DeriveSharedKey to use NewECDecryptor so TDF salt is propagated correctly through the shim - Add asym_decryption_test.go with coverage for new interface and methods - Update ec_key_pair_test.go and ec_decrypt_compressed_test.go to use non-deprecated APIs Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Dave Mihalcik <dmihalcik@virtru.com>
66d1f7d to
08055b2
Compare
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
|
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
sdk/tdf.go (1)
676-686:⚠️ Potential issue | 🟠 MajorDerive the wrap scheme from the parsed public key, not
kasInfo.Algorithm.Line 677 still trusts metadata to choose the EC/RSA path. If
kasInfo.Algorithmis empty or stale whilekasInfo.PublicKeyis actually EC, this falls into the RSA branch,generateWrapKeyWithRSAwill still encrypt with the EC encryptor returned byFromPublicPEM, and the manifest is emitted aswrappedwithout anEphemeralPublicKey. That key-access object is not rewrappable later.Please branch on the parsed key/encryptor type instead, or at least fail fast when the PEM type and declared algorithm disagree.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@sdk/tdf.go` around lines 676 - 686, The branch decision should be based on the parsed public key/encryptor type rather than trusting kasInfo.Algorithm; parse kasInfo.PublicKey (via the same flow that returns an encryptor from FromPublicPEM) and inspect its key type (EC vs RSA) before choosing generateWrapKeyWithEC or generateWrapKeyWithRSA, or else return an error if the parsed key type conflicts with kasInfo.Algorithm; update the logic around ktype/kasInfo.Algorithm, generateWrapKeyWithEC, generateWrapKeyWithRSA, FromPublicPEM, and the keyAccess fields (KeyType, WrappedKey, EphemeralPublicKey) so the manifest correctly records an EphemeralPublicKey for EC keys and prevents creating non-rewrappable entries when the PEM and declared algorithm disagree.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@lib/ocrypto/asym_decryption_test.go`:
- Around line 136-153: The TestAsymDecryptionKeyType loop should be converted to
subtests so failures show which RSA size failed: for each bits in {RSA2048Size,
RSA4096Size} call t.Run with a descriptive name (e.g. "RSA2048"/"RSA4096" or
fmt.Sprintf("%d", bits)), capture bits in the closure, and move t.Parallel()
into the subtest body; inside the subtest keep the existing logic
(NewRSAKeyPair, PrivateKeyInPemFormat, FromPrivatePEM, type assertion to
AsymDecryption and the KeyType checks against RSA2048Key/RSA4096Key).
In `@lib/ocrypto/asym_encryption.go`:
- Around line 288-293: Define a small interface (e.g.,
EphemeralPublicKeyExporter) that declares EphemeralPublicKeyInPemFormat()
(string, error) and have ECEncryptor implicitly implement it (the method already
exists), then update the public API to return or accept that interface where
callers need to export the ephemeral PEM instead of requiring a concrete
ocrypto.ECEncryptor; ensure references to PublicKeyEncryptor remain unchanged
unless the API should explicitly include the new interface, and update call
sites to use the new EphemeralPublicKeyExporter interface for PEM export without
downcasting to ECEncryptor.
In `@lib/ocrypto/ec_key_pair.go`:
- Around line 480-488: The two exported helpers (ECPrivateKeyInPemFormat,
ECPublicKeyInPemFormat) pass value-type ecdsa.PrivateKey/ecdsa.PublicKey to
privateKeyInPemFormat which only handles pointer cases, so zero-value structs
bypass nil checks; update privateKeyInPemFormat to add switch cases for the
value types (case ecdsa.PrivateKey: and case ecdsa.PublicKey:) and perform the
same validity checks as for pointers (e.g., for ecdsa.PrivateKey verify key.D !=
nil and for ecdsa.PublicKey verify key.X != nil && key.Y != nil), returning the
same error ("failed to generate PEM formatted private key"/appropriate public
key error) when those fields are nil to prevent falling through to
x509.MarshalPKCS8PrivateKey with a zero-value.
In `@sdk/experimental/tdf/key_access.go`:
- Around line 186-201: The EC wrapper currently returns the non-canonical string
"eccWrapped" from wrapKeyWithEC; update the return to use the canonical key type
"ec-wrapped" so it matches the rest of the codebase (see wrapKeyWithEC in
key_access.go and the canonical uses in sdk/tdf.go and
lib/ocrypto/asym_encryption.go), keeping the other return values and error
handling unchanged.
- Around line 165-182: The code currently chooses the wrapping branch using
pubKeyInfo.Algorithm (ktype) even after parsing the PEM; instead, base the
decision on the parsed kasPublicKey's runtime type: after
ocrypto.FromPublicPEM(...) check if kasPublicKey implements ocrypto.ECEncryptor
(or use a type switch) and call wrapKeyWithEC(ktype, epk, symKey) when it does,
otherwise call wrapKeyWithRSA(kasPublicKey, symKey); update error paths to
report a clear mismatch if the parsed key cannot perform the expected operations
and remove reliance on pubKeyInfo.Algorithm for branching.
In `@sdk/kas_client.go`:
- Around line 196-208: In handleECKeyResponse, explicitly validate that
response.GetSessionPublicKey() is not empty before calling pem.Decode so you can
return a clearer error; check the returned string from
kas.RewrapResponse.GetSessionPublicKey(), if it's empty return a descriptive
error like "empty KAS session public key" and only then call pem.Decode,
preserving the existing session key type assertion (k.sessionKey as
ocrypto.ECDecryptor) and the subsequent call to k.processECResponse with
block.Bytes.
---
Outside diff comments:
In `@sdk/tdf.go`:
- Around line 676-686: The branch decision should be based on the parsed public
key/encryptor type rather than trusting kasInfo.Algorithm; parse
kasInfo.PublicKey (via the same flow that returns an encryptor from
FromPublicPEM) and inspect its key type (EC vs RSA) before choosing
generateWrapKeyWithEC or generateWrapKeyWithRSA, or else return an error if the
parsed key type conflicts with kasInfo.Algorithm; update the logic around
ktype/kasInfo.Algorithm, generateWrapKeyWithEC, generateWrapKeyWithRSA,
FromPublicPEM, and the keyAccess fields (KeyType, WrappedKey,
EphemeralPublicKey) so the manifest correctly records an EphemeralPublicKey for
EC keys and prevents creating non-rewrappable entries when the PEM and declared
algorithm disagree.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 0a7022c7-f77c-4a72-8811-b1965976ce94
📒 Files selected for processing (21)
.gitignorelib/ocrypto/asym_decryption.golib/ocrypto/asym_decryption_test.golib/ocrypto/asym_encryption.golib/ocrypto/ec_decrypt_compressed_test.golib/ocrypto/ec_key_pair.golib/ocrypto/ec_key_pair_test.golib/ocrypto/rsa_key_pair.gosdk/codegen/runner/generate.gosdk/experimental/tdf/key_access.gosdk/experimental/tdf/key_access_test.gosdk/kas_client.gosdk/kas_client_test.gosdk/tdf.gosdk/tdf_config.gosdk/tdf_test.goservice/internal/security/basic_manager.goservice/internal/security/basic_manager_test.goservice/internal/security/in_process_provider.goservice/kas/access/rewrap.goservice/trust/key_manager.go
| func TestAsymDecryptionKeyType(t *testing.T) { | ||
| t.Parallel() | ||
|
|
||
| for _, bits := range []int{RSA2048Size, RSA4096Size} { | ||
| kp, err := NewRSAKeyPair(bits) | ||
| require.NoError(t, err) | ||
| privPEM, err := kp.PrivateKeyInPemFormat() | ||
| require.NoError(t, err) | ||
| d, err := FromPrivatePEM(privPEM) | ||
| require.NoError(t, err) | ||
| ad, ok := d.(AsymDecryption) | ||
| require.True(t, ok) | ||
| if bits == RSA2048Size { | ||
| require.Equal(t, RSA2048Key, ad.KeyType()) | ||
| } else { | ||
| require.Equal(t, RSA4096Key, ad.KeyType()) | ||
| } | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider using subtests for each RSA key size to improve test diagnostics.
The loop iterates over multiple RSA key sizes but doesn't wrap each iteration in t.Run(). This is inconsistent with the other tests in this file and makes it harder to identify which key size failed if there's an assertion error.
♻️ Proposed fix to add subtests
func TestAsymDecryptionKeyType(t *testing.T) {
t.Parallel()
- for _, bits := range []int{RSA2048Size, RSA4096Size} {
+ sizes := []struct {
+ bits int
+ kt KeyType
+ }{
+ {RSA2048Size, RSA2048Key},
+ {RSA4096Size, RSA4096Key},
+ }
+
+ for _, tc := range sizes {
+ t.Run(string(tc.kt), func(t *testing.T) {
+ t.Parallel()
+ kp, err := NewRSAKeyPair(tc.bits)
+ require.NoError(t, err)
+ privPEM, err := kp.PrivateKeyInPemFormat()
+ require.NoError(t, err)
+ d, err := FromPrivatePEM(privPEM)
+ require.NoError(t, err)
+ ad, ok := d.(AsymDecryption)
+ require.True(t, ok)
+ require.Equal(t, tc.kt, ad.KeyType())
+ })
+ }
- kp, err := NewRSAKeyPair(bits)
- require.NoError(t, err)
- privPEM, err := kp.PrivateKeyInPemFormat()
- require.NoError(t, err)
- d, err := FromPrivatePEM(privPEM)
- require.NoError(t, err)
- ad, ok := d.(AsymDecryption)
- require.True(t, ok)
- if bits == RSA2048Size {
- require.Equal(t, RSA2048Key, ad.KeyType())
- } else {
- require.Equal(t, RSA4096Key, ad.KeyType())
- }
- }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func TestAsymDecryptionKeyType(t *testing.T) { | |
| t.Parallel() | |
| for _, bits := range []int{RSA2048Size, RSA4096Size} { | |
| kp, err := NewRSAKeyPair(bits) | |
| require.NoError(t, err) | |
| privPEM, err := kp.PrivateKeyInPemFormat() | |
| require.NoError(t, err) | |
| d, err := FromPrivatePEM(privPEM) | |
| require.NoError(t, err) | |
| ad, ok := d.(AsymDecryption) | |
| require.True(t, ok) | |
| if bits == RSA2048Size { | |
| require.Equal(t, RSA2048Key, ad.KeyType()) | |
| } else { | |
| require.Equal(t, RSA4096Key, ad.KeyType()) | |
| } | |
| } | |
| func TestAsymDecryptionKeyType(t *testing.T) { | |
| t.Parallel() | |
| sizes := []struct { | |
| bits int | |
| kt KeyType | |
| }{ | |
| {RSA2048Size, RSA2048Key}, | |
| {RSA4096Size, RSA4096Key}, | |
| } | |
| for _, tc := range sizes { | |
| t.Run(string(tc.kt), func(t *testing.T) { | |
| t.Parallel() | |
| kp, err := NewRSAKeyPair(tc.bits) | |
| require.NoError(t, err) | |
| privPEM, err := kp.PrivateKeyInPemFormat() | |
| require.NoError(t, err) | |
| d, err := FromPrivatePEM(privPEM) | |
| require.NoError(t, err) | |
| ad, ok := d.(AsymDecryption) | |
| require.True(t, ok) | |
| require.Equal(t, tc.kt, ad.KeyType()) | |
| }) | |
| } | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@lib/ocrypto/asym_decryption_test.go` around lines 136 - 153, The
TestAsymDecryptionKeyType loop should be converted to subtests so failures show
which RSA size failed: for each bits in {RSA2048Size, RSA4096Size} call t.Run
with a descriptive name (e.g. "RSA2048"/"RSA4096" or fmt.Sprintf("%d", bits)),
capture bits in the closure, and move t.Parallel() into the subtest body; inside
the subtest keep the existing logic (NewRSAKeyPair, PrivateKeyInPemFormat,
FromPrivatePEM, type assertion to AsymDecryption and the KeyType checks against
RSA2048Key/RSA4096Key).
| func (e ECEncryptor) EphemeralPublicKeyInPemFormat() (string, error) { | ||
| if e.ek == nil { | ||
| return "", errors.New("failed to generate PEM formatted public key") | ||
| } | ||
|
|
||
| return publicKeyInPemFormat(e.ek.PublicKey()) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Expose this through an interface, not only the concrete ECEncryptor.
Callers now have to downcast PublicKeyEncryptor back to ocrypto.ECEncryptor in order to export the ephemeral PEM. That leaks the concrete implementation back into higher layers and makes the API fragile to future representation changes. A small EC-specific interface for this method would keep the new asym_* surface properly abstracted.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@lib/ocrypto/asym_encryption.go` around lines 288 - 293, Define a small
interface (e.g., EphemeralPublicKeyExporter) that declares
EphemeralPublicKeyInPemFormat() (string, error) and have ECEncryptor implicitly
implement it (the method already exists), then update the public API to return
or accept that interface where callers need to export the ephemeral PEM instead
of requiring a concrete ocrypto.ECEncryptor; ensure references to
PublicKeyEncryptor remain unchanged unless the API should explicitly include the
new interface, and update call sites to use the new EphemeralPublicKeyExporter
interface for PEM export without downcasting to ECEncryptor.
| // ECPrivateKeyInPemFormat Returns private key in pem format. | ||
| func ECPrivateKeyInPemFormat(privateKey ecdsa.PrivateKey) (string, error) { | ||
| privateKeyBytes, err := x509.MarshalPKCS8PrivateKey(privateKey) | ||
| if err != nil { | ||
| return "", fmt.Errorf("x509.MarshalPKCS8PrivateKey failed: %w", err) | ||
| } | ||
|
|
||
| privateKeyPem := pem.EncodeToMemory( | ||
| &pem.Block{ | ||
| Type: "PRIVATE KEY", | ||
| Bytes: privateKeyBytes, | ||
| }, | ||
| ) | ||
| return string(privateKeyPem), nil | ||
| return privateKeyInPemFormat(privateKey) | ||
| } | ||
|
|
||
| // ECPublicKeyInPemFormat Returns public key in pem format. | ||
| func ECPublicKeyInPemFormat(publicKey ecdsa.PublicKey) (string, error) { | ||
| pkb, err := x509.MarshalPKIXPublicKey(publicKey) | ||
| if err != nil { | ||
| return "", fmt.Errorf("x509.MarshalPKIXPublicKey failed: %w", err) | ||
| } | ||
|
|
||
| publicKeyPem := pem.EncodeToMemory( | ||
| &pem.Block{ | ||
| Type: "PUBLIC KEY", | ||
| Bytes: pkb, | ||
| }, | ||
| ) | ||
|
|
||
| return string(publicKeyPem), nil | ||
| return publicKeyInPemFormat(publicKey) | ||
| } |
There was a problem hiding this comment.
Value type passed to privateKeyInPemFormat bypasses nil checks.
ECPrivateKeyInPemFormat and ECPublicKeyInPemFormat accept value types (ecdsa.PrivateKey, ecdsa.PublicKey), but privateKeyInPemFormat only checks for pointer types in its switch. A zero-value struct will fall through to x509.MarshalPKCS8PrivateKey, potentially causing a less informative panic or error instead of the intended "failed to generate PEM formatted private key" message.
🛡️ Proposed fix to handle value types
// ECPrivateKeyInPemFormat Returns private key in pem format.
func ECPrivateKeyInPemFormat(privateKey ecdsa.PrivateKey) (string, error) {
- return privateKeyInPemFormat(privateKey)
+ return privateKeyInPemFormat(&privateKey)
}
// ECPublicKeyInPemFormat Returns public key in pem format.
func ECPublicKeyInPemFormat(publicKey ecdsa.PublicKey) (string, error) {
- return publicKeyInPemFormat(publicKey)
+ return publicKeyInPemFormat(&publicKey)
}Alternatively, add value-type cases to the switch in privateKeyInPemFormat:
case ecdsa.PrivateKey:
if key.D == nil {
return "", errors.New("failed to generate PEM formatted private key")
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@lib/ocrypto/ec_key_pair.go` around lines 480 - 488, The two exported helpers
(ECPrivateKeyInPemFormat, ECPublicKeyInPemFormat) pass value-type
ecdsa.PrivateKey/ecdsa.PublicKey to privateKeyInPemFormat which only handles
pointer cases, so zero-value structs bypass nil checks; update
privateKeyInPemFormat to add switch cases for the value types (case
ecdsa.PrivateKey: and case ecdsa.PublicKey:) and perform the same validity
checks as for pointers (e.g., for ecdsa.PrivateKey verify key.D != nil and for
ecdsa.PublicKey verify key.X != nil && key.Y != nil), returning the same error
("failed to generate PEM formatted private key"/appropriate public key error)
when those fields are nil to prevent falling through to
x509.MarshalPKCS8PrivateKey with a zero-value.
| // Determine key type based on algorithm | ||
| ktype := ocrypto.KeyType(pubKeyInfo.Algorithm) | ||
|
|
||
| kasPublicKey, err := ocrypto.FromPublicPEM(pubKeyInfo.PEM) | ||
| if err != nil { | ||
| return "", "", "", fmt.Errorf("failed to create parse KAS public key: %w", err) | ||
| } | ||
|
|
||
| if ocrypto.IsECKeyType(ktype) { | ||
| // Handle EC key wrapping | ||
| return wrapKeyWithEC(ktype, pubKeyInfo.PEM, symKey) | ||
| if epk, ok := kasPublicKey.(ocrypto.ECEncryptor); ok { | ||
| // Handle EC key wrapping | ||
| return wrapKeyWithEC(ktype, epk, symKey) | ||
| } | ||
| return "", "", "", fmt.Errorf("incorrect key type for %v", ktype) | ||
| } | ||
| // Handle RSA key wrapping | ||
| wrapped, err := wrapKeyWithRSA(pubKeyInfo.PEM, symKey) | ||
| wrapped, err := wrapKeyWithRSA(kasPublicKey, symKey) | ||
| return wrapped, "wrapped", "", err |
There was a problem hiding this comment.
Choose the wrapping flow from kasPublicKey, not pubKeyInfo.Algorithm.
You already parse the PEM on Line 168, but Line 173 still routes by the declared algorithm string. If that field is missing or stale for an EC key, the RSA branch will run, wrapKeyWithRSA will encrypt through the EC encryptor, and no ephemeral public key will be serialized.
🔧 Suggested fix
- // Determine key type based on algorithm
- ktype := ocrypto.KeyType(pubKeyInfo.Algorithm)
-
kasPublicKey, err := ocrypto.FromPublicPEM(pubKeyInfo.PEM)
if err != nil {
return "", "", "", fmt.Errorf("failed to create parse KAS public key: %w", err)
}
- if ocrypto.IsECKeyType(ktype) {
- if epk, ok := kasPublicKey.(ocrypto.ECEncryptor); ok {
- // Handle EC key wrapping
- return wrapKeyWithEC(ktype, epk, symKey)
- }
- return "", "", "", fmt.Errorf("incorrect key type for %v", ktype)
- }
- // Handle RSA key wrapping
- wrapped, err := wrapKeyWithRSA(kasPublicKey, symKey)
- return wrapped, "wrapped", "", err
+ switch kasPublicKey.Type() {
+ case ocrypto.EC:
+ epk, ok := kasPublicKey.(ocrypto.ECEncryptor)
+ if !ok {
+ return "", "", "", fmt.Errorf("unexpected encryptor type %T", kasPublicKey)
+ }
+ return wrapKeyWithEC(epk.KeyType(), epk, symKey)
+ case ocrypto.RSA:
+ wrapped, err := wrapKeyWithRSA(kasPublicKey, symKey)
+ return wrapped, "wrapped", "", err
+ default:
+ return "", "", "", fmt.Errorf("unsupported KAS public key type: %v", kasPublicKey.KeyType())
+ }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@sdk/experimental/tdf/key_access.go` around lines 165 - 182, The code
currently chooses the wrapping branch using pubKeyInfo.Algorithm (ktype) even
after parsing the PEM; instead, base the decision on the parsed kasPublicKey's
runtime type: after ocrypto.FromPublicPEM(...) check if kasPublicKey implements
ocrypto.ECEncryptor (or use a type switch) and call wrapKeyWithEC(ktype, epk,
symKey) when it does, otherwise call wrapKeyWithRSA(kasPublicKey, symKey);
update error paths to report a clear mismatch if the parsed key cannot perform
the expected operations and remove reliance on pubKeyInfo.Algorithm for
branching.
| func wrapKeyWithEC(keyType ocrypto.KeyType, kasPublicKey ocrypto.ECEncryptor, symKey []byte) (string, string, string, error) { | ||
| if !ocrypto.IsECKeyType(kasPublicKey.KeyType()) { | ||
| return "", "", "", fmt.Errorf("unexpected KAS public key type: %v", kasPublicKey.KeyType()) | ||
| } | ||
|
|
||
| // Get ephemeral private key | ||
| ephemeralPrivKey, err := ecKeyPair.PrivateKeyInPemFormat() | ||
| wrapped, err := kasPublicKey.Encrypt(symKey) | ||
| if err != nil { | ||
| return "", "", "", fmt.Errorf("failed to get ephemeral private key: %w", err) | ||
| return "", "", "", fmt.Errorf("failed to wrap with %v: %w", keyType, err) | ||
| } | ||
|
|
||
| // Compute ECDH shared secret | ||
| ecdhKey, err := ocrypto.ComputeECDHKey([]byte(ephemeralPrivKey), []byte(kasPublicKeyPEM)) | ||
| epk, err := kasPublicKey.EphemeralPublicKeyInPemFormat() | ||
| if err != nil { | ||
| return "", "", "", fmt.Errorf("failed to compute ECDH key: %w", err) | ||
| return "", "", "", fmt.Errorf("failed to export ephemeral public key: %w", err) | ||
| } | ||
|
|
||
| // Derive wrapping key using HKDF | ||
| salt := tdfSalt() | ||
| wrapKey, err := ocrypto.CalculateHKDF(salt, ecdhKey) | ||
| if err != nil { | ||
| return "", "", "", fmt.Errorf("failed to derive wrap key: %w", err) | ||
| } | ||
|
|
||
| // Ensure we have the right length for wrapping, trim if needed, or error if too short | ||
| if len(wrapKey) > len(symKey) { | ||
| wrapKey = wrapKey[:len(symKey)] | ||
| } else if len(wrapKey) < len(symKey) { | ||
| return "", "", "", fmt.Errorf("wrap key too short: got %d, expected at least %d", | ||
| len(wrapKey), len(symKey)) | ||
| } | ||
|
|
||
| wrapped := make([]byte, len(symKey)) | ||
| for i := range symKey { | ||
| wrapped[i] = symKey[i] ^ wrapKey[i] | ||
| } | ||
|
|
||
| return string(ocrypto.Base64Encode(wrapped)), "eccWrapped", ephemeralPubKey, nil | ||
| return string(ocrypto.Base64Encode(wrapped)), "eccWrapped", epk, nil |
There was a problem hiding this comment.
Emit the canonical EC wrapper key type.
Line 201 returns "eccWrapped", but the rest of the codebase serializes EC-wrapped entries as ec-wrapped (sdk/tdf.go Lines 44-45 and lib/ocrypto/asym_encryption.go Lines 25-28). This will make experimental manifests diverge from the main SDK/KAS contract.
🔧 Suggested fix
- return string(ocrypto.Base64Encode(wrapped)), "eccWrapped", epk, nil
+ return string(ocrypto.Base64Encode(wrapped)), "ec-wrapped", epk, nil🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@sdk/experimental/tdf/key_access.go` around lines 186 - 201, The EC wrapper
currently returns the non-canonical string "eccWrapped" from wrapKeyWithEC;
update the return to use the canonical key type "ec-wrapped" so it matches the
rest of the codebase (see wrapKeyWithEC in key_access.go and the canonical uses
in sdk/tdf.go and lib/ocrypto/asym_encryption.go), keeping the other return
values and error handling unchanged.
| func (k *KASClient) handleECKeyResponse(response *kas.RewrapResponse) (map[string][]kaoResult, error) { | ||
| kasEphemeralPublicKey := response.GetSessionPublicKey() | ||
| clientPrivateKey, err := k.sessionKey.PrivateKeyInPemFormat() | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to get private key: %w", err) | ||
| } | ||
| ecdhKey, err := ocrypto.ComputeECDHKey([]byte(clientPrivateKey), []byte(kasEphemeralPublicKey)) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("ocrypto.ComputeECDHKey failed: %w", err) | ||
| } | ||
|
|
||
| digest := sha256.New() | ||
| digest.Write([]byte("TDF")) | ||
| salt := digest.Sum(nil) | ||
| sessionKey, err := ocrypto.CalculateHKDF(salt, ecdhKey) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("ocrypto.CalculateHKDF failed: %w", err) | ||
| kasEphemeralPublicKeyPEM := response.GetSessionPublicKey() | ||
| block, _ := pem.Decode([]byte(kasEphemeralPublicKeyPEM)) | ||
| if block == nil { | ||
| return nil, errors.New("failed to decode KAS session public key PEM") | ||
| } | ||
|
|
||
| aesGcm, err := ocrypto.NewAESGcm(sessionKey) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("ocrypto.NewAESGcm failed: %w", err) | ||
| ecDecryptor, ok := k.sessionKey.(ocrypto.ECDecryptor) | ||
| if !ok { | ||
| return nil, errors.New("session key is not an EC decryptor") | ||
| } | ||
|
|
||
| return k.processECResponse(response, aesGcm) | ||
| return k.processECResponse(response, ecDecryptor, block.Bytes) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider adding validation for empty session public key.
If response.GetSessionPublicKey() returns an empty string, pem.Decode will return nil for the block, which is caught. However, adding an explicit check for empty input would provide a clearer error message.
💡 Optional: Add explicit empty check
func (k *KASClient) handleECKeyResponse(response *kas.RewrapResponse) (map[string][]kaoResult, error) {
kasEphemeralPublicKeyPEM := response.GetSessionPublicKey()
+ if kasEphemeralPublicKeyPEM == "" {
+ return nil, errors.New("KAS session public key is empty")
+ }
block, _ := pem.Decode([]byte(kasEphemeralPublicKeyPEM))
if block == nil {
return nil, errors.New("failed to decode KAS session public key PEM")
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func (k *KASClient) handleECKeyResponse(response *kas.RewrapResponse) (map[string][]kaoResult, error) { | |
| kasEphemeralPublicKey := response.GetSessionPublicKey() | |
| clientPrivateKey, err := k.sessionKey.PrivateKeyInPemFormat() | |
| if err != nil { | |
| return nil, fmt.Errorf("failed to get private key: %w", err) | |
| } | |
| ecdhKey, err := ocrypto.ComputeECDHKey([]byte(clientPrivateKey), []byte(kasEphemeralPublicKey)) | |
| if err != nil { | |
| return nil, fmt.Errorf("ocrypto.ComputeECDHKey failed: %w", err) | |
| } | |
| digest := sha256.New() | |
| digest.Write([]byte("TDF")) | |
| salt := digest.Sum(nil) | |
| sessionKey, err := ocrypto.CalculateHKDF(salt, ecdhKey) | |
| if err != nil { | |
| return nil, fmt.Errorf("ocrypto.CalculateHKDF failed: %w", err) | |
| kasEphemeralPublicKeyPEM := response.GetSessionPublicKey() | |
| block, _ := pem.Decode([]byte(kasEphemeralPublicKeyPEM)) | |
| if block == nil { | |
| return nil, errors.New("failed to decode KAS session public key PEM") | |
| } | |
| aesGcm, err := ocrypto.NewAESGcm(sessionKey) | |
| if err != nil { | |
| return nil, fmt.Errorf("ocrypto.NewAESGcm failed: %w", err) | |
| ecDecryptor, ok := k.sessionKey.(ocrypto.ECDecryptor) | |
| if !ok { | |
| return nil, errors.New("session key is not an EC decryptor") | |
| } | |
| return k.processECResponse(response, aesGcm) | |
| return k.processECResponse(response, ecDecryptor, block.Bytes) | |
| func (k *KASClient) handleECKeyResponse(response *kas.RewrapResponse) (map[string][]kaoResult, error) { | |
| kasEphemeralPublicKeyPEM := response.GetSessionPublicKey() | |
| if kasEphemeralPublicKeyPEM == "" { | |
| return nil, errors.New("KAS session public key is empty") | |
| } | |
| block, _ := pem.Decode([]byte(kasEphemeralPublicKeyPEM)) | |
| if block == nil { | |
| return nil, errors.New("failed to decode KAS session public key PEM") | |
| } | |
| ecDecryptor, ok := k.sessionKey.(ocrypto.ECDecryptor) | |
| if !ok { | |
| return nil, errors.New("session key is not an EC decryptor") | |
| } | |
| return k.processECResponse(response, ecDecryptor, block.Bytes) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@sdk/kas_client.go` around lines 196 - 208, In handleECKeyResponse, explicitly
validate that response.GetSessionPublicKey() is not empty before calling
pem.Decode so you can return a clearer error; check the returned string from
kas.RewrapResponse.GetSessionPublicKey(), if it's empty return a descriptive
error like "empty KAS session public key" and only then call pem.Decode,
preserving the existing session key type assertion (k.sessionKey as
ocrypto.ECDecryptor) and the subsequent call to k.processECResponse with
block.Bytes.
Proposed Changes
Checklist
Testing Instructions
Summary by CodeRabbit
Release Notes
New Features
Deprecations
PrivateKeyDecryptorinterfaceBug Fixes
Tests