| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
1 parent 4d818b9 commit 7b70fea
6 files changed
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -121,7 +121,7 @@ class Request(transport.Request): | |||
| 121 | 121 | .. automethod:: __call__ | |
| 122 | 122 | """ | |
| 123 | 123 | ||
| 124 | - def __init__(self, session: aiohttp.ClientSession = None): | ||
| 124 | + def __init__(self, session: Optional[aiohttp.ClientSession] = None): | ||
| 125 | 125 | self._session = session | |
| 126 | 126 | self._closed = False | |
| 127 | 127 | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -17,7 +17,12 @@ | |||
| 17 | 17 | """ | |
| 18 | 18 | ||
| 19 | 19 | import asyncio | |
| 20 | + import contextlib | ||
| 20 | 21 | import logging | |
| 22 | + import os | ||
| 23 | + import ssl | ||
| 24 | + import tempfile | ||
| 25 | + from typing import Optional | ||
| 21 | 26 | ||
| 22 | 27 | from google.auth import exceptions | |
| 23 | 28 | import google.auth.transport._mtls_helper | |
@@ -26,6 +31,61 @@ | |||
| 26 | 31 | _LOGGER = logging.getLogger(__name__) | |
| 27 | 32 | ||
| 28 | 33 | ||
| 34 | + @contextlib.contextmanager | ||
| 35 | + def _create_temp_file(content: bytes): | ||
| 36 | + """Creates a temporary file with the given content. | ||
| 37 | + | ||
| 38 | + Args: | ||
| 39 | + content (bytes): The content to write to the file. | ||
| 40 | + | ||
| 41 | + Yields: | ||
| 42 | + str: The path to the temporary file. | ||
| 43 | + """ | ||
| 44 | + # Create a temporary file that is readable only by the owner. | ||
| 45 | + fd, file_path = tempfile.mkstemp() | ||
| 46 | + try: | ||
| 47 | + with os.fdopen(fd, "wb") as f: | ||
| 48 | + f.write(content) | ||
| 49 | + yield file_path | ||
| 50 | + finally: | ||
| 51 | + # Securely delete the file after use. | ||
| 52 | + if os.path.exists(file_path): | ||
| 53 | + os.remove(file_path) | ||
| 54 | + | ||
| 55 | + | ||
| 56 | + def make_client_cert_ssl_context( | ||
| 57 | + cert_bytes: bytes, key_bytes: bytes, passphrase: Optional[bytes] = None | ||
| 58 | + ) -> ssl.SSLContext: | ||
| 59 | + """Creates an SSLContext with the given client certificate and key. | ||
| 60 | + This function writes the certificate and key to temporary files so that | ||
| 61 | + ssl.create_default_context can load them, as the ssl module requires | ||
| 62 | + file paths for client certificates. These temporary files are deleted | ||
| 63 | + immediately after the SSL context is created. | ||
| 64 | + Args: | ||
| 65 | + cert_bytes (bytes): The client certificate content in PEM format. | ||
| 66 | + key_bytes (bytes): The client private key content in PEM format. | ||
| 67 | + passphrase (Optional[bytes]): The passphrase for the private key, if any. | ||
| 68 | + Returns: | ||
| 69 | + ssl.SSLContext: The configured SSL context with client certificate. | ||
| 70 | + | ||
| 71 | + Raises: | ||
| 72 | + google.auth.exceptions.TransportError: If there is an error loading the certificate. | ||
| 73 | + """ | ||
| 74 | + with _create_temp_file(cert_bytes) as cert_path, _create_temp_file( | ||
| 75 | + key_bytes | ||
| 76 | + ) as key_path: | ||
| 77 | + try: | ||
| 78 | + context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) | ||
| 79 | + context.load_cert_chain( | ||
| 80 | + certfile=cert_path, keyfile=key_path, password=passphrase | ||
| 81 | + ) | ||
| 82 | + return context | ||
| 83 | + except (ssl.SSLError, OSError, IOError, ValueError, RuntimeError) as exc: | ||
| 84 | + raise exceptions.TransportError( | ||
| 85 | + "Failed to load client certificate and key for mTLS." | ||
| 86 | + ) from exc | ||
| 87 | + | ||
| 88 | + | ||
| 29 | 89 | async def _run_in_executor(func, *args): | |
| 30 | 90 | """Run a blocking function in an executor to avoid blocking the event loop. | |
| 31 | 91 | ||
@@ -44,7 +104,7 @@ def default_client_cert_source(): | |||
| 44 | 104 | """Get a callback which returns the default client SSL credentials. | |
| 45 | 105 | ||
| 46 | 106 | Returns: | |
| 47 | - Awaitable[Callable[[], [bytes, bytes]]]: A callback which returns the default | ||
| 107 | + Awaitable[Callable[[], Tuple[bytes, bytes]]]: A callback which returns the default | ||
| 48 | 108 | client certificate bytes and private key bytes, both in PEM format. | |
| 49 | 109 | ||
| 50 | 110 | Raises: | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -21,17 +21,22 @@ | |||
| 21 | 21 | from google.auth import _exponential_backoff, exceptions | |
| 22 | 22 | from google.auth.aio import transport | |
| 23 | 23 | from google.auth.aio.credentials import Credentials | |
| 24 | + from google.auth.aio.transport import mtls | ||
| 24 | 25 | from google.auth.exceptions import TimeoutError | |
| 26 | + import google.auth.transport._mtls_helper | ||
| 25 | 27 | ||
| 26 | 28 | if TYPE_CHECKING: # pragma: NO COVER | |
| 29 | + import aiohttp | ||
| 27 | 30 | from aiohttp import ClientTimeout # type: ignore | |
| 28 | 31 | ||
| 29 | 32 | else: | |
| 30 | 33 | try: | |
| 34 | + import aiohttp | ||
| 31 | 35 | from aiohttp import ClientTimeout | |
| 32 | 36 | except (ImportError, AttributeError): | |
| 33 | 37 | ClientTimeout = None | |
| 34 | 38 | ||
| 39 | + # Tracks the internal aiohttp installation and usage | ||
| 35 | 40 | try: | |
| 36 | 41 | from google.auth.aio.transport.aiohttp import Request as AiohttpRequest | |
| 37 | 42 | ||
@@ -133,12 +138,88 @@ def __init__( | |||
| 133 | 138 | _auth_request = auth_request | |
| 134 | 139 | if not _auth_request and AIOHTTP_INSTALLED: | |
| 135 | 140 | _auth_request = AiohttpRequest() | |
| 141 | + self._is_mtls = False | ||
| 142 | + self._mtls_init_task = None | ||
| 143 | + self._cached_cert = None | ||
| 136 | 144 | if _auth_request is None: | |
| 137 | 145 | raise exceptions.TransportError( | |
| 138 | 146 | "`auth_request` must either be configured or the external package `aiohttp` must be installed to use the default value." | |
| 139 | 147 | ) | |
| 140 | 148 | self._auth_request = _auth_request | |
| 141 | 149 | ||
| 150 | + async def configure_mtls_channel(self, client_cert_callback=None): | ||
| 151 | + """Configure the client certificate and key for SSL connection. | ||
| 152 | + | ||
| 153 | + The function does nothing unless `GOOGLE_API_USE_CLIENT_CERTIFICATE` is | ||
| 154 | + explicitly set to `true`. In this case if client certificate and key are | ||
| 155 | + successfully obtained (from the given client_cert_callback or from application | ||
| 156 | + default SSL credentials), the underlying transport will be reconfigured | ||
| 157 | + to use mTLS. | ||
| 158 | + Note: This function does nothing if the `aiohttp` library is not | ||
| 159 | + installed. | ||
| 160 | + Important: Calling this method will close any ongoing API requests associated | ||
| 161 | + with the current session. To ensure a smooth transition, it is recommended | ||
| 162 | + to call this during session initialization. | ||
| 163 | + | ||
| 164 | + Args: | ||
| 165 | + client_cert_callback (Optional[Callable[[], (bytes, bytes)]]): | ||
| 166 | + The optional callback returns the client certificate and private | ||
| 167 | + key bytes both in PEM format. | ||
| 168 | + If the callback is None, application default SSL credentials | ||
| 169 | + will be used. | ||
| 170 | + | ||
| 171 | + Raises: | ||
| 172 | + google.auth.exceptions.MutualTLSChannelError: If mutual TLS channel | ||
| 173 | + creation failed for any reason. | ||
| 174 | + """ | ||
| 175 | + if self._mtls_init_task is None: | ||
| 176 | + | ||
| 177 | + async def _do_configure(): | ||
| 178 | + # Run the blocking check in an executor | ||
| 179 | + use_client_cert = await mtls._run_in_executor( | ||
| 180 | + google.auth.transport._mtls_helper.check_use_client_cert | ||
| 181 | + ) | ||
| 182 | + if not use_client_cert: | ||
| 183 | + self._is_mtls = False | ||
| 184 | + return | ||
| 185 | + | ||
| 186 | + try: | ||
| 187 | + ( | ||
| 188 | + self._is_mtls, | ||
| 189 | + cert, | ||
| 190 | + key, | ||
| 191 | + ) = await mtls.get_client_cert_and_key(client_cert_callback) | ||
| 192 | + | ||
| 193 | + if self._is_mtls: | ||
| 194 | + self._cached_cert = cert | ||
| 195 | + ssl_context = await mtls._run_in_executor( | ||
| 196 | + mtls.make_client_cert_ssl_context, cert, key | ||
| 197 | + ) | ||
| 198 | + | ||
| 199 | + # Re-create the auth request with the new SSL context | ||
| 200 | + if AIOHTTP_INSTALLED and isinstance( | ||
| 201 | + self._auth_request, AiohttpRequest | ||
| 202 | + ): | ||
| 203 | + connector = aiohttp.TCPConnector(ssl=ssl_context) | ||
| 204 | + new_session = aiohttp.ClientSession(connector=connector) | ||
| 205 | + | ||
| 206 | + old_auth_request = self._auth_request | ||
| 207 | + self._auth_request = AiohttpRequest(session=new_session) | ||
| 208 | + | ||
| 209 | + await old_auth_request.close() | ||
| 210 | + | ||
| 211 | + except ( | ||
| 212 | + exceptions.ClientCertError, | ||
| 213 | + ImportError, | ||
| 214 | + OSError, | ||
| 215 | + ) as caught_exc: | ||
| 216 | + new_exc = exceptions.MutualTLSChannelError(caught_exc) | ||
| 217 | + raise new_exc from caught_exc | ||
| 218 | + | ||
| 219 | + self._mtls_init_task = asyncio.create_task(_do_configure()) | ||
| 220 | + | ||
| 221 | + return await self._mtls_init_task | ||
| 222 | + | ||
| 142 | 223 | async def request( | |
| 143 | 224 | self, | |
| 144 | 225 | method: str, | |
@@ -182,22 +263,37 @@ async def request( | |||
| 182 | 263 | the configured `max_allowed_time` or the request exceeds the configured | |
| 183 | 264 | `timeout`. | |
| 184 | 265 | """ | |
| 185 | - | ||
| 266 | + if self._mtls_init_task: | ||
| 267 | + try: | ||
| 268 | + await self._mtls_init_task | ||
| 269 | + except Exception: | ||
| 270 | + # Suppress all exceptions from the background mTLS initialization task, | ||
| 271 | + # allowing the request to fail naturally elsewhere. | ||
| 272 | + pass | ||
| 186 | 273 | retries = _exponential_backoff.AsyncExponentialBackoff( | |
| 187 | 274 | total_attempts=total_attempts, | |
| 188 | 275 | ) | |
| 276 | + if headers is None: | ||
| 277 | + headers = {} | ||
| 189 | 278 | async with timeout_guard(max_allowed_time) as with_timeout: | |
| 190 | 279 | await with_timeout( | |
| 191 | 280 | # Note: before_request will attempt to refresh credentials if expired. | |
| 192 | 281 | self._credentials.before_request( | |
| 193 | 282 | self._auth_request, method, url, headers | |
| 194 | 283 | ) | |
| 195 | 284 | ) | |
| 285 | + actual_timeout: float = 0.0 | ||
| 286 | + if ClientTimeout is not None and isinstance(timeout, ClientTimeout): | ||
| 287 | + actual_timeout = timeout.total if timeout.total is not None else 0.0 | ||
| 288 | + elif isinstance(timeout, (int, float)): | ||
| 289 | + actual_timeout = float(timeout) | ||
| 196 | 290 | # Workaround issue in python 3.9 related to code coverage by adding `# pragma: no branch` | |
| 197 | 291 | # See https://github.com/googleapis/gapic-generator-python/pull/1174#issuecomment-1025132372 | |
| 198 | 292 | async for _ in retries: # pragma: no branch | |
| 199 | 293 | response = await with_timeout( | |
| 200 | - self._auth_request(url, method, data, headers, timeout, **kwargs) | ||
| 294 | + self._auth_request( | ||
| 295 | + url, method, data, headers, actual_timeout, **kwargs | ||
| 296 | + ) | ||
| 201 | 297 | ) | |
| 202 | 298 | if response.status_code not in transport.DEFAULT_RETRYABLE_STATUS_CODES: | |
| 203 | 299 | break | |
@@ -468,6 +564,11 @@ async def delete( | |||
| 468 | 564 | **kwargs, | |
| 469 | 565 | ) | |
| 470 | 566 | ||
| 567 | + @property | ||
| 568 | + def is_mtls(self): | ||
| 569 | + """Indicates if mutual TLS is enabled.""" | ||
| 570 | + return self._is_mtls | ||
| 571 | + | ||
| 471 | 572 | async def close(self) -> None: | |
| 472 | 573 | """ | |
| 473 | 574 | Close the underlying auth request session. | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -91,7 +91,7 @@ def blacken(session): | |||
| 91 | 91 | @nox.session(python=DEFAULT_PYTHON_VERSION) | |
| 92 | 92 | def mypy(session): | |
| 93 | 93 | """Verify type hints are mypy compatible.""" | |
| 94 | - session.install("-e", ".") | ||
| 94 | + session.install("-e", ".[aiohttp]") | ||
| 95 | 95 | session.install( | |
| 96 | 96 | "mypy", | |
| 97 | 97 | "types-certifi", | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -0,0 +1,123 @@ | |||
| 1 | + # Copyright 2026 Google LLC | ||
| 2 | + # | ||
| 3 | + # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | + # you may not use this file except in compliance with the License. | ||
| 5 | + # You may obtain a copy of the License at | ||
| 6 | + # | ||
| 7 | + # http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | + # | ||
| 9 | + # Unless required by applicable law or agreed to in writing, software | ||
| 10 | + # distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | + # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | + # See the License for the specific language governing permissions and | ||
| 13 | + # limitations under the License. | ||
| 14 | + | ||
| 15 | + import json | ||
| 16 | + import os | ||
| 17 | + import ssl | ||
| 18 | + from unittest import mock | ||
| 19 | + | ||
| 20 | + import pytest | ||
| 21 | + | ||
| 22 | + from google.auth import exceptions | ||
| 23 | + from google.auth.aio import credentials | ||
| 24 | + from google.auth.aio.transport import sessions | ||
| 25 | + | ||
| 26 | + # This is the valid "workload" format the library expects | ||
| 27 | + VALID_WORKLOAD_CONFIG = { | ||
| 28 | + "version": 1, | ||
| 29 | + "cert_configs": { | ||
| 30 | + "workload": {"cert_path": "/tmp/mock_cert.pem", "key_path": "/tmp/mock_key.pem"} | ||
| 31 | + }, | ||
| 32 | + } | ||
| 33 | + | ||
| 34 | + | ||
| 35 | + class TestSessionsMtls: | ||
| 36 | + @pytest.mark.asyncio | ||
| 37 | + async def test_configure_mtls_channel(self): | ||
| 38 | + """ | ||
| 39 | + Tests that the mTLS channel configures correctly when a | ||
| 40 | + valid workload config is mocked. | ||
| 41 | + """ | ||
| 42 | + with mock.patch.dict( | ||
| 43 | + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} | ||
| 44 | + ), mock.patch("os.path.exists") as mock_exists, mock.patch( | ||
| 45 | + "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) | ||
| 46 | + ), mock.patch( | ||
| 47 | + "google.auth.aio.transport.mtls.get_client_cert_and_key" | ||
| 48 | + ) as mock_helper, mock.patch( | ||
| 49 | + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" | ||
| 50 | + ) as mock_make_context: | ||
| 51 | + mock_exists.return_value = True | ||
| 52 | + mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") | ||
| 53 | + | ||
| 54 | + mock_context = mock.Mock(spec=ssl.SSLContext) | ||
| 55 | + mock_make_context.return_value = mock_context | ||
| 56 | + | ||
| 57 | + mock_creds = mock.AsyncMock(spec=credentials.Credentials) | ||
| 58 | + session = sessions.AsyncAuthorizedSession(mock_creds) | ||
| 59 | + | ||
| 60 | + await session.configure_mtls_channel() | ||
| 61 | + | ||
| 62 | + assert session._is_mtls is True | ||
| 63 | + mock_make_context.assert_called_once_with( | ||
| 64 | + b"fake_cert_data", b"fake_key_data" | ||
| 65 | + ) | ||
| 66 | + | ||
| 67 | + @pytest.mark.asyncio | ||
| 68 | + async def test_configure_mtls_channel_disabled(self): | ||
| 69 | + """ | ||
| 70 | + Tests behavior when the config file does not exist. | ||
| 71 | + """ | ||
| 72 | + with mock.patch.dict( | ||
| 73 | + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} | ||
| 74 | + ), mock.patch("os.path.exists") as mock_exists: | ||
| 75 | + mock_exists.return_value = False | ||
| 76 | + mock_creds = mock.AsyncMock(spec=credentials.Credentials) | ||
| 77 | + session = sessions.AsyncAuthorizedSession(mock_creds) | ||
| 78 | + | ||
| 79 | + await session.configure_mtls_channel() | ||
| 80 | + | ||
| 81 | + # If the file doesn't exist, it shouldn't error; it just won't use mTLS | ||
| 82 | + assert session._is_mtls is False | ||
| 83 | + | ||
| 84 | + @pytest.mark.asyncio | ||
| 85 | + async def test_configure_mtls_channel_invalid_format(self): | ||
| 86 | + """ | ||
| 87 | + Verifies that the MutualTLSChannelError is raised for bad formats. | ||
| 88 | + """ | ||
| 89 | + with mock.patch.dict( | ||
| 90 | + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} | ||
| 91 | + ), mock.patch("os.path.exists") as mock_exists, mock.patch( | ||
| 92 | + "builtins.open", mock.mock_open(read_data='{"invalid": "format"}') | ||
| 93 | + ): | ||
| 94 | + mock_exists.return_value = True | ||
| 95 | + mock_creds = mock.AsyncMock(spec=credentials.Credentials) | ||
| 96 | + session = sessions.AsyncAuthorizedSession(mock_creds) | ||
| 97 | + | ||
| 98 | + with pytest.raises(exceptions.MutualTLSChannelError): | ||
| 99 | + await session.configure_mtls_channel() | ||
| 100 | + | ||
| 101 | + @pytest.mark.asyncio | ||
| 102 | + async def test_configure_mtls_channel_mock_callback(self): | ||
| 103 | + """ | ||
| 104 | + Tests mTLS configuration using bytes-returning callback. | ||
| 105 | + """ | ||
| 106 | + | ||
| 107 | + def mock_callback(): | ||
| 108 | + return (b"fake_cert_bytes", b"fake_key_bytes") | ||
| 109 | + | ||
| 110 | + with mock.patch.dict( | ||
| 111 | + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} | ||
| 112 | + ), mock.patch( | ||
| 113 | + "google.auth.transport.mtls.has_default_client_cert_source", | ||
| 114 | + return_value=True, | ||
| 115 | + ), mock.patch( | ||
| 116 | + "ssl.SSLContext.load_cert_chain" | ||
| 117 | + ): | ||
| 118 | + mock_creds = mock.AsyncMock(spec=credentials.Credentials) | ||
| 119 | + session = sessions.AsyncAuthorizedSession(mock_creds) | ||
| 120 | + | ||
| 121 | + await session.configure_mtls_channel(client_cert_callback=mock_callback) | ||
| 122 | + | ||
| 123 | + assert session._is_mtls is True | ||
| Back | FazBrowse Home | New Git URL |
0 commit comments