Repository navigation
Conversation
Adds `_MtlsCertManager` in `google.auth.transport._mtls_helper` to track the active mTLS client certificate, private key, leaf certificate SHA-256 fingerprint, and generation counter across certificate reloads. This is the first change in the gRPC mTLS client certificate rotation series. Subsequent PRs wire `_MtlsCertManager` into `requests.Request` and Compute Engine bound-token refreshes, `AuthMetadataPlugin`, `google.api_core.grpc_helpers.create_channel`, and `gapic-generator` so that TLS channel credentials and certificate-bound token refreshes stay on the same certificate generation.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces _MtlsCertManager, a thread-safe manager for handling mTLS client certificates, tracking their generation, and managing reloading. It also adds comprehensive unit tests for this new class. The feedback suggests catching all standard exceptions (Exception) during certificate reloads rather than a specific subset, ensuring that unexpected errors from custom certificate sources do not crash the application and are instead gracefully handled as transient reload errors.
| except ( | ||
| exceptions.ClientCertError, | ||
| exceptions.MutualTLSChannelError, | ||
| OSError, | ||
| ValueError, | ||
| ) as caught_exc: | ||
| if self._generation_id > 0: | ||
| _LOGGER.debug( | ||
| "Failed to reload mTLS client certificate; retaining " | ||
| "previous certificate (generation %d): %s", | ||
| self._generation_id, | ||
| caught_exc, | ||
| ) | ||
| return False | ||
| if isinstance(caught_exc, exceptions.ClientCertError): | ||
| raise | ||
| raise exceptions.ClientCertError(caught_exc) from caught_exc |
There was a problem hiding this comment.
The except block currently catches a specific list of exceptions (ClientCertError, MutualTLSChannelError, OSError, ValueError). However, when a custom client_cert_source callback is provided, it may raise other unexpected or custom exceptions (e.g., network errors, custom HSM/KMS exceptions, or even TypeError if the callback returns an unexpected format).
If such an exception is raised during a certificate reload (when self._generation_id > 0), it will propagate and crash the application/request instead of being gracefully handled as a transient reload error. To ensure maximum robustness and successfully retain the last known good certificate, we should catch all standard exceptions (Exception) during reload.
except Exception as caught_exc:
if self._generation_id > 0:
_LOGGER.debug(
"Failed to reload mTLS client certificate; retaining "
"previous certificate (generation %d): %s",
self._generation_id,
caught_exc,
)
return False
if isinstance(caught_exc, exceptions.ClientCertError):
raise
raise exceptions.ClientCertError(caught_exc) from caught_excReferences
- Do not replace historical graceful fallback behaviors (such as returning False/falling back to standard TLS) with exceptions if doing so would introduce breaking changes for downstream users and violate backwards compatibility.
lsirac
added
the
do not merge
| if cert_bytes is None: | ||
| return False | ||
|
|
||
| if fingerprint != self._cert_fingerprint: |
There was a problem hiding this comment.
I think we are introducing the same edge case here as we have in the past in both Python and Java. We need to verify that the key matches the cert before storing the pair (and retry if it doesn't). Otherwise we can take a snapshot of when one is updated, and the other isn't, and then we would fail until the next rotation.
Prev PRs:
| expected_generation is not None | ||
| and self._generation_id > expected_generation | ||
| ): | ||
| return True |
There was a problem hiding this comment.
refresh(expected_generation=...) only deduplicates concurrent callers when _generation_id advances. If concurrent requests fail with UNAUTHENTICATED or 401 while the cert on disk hasn't changed, or while a reload fails transiently, _generation_id stays equal to expected_generation and every waiting thread reloads from the cert source one after another under _lock.
There was a problem hiding this comment.
Good catch! I updated refresh() so threads waiting on the lock reuse the result of a reload that started after they arrived, including one that failed after startup and kept the old cert, instead of each reloading again. A burst of concurrent UNAUTHENTICATED/401 failures now causes at most two reloads: the one already in progress, plus one shared by the threads that arrived during it.
|
|
||
| def __deepcopy__(self, memo): | ||
| # Share the same thread-safe manager instance across deep copies. | ||
| return self |
There was a problem hiding this comment.
Should we also implement __copy__ returning self? With __getstate__ and __setstate__ defined below, copy.copy(manager) creates a separate manager with its own _lock and _generation_id instead of sharing the instance like copy.deepcopy(manager).
There was a problem hiding this comment.
Done in a05f953: added __copy__ returning self, so shallow and deep copies both share the same manager. The existing copy/pickle test now also asserts copy.copy(manager) is manager.
| return self | ||
|
|
||
| def __getstate__(self): | ||
| state = self.__dict__.copy() |
There was a problem hiding this comment.
__getstate__ keeps _client_cert_source in the state dict, so pickle.dumps() fails with AttributeError or PicklingError when client_cert_source is a closure like mtls.default_client_cert_source() or a lambda. Since GAPIC clients pass mtls.default_client_cert_source() by default, attaching this manager to Request will break pickle.dumps(request).
There was a problem hiding this comment.
Good catch. The manager now raises TypeError when pickled, since I realized that it also holds the private key. The later PR that attaches it to Request will leave the manager out of the Request's pickled state, so the restored Request rediscovers mTLS like it does today.
|
|
||
| expected_leaf_fp = _agent_identity_utils.calculate_certificate_fingerprint( | ||
| _agent_identity_utils.parse_certificate(TRUST_CHAIN_WITH_LEAF_BYTES) | ||
| ) |
There was a problem hiding this comment.
expected_leaf_fp runs calculate_certificate_fingerprint(parse_certificate(TRUST_CHAIN_WITH_LEAF_BYTES)), which is the exact same call _load_cert_and_fingerprint() makes internally, so assert manager._cert_fingerprint == expected_leaf_fp still passes even if parse_certificate hashes the root CA instead of the leaf cert pytest.public_cert_bytes.
There was a problem hiding this comment.
Good point. The expected fingerprint is now computed from the leaf cert (public_cert.pem) loaded on its own, so the test fails if the manager hashes the root or the wrong block of the chain.
|
|
||
| def worker(idx): | ||
| barrier.wait() | ||
| results[idx] = manager.refresh(expected_generation=1) |
There was a problem hiding this comment.
Because the mock source returns from memory in ~0.1ms without blocking, the first thread finishes refresh() within CPython's 5ms switch interval before the other threads enter refresh(), so this test still passes even if self._lock is removed.
| _, config_file = self._write_workload_config(tmp_path) | ||
|
|
||
| manager = _mtls_helper._MtlsCertManager( | ||
| certificate_config_path=str(config_file) |
There was a problem hiding this comment.
This test monkeypatches _cloud_sdk.get_config_path above, yet passing certificate_config_path=str(config_file) here skips get_config_path completely and leaves happy-path load and reload with certificate_config_path=None untested across TestMtlsCertManager.
…and cert source callback out of pickles
Adds
_MtlsCertManageringoogle.auth.transport._mtls_helperto track the active mTLS client certificate, private key, leaf certificate SHA-256 fingerprint, and generation counter across certificate reloads.This is the first change in the gRPC mTLS client certificate rotation series. Subsequent PRs wire
_MtlsCertManagerintorequests.Requestand Compute Engine bound-token refreshes,AuthMetadataPlugin,google.api_core.grpc_helpers.create_channel, andgapic-generatorso that TLS channel credentials and certificate-bound token refreshes stay on the same certificate generation.This is PR1 from go/cert-rotation-python-auth phase 1.