fix: use mTLS endpoint for Google OAuth2 token requests

GWeale · copybara-github · commit ffe41f050c14 · 2026-07-01T15:13:05.000-07:00
When a client certificate is configured, OAuth2 token exchange and refresh
now present the certificate and target the *.mtls.googleapis.com endpoint so
Context-Aware Access / token binding is honored. Gated to *.googleapis.com
token endpoints with a client cert available; third-party providers and
non-cert environments are unchanged.

Co-authored-by: George Weale &lt;gweale@google.com&gt;
PiperOrigin-RevId: 941319195
diff --git a/src/google/adk/auth/oauth2_credential_util.py b/src/google/adk/auth/oauth2_credential_util.py
@@ -22,6 +22,7 @@
 from authlib.oauth2.rfc6749 import OAuth2Token
 from fastapi.openapi.models import OAuth2
 
+from ..utils import _mtls_utils
 from ..utils.feature_decorator import experimental
 from .auth_credential import AuthCredential
 from .auth_schemes import AuthScheme
@@ -83,18 +84,28 @@ def create_oauth2_session(
 
   # Scope is intentionally omitted: token exchange and refresh don't require
   # it per RFC 6749, and some providers reject it on these requests.
-  return (
-      OAuth2Session(
-          auth_credential.oauth2.client_id,
-          auth_credential.oauth2.client_secret,
-          redirect_uri=auth_credential.oauth2.redirect_uri,
-          state=auth_credential.oauth2.state,
-          token_endpoint_auth_method=auth_credential.oauth2.token_endpoint_auth_method,
-          code_challenge_method=auth_credential.oauth2.code_challenge_method,
-      ),
-      token_endpoint,
+  session = OAuth2Session(
+      auth_credential.oauth2.client_id,
+      auth_credential.oauth2.client_secret,
+      redirect_uri=auth_credential.oauth2.redirect_uri,
+      state=auth_credential.oauth2.state,
+      token_endpoint_auth_method=auth_credential.oauth2.token_endpoint_auth_method,
+      code_challenge_method=auth_credential.oauth2.code_challenge_method,
   )
 
+  # When a client certificate is configured, route Google token requests through
+  # the mTLS endpoint and present the cert so Context-Aware Access / token
+  # binding is honored. Non-Google providers and non-cert environments keep the
+  # existing behavior.
+  if (
+      _mtls_utils.is_non_mtls_googleapis_endpoint(token_endpoint)
+      and _mtls_utils.use_client_cert_effective()
+  ):
+    if _mtls_utils.configure_session_for_mtls(session):
+      token_endpoint = _mtls_utils.effective_googleapis_endpoint(token_endpoint)
+
+  return session, token_endpoint
+
 
 @experimental
 def update_credential_with_tokens(
diff --git a/src/google/adk/utils/_mtls_utils.py b/src/google/adk/utils/_mtls_utils.py
@@ -17,10 +17,22 @@
 from __future__ import annotations
 
 import enum
+import logging
 import os
+from typing import TYPE_CHECKING
+from urllib.parse import urlsplit
+from urllib.parse import urlunsplit
 
 from google.auth.transport import mtls
 
+if TYPE_CHECKING:
+  import requests
+
+logger = logging.getLogger("google_adk." + __name__)
+
+_GOOGLEAPIS_SUFFIX = ".googleapis.com"
+_MTLS_GOOGLEAPIS_SUFFIX = ".mtls.googleapis.com"
+
 
 class MtlsEndpoint(enum.Enum):
   """Enum for the mTLS endpoint setting."""
@@ -30,10 +42,21 @@ class MtlsEndpoint(enum.Enum):
   NEVER = "never"
 
 
+def _mtls_endpoint_setting() -> MtlsEndpoint:
+  """Returns the GOOGLE_API_USE_MTLS_ENDPOINT setting, defaulting to AUTO."""
+  setting = os.getenv(
+      "GOOGLE_API_USE_MTLS_ENDPOINT", MtlsEndpoint.AUTO.value
+  ).lower()
+  try:
+    return MtlsEndpoint(setting)
+  except ValueError:
+    return MtlsEndpoint.AUTO
+
+
 def use_client_cert_effective() -> bool:
   """Returns whether client certificate should be used for mTLS."""
   try:
-    return mtls.should_use_client_cert()
+    return bool(mtls.should_use_client_cert())
   except (ImportError, AttributeError):
     return (
         os.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false").lower()
@@ -53,16 +76,76 @@ def get_api_endpoint(
       mtls_template: Template for mTLS regional endpoint (e.g.
         "secretmanager.{location}.rep.mtls.googleapis.com").
   """
-  use_mtls_endpoint_str = os.getenv(
-      "GOOGLE_API_USE_MTLS_ENDPOINT", MtlsEndpoint.AUTO.value
-  ).lower()
-  try:
-    use_mtls_endpoint = MtlsEndpoint(use_mtls_endpoint_str)
-  except ValueError:
-    use_mtls_endpoint = MtlsEndpoint.AUTO
-
+  use_mtls_endpoint = _mtls_endpoint_setting()
   if (use_mtls_endpoint == MtlsEndpoint.ALWAYS) or (
       use_mtls_endpoint == MtlsEndpoint.AUTO and use_client_cert_effective()
   ):
     return mtls_template.format(location=location)
   return default_template.format(location=location)
+
+
+def is_non_mtls_googleapis_endpoint(url: str) -> bool:
+  """Returns whether url points at a *.googleapis.com host without the mTLS infix."""
+  if not url:
+    return False
+  host = urlsplit(url).hostname or ""
+  return (
+      host.endswith(_GOOGLEAPIS_SUFFIX) and _MTLS_GOOGLEAPIS_SUFFIX not in host
+  )
+
+
+def effective_googleapis_endpoint(url: str) -> str:
+  """Rewrites a *.googleapis.com url to its .mtls.googleapis.com variant.
+
+  Honors GOOGLE_API_USE_MTLS_ENDPOINT=never as an opt-out. Hosts that are not
+  googleapis.com hosts, or are already mTLS hosts, are returned unchanged so
+  non-Google providers are never affected.
+  """
+  if not is_non_mtls_googleapis_endpoint(url):
+    return url
+  if _mtls_endpoint_setting() == MtlsEndpoint.NEVER:
+    return url
+  parsed = urlsplit(url)
+  host = parsed.hostname or ""
+  new_host = host[: -len(_GOOGLEAPIS_SUFFIX)] + _MTLS_GOOGLEAPIS_SUFFIX
+  netloc = f"{new_host}:{parsed.port}" if parsed.port else new_host
+  return urlunsplit(
+      (parsed.scheme, netloc, parsed.path, parsed.query, parsed.fragment)
+  )
+
+
+def configure_session_for_mtls(session: requests.Session) -> bool:
+  """Mounts a mutual-TLS adapter on a requests session when a client cert exists.
+
+  authlib's OAuth2Session is a requests.Session but not a google-auth
+  AuthorizedSession, so it lacks configure_mtls_channel(). This replicates that
+  method's effect: load the application-default client certificate and mount an
+  adapter that presents it on https connections.
+
+  Returns True if a client certificate was found and the adapter was mounted.
+  """
+  try:
+    from google.auth import exceptions as ga_exceptions
+    from google.auth.transport import _mtls_helper
+    from google.auth.transport.requests import _MutualTlsAdapter
+  except ImportError:
+    return False
+
+  cert_source = (
+      mtls.default_client_cert_source()
+      if mtls.has_default_client_cert_source()
+      else None
+  )
+  try:
+    is_mtls, cert, key = _mtls_helper.get_client_cert_and_key(cert_source)
+  except (ImportError, ga_exceptions.GoogleAuthError) as e:
+    logger.warning(
+        "Could not load client certificate for mTLS; falling back to non-mTLS"
+        " token request: %s",
+        e,
+    )
+    return False
+
+  if is_mtls:
+    session.mount("https://", _MutualTlsAdapter(cert, key))
+  return bool(is_mtls)
diff --git a/tests/unittests/auth/test_oauth2_credential_util.py b/tests/unittests/auth/test_oauth2_credential_util.py
@@ -15,6 +15,7 @@
 import time
 from typing import Optional
 from unittest.mock import Mock
+from unittest.mock import patch
 
 from authlib.oauth2.rfc6749 import OAuth2Token
 from fastapi.openapi.models import OAuth2
@@ -151,6 +152,104 @@ def test_create_oauth2_session_missing_credentials(self):
     assert client is None
     assert token_endpoint is None
 
+  def _google_openid_scheme(self) -> OpenIdConnectWithConfig:
+    """OpenID Connect scheme that uses Google's OAuth2 token endpoint."""
+    return OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://accounts.google.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://accounts.google.com/o/oauth2/v2/auth",
+        token_endpoint="https://oauth2.googleapis.com/token",
+        scopes=["openid"],
+    )
+
+  @patch.dict("os.environ", {}, clear=True)
+  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
+  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
+  def test_create_oauth2_session_google_endpoint_uses_mtls(
+      self, mock_use_cert, mock_configure
+  ):
+    """Google token endpoint is switched to mTLS when a cert is mounted."""
+    mock_use_cert.return_value = True
+    mock_configure.return_value = True
+    credential = create_oauth2_auth_credential(
+        auth_type=AuthCredentialTypes.OAUTH2
+    )
+
+    client, token_endpoint = create_oauth2_session(
+        self._google_openid_scheme(), credential
+    )
+
+    assert client is not None
+    assert token_endpoint == "https://oauth2.mtls.googleapis.com/token"
+    mock_configure.assert_called_once_with(client)
+
+  @patch.dict("os.environ", {}, clear=True)
+  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
+  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
+  def test_create_oauth2_session_google_endpoint_no_cert_keeps_plain(
+      self, mock_use_cert, mock_configure
+  ):
+    """Without a client cert the plain Google endpoint is kept."""
+    mock_use_cert.return_value = False
+    credential = create_oauth2_auth_credential(
+        auth_type=AuthCredentialTypes.OAUTH2
+    )
+
+    _, token_endpoint = create_oauth2_session(
+        self._google_openid_scheme(), credential
+    )
+
+    assert token_endpoint == "https://oauth2.googleapis.com/token"
+    mock_configure.assert_not_called()
+
+  @patch.dict("os.environ", {}, clear=True)
+  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
+  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
+  def test_create_oauth2_session_cert_unavailable_keeps_plain(
+      self, mock_use_cert, mock_configure
+  ):
+    """If the adapter cannot be mounted, the endpoint is not switched."""
+    mock_use_cert.return_value = True
+    mock_configure.return_value = False
+    credential = create_oauth2_auth_credential(
+        auth_type=AuthCredentialTypes.OAUTH2
+    )
+
+    client, token_endpoint = create_oauth2_session(
+        self._google_openid_scheme(), credential
+    )
+
+    assert token_endpoint == "https://oauth2.googleapis.com/token"
+    mock_configure.assert_called_once_with(client)
+
+  @patch.dict("os.environ", {}, clear=True)
+  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
+  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
+  def test_create_oauth2_session_non_google_endpoint_skips_mtls(
+      self, mock_use_cert, mock_configure
+  ):
+    """Non-Google providers are never switched to an mTLS endpoint."""
+    mock_use_cert.return_value = True
+    credential = create_oauth2_auth_credential(
+        auth_type=AuthCredentialTypes.OAUTH2
+    )
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+
+    _, token_endpoint = create_oauth2_session(scheme, credential)
+
+    assert token_endpoint == "https://example.com/token"
+    mock_configure.assert_not_called()
+
   @pytest.mark.parametrize(
       "token_endpoint_auth_method, expected_auth_method",
       [
diff --git a/tests/unittests/utils/test_mtls_utils.py b/tests/unittests/utils/test_mtls_utils.py
