feat: Adding Agent Identity bound token support and handling certificate mismatches with retries

[agrawalradhika-cell]

This PR includes adding changes which are for -

• Adding support for Agent Identity bound tokens which will be used for Agent Identity (feat: Add support for Agent Identity bound tokens #1821)
• Adding the retry logic when certificates mismatch for credentials used for Agent Identities on GKE and Cloud Run Workloads. (feat: Add retry logic when certificate mismatch for existing credentials & Agent Identity workloads #1841)

[gemini-code-assist]

Summary of Changes

Hello @agrawalradhika-cell, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the security and resilience of Agent Identity authentication by introducing certificate-bound tokens and robust retry mechanisms for certificate management. It ensures that access tokens are securely tied to specific client certificates and that the system can dynamically adapt to certificate rotations without manual intervention, particularly for workloads on GKE and Cloud Run.

Highlights

• Agent Identity Certificate Utilities: Introduced a new module, _agent_identity_utils.py, to centralize logic for handling Agent Identity certificates. This includes functions for parsing certificates, determining if a certificate is an Agent Identity type, calculating certificate fingerprints, and managing certificate file paths with retry mechanisms.
• Certificate-Bound Token Support: Implemented support for certificate-bound tokens for Agent Identities. This enhances security by binding access tokens to specific client certificates, preventing token sharing. This is integrated into compute_engine, external_account, and identity_pool credentials.
• Dynamic mTLS Channel Reconfiguration: Added retry logic to requests.py and urllib3.py transport layers. If an UNAUTHORIZED (401) response is received during an mTLS request and the client certificate has changed, the mTLS channel will attempt to reconfigure itself with the new certificate, improving resilience to certificate rotation.
• New Environment Variables: Defined two new environment variables: GOOGLE_API_CERTIFICATE_CONFIG to specify the path to the certificate configuration file, and GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES to allow users to opt-out of certificate-bound tokens.
• Cryptography as Optional Dependency: The cryptography library, essential for certificate operations, has been moved to an optional dependency, allowing for more flexible installations.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for Agent Identity bound tokens, which enhances security by binding access tokens to the client's mTLS certificate. It adds a new utility module, _agent_identity_utils.py, for handling agent identity certificates. The changes also include retry logic in the requests and urllib3 transports to handle certificate rotation by reconfiguring the mTLS channel upon receiving a 401 Unauthorized response. My review focuses on improving maintainability and robustness. I've suggested refactoring repeated code in _agent_identity_utils.py into a decorator and improving error logging in the certificate polling logic. I've also pointed out the use of overly broad exception handling in the transport layers, which could mask bugs, and recommended using more specific exceptions.

[gemini-code-assist]

Catching the base Exception class is too broad. It can mask unexpected errors and even catch system-exiting exceptions like KeyboardInterrupt. Since configure_mtls_channel is documented to raise google.auth.exceptions.MutualTLSChannelError on failure, it would be safer to catch that specific exception.

Suggested change

	except Exception as e:
	except exceptions.MutualTLSChannelError as e:

[gemini-code-assist]

Catching the base Exception class is too broad. It can mask unexpected errors and even catch system-exiting exceptions like KeyboardInterrupt. Since configure_mtls_channel is documented to raise google.auth.exceptions.MutualTLSChannelError on failure, it would be safer to catch that specific exception.

Suggested change

	except Exception as e:
	except exceptions.MutualTLSChannelError as e:

[gemini-code-assist]

The except (IOError, ValueError, KeyError): pass block is broad and silences all but the first error encountered during the polling loop. This can make debugging difficult. For instance, if the file is initially not found (IOError), a warning is logged. If on a subsequent attempt the file is found but is malformed (ValueError), this error will be silently ignored until the timeout. To improve debuggability, consider logging subsequent errors at a DEBUG level.

        except (IOError, ValueError, KeyError) as e:
            if not has_logged_warning:
                _LOGGER.warning(
                    "Certificate config file not found at %s (from %s environment "
                    "variable). Retrying for up to %s seconds.",
                    cert_config_path,
                    environment_vars.GOOGLE_API_CERTIFICATE_CONFIG,
                    _TOTAL_TIMEOUT,
                )
                has_logged_warning = True
            else:
                _LOGGER.debug("Polling for certificate failed again: %s", e)
            pass

[gemini-code-assist]

To improve maintainability and reduce code duplication, the repeated try...except ImportError block for the cryptography dependency in parse_certificate, _is_agent_identity_certificate, and calculate_certificate_fingerprint could be refactored into a decorator.

For example, you could define a decorator like this:

import functools

def _require_cryptography(func):
    """A decorator to handle the optional cryptography dependency."""
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except ImportError as e:
            raise ImportError(CRYPTOGRAPHY_NOT_FOUND_ERROR) from e
    return wrapper

And then apply it to the relevant functions:

@_require_cryptography
def parse_certificate(cert_bytes):
    from cryptography import x509
    return x509.load_pem_x509_certificate(cert_bytes)

@_require_cryptography
def _is_agent_identity_certificate(cert):
    # ... function body ...

@_require_cryptography
def calculate_certificate_fingerprint(cert):
    # ... function body ...

This would make the code cleaner and easier to maintain.

[daniel-sanche]

LGTM

This PR is from a stanging branch, which was made up of individually approved PRs