feat: add metadata exchange to Python

Author: jackwotherspoon

google/cloud/sql/connector/connector.py

@@ -21,9 +21,10 @@
 import logging
 import os
 import socket
+import struct
 from threading import Thread
 from types import TracebackType
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Optional, TYPE_CHECKING, Union
 
 import google.auth
 from google.auth.credentials import Credentials
@@ -44,11 +45,17 @@
 from google.cloud.sql.connector.resolver import DnsResolver
 from google.cloud.sql.connector.utils import format_database_user
 from google.cloud.sql.connector.utils import generate_keys
+import google.cloud.sql.proto.cloud_sql_metadata_exchange_pb2 as connectorspb
+
+if TYPE_CHECKING:
+    import ssl
 
 logger = logging.getLogger(name=__name__)
 
 ASYNC_DRIVERS = ["asyncpg"]
 SERVER_PROXY_PORT = 3307
+# the maximum amount of time to wait before aborting a metadata exchange
+IO_TIMEOUT = 30
 _DEFAULT_SCHEME = "https://"
 _DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
 _SQLADMIN_HOST_TEMPLATE = "sqladmin.{universe_domain}"
@@ -391,6 +398,9 @@ async def connect_async(
                 socket.create_connection((ip_address, SERVER_PROXY_PORT)),
                 server_hostname=ip_address,
             )
+            # Perform Metadata Exchange Protocol
+            metadata_partial = partial(self.metadata_exchange, sock)
+            sock = await self._loop.run_in_executor(None, metadata_partial)
             # If this connection was opened using a domain name, then store it
             # for later in case we need to forcibly close it on failover.
             if conn_info.conn_name.domain_name:
@@ -409,6 +419,86 @@ async def connect_async(
             await monitored_cache.force_refresh()
             raise
 
+    def metadata_exchange(self, sock: ssl.SSLSocket) -> ssl.SSLSocket:
+        """
+        Sends metadata about the connection prior to the database
+        protocol taking over.
+        The exchange consists of four steps:
+        1. Prepare a CloudSQLConnectRequest including the socket protocol and
+           the user agent.
+        2. Write the size of the message as a big endian uint32 (4 bytes) to
+           the server followed by the serialized message. The length does not
+           include the initial four bytes.
+        3. Read a big endian uint32 (4 bytes) from the server. This is the
+           CloudSQLConnectResponse message length and does not include the
+           initial four bytes.
+        4. Parse the response using the message length in step 3. If the
+           response is not OK, return the response's error. If there is no error,
+           the metadata exchange has succeeded and the connection is complete.
+        Args:
+            sock (ssl.SSLSocket): The mTLS/SSL socket to perform metadata
+                exchange on.
+        Returns:
+            sock (ssl.SSLSocket): mTLS/SSL socket connected to Cloud SQL Proxy
+                server.
+        """
+        # form metadata exchange request
+        req = connectorspb.CloudSQLConnectRequest(
+            user_agent=f"{self._client._user_agent}",  # type: ignore
+            protocol_type=connectorspb.CloudSQLConnectRequest.TCP,
+        )
+
+        # set I/O timeout
+        sock.settimeout(IO_TIMEOUT)
+
+        # pack big-endian unsigned integer (4 bytes)
+        packed_len = struct.pack(">I", req.ByteSize())
+
+        # send metadata message length and request message
+        sock.sendall(packed_len + req.SerializeToString())
+
+        # form metadata exchange response
+        resp = connectorspb.CloudSQLConnectResponse()
+
+        # read metadata message length (4 bytes)
+        message_len_buffer_size = struct.Struct(">I").size
+        message_len_buffer = b""
+        while message_len_buffer_size > 0:
+            chunk = sock.recv(message_len_buffer_size)
+            if not chunk:
+                raise RuntimeError(
+                    "Connection closed while getting metadata exchange length!"
+                )
+            message_len_buffer += chunk
+            message_len_buffer_size -= len(chunk)
+
+        (message_len,) = struct.unpack(">I", message_len_buffer)
+
+        # read metadata exchange message
+        buffer = b""
+        while message_len > 0:
+            chunk = sock.recv(message_len)
+            if not chunk:
+                raise RuntimeError(
+                    "Connection closed while performing metadata exchange!"
+                )
+            buffer += chunk
+            message_len -= len(chunk)
+
+        # parse metadata exchange response from buffer
+        resp.ParseFromString(buffer)
+
+        # reset socket back to blocking mode
+        sock.setblocking(True)
+
+        # validate metadata exchange response
+        if resp.response_code != connectorspb.CloudSQLConnectResponse.OK:
+            raise ValueError(
+                f"Metadata Exchange request has failed with error: {resp.error}"
+            )
+
+        return sock
+
     async def _remove_cached(
         self, instance_connection_string: str, enable_iam_auth: bool
     ) -> None:

google/cloud/sql/proto/cloud_sql_metadata_exchange_pb2.py

@@ -0,0 +1,48 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: google/cloud/sql/v1beta4/cloud_sql_metadata_exchange.proto
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
+    b'\n:google/cloud/sql/v1beta4/cloud_sql_metadata_exchange.proto\x12\x18google.cloud.sql.v1beta4"\x90\x02\n\x16\x43loudSQLConnectRequest\x12\x14\n\x07version\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nuser_agent\x18\x02 \x01(\tH\x01\x88\x01\x01\x12Y\n\rprotocol_type\x18\x03 \x01(\x0e\x32=.google.cloud.sql.v1beta4.CloudSQLConnectRequest.ProtocolTypeH\x02\x88\x01\x01"?\n\x0cProtocolType\x12\x1d\n\x19PROTOCOL_TYPE_UNSPECIFIED\x10\x00\x12\x07\n\x03TCP\x10\x01\x12\x07\n\x03UDS\x10\x02\x42\n\n\x08_versionB\r\n\x0b_user_agentB\x10\n\x0e_protocol_type"\x8d\x01\n\x17\x43loudSQLConnectResponse\x12\x42\n\rresponse_code\x18\x01 \x01(\x0e\x32&.google.cloud.sql.v1beta4.ResponseCodeH\x00\x88\x01\x01\x12\x12\n\x05\x65rror\x18\x02 \x01(\tH\x01\x88\x01\x01\x42\x10\n\x0e_response_codeB\x08\n\x06_error*@\n\x0cResponseCode\x12\x1d\n\x19RESPONSE_CODE_UNSPECIFIED\x10\x00\x12\x06\n\x02OK\x10\x01\x12\t\n\x05\x45RROR\x10\x02\x62\x06proto3'
+)
+
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
+_builder.BuildTopDescriptorsAndMessages(
+    DESCRIPTOR, "google.cloud.sql.v1beta4.cloud_sql_metadata_exchange_pb2", globals()
+)
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+    DESCRIPTOR._options = None
+    _RESPONSECODE._serialized_start = 507
+    _RESPONSECODE._serialized_end = 571
+    _CLOUDSQLCONNECTREQUEST._serialized_start = 89
+    _CLOUDSQLCONNECTREQUEST._serialized_end = 361
+    _CLOUDSQLCONNECTREQUEST_PROTOCOLTYPE._serialized_start = 253
+    _CLOUDSQLCONNECTREQUEST_PROTOCOLTYPE._serialized_end = 316
+    _CLOUDSQLCONNECTRESPONSE._serialized_start = 364
+    _CLOUDSQLCONNECTRESPONSE._serialized_end = 505
+# @@protoc_insertion_point(module_scope)

google/cloud/sql/proto/cloud_sql_metadata_exchange_pb2.pyi

@@ -0,0 +1,65 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import ClassVar as _ClassVar
+from typing import Optional as _Optional
+from typing import Union as _Union
+
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
+
+DESCRIPTOR: _descriptor.FileDescriptor
+ERROR: ResponseCode
+OK: ResponseCode
+RESPONSE_CODE_UNSPECIFIED: ResponseCode
+
+class CloudSQLConnectRequest(_message.Message):
+    __slots__ = ["protocol_type", "user_agent", "version"]
+
+    class ProtocolType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+        __slots__ = []
+
+    PROTOCOL_TYPE_FIELD_NUMBER: _ClassVar[int]
+    PROTOCOL_TYPE_UNSPECIFIED: CloudSQLConnectRequest.ProtocolType
+    TCP: CloudSQLConnectRequest.ProtocolType
+    UDS: CloudSQLConnectRequest.ProtocolType
+    USER_AGENT_FIELD_NUMBER: _ClassVar[int]
+    VERSION_FIELD_NUMBER: _ClassVar[int]
+    protocol_type: CloudSQLConnectRequest.ProtocolType
+    user_agent: str
+    version: str
+    def __init__(
+        self,
+        version: _Optional[str] = ...,
+        user_agent: _Optional[str] = ...,
+        protocol_type: _Optional[
+            _Union[CloudSQLConnectRequest.ProtocolType, str]
+        ] = ...,
+    ) -> None: ...
+
+class CloudSQLConnectResponse(_message.Message):
+    __slots__ = ["error", "response_code"]
+    ERROR_FIELD_NUMBER: _ClassVar[int]
+    RESPONSE_CODE_FIELD_NUMBER: _ClassVar[int]
+    error: str
+    response_code: ResponseCode
+    def __init__(
+        self,
+        response_code: _Optional[_Union[ResponseCode, str]] = ...,
+        error: _Optional[str] = ...,
+    ) -> None: ...
+
+class ResponseCode(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+    __slots__ = []