From 5c6cdcd197a6780fc86d9183fa208f78c8a975d9 Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Mon, 6 Oct 2025 11:27:41 -0700 Subject: [PATCH] feat: Support Oauth2 client credentials grant type PiperOrigin-RevId: 815813477 --- src/google/adk/auth/credential_manager.py | 43 ++++++- .../exchanger/oauth2_credential_exchanger.py | 101 ++++++++++++++- src/google/adk/auth/oauth2_credential_util.py | 37 +++--- .../test_oauth2_credential_exchanger.py | 116 ++++++++++++++++++ .../unittests/auth/test_credential_manager.py | 106 ++++++++++++++++ 5 files changed, 381 insertions(+), 22 deletions(-) diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index 5f51bae8..3503e1e4 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -25,11 +25,11 @@ from .auth_credential import AuthCredential from .auth_credential import AuthCredentialTypes from .auth_schemes import AuthSchemeType from .auth_schemes import ExtendedOAuth2 +from .auth_schemes import OpenIdConnectWithConfig from .auth_tool import AuthConfig from .exchanger.base_credential_exchanger import BaseCredentialExchanger from .exchanger.credential_exchanger_registry import CredentialExchangerRegistry from .oauth2_discovery import OAuth2DiscoveryManager -from .refresher.base_credential_refresher import BaseCredentialRefresher from .refresher.credential_refresher_registry import CredentialRefresherRegistry logger = logging.getLogger("google_adk." + __name__) @@ -85,8 +85,17 @@ class CredentialManager: # Register default exchangers and refreshers # TODO: support service account credential exchanger + from .exchanger.oauth2_credential_exchanger import OAuth2CredentialExchanger from .refresher.oauth2_credential_refresher import OAuth2CredentialRefresher + oauth2_exchanger = OAuth2CredentialExchanger() + self._exchanger_registry.register( + AuthCredentialTypes.OAUTH2, oauth2_exchanger + ) + self._exchanger_registry.register( + AuthCredentialTypes.OPEN_ID_CONNECT, oauth2_exchanger + ) + oauth2_refresher = OAuth2CredentialRefresher() self._refresher_registry.register( AuthCredentialTypes.OAUTH2, oauth2_refresher @@ -134,9 +143,14 @@ class CredentialManager: credential = await self._load_from_auth_response(callback_context) was_from_auth_response = True - # Step 5: If still no credential available, return None + # Step 5: If still no credential available, check if client credentials if not credential: - return None + # For client credentials flow, use raw credentials directly + if self._is_client_credentials_flow(): + credential = self._auth_config.raw_auth_credential + else: + # For authorization code flow, return None to trigger user authorization + return None # Step 6: Exchange credential if needed (e.g., service account to access token) credential, was_exchanged = await self._exchange_credential(credential) @@ -328,3 +342,26 @@ class CredentialManager: and not flows.authorizationCode.tokenUrl ) return False + + def _is_client_credentials_flow(self) -> bool: + """Check if the auth scheme uses client credentials flow. + + Supports both OAuth2 and OIDC schemes. + + Returns: + True if using client credentials flow, False otherwise. + """ + auth_scheme = self._auth_config.auth_scheme + + # Check OAuth2 schemes + if isinstance(auth_scheme, OAuth2) and auth_scheme.flows: + return auth_scheme.flows.clientCredentials is not None + + # Check OIDC schemes + if isinstance(auth_scheme, OpenIdConnectWithConfig): + return ( + auth_scheme.grant_types_supported is not None + and "client_credentials" in auth_scheme.grant_types_supported + ) + + return False diff --git a/src/google/adk/auth/exchanger/oauth2_credential_exchanger.py b/src/google/adk/auth/exchanger/oauth2_credential_exchanger.py index 4231a7c1..71c5f7c8 100644 --- a/src/google/adk/auth/exchanger/oauth2_credential_exchanger.py +++ b/src/google/adk/auth/exchanger/oauth2_credential_exchanger.py @@ -19,9 +19,11 @@ from __future__ import annotations import logging from typing import Optional +from fastapi.openapi.models import OAuth2 from google.adk.auth.auth_credential import AuthCredential from google.adk.auth.auth_schemes import AuthScheme from google.adk.auth.auth_schemes import OAuthGrantType +from google.adk.auth.auth_schemes import OpenIdConnectWithConfig from google.adk.auth.oauth2_credential_util import create_oauth2_session from google.adk.auth.oauth2_credential_util import update_credential_with_tokens from google.adk.utils.feature_decorator import experimental @@ -81,9 +83,100 @@ class OAuth2CredentialExchanger(BaseCredentialExchanger): if auth_credential.oauth2 and auth_credential.oauth2.access_token: return auth_credential + # Determine grant type from auth_scheme + grant_type = self._determine_grant_type(auth_scheme) + + if grant_type == OAuthGrantType.CLIENT_CREDENTIALS: + return await self._exchange_client_credentials( + auth_credential, auth_scheme + ) + elif grant_type == OAuthGrantType.AUTHORIZATION_CODE: + return await self._exchange_authorization_code( + auth_credential, auth_scheme + ) + else: + logger.warning("Unsupported OAuth2 grant type: %s", grant_type) + return auth_credential + + def _determine_grant_type( + self, auth_scheme: AuthScheme + ) -> Optional[OAuthGrantType]: + """Determine the OAuth2 grant type from the auth scheme. + + Args: + auth_scheme: The OAuth2 authentication scheme. + + Returns: + The OAuth2 grant type or None if cannot be determined. + """ + if isinstance(auth_scheme, OAuth2) and auth_scheme.flows: + return OAuthGrantType.from_flow(auth_scheme.flows) + elif isinstance(auth_scheme, OpenIdConnectWithConfig): + # Check supported grant types for OIDC + if ( + auth_scheme.grant_types_supported + and "client_credentials" in auth_scheme.grant_types_supported + ): + return OAuthGrantType.CLIENT_CREDENTIALS + else: + # Default to authorization code if client credentials not supported + return OAuthGrantType.AUTHORIZATION_CODE + + return None + + async def _exchange_client_credentials( + self, + auth_credential: AuthCredential, + auth_scheme: AuthScheme, + ) -> AuthCredential: + """Exchange client credentials for access token. + + Args: + auth_credential: The OAuth2 credential to exchange. + auth_scheme: The OAuth2 authentication scheme. + + Returns: + The credential with access token. + """ client, token_endpoint = create_oauth2_session(auth_scheme, auth_credential) if not client: - logger.warning("Could not create OAuth2 session for token exchange") + logger.warning( + "Could not create OAuth2 session for client credentials exchange" + ) + return auth_credential + + try: + tokens = client.fetch_token( + token_endpoint, + grant_type=OAuthGrantType.CLIENT_CREDENTIALS, + ) + update_credential_with_tokens(auth_credential, tokens) + logger.debug("Successfully exchanged client credentials for access token") + except Exception as e: + logger.error("Failed to exchange client credentials: %s", e) + return auth_credential + + return auth_credential + + async def _exchange_authorization_code( + self, + auth_credential: AuthCredential, + auth_scheme: AuthScheme, + ) -> AuthCredential: + """Exchange authorization code for access token. + + Args: + auth_credential: The OAuth2 credential to exchange. + auth_scheme: The OAuth2 authentication scheme. + + Returns: + The credential with access token. + """ + client, token_endpoint = create_oauth2_session(auth_scheme, auth_credential) + if not client: + logger.warning( + "Could not create OAuth2 session for authorization code exchange" + ) return auth_credential try: @@ -94,11 +187,9 @@ class OAuth2CredentialExchanger(BaseCredentialExchanger): grant_type=OAuthGrantType.AUTHORIZATION_CODE, ) update_credential_with_tokens(auth_credential, tokens) - logger.debug("Successfully exchanged OAuth2 tokens") + logger.debug("Successfully exchanged authorization code for access token") except Exception as e: - # TODO reconsider whether we should raise errors in this case - logger.error("Failed to exchange OAuth2 tokens: %s", e) - # Return original credential on failure + logger.error("Failed to exchange authorization code: %s", e) return auth_credential return auth_credential diff --git a/src/google/adk/auth/oauth2_credential_util.py b/src/google/adk/auth/oauth2_credential_util.py index cc315bd2..d90103a8 100644 --- a/src/google/adk/auth/oauth2_credential_util.py +++ b/src/google/adk/auth/oauth2_credential_util.py @@ -18,6 +18,8 @@ import logging from typing import Optional from typing import Tuple +from authlib.integrations.requests_client import OAuth2Session +from authlib.oauth2.rfc6749 import OAuth2Token from fastapi.openapi.models import OAuth2 from ..utils.feature_decorator import experimental @@ -25,15 +27,6 @@ from .auth_credential import AuthCredential from .auth_schemes import AuthScheme from .auth_schemes import OpenIdConnectWithConfig -try: - from authlib.integrations.requests_client import OAuth2Session - from authlib.oauth2.rfc6749 import OAuth2Token - - AUTHLIB_AVAILABLE = True -except ImportError: - AUTHLIB_AVAILABLE = False - - logger = logging.getLogger("google_adk." + __name__) @@ -53,18 +46,34 @@ def create_oauth2_session( """ if isinstance(auth_scheme, OpenIdConnectWithConfig): if not hasattr(auth_scheme, "token_endpoint"): + logger.warning("OpenIdConnect scheme missing token_endpoint") return None, None token_endpoint = auth_scheme.token_endpoint - scopes = auth_scheme.scopes + scopes = auth_scheme.scopes or [] elif isinstance(auth_scheme, OAuth2): + # Support both authorization code and client credentials flows if ( - not auth_scheme.flows.authorizationCode - or not auth_scheme.flows.authorizationCode.tokenUrl + auth_scheme.flows.authorizationCode + and auth_scheme.flows.authorizationCode.tokenUrl ): + token_endpoint = auth_scheme.flows.authorizationCode.tokenUrl + scopes = list(auth_scheme.flows.authorizationCode.scopes.keys()) + elif ( + auth_scheme.flows.clientCredentials + and auth_scheme.flows.clientCredentials.tokenUrl + ): + token_endpoint = auth_scheme.flows.clientCredentials.tokenUrl + scopes = list(auth_scheme.flows.clientCredentials.scopes.keys()) + else: + logger.warning( + "OAuth2 scheme missing required flow configuration. Expected either" + " authorizationCode.tokenUrl or clientCredentials.tokenUrl. Auth" + " scheme: %s", + auth_scheme, + ) return None, None - token_endpoint = auth_scheme.flows.authorizationCode.tokenUrl - scopes = list(auth_scheme.flows.authorizationCode.scopes.keys()) else: + logger.warning(f"Unsupported auth_scheme type: {type(auth_scheme)}") return None, None if ( diff --git a/tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py b/tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py index 31288511..bc5990c6 100644 --- a/tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py +++ b/tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py @@ -17,6 +17,9 @@ from unittest.mock import Mock from unittest.mock import patch from authlib.oauth2.rfc6749 import OAuth2Token +from fastapi.openapi.models import OAuth2 +from fastapi.openapi.models import OAuthFlowClientCredentials +from fastapi.openapi.models import OAuthFlows from google.adk.auth.auth_credential import AuthCredential from google.adk.auth.auth_credential import AuthCredentialTypes from google.adk.auth.auth_credential import OAuth2Auth @@ -218,3 +221,116 @@ class TestOAuth2CredentialExchanger: # Should return original credential when authlib is not available assert result == credential assert result.oauth2.access_token is None + + @patch("google.adk.auth.oauth2_credential_util.OAuth2Session") + @pytest.mark.asyncio + async def test_exchange_client_credentials_success(self, mock_oauth2_session): + """Test successful client credentials exchange.""" + # Setup mock + mock_client = Mock() + mock_oauth2_session.return_value = mock_client + mock_tokens = OAuth2Token({ + "access_token": "client_access_token", + "expires_at": int(time.time()) + 3600, + "expires_in": 3600, + }) + mock_client.fetch_token.return_value = mock_tokens + + # Create OAuth2 scheme with client credentials flow + flows = OAuthFlows( + clientCredentials=OAuthFlowClientCredentials( + tokenUrl="https://example.com/token", + scopes={"read": "Read access", "write": "Write access"}, + ) + ) + scheme = OAuth2(flows=flows) + + credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="test_client_id", + client_secret="test_client_secret", + ), + ) + + exchanger = OAuth2CredentialExchanger() + result = await exchanger.exchange(credential, scheme) + + # Verify client credentials exchange was successful + assert result.oauth2.access_token == "client_access_token" + mock_client.fetch_token.assert_called_once_with( + "https://example.com/token", + grant_type="client_credentials", + ) + + @patch("google.adk.auth.oauth2_credential_util.OAuth2Session") + @pytest.mark.asyncio + async def test_exchange_client_credentials_failure(self, mock_oauth2_session): + """Test client credentials exchange failure.""" + # Setup mock to raise exception during fetch_token + mock_client = Mock() + mock_oauth2_session.return_value = mock_client + mock_client.fetch_token.side_effect = Exception( + "Client credentials fetch failed" + ) + + # Create OAuth2 scheme with client credentials flow + flows = OAuthFlows( + clientCredentials=OAuthFlowClientCredentials( + tokenUrl="https://example.com/token", scopes={"read": "Read access"} + ) + ) + scheme = OAuth2(flows=flows) + + credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id="test_client_id", + client_secret="test_client_secret", + ), + ) + + exchanger = OAuth2CredentialExchanger() + result = await exchanger.exchange(credential, scheme) + + # Should return original credential when client credentials exchange fails + assert result == credential + assert result.oauth2.access_token is None + mock_client.fetch_token.assert_called_once() + + @pytest.mark.asyncio + async def test_determine_grant_type_client_credentials(self): + """Test grant type determination for client credentials.""" + flows = OAuthFlows( + clientCredentials=OAuthFlowClientCredentials( + tokenUrl="https://example.com/token", scopes={"read": "Read access"} + ) + ) + scheme = OAuth2(flows=flows) + + exchanger = OAuth2CredentialExchanger() + grant_type = exchanger._determine_grant_type(scheme) + + from google.adk.auth.auth_schemes import OAuthGrantType + + assert grant_type == OAuthGrantType.CLIENT_CREDENTIALS + + @pytest.mark.asyncio + async def test_determine_grant_type_openid_connect(self): + """Test grant type determination for OpenID Connect (defaults to auth code).""" + 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"], + ) + + exchanger = OAuth2CredentialExchanger() + grant_type = exchanger._determine_grant_type(scheme) + + from google.adk.auth.auth_schemes import OAuthGrantType + + assert grant_type == OAuthGrantType.AUTHORIZATION_CODE diff --git a/tests/unittests/auth/test_credential_manager.py b/tests/unittests/auth/test_credential_manager.py index 8ecf6946..da3ea542 100644 --- a/tests/unittests/auth/test_credential_manager.py +++ b/tests/unittests/auth/test_credential_manager.py @@ -105,6 +105,9 @@ class TestCredentialManager: auth_config = Mock(spec=AuthConfig) auth_config.raw_auth_credential = None auth_config.exchanged_auth_credential = None + # Add auth_scheme for the _is_client_credentials_flow method + auth_config.auth_scheme = Mock() + auth_config.auth_scheme.flows = None callback_context = Mock() @@ -562,6 +565,109 @@ class TestCredentialManager: assert manager._auth_config.auth_scheme == implicit_oauth2_scheme + def test_is_client_credentials_flow_oauth2_with_client_credentials(self): + """Test _is_client_credentials_flow returns True for OAuth2 with client credentials.""" + from fastapi.openapi.models import OAuth2 + from fastapi.openapi.models import OAuthFlowClientCredentials + from fastapi.openapi.models import OAuthFlows + + # Create OAuth2 scheme with client credentials flow + auth_scheme = OAuth2( + flows=OAuthFlows( + clientCredentials=OAuthFlowClientCredentials( + tokenUrl="https://example.com/token" + ) + ) + ) + + auth_config = Mock(spec=AuthConfig) + auth_config.auth_scheme = auth_scheme + auth_config.raw_auth_credential = None + auth_config.exchanged_auth_credential = None + + manager = CredentialManager(auth_config) + + assert manager._is_client_credentials_flow() is True + + def test_is_client_credentials_flow_oauth2_without_client_credentials(self): + """Test _is_client_credentials_flow returns False for OAuth2 without client credentials.""" + from fastapi.openapi.models import OAuth2 + from fastapi.openapi.models import OAuthFlowAuthorizationCode + from fastapi.openapi.models import OAuthFlows + + # Create OAuth2 scheme with authorization code flow only + auth_scheme = OAuth2( + flows=OAuthFlows( + authorizationCode=OAuthFlowAuthorizationCode( + authorizationUrl="https://example.com/auth", + tokenUrl="https://example.com/token", + ) + ) + ) + + auth_config = Mock(spec=AuthConfig) + auth_config.auth_scheme = auth_scheme + auth_config.raw_auth_credential = None + auth_config.exchanged_auth_credential = None + + manager = CredentialManager(auth_config) + + assert manager._is_client_credentials_flow() is False + + def test_is_client_credentials_flow_oidc_with_client_credentials(self): + """Test _is_client_credentials_flow returns True for OIDC with client credentials.""" + from google.adk.auth.auth_schemes import OpenIdConnectWithConfig + + # Create OIDC scheme with client credentials support + auth_scheme = OpenIdConnectWithConfig( + authorization_endpoint="https://example.com/auth", + token_endpoint="https://example.com/token", + grant_types_supported=["authorization_code", "client_credentials"], + ) + + auth_config = Mock(spec=AuthConfig) + auth_config.auth_scheme = auth_scheme + auth_config.raw_auth_credential = None + auth_config.exchanged_auth_credential = None + + manager = CredentialManager(auth_config) + + assert manager._is_client_credentials_flow() is True + + def test_is_client_credentials_flow_oidc_without_client_credentials(self): + """Test _is_client_credentials_flow returns False for OIDC without client credentials.""" + from google.adk.auth.auth_schemes import OpenIdConnectWithConfig + + # Create OIDC scheme without client credentials support + auth_scheme = OpenIdConnectWithConfig( + authorization_endpoint="https://example.com/auth", + token_endpoint="https://example.com/token", + grant_types_supported=["authorization_code"], + ) + + auth_config = Mock(spec=AuthConfig) + auth_config.auth_scheme = auth_scheme + auth_config.raw_auth_credential = None + auth_config.exchanged_auth_credential = None + + manager = CredentialManager(auth_config) + + assert manager._is_client_credentials_flow() is False + + def test_is_client_credentials_flow_other_scheme(self): + """Test _is_client_credentials_flow returns False for other auth schemes.""" + # Create a non-OAuth2/OIDC scheme + auth_scheme = Mock() + + auth_config = Mock(spec=AuthConfig) + auth_config.auth_scheme = auth_scheme + auth_config.raw_auth_credential = None + auth_config.exchanged_auth_credential = None + + manager = CredentialManager(auth_config) + + assert manager._is_client_credentials_flow() is False + @pytest.fixture def oauth2_auth_scheme():