feat: Add retry logic when certificate mismatch for existing credentials & Agent Identity workloads (#1841)

feat: Add retry logic when certificate mismatch for existing credentials
& Agent Identity workloads

This change introduces retry support when requests are created for
existing credentials and Agent Identities on GKE and Cloud Run
Workloads. When 401(Unauthorized) error is created, due to certificate
at time of configuration of mTLS channel being different from the
current certificate, a retry is added to the request by configuring the
mTLS channel with the current certificate.

---------

Signed-off-by: Radhika Agrawal <[email protected]>
Co-authored-by: nbayati <[email protected]>
Co-authored-by: Andy Zhao <[email protected]>

Author: 3 people

google/auth/_agent_identity_utils.py

@@ -20,10 +20,11 @@
 import os
 import re
 import time
-from urllib.parse import urlparse, quote
+from urllib.parse import quote, urlparse
 
 from google.auth import environment_vars
 from google.auth import exceptions
+from google.auth.transport import _mtls_helper
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -260,3 +261,21 @@ def should_request_bound_token(cert):
         == "true"
     )
     return is_agent_cert and is_opted_in
+
+
+def call_client_cert_callback():
+    """Calls the client cert callback and returns the certificate and key."""
+    _, cert_bytes, key_bytes, passphrase = _mtls_helper.get_client_ssl_credentials(
+        generate_encrypted_key=True
+    )
+    return cert_bytes, key_bytes
+
+
+def get_cached_cert_fingerprint(cached_cert):
+    """Returns the fingerprint of the cached certificate."""
+    if cached_cert:
+        cert_obj = parse_certificate(cached_cert)
+        cached_cert_fingerprint = calculate_certificate_fingerprint(cert_obj)
+    else:
+        raise ValueError("mTLS connection is not configured.")
+    return cached_cert_fingerprint

google/auth/transport/_mtls_helper.py

@@ -20,6 +20,7 @@
 import re
 import subprocess
 
+from google.auth import _agent_identity_utils
 from google.auth import environment_vars
 from google.auth import exceptions
 
@@ -474,3 +475,29 @@ def check_use_client_cert():
             ) as e:
                 _LOGGER.debug("error decoding certificate: %s", e)
         return False
+
+
+def check_parameters_for_unauthorized_response(cached_cert):
+    """Returns the cached and current cert fingerprint for reconfiguring mTLS.
+
+    Args:
+        cached_cert(bytes): The cached client certificate.
+
+    Returns:
+        bytes: The client callback cert bytes.
+        bytes: The client callback key bytes.
+        str: The base64-encoded SHA256 cached fingerprint.
+        str: The base64-encoded SHA256 current cert fingerprint.
+    """
+    call_cert_bytes, call_key_bytes = _agent_identity_utils.call_client_cert_callback()
+    cert_obj = _agent_identity_utils.parse_certificate(call_cert_bytes)
+    current_cert_fingerprint = _agent_identity_utils.calculate_certificate_fingerprint(
+        cert_obj
+    )
+    if cached_cert:
+        cached_fingerprint = _agent_identity_utils.get_cached_cert_fingerprint(
+            cached_cert
+        )
+    else:
+        cached_fingerprint = current_cert_fingerprint
+    return call_cert_bytes, call_key_bytes, cached_fingerprint, current_cert_fingerprint

google/auth/transport/requests.py

@@ -17,6 +17,7 @@
 from __future__ import absolute_import
 
 import functools
+import http.client as http_client
 import logging
 import numbers
 import time
@@ -36,6 +37,7 @@
 from google.auth import _helpers
 from google.auth import exceptions
 from google.auth import transport
+from google.auth.transport import _mtls_helper
 import google.auth.transport._mtls_helper
 from google.oauth2 import service_account
 
@@ -463,6 +465,7 @@ def configure_mtls_channel(self, client_cert_callback=None):
 
             if self._is_mtls:
                 mtls_adapter = _MutualTlsAdapter(cert, key)
+                self._cached_cert = cert
                 self.mount("https://", mtls_adapter)
         except (
             exceptions.ClientCertError,
@@ -502,6 +505,10 @@ def request(
                 itself does not timeout, e.g. if a large file is being
                 transmitted. The timout error will be raised after such
                 request completes.
+        Raises:
+            google.auth.exceptions.MutualTLSChannelError: If mutual TLS
+                channel creation fails for any reason.
+            ValueError: If the client certificate is invalid.
         """
         # pylint: disable=arguments-differ
         # Requests has a ton of arguments to request, but only two
@@ -551,7 +558,31 @@ def request(
             response.status_code in self._refresh_status_codes
             and _credential_refresh_attempt < self._max_refresh_attempts
         ):
-
+            # Handle unauthorized permission error(401 status code)
+            if response.status_code == http_client.UNAUTHORIZED:
+                if self.is_mtls:
+                    call_cert_bytes, call_key_bytes, cached_fingerprint, current_cert_fingerprint = _mtls_helper.check_parameters_for_unauthorized_response(
+                        self._cached_cert
+                    )
+                    if cached_fingerprint != current_cert_fingerprint:
+                        try:
+                            _LOGGER.info(
+                                "Client certificate has changed, reconfiguring mTLS "
+                                "channel."
+                            )
+                            self.configure_mtls_channel(
+                                lambda: (call_cert_bytes, call_key_bytes)
+                            )
+                        except Exception as e:
+                            _LOGGER.error("Failed to reconfigure mTLS channel: %s", e)
+                            raise exceptions.MutualTLSChannelError(
+                                "Failed to reconfigure mTLS channel"
+                            ) from e
+                    else:
+                        _LOGGER.info(
+                            "Skipping reconfiguration of mTLS channel because the client"
+                            " certificate has not changed."
+                        )
             _LOGGER.info(
                 "Refreshing credentials due to a %s response. Attempt %s/%s.",
                 response.status_code,

google/auth/transport/urllib3.py

@@ -16,6 +16,7 @@
 
 from __future__ import absolute_import
 
+import http.client as http_client
 import logging
 import warnings
 
@@ -52,6 +53,7 @@
 from google.auth import _helpers
 from google.auth import exceptions
 from google.auth import transport
+from google.auth.transport import _mtls_helper
 from google.oauth2 import service_account
 
 if version.parse(urllib3.__version__) >= version.parse("2.0.0"):  # pragma: NO COVER
@@ -299,6 +301,7 @@ def __init__(
         # Request instance used by internal methods (for example,
         # credentials.refresh).
         self._request = Request(self.http)
+        self._is_mtls = False
 
         # https://google.aip.dev/auth/4111
         # Attempt to use self-signed JWTs when a service account is used.
@@ -335,7 +338,10 @@ def configure_mtls_channel(self, client_cert_callback=None):
         """
         use_client_cert = transport._mtls_helper.check_use_client_cert()
         if not use_client_cert:
+            self._is_mtls = False
             return False
+        else:
+            self._is_mtls = True
         try:
             import OpenSSL
         except ImportError as caught_exc:
@@ -349,6 +355,7 @@ def configure_mtls_channel(self, client_cert_callback=None):
 
             if found_cert_key:
                 self.http = _make_mutual_tls_http(cert, key)
+                self._cached_cert = cert
             else:
                 self.http = _make_default_http()
         except (
@@ -381,6 +388,11 @@ def urlopen(self, method, url, body=None, headers=None, **kwargs):
         if headers is None:
             headers = self.headers
 
+        use_mtls = False
+        if self._is_mtls:
+            MTLS_URL_PREFIXES = ["mtls.googleapis.com", "mtls.sandbox.googleapis.com"]
+            use_mtls = any([prefix in url for prefix in MTLS_URL_PREFIXES])
+
         # Make a copy of the headers. They will be modified by the credentials
         # and we want to pass the original headers if we recurse.
         request_headers = headers.copy()
@@ -402,6 +414,34 @@ def urlopen(self, method, url, body=None, headers=None, **kwargs):
             response.status in self._refresh_status_codes
             and _credential_refresh_attempt < self._max_refresh_attempts
         ):
+            if response.status == http_client.UNAUTHORIZED:
+                if use_mtls:
+                    call_cert_bytes, call_key_bytes, cached_fingerprint, current_cert_fingerprint = _mtls_helper.check_parameters_for_unauthorized_response(
+                        self._cached_cert
+                    )
+                    if cached_fingerprint != current_cert_fingerprint:
+                        try:
+                            _LOGGER.info(
+                                "Client certificate has changed, reconfiguring mTLS "
+                                "channel."
+                            )
+                            self.configure_mtls_channel(
+                                client_cert_callback=lambda: (
+                                    call_cert_bytes,
+                                    call_key_bytes,
+                                )
+                            )
+                        except Exception as e:
+                            _LOGGER.error("Failed to reconfigure mTLS channel: %s", e)
+                            raise exceptions.MutualTLSChannelError(
+                                "Failed to reconfigure mTLS channel"
+                            ) from e
+
+                    else:
+                        _LOGGER.info(
+                            "Skipping reconfiguration of mTLS channel because the "
+                            "client certificate has not changed."
+                        )
 
             _LOGGER.info(
                 "Refreshing credentials due to a %s response. Attempt %s/%s.",

noxfile.py

@@ -105,7 +105,7 @@ def mypy(session):
         "types-requests",
         "types-setuptools",
         "types-mock",
-        "pytest",
+        "pytest<8.0.0",
     )
     session.run("mypy", "-p", "google", "-p", "tests", "-p", "tests_async")
 
@@ -130,6 +130,7 @@ def unit(session):
 
 @nox.session(python=DEFAULT_PYTHON_VERSION)
 def cover(session):
+    session.env["PIP_EXTRA_INDEX_URL"] = "https://pypi.org/simple"
     session.install("-e", ".[testing]")
     session.run(
         "pytest",

tests/test_agent_identity_utils.py

@@ -15,7 +15,6 @@
 import base64
 import hashlib
 import json
-
 import urllib.parse
 
 from cryptography import x509
@@ -104,7 +103,9 @@ def test_calculate_certificate_fingerprint(self):
         mock_cert.public_bytes.return_value = b"der-bytes"
 
         # Expected: base64 (standard), unpadded, then URL-encoded
-        base64_fingerprint = base64.b64encode(hashlib.sha256(b"der-bytes").digest()).decode("utf-8")
+        base64_fingerprint = base64.b64encode(
+            hashlib.sha256(b"der-bytes").digest()
+        ).decode("utf-8")
         unpadded_base64_fingerprint = base64_fingerprint.rstrip("=")
         expected_fingerprint = urllib.parse.quote(unpadded_base64_fingerprint)
 
@@ -260,6 +261,54 @@ def test_get_and_parse_agent_identity_certificate_success(
         mock_parse_certificate.assert_called_once_with(b"cert_bytes")
         assert result == mock_parse_certificate.return_value
 
+    @mock.patch("time.sleep", return_value=None)
+    @mock.patch("google.auth._agent_identity_utils._is_certificate_file_ready")
+    def test_get_agent_identity_certificate_path_fallback_to_well_known_path(
+        self, mock_is_ready, mock_sleep, monkeypatch
+    ):
+        # Set a dummy config path that won't be found.
+        monkeypatch.setenv(
+            environment_vars.GOOGLE_API_CERTIFICATE_CONFIG, "/dummy/config.json"
+        )
+
+        # First, the primary path from the (mocked) config is not ready.
+        # Then, the fallback well-known path is ready.
+        mock_is_ready.side_effect = [False, True]
+
+        result = _agent_identity_utils.get_agent_identity_certificate_path()
+
+        assert result == _agent_identity_utils._WELL_KNOWN_CERT_PATH
+        # The sleep should have been called once before the fallback is checked.
+        mock_sleep.assert_called_once()
+        assert mock_is_ready.call_count == 2
+
+    @mock.patch("google.auth.transport._mtls_helper.get_client_ssl_credentials")
+    def test_call_client_cert_callback(self, mock_get_client_ssl_credentials):
+        mock_get_client_ssl_credentials.return_value = (
+            True,
+            b"cert_bytes",
+            b"key_bytes",
+            b"passphrase",
+        )
+
+        cert, key = _agent_identity_utils.call_client_cert_callback()
+
+        assert cert == b"cert_bytes"
+        assert key == b"key_bytes"
+        mock_get_client_ssl_credentials.assert_called_once_with(
+            generate_encrypted_key=True
+        )
+
+    def test_get_cached_cert_fingerprint_no_cert(self):
+        with pytest.raises(ValueError, match="mTLS connection is not configured."):
+            _agent_identity_utils.get_cached_cert_fingerprint(None)
+
+    def test_get_cached_cert_fingerprint_with_cert(self):
+        fingerprint = _agent_identity_utils.get_cached_cert_fingerprint(
+            NON_AGENT_IDENTITY_CERT_BYTES
+        )
+        assert isinstance(fingerprint, str)
+
 
 class TestAgentIdentityUtilsNoCryptography:
     @pytest.fixture(autouse=True)

tests/transport/test__mtls_helper.py

@@ -23,8 +23,9 @@
 from google.auth import exceptions
 from google.auth.transport import _mtls_helper
 
+CERT_MOCK_VAL = b"cert"
+KEY_MOCK_VAL = b"key"
 CONTEXT_AWARE_METADATA = {"cert_provider_command": ["some command"]}
-
 ENCRYPTED_EC_PRIVATE_KEY = b"""-----BEGIN ENCRYPTED PRIVATE KEY-----
 MIHkME8GCSqGSIb3DQEFDTBCMCkGCSqGSIb3DQEFDDAcBAgl2/yVgs1h3QICCAAw
 DAYIKoZIhvcNAgkFADAVBgkrBgEEAZdVAQIECJk2GRrvxOaJBIGQXIBnMU4wmciT
@@ -813,3 +814,64 @@ def test_check_use_client_cert_when_file_does_not_exist(self, monkeypatch):
         monkeypatch.setenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "")
         use_client_cert = _mtls_helper.check_use_client_cert()
         assert use_client_cert is False
+
+
+class TestMtlsHelper:
+    @mock.patch("google.auth.transport._mtls_helper._agent_identity_utils")
+    def test_check_parameters_for_unauthorized_response_with_cached_cert(
+        self, mock_agent_identity_utils
+    ):
+        mock_agent_identity_utils.call_client_cert_callback.return_value = (
+            CERT_MOCK_VAL,
+            KEY_MOCK_VAL,
+        )
+        mock_agent_identity_utils.get_cached_cert_fingerprint.return_value = (
+            "cached_fingerprint"
+        )
+        mock_agent_identity_utils.calculate_certificate_fingerprint.return_value = (
+            "current_fingerprint"
+        )
+
+        (
+            cert,
+            key,
+            cached_fingerprint,
+            current_fingerprint,
+        ) = _mtls_helper.check_parameters_for_unauthorized_response(
+            cached_cert=b"cached_cert_bytes"
+        )
+
+        assert cert == CERT_MOCK_VAL
+        assert key == KEY_MOCK_VAL
+        assert cached_fingerprint == "cached_fingerprint"
+        assert current_fingerprint == "current_fingerprint"
+        mock_agent_identity_utils.call_client_cert_callback.assert_called_once()
+        mock_agent_identity_utils.get_cached_cert_fingerprint.assert_called_once_with(
+            b"cached_cert_bytes"
+        )
+
+    @mock.patch("google.auth.transport._mtls_helper._agent_identity_utils")
+    def test_check_parameters_for_unauthorized_response_without_cached_cert(
+        self, mock_agent_identity_utils
+    ):
+        mock_agent_identity_utils.call_client_cert_callback.return_value = (
+            CERT_MOCK_VAL,
+            KEY_MOCK_VAL,
+        )
+        mock_agent_identity_utils.calculate_certificate_fingerprint.return_value = (
+            "current_fingerprint"
+        )
+
+        (
+            cert,
+            key,
+            cached_fingerprint,
+            current_fingerprint,
+        ) = _mtls_helper.check_parameters_for_unauthorized_response(cached_cert=None)
+
+        assert cert == CERT_MOCK_VAL
+        assert key == KEY_MOCK_VAL
+        assert cached_fingerprint == "current_fingerprint"
+        assert current_fingerprint == "current_fingerprint"
+        mock_agent_identity_utils.call_client_cert_callback.assert_called_once()
+        mock_agent_identity_utils.get_cached_cert_fingerprint.assert_not_called()