fix: send PostgreSQL v3 StartupMessage in proactive IAM refresh probe for MCP by kgala2 · Pull Request #1167 · GoogleCloudPlatform/cloud-sql-go-connector

kgala2 · 2026-10-01T18:50:36Z

Fixes GoogleCloudPlatform/cloud-sql-proxy#2553
Builds on #1150

Summary

When using Managed Connection Pooling (MCP / PgBouncer) with Auto IAM authentication on PostgreSQL, PgBouncer caches the client's 1-hour IAM OAuth2 token in its per-(database, username) credential store (db->user_tree) and uses that cached token whenever it opens a new backend connection to PostgreSQL.

A TLS-only probe on port 3307 completes the TLS handshake, which extracts the X.509 IAM token into a temporary per-socket credential struct, but PgBouncer only copies those credentials into the shared (database, username) pool after receiving a PostgreSQL v3.0 StartupMessage (0x00030000 with user and database) and completing the client pool/authentication setup. Closing the TLS socket immediately after the handshake frees the temporary per-socket credentials without updating the pool's cached token.

Changes

Proactive refresh probe (RefreshAheadCache.probeConnection):
- After each background certificate refresh when Auto IAM auth is enabled, probe the instance on port 3307 respecting the configured IP type (WithRefreshAheadIPType) and refreshTimeout.
Capture active PostgreSQL IAM (user, database) principals (NewPostgresStartupSnifferConn):
- Wrap IAM PostgreSQL connections in connectInstanceIP with a lightweight first-write sniffer (NewPostgresStartupSnifferConn) that extracts (user, database) from the client's initial PostgreSQL v3 StartupMessage (skipping an 8-byte SSLRequest if present) and passes all subsequent writes through directly.
Complete PostgreSQL Startup + Terminate on probe:
- For each recorded (user, database) principal, send a PostgreSQL v3.0 StartupMessage over the probe TLS connection, read the server's initial authentication response (AuthenticationOk 'R') so PgBouncer updates the (database, username) pool credentials, and send Terminate ('X') before closing.

…kens

Code review comments addressed: - Respect the IP settings in connection configuration (WithRefreshAheadIPType / connectionType) rather than probing all available IPs - Use configured refreshTimeout for probe context timeout

hessjcg · 2026-10-01T20:07:44Z

 }

+func (c *monitoredCache) recordIAMPrincipal(user, database string) {
+	if r, ok := c.connectionInfoCache.(interface{ RecordIAMPrincipal(string, string) }); ok {


Add RecordIAMPrincipal() to the interface for ConnectionInfoCache.

Done in f08e885: added RecordIAMPrincipal(user, database string) directly to the connectionInfoCache interface in dialer.go, implemented it on both RefreshAheadCache and LazyRefreshCache, and removed the recordIAMPrincipal type-assertion wrapper from monitoredCache (since monitoredCache embeds connectionInfoCache and promotes RecordIAMPrincipal directly).

hessjcg · 2026-10-01T20:11:28Z

 	tp auth.TokenProvider,
 	dialerID string,
 	useIAMAuthNDial bool,
+	opts ...RefreshAheadOption,


Don't use opts, pass explicit parameters.

Done in f08e885: removed RefreshAheadOption, WithRefreshAheadDialFunc, and WithRefreshAheadIPType, and updated NewRefreshAheadCache and NewLazyRefreshCache to accept dialFunc and ipType as explicit parameters.

hessjcg · 2026-10-01T20:11:50Z

+// RecordIAMPrincipal records a PostgreSQL (user, database) pair observed on an
+// active Auto-IAM connection so that subsequent background refresh probes can
+// send a StartupMessage and update PgBouncer's cached pool credentials.
+func (i *RefreshAheadCache) RecordIAMPrincipal(user, database string) {


What about the lazy cache?

Done in f08e885: added RecordIAMPrincipal(user, database string) and proactive IAM refresh probing (probeInstanceConnection) to LazyRefreshCache in internal/cloudsql/lazy.go (sharing appendIAMPrincipal and probeInstanceConnection with RefreshAheadCache), and added TestLazyRefreshCache_ProbeConnection_PostgresStartupPacket in internal/cloudsql/lazy_test.go.

hessjcg · 2026-10-01T20:13:13Z

 			r.result, r.err = i.r.ConnectionInfo(
 				ctx, i.connName, useIAMAuthN,
 			)
+			if r.err == nil && useIAMAuthN {


See if this can be cleaned up. Not idiomatic go. Use the switch statement below if possible.

Done in f08e885: moved if i.useIAMAuthNDial { i.probeConnection(ctx, r.result) } inside case nil: of the switch r.err statement below.

hessjcg · 2026-10-01T20:16:04Z

+}
+
+func (s *pgStartupSnifferConn) Write(b []byte) (int, error) {
+	s.mu.Lock()


s.mu.Lock() if s.done { s.mu.Unlock() return s.Conn.Write(b) } s.buf = append(s.buf, b...) if user, db, complete := parsePostgresStartupPacket(s.buf); complete { s.done = true s.buf = nil if user != "" { s.onStartup(user, db) } } else if len(s.buf) > maxPgStartupPacketLen+8 { s.done = true s.buf = nil } s.mu.Unlock() return s.Conn.Write(b)

Done in f08e885: updated pgStartupSnifferConn.Write to use the early-return guard (if s.done { s.mu.Unlock(); return s.Conn.Write(b) }).

…e, explicit params, interface method)

hessjcg

Approved after pending comment is addressed.

hessjcg · 2026-10-05T15:22:00Z

+	maxPgStartupPacketLen = 10000
+)
+
+func (i *RefreshAheadCache) probeConnection(ctx context.Context, ci ConnectionInfo) error {


Up to this point, we had kept all of the actual "make a connection" logic in dialer.go for a separation of concerns. Now we have redundant "make a connection" code here in instance.go. Can you figure out a way to reuse the dialer to make a connection for this probe?

hessjcg · 2026-10-05T15:24:20Z

+// NewPostgresStartupSnifferConn wraps conn so that the first PostgreSQL v3
+// StartupMessage written by the client is inspected to extract (user, database)
+// and passed to onStartup. Subsequent writes pass through with zero overhead.
+func NewPostgresStartupSnifferConn(conn net.Conn, onStartup func(user, database string)) net.Conn {


Let's rename "pgStartupSnifferConn" to "postgresConn". I want the name to cover more thank just the startup. We will likely add more features that require client-side interpretation of the postgres database protocol.

hessjcg · 2026-10-05T15:25:55Z

+
+type pgStartupSnifferConn struct {
+	net.Conn
+	done      bool


Rename this "startupMessageSent"

hessjcg · 2026-10-05T15:27:06Z

+	return s.Conn.Write(b)
+}
+
+func parsePostgresStartupPacket(buf []byte) (user, database string, complete bool) {


This assumes that the startup message is always sent with a single call to Conn.Write(). This may not be true.You will to make this work when the startup message is split across two or more Conn.Write() operations.

hessjcg · 2026-10-05T15:31:58Z

+	if err != nil {
+		t.Fatalf("encode StartupMessage failed: %v", err)
+	}
+	if _, err := sniffer.Write(startupPkt); err != nil {


Break the initial startup request into two Write() calls to ensure that the postgres protocol aware Conn can handle connection state across IO operations:

// Write the first byte _, err := sniffer.Write(startupPkt[0:1]) if err != nil { //... } _, err := sniffer.Write(startupPkt[1:]) if err != nil { //... }

hessjcg · 2026-10-05T15:33:42Z

+				backend := pgproto3.NewBackend(c, c)
+				msg, err := backend.ReceiveStartupMessage()
+				if err != nil {
+					return


Do close(startupProbeCh) ; return here so that you don't have to wait 5 seconds to check this channel for a result. Same on line 689

hessjcg and others added 3 commits October 1, 2026 18:48

feat: proactively probe database on Auto-IAM refresh to update MCP to…

ad43966

…kens

fix: address code review comments on proactive IAM token probe

df8bcaa

Code review comments addressed: - Respect the IP settings in connection configuration (WithRefreshAheadIPType / connectionType) rather than probing all available IPs - Use configured refreshTimeout for probe context timeout

fix: send PostgreSQL v3 StartupMessage in proactive IAM refresh probe

b7f20f5

hessjcg reviewed Oct 1, 2026

View reviewed changes

refactor: address PR #1167 review comments (pgproto3, LazyRefreshCach…

f08e885

…e, explicit params, interface method)

kgala2 marked this pull request as ready for review October 2, 2026 19:03

kgala2 requested a review from a team as a code owner October 2, 2026 19:03

blunderbuss-gcf Bot assigned hessjcg Oct 2, 2026

hessjcg previously approved these changes Oct 2, 2026

View reviewed changes

Comment thread internal/cloudsql/instance.go Outdated

refactor: remove unnecessary mutex from pgStartupSnifferConn

46d232c

kgala2 dismissed hessjcg’s stale review via 46d232c October 2, 2026 21:10

fix: acquire mutex before closing ready channel in scheduleRefresh

a081464

hessjcg approved these changes Oct 5, 2026

View reviewed changes

hessjcg requested changes Oct 5, 2026

View reviewed changes

hessjcg reviewed Oct 5, 2026

View reviewed changes

Conversation

kgala2 commented Oct 1, 2026

Summary

Changes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hessjcg left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants