fix(wire): decrypt reads on pgjdbc describe-statement and prepared-reuse flows

cmd/proxy/server.go

@@ -128,12 +128,16 @@ func (s *server) addr() string {
 // relays — Accept blocks on a semaphore so the kernel queue absorbs the
 // backpressure instead of the process opening unbounded sessions.
 func (s *server) serve() error {
+	s.mu.Lock()
+	ln := s.listener
+	s.mu.Unlock()
+
 	var sem chan struct{}
 	if s.cfg.MaxConnections > 0 {
 		sem = make(chan struct{}, s.cfg.MaxConnections)
 	}
 	for {
-		conn, err := s.listener.Accept()
+		conn, err := ln.Accept()
 		if err != nil {
 			if errors.Is(err, net.ErrClosed) {
 				return nil
@@ -147,7 +151,22 @@ func (s *server) serve() error {
 		if s.cfg.HandshakeTimeout > 0 {
 			_ = conn.SetDeadline(time.Now().Add(s.cfg.HandshakeTimeout))
 		}
-		s.wg.Add(1)
+		// Register the connection under the lock close() takes before
+		// waiting, so a conn accepted concurrently with close either
+		// increments the WaitGroup before close waits or is dropped.
+		s.mu.Lock()
+		closed := s.listener == nil
+		if !closed {
+			s.wg.Add(1)
+		}
+		s.mu.Unlock()
+		if closed {
+			_ = conn.Close()
+			if sem != nil {
+				<-sem
+			}
+			return nil
+		}
 		go func() {
 			defer s.wg.Done()
 			defer func() {
@@ -164,6 +183,7 @@ func (s *server) serve() error {
 func (s *server) close() error {
 	s.mu.Lock()
 	ln := s.listener
+	s.listener = nil
 	s.mu.Unlock()
 	if ln == nil {
 		return nil

go.mod

@@ -19,6 +19,7 @@ require (
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
 	github.com/kr/text v0.2.0 // indirect
+	github.com/kylelemons/godebug v1.1.0 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/prometheus/client_model v0.6.2 // indirect
 	github.com/prometheus/common v0.66.1 // indirect

internal/observe/metrics.go

@@ -45,11 +45,30 @@ var (
 		Name: "kkp_kms_calls_total",
 		Help: "Calls to the KMS (Encrypt/Decrypt of DEKs), by op and outcome.",
 	}, []string{"op", "outcome"})
+
+	// CiphertextPassthrough counts DataRows that left the proxy still
+	// carrying a ciphertext envelope because no decrypt plan matched, by
+	// reason. Any non-zero value is an incident: the client saw raw $KKP$
+	// blobs (the verify-email failure class).
+	CiphertextPassthrough = prometheus.NewCounterVec(prometheus.CounterOpts{
+		Name: "kkp_ciphertext_passthrough_total",
+		Help: "DataRows passed through undecrypted despite carrying a ciphertext envelope, by reason.",
+	}, []string{"reason"})
+
+	// DoubleEncrypted counts decrypted values that still carry a ciphertext
+	// envelope: rows whose stored value IS a leaked ciphertext (written back
+	// by the client during a passthrough window and encrypted again). Used
+	// to inventory corrupted rows for cleanup.
+	DoubleEncrypted = prometheus.NewCounterVec(prometheus.CounterOpts{
+		Name: "kkp_double_encrypted_total",
+		Help: "Decrypted values still carrying a ciphertext envelope (double-encrypted rows), by table/column.",
+	}, []string{"table", "column"})
 )
 
 func init() {
 	prometheus.MustRegister(
 		DecryptTotal, DecryptFailures, EncryptTotal,
-		UnrecognizedPIISQL, KMSCallTotal,
+		UnrecognizedPIISQL, KMSCallTotal, CiphertextPassthrough,
+		DoubleEncrypted,
 	)
 }

internal/rewrite/analyze.go

@@ -49,6 +49,10 @@ type ColumnParam struct {
 // to the filtered columns on the read/search path.
 type Analysis struct {
 	Kind StmtKind
+	// SQL is this statement's own text within a multi-statement
+	// simple-protocol query string. Set by AnalyzeAll only; empty for
+	// Analyze, whose callers already hold the (single) statement text.
+	SQL string
 	// Table is the target table (INSERT/UPDATE/DELETE) or the first FROM table
 	// (SELECT). Empty if no single table applies.
 	Table string
@@ -76,19 +80,56 @@ func Analyze(sql string) (*Analysis, error) {
 	if len(stmts) != 1 {
 		return nil, fmt.Errorf("rewrite: expected exactly one statement, got %d", len(stmts))
 	}
+	return analyzeNode(stmts[0].GetStmt()), nil
+}
+
+// AnalyzeAll parses a simple-protocol query string, which may carry several
+// semicolon-separated statements, and returns one Analysis per statement in
+// order, each carrying its own statement text. The backend answers each
+// statement with its own result cycle, so the wire layer needs one
+// result-queue entry per statement to stay in sync.
+func AnalyzeAll(sql string) ([]*Analysis, error) {
+	result, err := pg.Parse(sql)
+	if err != nil {
+		return nil, fmt.Errorf("rewrite: parse: %w", err)
+	}
+	stmts := result.GetStmts()
+	analyses := make([]*Analysis, 0, len(stmts))
+	for _, st := range stmts {
+		a := analyzeNode(st.GetStmt())
+		a.SQL = statementText(sql, st)
+		analyses = append(analyses, a)
+	}
+	return analyses, nil
+}
+
+// statementText slices one statement's own text out of a multi-statement
+// query string using the parser's location/length, so per-statement logging
+// and PII heuristics do not see the neighbouring statements.
+func statementText(sql string, st *pg.RawStmt) string {
+	start := int(st.GetStmtLocation())
+	if start < 0 || start > len(sql) {
+		return sql
+	}
+	end := len(sql)
+	if l := int(st.GetStmtLen()); l > 0 && start+l <= len(sql) {
+		end = start + l
+	}
+	return strings.TrimLeft(sql[start:end], " \t\r\n;")
+}
 
-	node := stmts[0].GetStmt()
+func analyzeNode(node *pg.Node) *Analysis {
 	switch {
 	case node.GetInsertStmt() != nil:
-		return analyzeInsert(node.GetInsertStmt()), nil
+		return analyzeInsert(node.GetInsertStmt())
 	case node.GetUpdateStmt() != nil:
-		return analyzeUpdate(node.GetUpdateStmt()), nil
+		return analyzeUpdate(node.GetUpdateStmt())
 	case node.GetSelectStmt() != nil:
-		return analyzeSelect(node.GetSelectStmt()), nil
+		return analyzeSelect(node.GetSelectStmt())
 	case node.GetDeleteStmt() != nil:
-		return analyzeDelete(node.GetDeleteStmt()), nil
+		return analyzeDelete(node.GetDeleteStmt())
 	default:
-		return &Analysis{Kind: KindOther}, nil
+		return &Analysis{Kind: KindOther}
 	}
 }
 

internal/rewrite/analyze_test.go

@@ -163,3 +163,24 @@ func TestAnalyzeErrors(t *testing.T) {
 		t.Error("Analyze accepted multiple statements")
 	}
 }
+
+func TestAnalyzeAllPerStatementSQL(t *testing.T) {
+	t.Parallel()
+
+	analyses, err := AnalyzeAll("select id from realm; select email from user_entity where id = 'u1'")
+	if err != nil {
+		t.Fatalf("AnalyzeAll: %v", err)
+	}
+	if len(analyses) != 2 {
+		t.Fatalf("statements = %d, want 2", len(analyses))
+	}
+	if analyses[0].SQL != "select id from realm" {
+		t.Errorf("first statement SQL = %q", analyses[0].SQL)
+	}
+	if analyses[1].SQL != "select email from user_entity where id = 'u1'" {
+		t.Errorf("second statement SQL = %q", analyses[1].SQL)
+	}
+	if analyses[0].Table != "realm" || analyses[1].Table != "user_entity" {
+		t.Errorf("tables = %q, %q", analyses[0].Table, analyses[1].Table)
+	}
+}