feat(policy): Add UnsafeUpdateKey implementation#3731
Conversation
Signed-off-by: Chris Reed <creed@virtru.com>
Signed-off-by: Chris Reed <creed@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 |
github-actions
Bot
added
comp:db
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 an 'UnsafeUpdateKey' operation to the KAS registry, enabling administrative updates to key modes and provider configurations. This functionality is designed for support scenarios where keys need to be reconfigured without deleting and recreating them, ensuring existing TDF decryptability is maintained where applicable. The changes include full gRPC service definitions, database layer updates, and robust validation to prevent invalid state transitions. 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. A key once set in stone and cold, Now shifts its mode as we are told. With unsafe calls and careful care, We update keys with grace and flair. Footnotes
|
![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 an unsafe key update operation (UnsafeUpdateKey) to the policy service, allowing support-only switching between REMOTE and PUBLIC_KEY_ONLY key modes or updating the provider configuration for existing remote keys. The changes span protocol definitions, database-layer implementation, audit logging, and comprehensive integration and unit tests. The code reviewer provided several valuable suggestions to enhance the robustness and cleanliness of the implementation, including strengthening database-level validation of provider_config_id based on the target key mode, adding defensive nil checks to prevent potential panics, returning nil instead of empty structs in audit logs to reduce noise, and adding test cases to verify invalid provider configuration ID validation.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| func validateUnsafeUpdateKey(existing *policy.KasKey, r *unsafe.UnsafeUpdateKeyRequest) (unsafeUpdateKeyParams, error) { | ||
| existingMode := existing.GetKey().GetKeyMode() | ||
| params := unsafeUpdateKeyParams{ | ||
| ID: r.GetId(), | ||
| ProviderConfigID: pgtypeUUID(r.GetProviderConfigId()), | ||
| } | ||
|
|
||
| if existingMode != policy.KeyMode_KEY_MODE_PUBLIC_KEY_ONLY && existingMode != policy.KeyMode_KEY_MODE_REMOTE { | ||
| return params, ErrUnsafeUpdateKeyExistingModeUnsupported | ||
| } | ||
|
|
||
| if r.GetKeyMode() == policy.KeyMode_KEY_MODE_UNSPECIFIED && existingMode != policy.KeyMode_KEY_MODE_REMOTE { | ||
| return params, ErrUnsafeUpdateKeyProviderConfigExistingMode | ||
| } | ||
|
|
||
| if r.GetKeyMode() != policy.KeyMode_KEY_MODE_UNSPECIFIED { | ||
| params.KeyMode = pgtypeInt4(int32(r.GetKeyMode()), true) | ||
| } | ||
|
|
||
| return params, nil | ||
| } |
There was a problem hiding this comment.
The database layer should robustly validate the provider_config_id based on the target key mode rather than relying solely on the proto-level CEL validation. This prevents inconsistent database states (e.g., a REMOTE key with an empty/invalid provider_config_id, or a PUBLIC_KEY_ONLY key with a provider_config_id set) if the gRPC validation is bypassed or when called programmatically.
func validateUnsafeUpdateKey(existing *policy.KasKey, r *unsafe.UnsafeUpdateKeyRequest) (unsafeUpdateKeyParams, error) {
existingMode := existing.GetKey().GetKeyMode()
params := unsafeUpdateKeyParams{
ID: r.GetId(),
ProviderConfigID: pgtypeUUID(r.GetProviderConfigId()),
}
if existingMode != policy.KeyMode_KEY_MODE_PUBLIC_KEY_ONLY && existingMode != policy.KeyMode_KEY_MODE_REMOTE {
return params, ErrUnsafeUpdateKeyExistingModeUnsupported
}
if r.GetKeyMode() == policy.KeyMode_KEY_MODE_UNSPECIFIED && existingMode != policy.KeyMode_KEY_MODE_REMOTE {
return params, ErrUnsafeUpdateKeyProviderConfigExistingMode
}
targetMode := r.GetKeyMode()
if targetMode == policy.KeyMode_KEY_MODE_UNSPECIFIED {
targetMode = existingMode
}
if targetMode == policy.KeyMode_KEY_MODE_REMOTE {
if r.GetProviderConfigId() == "" || !params.ProviderConfigID.Valid {
return params, db.ErrUUIDInvalid
}
} else if targetMode == policy.KeyMode_KEY_MODE_PUBLIC_KEY_ONLY {
if r.GetProviderConfigId() != "" {
return params, errors.Join(db.ErrSelectIdentifierInvalid, fmt.Errorf("provider_config_id must be empty for PUBLIC_KEY_ONLY mode"))
}
}
if r.GetKeyMode() != policy.KeyMode_KEY_MODE_UNSPECIFIED {
params.KeyMode = pgtypeInt4(int32(r.GetKeyMode()), true)
}
return params, nil
}| import ( | ||
| "testing" | ||
|
|
||
| "github.com/jackc/pgx/v5/pgtype" | ||
| "github.com/opentdf/platform/protocol/go/policy" | ||
| "github.com/opentdf/platform/protocol/go/policy/unsafe" | ||
| "github.com/stretchr/testify/require" | ||
| ) |
There was a problem hiding this comment.
Import the service/pkg/db package to allow using db.ErrUUIDInvalid in the new test case verifying invalid provider_config_id validation.
| import ( | |
| "testing" | |
| "github.com/jackc/pgx/v5/pgtype" | |
| "github.com/opentdf/platform/protocol/go/policy" | |
| "github.com/opentdf/platform/protocol/go/policy/unsafe" | |
| "github.com/stretchr/testify/require" | |
| ) | |
| import ( | |
| "testing" | |
| "github.com/jackc/pgx/v5/pgtype" | |
| "github.com/opentdf/platform/protocol/go/policy" | |
| "github.com/opentdf/platform/protocol/go/policy/unsafe" | |
| "github.com/opentdf/platform/service/pkg/db" | |
| "github.com/stretchr/testify/require" | |
| ) |
| { | ||
| name: "existing config root key rejected before request validation", | ||
| existingMode: policy.KeyMode_KEY_MODE_CONFIG_ROOT_KEY, | ||
| requestMode: policy.KeyMode_KEY_MODE_REMOTE, | ||
| providerConfig: unsafeUpdateKeyTestUUID, | ||
| wantErr: ErrUnsafeUpdateKeyExistingModeUnsupported, | ||
| }, | ||
| } |
There was a problem hiding this comment.
Add a test case to verify that an invalid provider_config_id is correctly rejected when the target key mode is KEY_MODE_REMOTE.
{
name: "existing config root key rejected before request validation",
existingMode: policy.KeyMode_KEY_MODE_CONFIG_ROOT_KEY,
requestMode: policy.KeyMode_KEY_MODE_REMOTE,
providerConfig: unsafeUpdateKeyTestUUID,
wantErr: ErrUnsafeUpdateKeyExistingModeUnsupported,
},
{
name: "invalid provider config id for remote mode",
existingMode: policy.KeyMode_KEY_MODE_REMOTE,
requestMode: policy.KeyMode_KEY_MODE_REMOTE,
providerConfig: "invalid-uuid",
wantErr: db.ErrUUIDInvalid,
},
}| func unsafeUpdateKeyAuditValue(kasKey *policy.KasKey) *policy.KasKey { | ||
| if kasKey.GetKey() == nil { | ||
| return nil | ||
| } |
There was a problem hiding this comment.
Add an explicit nil check for kasKey to ensure robust defensive programming and prevent potential nil pointer dereference panics.
| func unsafeUpdateKeyAuditValue(kasKey *policy.KasKey) *policy.KasKey { | |
| if kasKey.GetKey() == nil { | |
| return nil | |
| } | |
| func unsafeUpdateKeyAuditValue(kasKey *policy.KasKey) *policy.KasKey { | |
| if kasKey == nil || kasKey.GetKey() == nil { | |
| return nil | |
| } |
| func unsafeUpdateKeyAuditProviderConfig(providerConfig *policy.KeyProviderConfig) *policy.KeyProviderConfig { | ||
| if providerConfig == nil { | ||
| return &policy.KeyProviderConfig{} | ||
| } |
There was a problem hiding this comment.
When providerConfig is nil, return nil instead of an empty &policy.KeyProviderConfig{}. Returning an empty struct results in empty objects in the serialized audit logs (e.g., "providerConfig": {}), which is noisy and misleading.
| func unsafeUpdateKeyAuditProviderConfig(providerConfig *policy.KeyProviderConfig) *policy.KeyProviderConfig { | |
| if providerConfig == nil { | |
| return &policy.KeyProviderConfig{} | |
| } | |
| func unsafeUpdateKeyAuditProviderConfig(providerConfig *policy.KeyProviderConfig) *policy.KeyProviderConfig { | |
| if providerConfig == nil { | |
| return nil | |
| } |
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
Signed-off-by: Chris Reed <creed@virtru.com>
Signed-off-by: Chris Reed <creed@virtru.com>
|
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 UnsafeUpdateKey for KAS registry keys.
Changes
Added DB support for unsafe key updates.
Supports:
Added validation for unsupported existing key modes.
Added typed errors for unsafe update validation failures.
Added integration coverage in kas_registry_key_test.go.
Added focused unit coverage for unsafe key update validation.
Depends on: