Skip to content

fix: send PostgreSQL v3 StartupMessage in proactive IAM refresh probe for MCP - #1480

Open
kgala2 wants to merge 5 commits into
mainfrom
draft-mcp-iam-token-refresh
Open

kgala2 wants to merge 5 commits into
mainfrom
draft-mcp-iam-token-refresh

Conversation

@kgala2

@kgala2 kgala2 commented Oct 1, 2026

Copy link
Copy Markdown
Collaborator

Summary

Supersedes / extends #1467 to ensure proactive IAM refresh probes actually update the cached IAM access token in Managed Connection Pooling (MCP) PgBouncer for long-lived persistent connections (GoogleCloudPlatform/cloud-sql-proxy#2553).

Why a TLS-only probe is insufficient

When Managed Connection Pooling (PgBouncer) accepts a TLS connection on port 3307:

  1. During SBUF_EV_TLS_READY, AutoIAMInit(sbuf) extracts the IAM access token from the ephemeral X.509 certificate extension and stores it in the per-socket client->auto_iam_login_user_credentials->passwd (client->db and client->login_user_credentials are still NULL because the X.509 client certificate does not contain the PostgreSQL database or user).
  2. PgBouncer only copies client->auto_iam_login_user_credentials->passwd into the shared pool's client->login_user_credentials->passwd (db->user_tree, which is keyed per (database, username)) inside finish_set_pool(), which is reached only after reading a PostgreSQL v3 StartupMessage (PKT_STARTUP_V3 -> decide_startup_pool -> set_pool(client, dbname, username, ...)).
  3. If the probe closes the TLS socket immediately after the TLS handshake without sending a StartupMessage, PgBouncer frees client->auto_iam_login_user_credentials on disconnect without ever updating pool->user_credentials->passwd, causing idle persistent connections to fail after ~1 hour once the initial IAM token expires.

What this PR changes

  • Adds @abc.abstractmethod def record_principal(self, user: str, database: str) -> None to ConnectionInfoCache (RefreshAheadCache, LazyRefreshCache, and MonitoredCache).
  • Records the (user, db) pair in Connector.connect_async when enable_iam_auth=True.
  • Updates _probe_instance_connection (shared by RefreshAheadCache and LazyRefreshCache) to send a PostgreSQL v3 StartupMessage ([int32 length][int32 196608]["user\0<user>\0database\0<database>\0\0"]) for each recorded (user, database) pair over the TLS probe connection, read the initial server authentication response, and send a clean PostgreSQL Terminate packet ('X') before closing.

hessjcg and others added 5 commits October 1, 2026 22:51
Code review comments addressed:
- Respect the IP settings in connection configuration (self._ip_type) rather than probing all available IPs
- Use SERVER_PROXY_PORT constant (3307)
- Use DEFAULT_CONNECT_TIMEOUT constant (30) / configured timeout
@kgala2
kgala2 marked this pull request as ready for review October 2, 2026 19:03
@kgala2
kgala2 requested a review from a team as a code owner October 2, 2026 19:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants