feat(policy): Add UnsafeUpdateKey implementation#3731
Conversation
Signed-off-by: Chris Reed <creed@virtru.com>
Signed-off-by: Chris Reed <creed@virtru.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: ASSERTIVE Plan: Pro Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (5)
📝 WalkthroughWalkthroughImplements the ChangesUnsafe key update
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant UnsafeService
participant PolicyDBClient
participant Database
participant AuditLogger
Client->>UnsafeService: UnsafeUpdateKey request
UnsafeService->>PolicyDBClient: Load existing key
PolicyDBClient->>Database: Validate and update key
Database-->>PolicyDBClient: Updated key
PolicyDBClient-->>UnsafeService: Return updated key
UnsafeService->>AuditLogger: Record original and updated values
UnsafeService-->>Client: UnsafeUpdateKey response
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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.
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:
|
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
service/policy/kasregistry/key_access_server_registry.go (1)
299-324: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winDouble-
StatusifyErrorwrapping silently downgrades allUpdateKeyfailures toCodeInternal.Line 303-304 now statusifies the error inside the
RunInTxcallback, but line 322-324 (unchanged) statusifies it again outside.StatusifyError's first call wraps the error inconnect.NewError(code, errors.New(ErrTextXxx))— a brand-new error unrelated to the original sentinel (e.g.db.ErrCannotUpdateToUnspecified).connect.Error.Unwrap()returns that fresh inner error, so on the second pass none of theerrors.Is(err, ErrXxx)checks inStatusifyErrormatch, and it falls through to the genericconnect.CodeInternalfallback — discarding the correctly classified code and message from the first call.This affects every failure path of the pre-existing
UpdateKeyRPC (e.g.ErrCannotUpdateToUnspecifiedshould map toCodeInvalidArgument, but will now returnCodeInternal).service/policy/unsafe/unsafe.go's newUnsafeUpdateKeyshows the correct pattern: statusify once inside the callback, then propagate the error unchanged outside.🐛 Proposed fix: remove the redundant outer statusify
if err != nil { - return nil, db.StatusifyError(ctx, s.logger, err, db.ErrTextUpdateFailed, slog.String("key_access_server_keys", req.Msg.GetId())) + return nil, err } return connect.NewResponse(rsp), nil }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@service/policy/kasregistry/key_access_server_registry.go` around lines 299 - 324, Remove the redundant outer StatusifyError call after RunInTx in the UpdateKey flow. Return the transaction error unchanged so the statusified error produced inside the callback retains its original Connect code and message, matching the propagation pattern used by UnsafeUpdateKey.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@service/policy/db/key_access_server_registry.go`:
- Around line 623-630: In the KeyMode_KEY_MODE_UNSPECIFIED branch of the key
access server registry update logic, validate newProviderConfiguration.Valid
before checking existingMode so an empty provider configuration returns
ErrUnsafeUpdateKeyProviderConfigRequired rather than the existing-mode error.
Preserve the current behavior for valid provider configurations and add coverage
for existingMode PUBLIC_KEY_ONLY, requestMode UNSPECIFIED, and an empty provider
configuration.
---
Outside diff comments:
In `@service/policy/kasregistry/key_access_server_registry.go`:
- Around line 299-324: Remove the redundant outer StatusifyError call after
RunInTx in the UpdateKey flow. Return the transaction error unchanged so the
statusified error produced inside the callback retains its original Connect code
and message, matching the propagation pattern used by UnsafeUpdateKey.
🪄 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: 2aa85237-02e8-422e-a405-af339e236edf
📒 Files selected for processing (10)
service/integration/kas_registry_key_test.goservice/pkg/db/errors.goservice/pkg/db/errors_test.goservice/policy/db/key_access_server_registry.goservice/policy/db/key_access_server_registry.sql.goservice/policy/db/key_access_server_registry_test.goservice/policy/db/queries/key_access_server_registry.sqlservice/policy/kasregistry/key_access_server_registry.goservice/policy/unsafe/unsafe.goservice/policy/unsafe/unsafe_test.go
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:
|
|
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:
Summary by CodeRabbit
New Features
Bug Fixes
Tests