feat: add support for auto IAM authentication to Connector

google/api/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 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.

google/api/field_behavior_pb2.py

@@ -0,0 +1,52 @@
+# Copyright 2024 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.
+
+# type: ignore
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: google/api/field_behavior.proto
+# isort: skip_file
+"""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()
+
+
+from google.protobuf import descriptor_pb2 as google_dot_protobuf_dot_descriptor__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
+    b"\n\x1fgoogle/api/field_behavior.proto\x12\ngoogle.api\x1a google/protobuf/descriptor.proto*\xb6\x01\n\rFieldBehavior\x12\x1e\n\x1a\x46IELD_BEHAVIOR_UNSPECIFIED\x10\x00\x12\x0c\n\x08OPTIONAL\x10\x01\x12\x0c\n\x08REQUIRED\x10\x02\x12\x0f\n\x0bOUTPUT_ONLY\x10\x03\x12\x0e\n\nINPUT_ONLY\x10\x04\x12\r\n\tIMMUTABLE\x10\x05\x12\x12\n\x0eUNORDERED_LIST\x10\x06\x12\x15\n\x11NON_EMPTY_DEFAULT\x10\x07\x12\x0e\n\nIDENTIFIER\x10\x08:Q\n\x0e\x66ield_behavior\x12\x1d.google.protobuf.FieldOptions\x18\x9c\x08 \x03(\x0e\x32\x19.google.api.FieldBehaviorBp\n\x0e\x63om.google.apiB\x12\x46ieldBehaviorProtoP\x01ZAgoogle.golang.org/genproto/googleapis/api/annotations;annotations\xa2\x02\x04GAPIb\x06proto3"
+)
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(
+    DESCRIPTOR, "google.api.field_behavior_pb2", _globals
+)
+if _descriptor._USE_C_DESCRIPTORS == False:
+    google_dot_protobuf_dot_descriptor__pb2.FieldOptions.RegisterExtension(
+        field_behavior
+    )
+
+    DESCRIPTOR._options = None
+    DESCRIPTOR._serialized_options = b"\n\016com.google.apiB\022FieldBehaviorProtoP\001ZAgoogle.golang.org/genproto/googleapis/api/annotations;annotations\242\002\004GAPI"
+    _globals["_FIELDBEHAVIOR"]._serialized_start = 82
+    _globals["_FIELDBEHAVIOR"]._serialized_end = 264
+# @@protoc_insertion_point(module_scope)

google/cloud/alloydb/connector/client.py

@@ -67,6 +67,7 @@ def __init__(
         self._client = client if client else aiohttp.ClientSession(headers=headers)
         self._credentials = credentials
         self._alloydb_api_endpoint = alloydb_api_endpoint
+        self._user_agent = USER_AGENT
 
     async def _get_metadata(
         self,
@@ -149,6 +150,7 @@ async def _get_client_certificate(
         data = {
             "publicKey": pub_key,
             "certDuration": "3600s",
+            "useMetadataExchange": True,
         }
 
         resp = await self._client.post(

google/cloud/alloydb/connector/connector.py

@@ -16,6 +16,8 @@
 
 import asyncio
 from functools import partial
+import socket
+import struct
 from threading import Thread
 from types import TracebackType
 from typing import Any, Dict, Optional, Type, TYPE_CHECKING
@@ -27,10 +29,15 @@
 from google.cloud.alloydb.connector.instance import Instance
 import google.cloud.alloydb.connector.pg8000 as pg8000
 from google.cloud.alloydb.connector.utils import generate_keys
+import google.cloud.alloydb_connectors_v1.proto.resources_pb2 as connectorspb
 
 if TYPE_CHECKING:
+    import ssl
+
     from google.auth.credentials import Credentials
 
+SERVER_PROXY_PORT = 5433
+
 
 class Connector:
     """A class to configure and create connections to Cloud SQL instances.
@@ -45,13 +52,15 @@ class Connector:
             Defaults to None, picking up project from environment.
         alloydb_api_endpoint (str): Base URL to use when calling
             the AlloyDB API endpoint. Defaults to "https://alloydb.googleapis.com".
+        enable_iam_auth (bool): Enables automatic IAM database authentication.
     """
 
     def __init__(
         self,
         credentials: Optional[Credentials] = None,
         quota_project: Optional[str] = None,
         alloydb_api_endpoint: str = "https://alloydb.googleapis.com",
+        enable_iam_auth: bool = False,
     ) -> None:
         # create event loop and start it in background thread
         self._loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
@@ -61,6 +70,7 @@ def __init__(
         # initialize default params
         self._quota_project = quota_project
         self._alloydb_api_endpoint = alloydb_api_endpoint
+        self._enable_iam_auth = enable_iam_auth
         # initialize credentials
         scopes = ["https://www.googleapis.com/auth/cloud-platform"]
         if credentials:
@@ -123,6 +133,7 @@ async def connect_async(self, instance_uri: str, driver: str, **kwargs: Any) ->
             self._client = AlloyDBClient(
                 self._alloydb_api_endpoint, self._quota_project, self._credentials
             )
+        enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth)
         # use existing connection info if possible
         if instance_uri in self._instances:
             instance = self._instances[instance_uri]
@@ -150,13 +161,117 @@ async def connect_async(self, instance_uri: str, driver: str, **kwargs: Any) ->
 
         # synchronous drivers are blocking and run using executor
         try:
-            connect_partial = partial(connector, ip_address, context, **kwargs)
+            metadata_partial = partial(
+                self.metadata_exchange, ip_address, context, enable_iam_auth, driver
+            )
+            sock = await self._loop.run_in_executor(None, metadata_partial)
+            connect_partial = partial(connector, sock, **kwargs)
             return await self._loop.run_in_executor(None, connect_partial)
         except Exception:
             # we attempt a force refresh, then throw the error
             await instance.force_refresh()
             raise
 
+    def metadata_exchange(
+        self, ip_address: str, ctx: ssl.SSLContext, enable_iam_auth: bool, driver: str
+    ) -> ssl.SSLSocket:
+        """
+        Sends metadata about the connection prior to the database
+        protocol taking over.
+
+        The exchange consists of four steps:
+
+        1. Prepare a MetadataExchangeRequest including the IAM Principal's
+           OAuth2 token, the user agent, and the requested authentication type.
+
+        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
+           MetadataExchangeResponse 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:
+            ip_address (str): IP address of AlloyDB instance to connect to.
+            ctx (ssl.SSLContext): Context used to create a TLS connection
+                with AlloyDB instance ssl certificates.
+            enable_iam_auth (bool): Flag to enable IAM database authentication.
+            driver (str): A string representing the database driver to connect with.
+                Supported drivers are pg8000.
+
+        Returns:
+            sock (ssl.SSLSocket): mTLS/SSL socket connected to AlloyDB Proxy server.
+        """
+        # Create socket and wrap with SSL/TLS context
+        sock = ctx.wrap_socket(
+            socket.create_connection((ip_address, SERVER_PROXY_PORT)),
+            server_hostname=ip_address,
+        )
+        # set auth type for metadata exchange
+        auth_type = connectorspb.MetadataExchangeRequest.DB_NATIVE
+        if enable_iam_auth:
+            auth_type = connectorspb.MetadataExchangeRequest.AUTO_IAM
+
+        # form metadata exchange request
+        req = connectorspb.MetadataExchangeRequest(
+            user_agent=f"{self._client._user_agent}+{driver}",  # type: ignore
+            auth_type=auth_type,
+            oauth2_token=self._credentials.token,
+        )
+
+        # set I/O timeout
+        sock.settimeout(30)
+
+        # pack big-endian unsigned integer (4 bytes)
+        packed_len = struct.pack(">I", req.ByteSize())
+
+        # send metadata message length
+        sock.sendall(packed_len)
+        # send metadata request message
+        sock.sendall(req.SerializeToString())
+
+        # form metadata exchange response
+        resp = connectorspb.MetadataExchangeResponse()
+
+        # 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)
+
+        # validate metadata exchange response
+        if resp.response_code != connectorspb.MetadataExchangeResponse.OK:
+            raise ValueError("Metadata Exchange request has failed!")
+
+        return sock
+
     def __enter__(self) -> "Connector":
         """Enter context manager by returning Connector object"""
         return self

google/cloud/alloydb/connector/pg8000.py

@@ -11,49 +11,35 @@
 # 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.
-import socket
-import ssl
 from typing import Any, TYPE_CHECKING
 
-SERVER_PROXY_PORT = 5433
-
 if TYPE_CHECKING:
+    import ssl
+
     import pg8000
 
 
-def connect(
-    ip_address: str, ctx: ssl.SSLContext, **kwargs: Any
-) -> "pg8000.dbapi.Connection":
+def connect(sock: "ssl.SSLSocket", **kwargs: Any) -> "pg8000.dbapi.Connection":
     """Create a pg8000 DBAPI connection object.
 
     Args:
-        ip_address (str): IP address of AlloyDB instance to connect to.
-        ctx (ssl.SSLContext): Context used to create a TLS connection
-            with AlloyDB instance ssl certificates.
+        sock (ssl.SSLSocket): SSL/TLS secure socket stream connected to the
+            AlloyDB proxy server.
 
     Returns:
         pg8000.dbapi.Connection: A pg8000 Connection object for
         the AlloyDB instance.
     """
-    # Connecting through pg8000 is done by passing in an SSL Context and setting the
-    # "request_ssl" attr to false. This works because when "request_ssl" is false,
-    # the driver skips the database level SSL/TLS exchange, but still uses the
-    # ssl_context (if it is not None) to create the connection.
     try:
         import pg8000
     except ImportError:
         raise ImportError(
             'Unable to import module "pg8000." Please install and try again.'
         )
-    # Create socket and wrap with context.
-    sock = ctx.wrap_socket(
-        socket.create_connection((ip_address, SERVER_PROXY_PORT)),
-        server_hostname=ip_address,
-    )
 
     user = kwargs.pop("user")
     db = kwargs.pop("db")
-    passwd = kwargs.pop("password")
+    passwd = kwargs.pop("password", None)
     return pg8000.dbapi.connect(
         user,
         database=db,

google/cloud/alloydb_connectors_v1/proto/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 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.

google/cloud/alloydb_connectors_v1/proto/resources_pb2.py

@@ -0,0 +1,61 @@
+# Copyright 2024 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/alloydb_connectors_v1/proto/resources.proto
+# isort: skip_file
+"""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()
+
+
+from google.api import field_behavior_pb2 as google_dot_api_dot_field__behavior__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
+    b'\n8google/cloud/alloydb_connectors_v1/proto/resources.proto\x12"google.cloud.alloydb.connectors.v1\x1a\x1fgoogle/api/field_behavior.proto"\xe6\x01\n\x17MetadataExchangeRequest\x12\x18\n\nuser_agent\x18\x01 \x01(\tB\x04\xe2\x41\x01\x01\x12W\n\tauth_type\x18\x02 \x01(\x0e\x32\x44.google.cloud.alloydb.connectors.v1.MetadataExchangeRequest.AuthType\x12\x14\n\x0coauth2_token\x18\x03 \x01(\t"B\n\x08\x41uthType\x12\x19\n\x15\x41UTH_TYPE_UNSPECIFIED\x10\x00\x12\r\n\tDB_NATIVE\x10\x01\x12\x0c\n\x08\x41UTO_IAM\x10\x02"\xd3\x01\n\x18MetadataExchangeResponse\x12`\n\rresponse_code\x18\x01 \x01(\x0e\x32I.google.cloud.alloydb.connectors.v1.MetadataExchangeResponse.ResponseCode\x12\x13\n\x05\x65rror\x18\x02 \x01(\tB\x04\xe2\x41\x01\x01"@\n\x0cResponseCode\x12\x1d\n\x19RESPONSE_CODE_UNSPECIFIED\x10\x00\x12\x06\n\x02OK\x10\x01\x12\t\n\x05\x45RROR\x10\x02\x42\xf5\x01\n&com.google.cloud.alloydb.connectors.v1B\x0eResourcesProtoP\x01ZFcloud.google.com/go/alloydb/connectors/apiv1/connectorspb;connectorspb\xaa\x02"Google.Cloud.AlloyDb.Connectors.V1\xca\x02"Google\\Cloud\\AlloyDb\\Connectors\\V1\xea\x02&Google::Cloud::AlloyDb::Connectors::V1b\x06proto3'
+)
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(
+    DESCRIPTOR, "google.cloud.alloydb_connectors_v1.proto.resources_pb2", _globals
+)
+if _descriptor._USE_C_DESCRIPTORS == False:
+    DESCRIPTOR._options = None
+    DESCRIPTOR._serialized_options = b'\n&com.google.cloud.alloydb.connectors.v1B\016ResourcesProtoP\001ZFcloud.google.com/go/alloydb/connectors/apiv1/connectorspb;connectorspb\252\002"Google.Cloud.AlloyDb.Connectors.V1\312\002"Google\\Cloud\\AlloyDb\\Connectors\\V1\352\002&Google::Cloud::AlloyDb::Connectors::V1'
+    _METADATAEXCHANGEREQUEST.fields_by_name["user_agent"]._options = None
+    _METADATAEXCHANGEREQUEST.fields_by_name[
+        "user_agent"
+    ]._serialized_options = b"\342A\001\001"
+    _METADATAEXCHANGERESPONSE.fields_by_name["error"]._options = None
+    _METADATAEXCHANGERESPONSE.fields_by_name[
+        "error"
+    ]._serialized_options = b"\342A\001\001"
+    _globals["_METADATAEXCHANGEREQUEST"]._serialized_start = 130
+    _globals["_METADATAEXCHANGEREQUEST"]._serialized_end = 360
+    _globals["_METADATAEXCHANGEREQUEST_AUTHTYPE"]._serialized_start = 294
+    _globals["_METADATAEXCHANGEREQUEST_AUTHTYPE"]._serialized_end = 360
+    _globals["_METADATAEXCHANGERESPONSE"]._serialized_start = 363
+    _globals["_METADATAEXCHANGERESPONSE"]._serialized_end = 574
+    _globals["_METADATAEXCHANGERESPONSE_RESPONSECODE"]._serialized_start = 510
+    _globals["_METADATAEXCHANGERESPONSE_RESPONSECODE"]._serialized_end = 574
+# @@protoc_insertion_point(module_scope)

requirements.txt

@@ -2,3 +2,4 @@ aiohttp==3.9.1
 cryptography==41.0.7
 google-auth==2.26.2
 requests==2.31.0
+protobuf==4.25.1