Update the retry logic and refactor cert handling.

Author: nbayati

google/auth/_agent_identity_utils.py

@@ -18,8 +18,8 @@
 import hashlib
 import os
 import re
+import time
 
-from google.auth import _exponential_backoff
 from google.auth import environment_vars
 from google.auth import exceptions
 
@@ -51,11 +51,10 @@ def get_agent_identity_certificate_path():
     if not cert_config_path:
         return None
 
-    # Use exponential backoff to retry loading the certificate config file.
-    # This is to handle the race condition where the env var is set before the file exists.
-    backoff = _exponential_backoff.ExponentialBackoff(total_attempts=5)
-
-    for _ in backoff:
+    # Poll for the config file and the certificate file to be available.
+    # Phase 1: Poll rapidly for 5 seconds (50 * 0.1s).
+    # Phase 2: Slow down polling for the next 25 seconds (50 * 0.5s).
+    for i in range(100):
         if os.path.exists(cert_config_path):
             with open(cert_config_path, "r") as f:
                 cert_config = json.load(f)
@@ -66,20 +65,41 @@ def get_agent_identity_certificate_path():
                 )
                 if cert_path and os.path.exists(cert_path):
                     return cert_path
+        if i < 50:
+            time.sleep(0.1)
+        else:
+            time.sleep(0.5)
 
     raise exceptions.RefreshError(
-        "Certificate config or certificate file not found after multiple retries."
+        "Certificate config or certificate file not found after multiple retries. "
+        f"If you are using Agent Engine, you can export "
+        f"{environment_vars.GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES} to false to "
+        "disable cert bound tokens to fall back to unbound tokens."
     )
 
 
-def _is_agent_identity_certificate(cert_bytes):
+def parse_certificate(cert_bytes):
+    """Parses a PEM-encoded certificate.
+
+    Args:
+        cert_bytes (bytes): The PEM-encoded certificate bytes.
+
+    Returns:
+        cryptography.x509.Certificate: The parsed certificate object.
+    """
+    from cryptography import x509
+
+    return x509.load_pem_x509_certificate(cert_bytes)
+
+
+def _is_agent_identity_certificate(cert):
     """Checks if a certificate is an Agent Identity certificate.
 
     This is determined by checking the Subject Alternative Name (SAN) for a
     SPIFFE ID with a trust domain matching Agent Identity patterns.
 
     Args:
-        cert_bytes (bytes): The PEM-encoded certificate bytes.
+        cert (cryptography.x509.Certificate): The parsed certificate object.
 
     Returns:
         bool: True if the certificate is an Agent Identity certificate,
@@ -88,7 +108,6 @@ def _is_agent_identity_certificate(cert_bytes):
     from cryptography import x509
     from cryptography.x509.oid import ExtensionOID
 
-    cert = x509.load_pem_x509_certificate(cert_bytes)
     try:
         ext = cert.extensions.get_extension_for_oid(
             ExtensionOID.SUBJECT_ALTERNATIVE_NAME
@@ -107,37 +126,39 @@ def _is_agent_identity_certificate(cert_bytes):
     return False
 
 
-def calculate_certificate_fingerprint(cert_bytes):
+def calculate_certificate_fingerprint(cert):
     """Calculates the base64-encoded SHA256 hash of a DER-encoded certificate.
 
     Args:
-        cert_bytes (bytes): The PEM-encoded certificate bytes.
+        cert (cryptography.x509.Certificate): The parsed certificate object.
 
     Returns:
         str: The base64-encoded SHA256 fingerprint.
     """
-    from cryptography import x509
     from cryptography.hazmat.primitives import serialization
 
-    cert = x509.load_pem_x509_certificate(cert_bytes)
     der_cert = cert.public_bytes(serialization.Encoding.DER)
     fingerprint = hashlib.sha256(der_cert).digest()
-    return base64.b64encode(fingerprint).decode("utf-8")
+    return base64.urlsafe_b64encode(fingerprint).rstrip(b"=").decode("utf-8")
 
 
-def should_request_bound_token(cert_bytes):
+def should_request_bound_token(cert):
     """Determines if a bound token should be requested.
 
     This is based on the GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES
     environment variable and whether the certificate is an agent identity cert.
 
+    Args:
+        cert (cryptography.x509.Certificate): The parsed certificate object.
+
     Returns:
         bool: True if a bound token should be requested, False otherwise.
     """
-    is_agent_cert = _is_agent_identity_certificate(cert_bytes)
+    is_agent_cert = _is_agent_identity_certificate(cert)
     is_opted_in = (
         os.environ.get(
-            "GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES", "true"
+            environment_vars.GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES,
+            "true",
         ).lower()
         == "true"
     )

google/auth/compute_engine/_metadata.py

@@ -368,18 +368,14 @@ def get_service_account_token(request, service_account="default", scopes=None):
             scopes = ",".join(scopes)
         params["scopes"] = scopes
 
-    try:
-        cert_path = _agent_identity_utils.get_agent_identity_certificate_path()
-        if cert_path:
-            with open(cert_path, "rb") as cert_file:
-                cert_bytes = cert_file.read()
-            if _agent_identity_utils.should_request_bound_token(cert_bytes):
-                fingerprint = _agent_identity_utils.calculate_certificate_fingerprint(
-                    cert_bytes
-                )
-                params["bindCertificateFingerprint"] = fingerprint
-    except exceptions.RefreshError as e:
-        _LOGGER.warning("Could not load agent identity certificate: %s", e)
+    cert_path = _agent_identity_utils.get_agent_identity_certificate_path()
+    if cert_path:
+        with open(cert_path, "rb") as cert_file:
+            cert_bytes = cert_file.read()
+        cert = _agent_identity_utils.parse_certificate(cert_bytes)
+        if _agent_identity_utils.should_request_bound_token(cert):
+            fingerprint = _agent_identity_utils.calculate_certificate_fingerprint(cert)
+            params["bindCertificateFingerprint"] = fingerprint
 
     metrics_header = {
         metrics.API_CLIENT_HEADER: metrics.token_request_access_token_mds()

google/auth/environment_vars.py

@@ -90,3 +90,8 @@
 GOOGLE_API_CERTIFICATE_CONFIG = "GOOGLE_API_CERTIFICATE_CONFIG"
 """Environment variable defining the location of Google API certificate config
 file."""
+
+GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES = (
+    "GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES"
+)
+"""Environment variable to prevent agent token sharing for GCP services."""

google/auth/identity_pool.py

@@ -564,9 +564,10 @@ def refresh(self, request):
         # Check if the credential is X.509 based.
         if self._credential_source_certificate is not None:
             cert_bytes = self._get_cert_bytes()
-            if _agent_identity_utils.should_request_bound_token(cert_bytes):
+            cert = _agent_identity_utils.parse_certificate(cert_bytes)
+            if _agent_identity_utils.should_request_bound_token(cert):
                 cert_fingerprint = _agent_identity_utils.calculate_certificate_fingerprint(
-                    cert_bytes
+                    cert
                 )
 
         self._refresh_token(request, cert_fingerprint=cert_fingerprint)

tests/compute_engine/test__metadata.py

@@ -21,6 +21,7 @@
 import mock
 import pytest  # type: ignore
 
+from google.auth import _agent_identity_utils
 from google.auth import _helpers
 from google.auth import environment_vars
 from google.auth import exceptions
@@ -38,6 +39,29 @@
     DATA_DIR, "smbios_product_name_non_google"
 )
 
+# A mock PEM-encoded certificate without an Agent Identity SPIFFE ID.
+NON_AGENT_IDENTITY_CERT_BYTES = (
+    b"-----BEGIN CERTIFICATE-----\n"
+    b"MIIDIzCCAgugAwIBAgIJAMfISuBQ5m+5MA0GCSqGSIb3DQEBBQUAMBUxEzARBgNV\n"
+    b"BAMTCnVuaXQtdGVzdHMwHhcNMTExMjA2MTYyNjAyWhcNMjExMjAzMTYyNjAyWjAV\n"
+    b"MRMwEQYDVQQDEwp1bml0LXRlc3RzMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIB\n"
+    b"CgKCAQEA4ej0p7bQ7L/r4rVGUz9RN4VQWoej1Bg1mYWIDYslvKrk1gpj7wZgkdmM\n"
+    b"7oVK2OfgrSj/FCTkInKPqaCR0gD7K80q+mLBrN3PUkDrJQZpvRZIff3/xmVU1Wer\n"
+    b"uQLFJjnFb2dqu0s/FY/2kWiJtBCakXvXEOb7zfbINuayL+MSsCGSdVYsSliS5qQp\n"
+    b"gyDap+8b5fpXZVJkq92hrcNtbkg7hCYUJczt8n9hcCTJCfUpApvaFQ18pe+zpyl4\n"
+    b"+WzkP66I28hniMQyUlA1hBiskT7qiouq0m8IOodhv2fagSZKjOTTU2xkSBc//fy3\n"
+    b"ZpsL7WqgsZS7Q+0VRK8gKfqkxg5OYQIDAQABo3YwdDAdBgNVHQ4EFgQU2RQ8yO+O\n"
+    b"gN8oVW2SW7RLrfYd9jEwRQYDVR0jBD4wPIAU2RQ8yO+OgN8oVW2SW7RLrfYd9jGh\n"
+    b"GaQXMBUxEzARBgNVBAMTCnVuaXQtdGVzdHOCCQDHyErgUOZvuTAMBgNVHRMEBTAD\n"
+    b"AQH/MA0GCSqGSIb3DQEBBQUAA4IBAQBRv+M/6+FiVu7KXNjFI5pSN17OcW5QUtPr\n"
+    b"odJMlWrJBtynn/TA1oJlYu3yV5clc/71Vr/AxuX5xGP+IXL32YDF9lTUJXG/uUGk\n"
+    b"+JETpKmQviPbRsvzYhz4pf6ZIOZMc3/GIcNq92ECbseGO+yAgyWUVKMmZM0HqXC9\n"
+    b"ovNslqe0M8C1sLm1zAR5z/h/litE7/8O2ietija3Q/qtl2TOXJdCA6sgjJX2WUql\n"
+    b"ybrC55ct18NKf3qhpcEkGQvFU40rVYApJpi98DiZPYFdx1oBDp/f4uZ3ojpxRVFT\n"
+    b"cDwcJLfNRCPUhormsY7fDS9xSyThiHsW9mjJYdcaKQkwYZ0F11yB\n"
+    b"-----END CERTIFICATE-----\n"
+)
+
 ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE = (
     "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/mds"
 )
@@ -641,21 +665,40 @@ def test_get_service_account_token_with_scopes_string(
     assert expiry == utcnow() + datetime.timedelta(seconds=ttl)
 
 
-@mock.patch("google.auth.compute_engine._metadata.get")
-@mock.patch("google.auth.compute_engine._metadata._LOGGER")
-@mock.patch("google.auth._agent_identity_utils.get_agent_identity_certificate_path")
-def test_get_service_account_token_agent_identity_refresh_error(
-    mock_get_agent_path, mock_logger, mock_get
-):
-    mock_get_agent_path.side_effect = exceptions.RefreshError("Test error")
-    mock_get.return_value = {"access_token": "token", "expires_in": 500}
+@mock.patch(
+    "google.auth._agent_identity_utils._is_agent_identity_certificate",
+    return_value=True,
+)
+def test_get_service_account_token_with_bound_token(mock_is_agent, tmpdir, monkeypatch):
+    # Create a mock certificate and a config file that points to it.
+    cert_path = tmpdir.join("cert.pem")
+    cert_path.write_binary(NON_AGENT_IDENTITY_CERT_BYTES)
+    config_path = tmpdir.join("config.json")
+    config_path.write(
+        json.dumps({"cert_configs": {"workload": {"cert_path": str(cert_path)}}})
+    )
+    monkeypatch.setenv(environment_vars.GOOGLE_API_CERTIFICATE_CONFIG, str(config_path))
+    monkeypatch.setenv(
+        environment_vars.GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES, "true"
+    )
 
-    _metadata.get_service_account_token(mock.sentinel.request)
+    # Create a mock request to simulate the metadata server response.
+    token_response = json.dumps({"access_token": "token", "expires_in": 3600})
+    request = make_request(token_response, headers={"content-type": "application/json"})
 
-    mock_logger.warning.assert_called_once_with(
-        "Could not load agent identity certificate: %s", mock.ANY
-    )
-    mock_get.assert_called_once()
+    # Call the function under test.
+    _metadata.get_service_account_token(request)
+
+    # Verify the request URL contains the correct fingerprint.
+    request.assert_called_once()
+    _, kwargs = request.call_args
+    url = kwargs["url"]
+
+    # Calculate the expected fingerprint.
+    cert = _agent_identity_utils.parse_certificate(NON_AGENT_IDENTITY_CERT_BYTES)
+    expected_fingerprint = _agent_identity_utils.calculate_certificate_fingerprint(cert)
+
+    assert f"bindCertificateFingerprint={expected_fingerprint}" in url
 
 
 def test_get_service_account_info():

tests/compute_engine/test_credentials.py

@@ -660,6 +660,7 @@ def test_build_trust_boundary_lookup_url_no_email(
 
     @mock.patch("google.auth.compute_engine._metadata.get")
     @mock.patch("google.auth._agent_identity_utils.get_agent_identity_certificate_path")
+    @mock.patch("google.auth._agent_identity_utils.parse_certificate")
     @mock.patch(
         "google.auth._agent_identity_utils.should_request_bound_token",
         return_value=True,
@@ -672,6 +673,7 @@ def test_refresh_with_agent_identity(
         self,
         mock_calculate_fingerprint,
         mock_should_request,
+        mock_parse_certificate,
         mock_get_path,
         mock_metadata_get,
         tmpdir,
@@ -688,7 +690,8 @@ def test_refresh_with_agent_identity(
         self.credentials.refresh(None)
 
         assert self.credentials.token == "token"
-        mock_should_request.assert_called_once_with(b"cert_content")
+        mock_parse_certificate.assert_called_once_with(b"cert_content")
+        mock_should_request.assert_called_once_with(mock_parse_certificate.return_value)
         kwargs = mock_metadata_get.call_args[1]
         assert kwargs["params"] == {
             "scopes": "one,two",
@@ -697,12 +700,18 @@ def test_refresh_with_agent_identity(
 
     @mock.patch("google.auth.compute_engine._metadata.get")
     @mock.patch("google.auth._agent_identity_utils.get_agent_identity_certificate_path")
+    @mock.patch("google.auth._agent_identity_utils.parse_certificate")
     @mock.patch(
         "google.auth._agent_identity_utils.should_request_bound_token",
         return_value=False,
     )
     def test_refresh_with_agent_identity_opt_out_or_not_agent(
-        self, mock_should_request, mock_get_path, mock_metadata_get, tmpdir
+        self,
+        mock_should_request,
+        mock_parse_certificate,
+        mock_get_path,
+        mock_metadata_get,
+        tmpdir,
     ):
         cert_path = tmpdir.join("cert.pem")
         cert_path.write(b"cert_content")
@@ -716,7 +725,8 @@ def test_refresh_with_agent_identity_opt_out_or_not_agent(
         self.credentials.refresh(None)
 
         assert self.credentials.token == "token"
-        mock_should_request.assert_called_once_with(b"cert_content")
+        mock_parse_certificate.assert_called_once_with(b"cert_content")
+        mock_should_request.assert_called_once_with(mock_parse_certificate.return_value)
         kwargs = mock_metadata_get.call_args[1]
         assert "bindCertificateFingerprint" not in kwargs.get("params", {})