|
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. |
|
0 commit comments