From d62acb6d5bb58f97e6de7aa7976067aacda120e9 Mon Sep 17 00:00:00 2001 From: Negar Bayati Date: Fri, 30 Oct 2026 06:23:59 +0000 Subject: [PATCH 1/2] feat(auth): add _MtlsCertManager for mTLS client certificate rotation 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. --- .../google/auth/transport/_mtls_helper.py | 141 +++++++++- .../tests/transport/test__mtls_helper.py | 252 +++++++++++++++++- 2 files changed, 391 insertions(+), 2 deletions(-) diff --git a/packages/google-auth/google/auth/transport/_mtls_helper.py b/packages/google-auth/google/auth/transport/_mtls_helper.py index 7ff5b2a81d81..4c49d9524e63 100644 --- a/packages/google-auth/google/auth/transport/_mtls_helper.py +++ b/packages/google-auth/google/auth/transport/_mtls_helper.py @@ -22,8 +22,9 @@ import subprocess import sys import tempfile +import threading from os import environ, getenv, path -from typing import Generator, List, Optional, Tuple, Union, cast +from typing import Callable, Generator, List, Optional, Tuple, Union, cast from urllib.parse import urlsplit from google.auth import _agent_identity_utils, _cloud_sdk, environment_vars, exceptions @@ -967,3 +968,141 @@ def is_mtls_endpoint(url: Optional[Union[str, bytes, object]]) -> bool: return False return hostname in _MTLS_EXACT_HOSTS or hostname.endswith(_MTLS_HOST_SUFFIXES) + + +class _MtlsCertManager: + """Thread-safe manager for mTLS client certificates and generation tracking. + + Tracks the active client certificate bytes, private key bytes, leaf + certificate SHA-256 fingerprint, and a ``generation_id`` counter that + increments only when the loaded certificate changes. + + Args: + client_cert_source (Optional[Callable[[], Tuple[bytes, bytes]]]): + Optional callback returning ``(cert_bytes, key_bytes)`` in PEM + format. When ``None``, certificates are discovered via + :func:`get_client_ssl_credentials`. + certificate_config_path (Optional[str]): Optional explicit path to a + ``certificate_config.json`` file used when ``client_cert_source`` + is ``None``. + """ + + def __init__( + self, + client_cert_source: Optional[Callable[[], Tuple[bytes, bytes]]] = None, + certificate_config_path: Optional[str] = None, + ) -> None: + self._client_cert_source = client_cert_source + self._certificate_config_path = certificate_config_path + self._cert: Optional[bytes] = None + self._key: Optional[bytes] = None + self._cert_fingerprint: Optional[str] = None + self._generation_id: int = 0 + self._lock = threading.Lock() + + @property + def generation_id(self) -> int: + return self._generation_id + + def get_credential_snapshot( + self, + ) -> Tuple[Optional[bytes], Optional[bytes], int]: + """Returns the current ``(cert_bytes, key_bytes, generation_id)`` under the lock.""" + with self._lock: + return self._cert, self._key, self._generation_id + + def _load_cert_and_fingerprint( + self, + ) -> Tuple[Optional[bytes], Optional[bytes], Optional[str]]: + """Loads the certificate, private key, and leaf SHA-256 fingerprint.""" + if self._client_cert_source is not None: + cert_bytes, key_bytes = self._client_cert_source() + else: + has_cert, cert_bytes, key_bytes, _ = get_client_ssl_credentials( + generate_encrypted_key=False, + certificate_config_path=self._certificate_config_path, + ) + if not has_cert and self._generation_id == 0: + return None, None, None + + if not cert_bytes or not key_bytes: + raise exceptions.ClientCertError( + "Client certificate or private key is empty." + ) + + cert_obj = _agent_identity_utils.parse_certificate(cert_bytes) + fingerprint = _agent_identity_utils.calculate_certificate_fingerprint(cert_obj) + + return cert_bytes, key_bytes, fingerprint + + def refresh(self, expected_generation: Optional[int] = None) -> bool: + """Loads the client certificate and increments ``generation_id`` if it changed. + + Args: + expected_generation (Optional[int]): If provided and the manager's + ``generation_id`` is already greater than ``expected_generation``, + returns ``True`` immediately without reading the certificate + source again. This can be set to the ``generation_id`` a caller + was using when a request failed so concurrent callers do not + load the certificate multiple times. + + Returns: + bool: ``True`` if the certificate changed (or ``generation_id`` is + already greater than ``expected_generation``), ``False`` + otherwise. + + Raises: + google.auth.exceptions.ClientCertError: If loading or parsing the + certificate fails on the initial load (``generation_id == 0``). + """ + with self._lock: + if ( + expected_generation is not None + and self._generation_id > expected_generation + ): + return True + + try: + cert_bytes, key_bytes, fingerprint = self._load_cert_and_fingerprint() + 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 + + if cert_bytes is None: + return False + + if fingerprint != self._cert_fingerprint: + self._cert = cert_bytes + self._key = key_bytes + self._cert_fingerprint = fingerprint + self._generation_id += 1 + return True + + return False + + def __deepcopy__(self, memo): + # Share the same thread-safe manager instance across deep copies. + return self + + def __getstate__(self): + state = self.__dict__.copy() + state["_lock"] = None + return state + + def __setstate__(self, state): + self.__dict__.update(state) + self._lock = threading.Lock() diff --git a/packages/google-auth/tests/transport/test__mtls_helper.py b/packages/google-auth/tests/transport/test__mtls_helper.py index c4e5b52c5973..8bc04437c8fb 100644 --- a/packages/google-auth/tests/transport/test__mtls_helper.py +++ b/packages/google-auth/tests/transport/test__mtls_helper.py @@ -12,10 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy +import json +import logging import os +import pickle import re import sys import tempfile +import threading from unittest import mock import pytest # type: ignore @@ -23,7 +28,7 @@ from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import ec -from google.auth import environment_vars, exceptions +from google.auth import _agent_identity_utils, environment_vars, exceptions from google.auth.transport import _mtls_helper if not hasattr(os, "MFD_CLOEXEC"): @@ -51,6 +56,14 @@ -----END PASSPHRASE-----""" PASSPHRASE_VALUE = b"password" +DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data") + +with open(os.path.join(DATA_DIR, "other_cert.pem"), "rb") as _fh: + OTHER_CERT_BYTES = _fh.read() + +with open(os.path.join(DATA_DIR, "trust_chain_with_leaf.pem"), "rb") as _fh: + TRUST_CHAIN_WITH_LEAF_BYTES = _fh.read() + def check_cert_and_key(content, expected_cert, expected_key): success = True @@ -2313,3 +2326,240 @@ def test_no_gke_fallback_when_implicit_config_has_no_workload( assert _mtls_helper._get_workload_cert_and_key() == (None, None) assert _mtls_helper.check_use_client_cert() is False mock_is_ready.assert_not_called() + + +class TestMtlsCertManager(object): + def _write_workload_config(self, tmp_path): + cert_file = tmp_path / "cert.pem" + key_file = tmp_path / "key.pem" + config_file = tmp_path / "certificate_config.json" + + cert_file.write_bytes(pytest.public_cert_bytes) + key_file.write_bytes(pytest.private_key_bytes) + config_file.write_text( + json.dumps( + { + "cert_configs": { + "workload": { + "cert_path": str(cert_file), + "key_path": str(key_file), + } + } + } + ) + ) + return cert_file, config_file + + def test_default_discovery_and_rotation_from_config_file(self, tmp_path): + cert_file, config_file = self._write_workload_config(tmp_path) + manager = _mtls_helper._MtlsCertManager( + certificate_config_path=str(config_file) + ) + assert manager.get_credential_snapshot() == (None, None, 0) + + assert manager.refresh() is True + assert manager.generation_id == 1 + assert manager.get_credential_snapshot() == ( + pytest.public_cert_bytes, + pytest.private_key_bytes, + 1, + ) + + assert manager.refresh() is False + assert manager.generation_id == 1 + + cert_file.write_bytes(OTHER_CERT_BYTES) + assert manager.refresh() is True + assert manager.generation_id == 2 + assert manager.get_credential_snapshot() == ( + OTHER_CERT_BYTES, + pytest.private_key_bytes, + 2, + ) + + def test_client_cert_source_and_pem_chain_leaf_fingerprint(self): + source = mock.Mock( + side_effect=[ + (TRUST_CHAIN_WITH_LEAF_BYTES, pytest.private_key_bytes), + (TRUST_CHAIN_WITH_LEAF_BYTES, pytest.private_key_bytes), + (OTHER_CERT_BYTES, pytest.private_key_bytes), + ] + ) + manager = _mtls_helper._MtlsCertManager(client_cert_source=source) + + expected_leaf_fp = _agent_identity_utils.calculate_certificate_fingerprint( + _agent_identity_utils.parse_certificate(TRUST_CHAIN_WITH_LEAF_BYTES) + ) + assert manager.refresh() is True + assert manager.generation_id == 1 + assert manager.get_credential_snapshot() == ( + TRUST_CHAIN_WITH_LEAF_BYTES, + pytest.private_key_bytes, + 1, + ) + assert manager._cert_fingerprint == expected_leaf_fp + + assert manager.refresh() is False + assert manager.generation_id == 1 + + assert manager.refresh() is True + assert manager.generation_id == 2 + assert manager.get_credential_snapshot() == ( + OTHER_CERT_BYTES, + pytest.private_key_bytes, + 2, + ) + assert source.call_count == 3 + + def test_expected_generation_coalesces_concurrent_refreshes(self): + barrier = threading.Barrier(4) + source = mock.Mock( + return_value=(pytest.public_cert_bytes, pytest.private_key_bytes) + ) + + manager = _mtls_helper._MtlsCertManager(client_cert_source=source) + assert manager.refresh() is True + assert manager.generation_id == 1 + assert source.call_count == 1 + + assert manager.refresh(expected_generation=0) is True + assert source.call_count == 1 + + source.return_value = (OTHER_CERT_BYTES, pytest.private_key_bytes) + results = [None] * 4 + + def worker(idx): + barrier.wait() + results[idx] = manager.refresh(expected_generation=1) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(4)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert results == [True, True, True, True] + assert manager.generation_id == 2 + assert source.call_count == 2 + + def test_initial_load_no_default_cert_returns_false(self, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + monkeypatch.setattr( + _mtls_helper._cloud_sdk, "get_config_path", lambda: str(tmp_path) + ) + manager = _mtls_helper._MtlsCertManager() + assert manager.refresh() is False + assert manager.generation_id == 0 + assert manager.get_credential_snapshot() == (None, None, 0) + + @pytest.mark.parametrize( + "side_effect", + [ + exceptions.ClientCertError("bad cert config"), + exceptions.MutualTLSChannelError("channel cert error"), + OSError("file missing"), + [(b"not a valid pem cert", pytest.private_key_bytes)], + [(b"", pytest.private_key_bytes)], + [(pytest.public_cert_bytes, b"")], + ], + ) + def test_initial_load_error_raises_client_cert_error(self, side_effect): + source = mock.Mock(side_effect=side_effect) + manager = _mtls_helper._MtlsCertManager(client_cert_source=source) + + with pytest.raises(exceptions.ClientCertError): + manager.refresh() + + assert manager.generation_id == 0 + assert manager.get_credential_snapshot() == (None, None, 0) + + @pytest.mark.parametrize( + "reload_side_effect", + [ + exceptions.ClientCertError("transient cert error"), + exceptions.MutualTLSChannelError("transient mtls error"), + OSError("transient read error"), + [(b"-----BEGIN CERTIFICATE-----\ntruncated", pytest.private_key_bytes)], + [(b"", pytest.private_key_bytes)], + [(pytest.public_cert_bytes, b"")], + ], + ) + def test_transient_reload_error_retains_last_known_good_cert( + self, reload_side_effect, caplog + ): + source = mock.Mock( + return_value=(pytest.public_cert_bytes, pytest.private_key_bytes) + ) + manager = _mtls_helper._MtlsCertManager(client_cert_source=source) + assert manager.refresh() is True + assert manager.generation_id == 1 + + source.side_effect = reload_side_effect + with caplog.at_level(logging.DEBUG, logger=_mtls_helper.__name__): + assert manager.refresh() is False + + assert "retaining previous certificate (generation 1)" in caplog.text + assert manager.generation_id == 1 + assert manager.get_credential_snapshot() == ( + pytest.public_cert_bytes, + pytest.private_key_bytes, + 1, + ) + + def test_transient_missing_default_cert_on_reload_retains_last_known_good( + self, tmp_path, monkeypatch, caplog + ): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + monkeypatch.setattr( + _mtls_helper._cloud_sdk, "get_config_path", lambda: str(tmp_path) + ) + _, config_file = self._write_workload_config(tmp_path) + + manager = _mtls_helper._MtlsCertManager( + certificate_config_path=str(config_file) + ) + assert manager.refresh() is True + + config_file.unlink() + with caplog.at_level(logging.DEBUG, logger=_mtls_helper.__name__): + assert manager.refresh() is False + + assert "retaining previous certificate (generation 1)" in caplog.text + assert manager.generation_id == 1 + assert manager.get_credential_snapshot() == ( + pytest.public_cert_bytes, + pytest.private_key_bytes, + 1, + ) + + def test_deepcopy_shares_instance_and_pickle_roundtrip_recreates_lock( + self, tmp_path + ): + cert_file, config_file = self._write_workload_config(tmp_path) + manager = _mtls_helper._MtlsCertManager( + certificate_config_path=str(config_file) + ) + assert manager.refresh() is True + + copied = copy.deepcopy(manager) + assert copied is manager + + restored = pickle.loads(pickle.dumps(manager)) + assert restored is not manager + assert restored.get_credential_snapshot() == ( + pytest.public_cert_bytes, + pytest.private_key_bytes, + 1, + ) + assert restored.generation_id == 1 + + cert_file.write_bytes(OTHER_CERT_BYTES) + assert restored.refresh() is True + assert restored.generation_id == 2 + assert restored.get_credential_snapshot() == ( + OTHER_CERT_BYTES, + pytest.private_key_bytes, + 2, + ) From 5bf98550bb7217c8b5f4eec0dbeda77a1c9eae87 Mon Sep 17 00:00:00 2001 From: Negar Bayati Date: Fri, 30 Oct 2026 21:29:13 +0000 Subject: [PATCH 2/2] test: fix failing test by mocking module logger instead of caplog --- .../tests/transport/test__mtls_helper.py | 51 +++++++------------ 1 file changed, 19 insertions(+), 32 deletions(-) diff --git a/packages/google-auth/tests/transport/test__mtls_helper.py b/packages/google-auth/tests/transport/test__mtls_helper.py index 8bc04437c8fb..d98dcdd9c8b2 100644 --- a/packages/google-auth/tests/transport/test__mtls_helper.py +++ b/packages/google-auth/tests/transport/test__mtls_helper.py @@ -14,7 +14,6 @@ import copy import json -import logging import os import pickle import re @@ -2454,19 +2453,18 @@ def test_initial_load_no_default_cert_returns_false(self, tmp_path, monkeypatch) assert manager.get_credential_snapshot() == (None, None, 0) @pytest.mark.parametrize( - "side_effect", + "cert_source", [ - exceptions.ClientCertError("bad cert config"), - exceptions.MutualTLSChannelError("channel cert error"), - OSError("file missing"), - [(b"not a valid pem cert", pytest.private_key_bytes)], - [(b"", pytest.private_key_bytes)], - [(pytest.public_cert_bytes, b"")], + mock.Mock(side_effect=exceptions.ClientCertError("bad cert config")), + mock.Mock(side_effect=exceptions.MutualTLSChannelError("channel error")), + mock.Mock(side_effect=OSError("file missing")), + mock.Mock(return_value=(b"not a valid pem cert", pytest.private_key_bytes)), + mock.Mock(return_value=(b"", pytest.private_key_bytes)), + mock.Mock(return_value=(pytest.public_cert_bytes, b"")), ], ) - def test_initial_load_error_raises_client_cert_error(self, side_effect): - source = mock.Mock(side_effect=side_effect) - manager = _mtls_helper._MtlsCertManager(client_cert_source=source) + def test_initial_load_error_raises_client_cert_error(self, cert_source): + manager = _mtls_helper._MtlsCertManager(client_cert_source=cert_source) with pytest.raises(exceptions.ClientCertError): manager.refresh() @@ -2474,32 +2472,21 @@ def test_initial_load_error_raises_client_cert_error(self, side_effect): assert manager.generation_id == 0 assert manager.get_credential_snapshot() == (None, None, 0) - @pytest.mark.parametrize( - "reload_side_effect", - [ - exceptions.ClientCertError("transient cert error"), - exceptions.MutualTLSChannelError("transient mtls error"), - OSError("transient read error"), - [(b"-----BEGIN CERTIFICATE-----\ntruncated", pytest.private_key_bytes)], - [(b"", pytest.private_key_bytes)], - [(pytest.public_cert_bytes, b"")], - ], - ) - def test_transient_reload_error_retains_last_known_good_cert( - self, reload_side_effect, caplog - ): + def test_transient_reload_error_retains_last_known_good_cert(self): source = mock.Mock( - return_value=(pytest.public_cert_bytes, pytest.private_key_bytes) + side_effect=[ + (pytest.public_cert_bytes, pytest.private_key_bytes), + OSError("transient read error"), + ] ) manager = _mtls_helper._MtlsCertManager(client_cert_source=source) assert manager.refresh() is True assert manager.generation_id == 1 - source.side_effect = reload_side_effect - with caplog.at_level(logging.DEBUG, logger=_mtls_helper.__name__): + with mock.patch.object(_mtls_helper._LOGGER, "debug") as mock_debug: assert manager.refresh() is False - assert "retaining previous certificate (generation 1)" in caplog.text + mock_debug.assert_called_once() assert manager.generation_id == 1 assert manager.get_credential_snapshot() == ( pytest.public_cert_bytes, @@ -2508,7 +2495,7 @@ def test_transient_reload_error_retains_last_known_good_cert( ) def test_transient_missing_default_cert_on_reload_retains_last_known_good( - self, tmp_path, monkeypatch, caplog + self, tmp_path, monkeypatch ): monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.setenv("USERPROFILE", str(tmp_path)) @@ -2523,10 +2510,10 @@ def test_transient_missing_default_cert_on_reload_retains_last_known_good( assert manager.refresh() is True config_file.unlink() - with caplog.at_level(logging.DEBUG, logger=_mtls_helper.__name__): + with mock.patch.object(_mtls_helper._LOGGER, "debug") as mock_debug: assert manager.refresh() is False - assert "retaining previous certificate (generation 1)" in caplog.text + assert "retaining previous certificate" in mock_debug.call_args[0][0] assert manager.generation_id == 1 assert manager.get_credential_snapshot() == ( pytest.public_cert_bytes,