fix(wire): decrypt reads on pgjdbc describe-statement and prepared-reuse flows - #2
Conversation
A simple-protocol query ('Q') may carry several semicolon-separated
statements, each answered by its own result cycle. Expose a per-statement
analysis so the wire layer can track one result-queue entry per statement.
Assisted-By: Claude
Signed-off-by: Andrei Kvapil <andrei.kvapil@aenix.io>
The decrypt-on-DataRow plan was only built when a RowDescription arrived
while its portal was executing. pgjdbc violates that assumption twice:
once a statement is server-prepared (prepareThreshold) it is described as
a statement — the RowDescription arrives outside any Execute and was
dropped — and warm re-executions send bare Bind/Execute with no Describe
at all, so no RowDescription flows. Both left the portal without a read
plan and raw ciphertext envelopes went to Keycloak (the stage
verify-email incident: EmailException on a $KKP$ blob in the address,
and the same blob surfacing in id_token email claims).
Track frontend Describes in a FIFO and attribute each backend
RowDescription/NoData answer to the described statement or portal,
caching the result columns on the statement; derive a portal's plan from
that cache when the portal itself was never described.
Track simple-protocol queries ('Q') as synthetic result-queue entries —
their untracked CommandComplete used to pop a pending extended-protocol
portal and desync every later result on the connection — planning them
like OnParse (fail-loud on literal PII writes, decrypt on reads).
Track Sync boundaries in both queues: clients pipeline the next batch
before consuming the previous batch's ReadyForQuery, so ErrorResponse
drops only the failing batch's entries (the backend skips to Sync) and
ReadyForQuery closes exactly one batch, dropping dead entries that never
produced results. Advance past PortalSuspended, and fail loudly —
counter plus WARN — whenever a DataRow that no plan matched still
carries a ciphertext envelope. PII-table WARN matching now respects
identifier boundaries so RESET_CREDENTIALS_FLOW no longer counts as
CREDENTIAL.
Assisted-By: Claude
Signed-off-by: Andrei Kvapil <andrei.kvapil@aenix.io>
Cover the pgjdbc describe-statement and server-prepared reuse flows, batches pipelined across an unconsumed ReadyForQuery, error recovery that keeps the pipelined batch, the simple-protocol select/COMMIT/multi-statement cycles and queue sync, NoData describe attribution, PortalSuspended resume, and the PII-warn identifier-boundary false positive. Assisted-By: Claude Signed-off-by: Andrei Kvapil <andrei.kvapil@aenix.io>
A value that decrypts into another ciphertext envelope is a corrupted row: ciphertext leaked to the client during a passthrough window and was written back as the value (Keycloak lowercases username/email on write, so the inner envelope is case-mangled and unrecoverable from that column). The proxy keeps decrypting one layer only and now counts such rows in kkp_double_encrypted_total with a WARN, so operators can find and clean them up. Assisted-By: Claude Signed-off-by: Andrei Kvapil <andrei.kvapil@aenix.io>
📝 WalkthroughWalkthroughThis PR fixes a shutdown race in the proxy's accept loop, adds multi-statement SQL analysis ( ChangesProxy Decrypt Pipeline and Protocol Correctness
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant Frontend
participant Session
participant Backend
Frontend->>Session: OnDescribe(statement/portal)
Session->>Session: enqueue pendingDescribe
Frontend->>Session: OnQuery(sql)
Session->>Session: AnalyzeAll(sql), enqueue synthetic entries
Frontend->>Session: OnSync
Session->>Session: push syncBoundary onto execQueue/describes
Backend->>Session: OnRowDescription
Session->>Session: pop pendingDescribe or fallback to executing portal
Session->>Session: learnColumns, ensureReadPlan
Backend->>Session: DataRow
Session->>Session: DecryptDataRow via ensureReadPlan
alt leak or double-encrypted
Session->>Session: flagCiphertextLeak / increment DoubleEncrypted
end
sequenceDiagram
participant Client
participant serve
participant close
close->>serve: s.listener = nil (under s.mu)
serve->>serve: Accept() on snapshotted listener
serve->>serve: check s.listener == nil
alt shutdown detected
serve->>Client: close accepted connection
else active
serve->>serve: register connection, s.wg.Add(1)
end
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
![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 tracking simple-protocol queries and server-prepared statement reuse in the KMS proxy, along with new metrics for tracking ciphertext passthrough and double-encrypted rows, and comprehensive regression tests. The review feedback highlights critical memory leaks across several queue-popping operations in session.go where sliced elements containing pointers are not zeroed out, preventing garbage collection of PreparedStatement and Portal objects. Additionally, it is recommended to extract individual statement SQL in multi-statement simple queries to prevent false positives in PII detection.
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.
| pd := s.describes[0] | ||
| s.describes = s.describes[1:] |
There was a problem hiding this comment.
Popping an element from s.describes by slicing (s.describes[1:]) leaves the popped element's pointers (stmt and portal) in the underlying array. This prevents the garbage collector from reclaiming the referenced PreparedStatement and Portal objects, leading to memory leaks in long-running proxy sessions. Explicitly zero out the popped element before slicing to allow GC to reclaim the memory.
pd := s.describes[0]
s.describes[0] = pendingDescribe{}
s.describes = s.describes[1:]| if len(s.describes) > 0 && !s.describes[0].boundary { | ||
| s.describes = s.describes[1:] | ||
| } |
There was a problem hiding this comment.
Popping an element from s.describes by slicing without zeroing out the popped element leaves the pointers in the underlying array, causing a memory leak. Explicitly zero out the popped element before slicing.
if len(s.describes) > 0 && !s.describes[0].boundary {
s.describes[0] = pendingDescribe{}
s.describes = s.describes[1:]
}| func (s *Session) OnErrorResponse() { | ||
| for len(s.execQueue) > 0 && s.execQueue[0] != syncBoundary { | ||
| s.execQueue = s.execQueue[1:] | ||
| } | ||
| for len(s.describes) > 0 && !s.describes[0].boundary { | ||
| s.describes = s.describes[1:] | ||
| } | ||
| } |
There was a problem hiding this comment.
When clearing the queues on error, slicing s.execQueue and s.describes without zeroing out the popped elements leaves references to Portal and PreparedStatement objects in the underlying arrays, causing memory leaks. Explicitly set them to nil or zero value before slicing.
| func (s *Session) OnErrorResponse() { | |
| for len(s.execQueue) > 0 && s.execQueue[0] != syncBoundary { | |
| s.execQueue = s.execQueue[1:] | |
| } | |
| for len(s.describes) > 0 && !s.describes[0].boundary { | |
| s.describes = s.describes[1:] | |
| } | |
| } | |
| func (s *Session) OnErrorResponse() { | |
| for len(s.execQueue) > 0 && s.execQueue[0] != syncBoundary { | |
| s.execQueue[0] = nil | |
| s.execQueue = s.execQueue[1:] | |
| } | |
| for len(s.describes) > 0 && !s.describes[0].boundary { | |
| s.describes[0] = pendingDescribe{} | |
| s.describes = s.describes[1:] | |
| } | |
| } |
| func (s *Session) OnReadyForQuery() { | ||
| dead := 0 | ||
| for len(s.execQueue) > 0 { | ||
| head := s.execQueue[0] | ||
| s.execQueue = s.execQueue[1:] | ||
| if head == syncBoundary { | ||
| break | ||
| } | ||
| dead++ | ||
| } | ||
| for len(s.describes) > 0 { | ||
| head := s.describes[0] | ||
| s.describes = s.describes[1:] | ||
| if head.boundary { | ||
| break | ||
| } | ||
| dead++ | ||
| } |
There was a problem hiding this comment.
When consuming the sync boundary and dead entries on ReadyForQuery, slicing s.execQueue and s.describes without zeroing out the popped elements causes memory leaks. Explicitly set them to nil or zero value before slicing.
func (s *Session) OnReadyForQuery() {
dead := 0
for len(s.execQueue) > 0 {
head := s.execQueue[0]
s.execQueue[0] = nil
s.execQueue = s.execQueue[1:]
if head == syncBoundary {
break
}
dead++
}
for len(s.describes) > 0 {
head := s.describes[0]
s.describes[0] = pendingDescribe{}
s.describes = s.describes[1:]
if head.boundary {
break
}
dead++
}| func (s *Session) popExecuting() { | ||
| if len(s.execQueue) > 0 { | ||
| if len(s.execQueue) > 0 && s.execQueue[0] != syncBoundary { | ||
| s.execQueue = s.execQueue[1:] | ||
| } | ||
| } |
There was a problem hiding this comment.
Popping the executing portal by slicing s.execQueue without zeroing out the popped element leaves a reference to the Portal in the underlying array, causing a memory leak. Explicitly set the popped element to nil before slicing.
| func (s *Session) popExecuting() { | |
| if len(s.execQueue) > 0 { | |
| if len(s.execQueue) > 0 && s.execQueue[0] != syncBoundary { | |
| s.execQueue = s.execQueue[1:] | |
| } | |
| } | |
| func (s *Session) popExecuting() { | |
| if len(s.execQueue) > 0 && s.execQueue[0] != syncBoundary { | |
| s.execQueue[0] = nil | |
| s.execQueue = s.execQueue[1:] | |
| } | |
| } |
| s.execQueue = append(s.execQueue, &Portal{ | ||
| Name: simplePortalName, | ||
| Stmt: &PreparedStatement{SQL: q.String, Analysis: a, WritePlan: plan}, | ||
| }) |
There was a problem hiding this comment.
In a multi-statement simple query, using the entire query string q.String as the SQL for every statement's PreparedStatement can cause false positives in piiTouchedButNotPlanned and inaccurate log messages. For example, if one statement in the batch touches PII but another does not, the non-PII statement's check will still see the PII table name in the full query string. Consider extracting the individual statement's SQL from q.String using the statement's location and length from the parser.
In a multi-statement simple query every synthetic entry carried the full query string, so per-statement logging and the PII-warn heuristic saw the neighbouring statements too. Slice each statement's own text out of the query string via the parser's location/length. Assisted-By: Claude Signed-off-by: Andrei Kvapil <andrei.kvapil@aenix.io>
Popping by reslicing left the popped Portal/PreparedStatement pointers in the backing array, pinning them from GC for the life of the connection. Route every pop through helpers that zero the slot first. Assisted-By: Claude Signed-off-by: Andrei Kvapil <andrei.kvapil@aenix.io>
serve() read s.listener without the lock, and a connection accepted concurrently with close() could wg.Add(1) after close() had started wg.Wait() — a data race the race detector flagged intermittently in TestServerStartAcceptStop. Snapshot the listener under the lock, register accepted connections under the same lock close() takes before waiting, and drop a connection that loses the race with shutdown. Assisted-By: Claude Signed-off-by: Andrei Kvapil <andrei.kvapil@aenix.io>
Kirill Ilin (sircthulhu)
left a comment
Kirill Ilin (sircthulhu)
left a comment
There was a problem hiding this comment.
LGTM.
Focused fix for a real incident (raw $KKP$ ciphertext reaching Keycloak's verifyEmail and id_token claims). Reviewed the diff against the merge-base and validated in a throwaway worktree.
Validation
go build ./...— cleango test ./...— all packages passgo test -race -count=3 ./cmd/proxy/... ./internal/wire/...— clean (matters, since this fixes a data race)make lint(pinned golangci-lint) — 0 issues
Why it's sound
The reworked two-queue wire state machine (execQueue + describes, each with syncBoundary markers) holds up against protocol-realistic flows:
- Describe-statement + warm reuse: columns cached on the
PreparedStatementat Describe time;ensureReadPlanlazily derives the portal plan on bare Bind/Execute cycles — the actual verify-email fix. - FIFO attribution of RowDescription/NoData to the head non-boundary describe is valid (Execute never emits RowDescription).
- Pipelining/boundaries: each Sync pushes exactly one boundary to both queues, each ReadyForQuery pops one — the queues stay in lockstep across cross-batch pipelining, error-in-batch-1, prefetch-then-execute, and interleaved simple/extended traffic.
- Fail-loud consistency: simple-protocol queries are now planned identically to
OnParse— a literal PII write or a literal deterministic-PII filter returnsErrUnencryptablePIIand tears down the connection rather than passing ciphertext or returning wrong rows.
server.go race fix is correct: serve() snapshots the listener under the lock, and wg.Add(1) is atomic with the listener == nil check under the same mutex close() takes before wg.Wait(). Semaphore token is balanced on the drop path. Race detector confirms.
Observability additions (kkp_ciphertext_passthrough_total, kkp_double_encrypted_total) are low-cardinality and correctly wired.
Non-blocking suggestions
DecryptDataRowlogs a WARN per double-encrypted field per row — during an incident where Keycloak polls a corrupted column this can flood logs. The metric alone suffices; consider dropping to debug or rate-limiting.- The literal-write fail-loud is tested (
TestSimpleQueryLiteralPIIWriteFailsLoud), but the literal-filter fail-loud (WHERE email = 'literal'→ErrUnencryptablePII) is newly reachable for simple queries and isn't directly covered. Low risk (sharedPlanWritepath), but a one-line test would document the stricter behavior.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
internal/wire/transform.go (1)
250-274: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
flagCiphertextLeakadds a per-value string conversion to the hot path for every non-PII passthrough row.
flagCiphertextLeakis called for everyDataRowon the "empty-plan" path — i.e., every row of every query against a table with no PII columns, which is likely the majority of proxy traffic. For each column value it does a fullstring(v)copy beforecrypto.Parseeven checks the cheap prefix. This adds allocation/copy overhead to the common case rather than only the rare "actually leaking ciphertext" case.Consider a byte-level prefix check before allocating a string, so only genuinely suspicious values pay the
crypto.Parsecost:♻️ Suggested optimization
+// (in internal/crypto, export the sentinel bytes or add a cheap check) +// func HasEnvelopePrefix(v []byte) bool { return bytes.HasPrefix(v, sentinelBytes) } + func flagCiphertextLeak(dr *pgproto3.DataRow, reason string) { for _, v := range dr.Values { - if v == nil { + if v == nil || !crypto.HasEnvelopePrefix(v) { continue } if _, ok, err := crypto.Parse(string(v)); ok || err != nil { observe.CiphertextPassthrough.WithLabelValues(reason).Inc() log.Printf("kkp: WARN ciphertext passed through undecrypted (%s) — read-path gap", reason) return } } }Also applies to: 304-319
🤖 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 `@internal/wire/transform.go` around lines 250 - 274, The hot path in Session.DecryptDataRow is paying an unnecessary allocation cost through flagCiphertextLeak for every passthrough DataRow, especially on the empty-plan branch. Update flagCiphertextLeak (and any shared callers like the no-portal path) to do a cheap byte-level ciphertext prefix check before converting each value to a string or calling crypto.Parse, so only genuinely suspicious values incur the parsing/allocation overhead.
🤖 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.
Nitpick comments:
In `@internal/wire/transform.go`:
- Around line 250-274: The hot path in Session.DecryptDataRow is paying an
unnecessary allocation cost through flagCiphertextLeak for every passthrough
DataRow, especially on the empty-plan branch. Update flagCiphertextLeak (and any
shared callers like the no-portal path) to do a cheap byte-level ciphertext
prefix check before converting each value to a string or calling crypto.Parse,
so only genuinely suspicious values incur the parsing/allocation overhead.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 966da047-9cca-448f-83c6-904340e501dc
📒 Files selected for processing (9)
cmd/proxy/server.gogo.modinternal/observe/metrics.gointernal/rewrite/analyze.gointernal/rewrite/analyze_test.gointernal/wire/pgjdbc_flow_test.gointernal/wire/relay.gointernal/wire/session.gointernal/wire/transform.go
2 participants
Problem
On stage, Keycloak 26.6.3 behind proxy v0.2.0 received raw
$KKP$…ciphertext instead of decrypted PII on warm connections:verifyEmailfailed withEmailException: Invalid address '$KKP$…', and the same blobs surfaced inid_tokenemail claims, breaking OIDC logins.Root cause
The decrypt-on-DataRow plan was built only when a RowDescription arrived while its portal was executing. Real pgjdbc traffic violates that: a server-prepared statement (past
prepareThreshold) is described as a statement — its RowDescription arrives outside any Execute and was dropped — and warm re-executions send bare Bind/Execute with no Describe, so no RowDescription flows at all. Simple-protocol queries were untracked, so their CommandComplete desynced the result queue, and clients pipeline the next batch before consuming the previous ReadyForQuery.Fix
kkp_ciphertext_passthrough_totalcounter + WARN) when an unplanned DataRow still carries a ciphertext envelope; PII-table WARN matching now respects identifier boundaries.kkp_double_encrypted_total+ WARN on decrypt.Testing
make lint/make testgreen; regression tests cover the describe-statement flow, server-prepared reuse, pipelined batches across ReadyForQuery, error recovery, simple-query cycles, NoData attribution, PortalSuspended, and double-encrypted row detection.Summary by CodeRabbit
New Features
Bug Fixes