diff --git a/src/google/adk/auth/auth_credential.py b/src/google/adk/auth/auth_credential.py index a71ae317..4a2add82 100644 --- a/src/google/adk/auth/auth_credential.py +++ b/src/google/adk/auth/auth_credential.py @@ -19,7 +19,6 @@ from typing import Any from typing import Dict from typing import List from typing import Literal -from typing import Optional from pydantic import alias_generators from pydantic import BaseModel @@ -40,9 +39,9 @@ class BaseModelWithConfig(BaseModel): class HttpCredentials(BaseModelWithConfig): """Represents the secret token value for HTTP authentication, like user name, password, oauth token, etc.""" - username: Optional[str] = None - password: Optional[str] = None - token: Optional[str] = None + username: str | None = None + password: str | None = None + token: str | None = None @classmethod def model_validate(cls, data: Dict[str, Any]) -> "HttpCredentials": @@ -62,40 +61,43 @@ class HttpAuth(BaseModelWithConfig): # Examples: 'basic', 'bearer' scheme: str credentials: HttpCredentials - additional_headers: Optional[Dict[str, str]] = None + additional_headers: Dict[str, str] | None = None class OAuth2Auth(BaseModelWithConfig): """Represents credential value and its metadata for a OAuth2 credential.""" - client_id: Optional[str] = None - client_secret: Optional[str] = None + client_id: str | None = None + client_secret: str | None = None # tool or adk can generate the auth_uri with the state info thus client # can verify the state - auth_uri: Optional[str] = None + auth_uri: str | None = None # A unique value generated at the start of the OAuth flow to bind the user's # session to the authorization request. This value is typically stored with # user session and passed to backend for validation. - nonce: Optional[str] = None - state: Optional[str] = None + nonce: str | None = None + state: str | None = None # tool or adk can decide the redirect_uri if they don't want client to decide - redirect_uri: Optional[str] = None - auth_response_uri: Optional[str] = None - auth_code: Optional[str] = None - access_token: Optional[str] = None - refresh_token: Optional[str] = None - id_token: Optional[str] = None - expires_at: Optional[int] = None - expires_in: Optional[int] = None - audience: Optional[str] = None - token_endpoint_auth_method: Optional[ + redirect_uri: str | None = None + auth_response_uri: str | None = None + auth_code: str | None = None + access_token: str | None = None + refresh_token: str | None = None + id_token: str | None = None + expires_at: int | None = None + expires_in: int | None = None + audience: str | None = None + code_verifier: str | None = None + code_challenge_method: str | None = None + token_endpoint_auth_method: ( Literal[ "client_secret_basic", "client_secret_post", "client_secret_jwt", "private_key_jwt", ] - ] = "client_secret_basic" + | None + ) = "client_secret_basic" class ServiceAccountCredential(BaseModelWithConfig): @@ -166,11 +168,11 @@ class ServiceAccount(BaseModelWithConfig): when ``use_id_token`` is True. """ - service_account_credential: Optional[ServiceAccountCredential] = None - scopes: Optional[List[str]] = None - use_default_credential: Optional[bool] = False - use_id_token: Optional[bool] = False - audience: Optional[str] = None + service_account_credential: ServiceAccountCredential | None = None + scopes: List[str] | None = None + use_default_credential: bool | None = False + use_id_token: bool | None = False + audience: str | None = None @model_validator(mode="after") def _validate_config(self) -> ServiceAccount: @@ -275,9 +277,9 @@ class AuthCredential(BaseModelWithConfig): auth_type: AuthCredentialTypes # Resource reference for the credential. # This will be supported in the future. - resource_ref: Optional[str] = None + resource_ref: str | None = None - api_key: Optional[str] = None - http: Optional[HttpAuth] = None - service_account: Optional[ServiceAccount] = None - oauth2: Optional[OAuth2Auth] = None + api_key: str | None = None + http: HttpAuth | None = None + service_account: ServiceAccount | None = None + oauth2: OAuth2Auth | None = None diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index ec7c7571..8e8f5d34 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -28,6 +28,7 @@ if TYPE_CHECKING: from ..sessions.state import State try: + from authlib.common.security import generate_token from authlib.integrations.requests_client import OAuth2Session AUTHLIB_AVAILABLE = True @@ -158,6 +159,8 @@ class AuthHandler: auth_scheme = self.auth_config.auth_scheme auth_credential = self.auth_config.raw_auth_credential + if not auth_credential or not auth_credential.oauth2: + raise ValueError("raw_auth_credential or oauth2 is empty") if isinstance(auth_scheme, OpenIdConnectWithConfig): authorization_endpoint = auth_scheme.authorization_endpoint @@ -190,6 +193,7 @@ class AuthHandler: auth_credential.oauth2.client_secret, scope=" ".join(scopes), redirect_uri=auth_credential.oauth2.redirect_uri, + code_challenge_method=auth_credential.oauth2.code_challenge_method, ) params = { "access_type": "offline", @@ -197,12 +201,30 @@ class AuthHandler: } if auth_credential.oauth2.audience: params["audience"] = auth_credential.oauth2.audience + + # If using PKCE with S256, ensure a code_verifier exists. + # If not provided in the credential, generate a cryptographically secure + # random token of 48 characters (OAuth2 recommends 43-128 characters). + code_verifier = auth_credential.oauth2.code_verifier + method = auth_credential.oauth2.code_challenge_method + + if method: + if method != "S256": + raise ValueError( + f"Unsupported code_challenge_method: {method}. Only 'S256' is" + " supported." + ) + if not code_verifier: + code_verifier = generate_token(48) + uri, state = client.create_authorization_url( - url=authorization_endpoint, **params + url=authorization_endpoint, code_verifier=code_verifier, **params ) exchanged_auth_credential = auth_credential.model_copy(deep=True) exchanged_auth_credential.oauth2.auth_uri = uri exchanged_auth_credential.oauth2.state = state + if code_verifier: + exchanged_auth_credential.oauth2.code_verifier = code_verifier return exchanged_auth_credential diff --git a/src/google/adk/auth/exchanger/oauth2_credential_exchanger.py b/src/google/adk/auth/exchanger/oauth2_credential_exchanger.py index 76f0c678..d3504bff 100644 --- a/src/google/adk/auth/exchanger/oauth2_credential_exchanger.py +++ b/src/google/adk/auth/exchanger/oauth2_credential_exchanger.py @@ -193,6 +193,12 @@ class OAuth2CredentialExchanger(BaseCredentialExchanger): return ExchangeResult(auth_credential, False) try: + kwargs = {} + # If a code_verifier is available (e.g. from PKCE), include it in the + # token exchange request. + if auth_credential.oauth2 and auth_credential.oauth2.code_verifier: + kwargs["code_verifier"] = auth_credential.oauth2.code_verifier + # Authlib already injects client_id for body-based client auth flows such # as client_secret_post, so passing it here would duplicate the field. tokens = client.fetch_token( @@ -202,6 +208,7 @@ class OAuth2CredentialExchanger(BaseCredentialExchanger): ), code=auth_credential.oauth2.auth_code, grant_type=OAuthGrantType.AUTHORIZATION_CODE, + **kwargs, ) update_credential_with_tokens(auth_credential, tokens) logger.debug("Successfully exchanged authorization code for access token") diff --git a/src/google/adk/auth/oauth2_credential_util.py b/src/google/adk/auth/oauth2_credential_util.py index df2f26c0..d0d1255f 100644 --- a/src/google/adk/auth/oauth2_credential_util.py +++ b/src/google/adk/auth/oauth2_credential_util.py @@ -92,6 +92,7 @@ def create_oauth2_session( 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, ) diff --git a/tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py b/tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py index 3a0a5647..25f92674 100644 --- a/tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py +++ b/tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py @@ -135,6 +135,57 @@ class TestOAuth2CredentialExchanger: assert exchange_result.was_exchanged mock_client.fetch_token.assert_called_once() + @patch("google.adk.auth.oauth2_credential_util.OAuth2Session") + async def test_exchange_success_pkce(self, mock_oauth2_session): + """Test successful token exchange with PKCE.""" + # Setup mock + mock_client = Mock() + mock_oauth2_session.return_value = mock_client + mock_tokens = OAuth2Token({ + "access_token": "new_access_token", + "refresh_token": "new_refresh_token", + "expires_at": int(time.time()) + 3600, + "expires_in": 3600, + }) + mock_client.fetch_token.return_value = mock_tokens + + 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"], + ) + credential = AuthCredential( + auth_type=AuthCredentialTypes.OPEN_ID_CONNECT, + oauth2=OAuth2Auth( + client_id="test_client_id", + client_secret="test_client_secret", + auth_response_uri="https://example.com/callback?code=auth_code", + auth_code="auth_code", + code_verifier="mock_code_verifier", + ), + ) + + exchanger = OAuth2CredentialExchanger() + exchange_result = await exchanger.exchange(credential, scheme) + + # Verify token exchange was successful + assert exchange_result.credential.oauth2.access_token == "new_access_token" + assert ( + exchange_result.credential.oauth2.refresh_token == "new_refresh_token" + ) + assert exchange_result.was_exchanged + mock_client.fetch_token.assert_called_once_with( + "https://example.com/token", + authorization_response="https://example.com/callback?code=auth_code", + code="auth_code", + grant_type=OAuthGrantType.AUTHORIZATION_CODE, + code_verifier="mock_code_verifier", + ) + async def test_exchange_missing_auth_scheme(self): """Test exchange with missing auth_scheme raises ValueError.""" credential = AuthCredential( diff --git a/tests/unittests/auth/test_auth_handler.py b/tests/unittests/auth/test_auth_handler.py index 2faeeb15..c19a5d93 100644 --- a/tests/unittests/auth/test_auth_handler.py +++ b/tests/unittests/auth/test_auth_handler.py @@ -53,12 +53,14 @@ class MockOAuth2Session: scope=None, redirect_uri=None, state=None, + **kwargs, ): self.client_id = client_id self.client_secret = client_secret self.scope = scope self.redirect_uri = redirect_uri self.state = state + self.extra_kwargs = kwargs def create_authorization_url(self, url, **kwargs): params = f"client_id={self.client_id}&scope={self.scope}" @@ -271,6 +273,54 @@ class TestGenerateAuthUri: assert "client_id=mock_client_id" in result.oauth2.auth_uri assert result.oauth2.state == "mock_state" + @patch("google.adk.auth.auth_handler.OAuth2Session") + def test_generate_auth_uri_pkce( + self, mock_oauth2_session, oauth2_auth_scheme, oauth2_credentials + ): + """Test generating an auth URI with PKCE.""" + oauth2_credentials.oauth2.code_challenge_method = "S256" + exchanged = oauth2_credentials.model_copy(deep=True) + + config = AuthConfig( + auth_scheme=oauth2_auth_scheme, + raw_auth_credential=oauth2_credentials, + exchanged_auth_credential=exchanged, + ) + + mock_client = Mock() + mock_oauth2_session.return_value = mock_client + mock_client.create_authorization_url.return_value = ( + "https://example.com/oauth2/authorize?code_challenge=...&code_challenge_method=S256", + "mock_state", + ) + + handler = AuthHandler(config) + result = handler.generate_auth_uri() + + assert result.oauth2.code_verifier is not None + assert len(result.oauth2.code_verifier) == 48 + mock_client.create_authorization_url.assert_called_once() + _, kwargs = mock_client.create_authorization_url.call_args + assert "code_verifier" in kwargs + assert kwargs["code_verifier"] == result.oauth2.code_verifier + + def test_generate_auth_uri_unsupported_pkce_method( + self, oauth2_auth_scheme, oauth2_credentials + ): + """Test generating an auth URI with unsupported PKCE method.""" + oauth2_credentials.oauth2.code_challenge_method = "plain" + exchanged = oauth2_credentials.model_copy(deep=True) + + config = AuthConfig( + auth_scheme=oauth2_auth_scheme, + raw_auth_credential=oauth2_credentials, + exchanged_auth_credential=exchanged, + ) + + handler = AuthHandler(config) + with pytest.raises(ValueError, match="Unsupported code_challenge_method"): + handler.generate_auth_uri() + class TestGenerateAuthRequest: """Tests for the generate_auth_request method."""