Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 35 additions & 6 deletions src/google/adk/tools/openapi_tool/auth/auth_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@
from fastapi.openapi.models import HTTPBase
from fastapi.openapi.models import HTTPBearer
from fastapi.openapi.models import OAuth2
from fastapi.openapi.models import OAuthFlowClientCredentials
from fastapi.openapi.models import OAuthFlows
from fastapi.openapi.models import OpenIdConnect
from fastapi.openapi.models import Schema
import httpx
Expand Down Expand Up @@ -152,13 +154,40 @@ def token_to_scheme_credential(
raise ValueError(f"Invalid security scheme type: {type}")


def _service_account_auth_scheme() -> OAuth2:
"""Auth scheme for Google Service Account credentials.

CredentialManager only auto-loads raw non-interactive credentials when the
scheme is an OAuth2/OIDC client-credentials flow. An HTTPBearer scheme makes
``_is_client_credentials_flow`` return False, so ``get_auth_credential``
returns None and the tool falls back to ``adk_request_credential`` instead of
exchanging the service account for a token.

The token URL is unused by ServiceAccountCredentialExchanger (ADC / JWT
assertion), but is required by the OAuth2 client-credentials model.
"""
return OAuth2(
flows=OAuthFlows(
clientCredentials=OAuthFlowClientCredentials(
# Placeholder only; SA exchange does not call this endpoint.
# Use the mTLS host form for compliance with Google API endpoint
# requirements.
tokenUrl="https://oauth2.mtls.googleapis.com/token",
scopes={},
)
)
)


def service_account_dict_to_scheme_credential(
config: Dict[str, Any],
scopes: List[str],
) -> Tuple[AuthScheme, AuthCredential]:
"""Creates AuthScheme and AuthCredential for Google Service Account.

Returns a bearer token scheme, and a service account credential.
Returns an OAuth2 client-credentials scheme (so CredentialManager can
exchange the service account) and a service account credential. After
exchange the credential is an HTTP bearer token.

Args:
config: A ServiceAccount object containing the Google Service Account
Expand All @@ -168,7 +197,6 @@ def service_account_dict_to_scheme_credential(
Returns:
Tuple: (AuthScheme, AuthCredential)
"""
auth_scheme = HTTPBearer(bearerFormat="JWT")
service_account = ServiceAccount(
service_account_credential=ServiceAccountCredential.model_construct(
**config
Expand All @@ -179,15 +207,17 @@ def service_account_dict_to_scheme_credential(
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
service_account=service_account,
)
return auth_scheme, auth_credential
return _service_account_auth_scheme(), auth_credential


def service_account_scheme_credential(
config: ServiceAccount,
) -> Tuple[AuthScheme, AuthCredential]:
"""Creates AuthScheme and AuthCredential for Google Service Account.

Returns a bearer token scheme, and a service account credential.
Returns an OAuth2 client-credentials scheme (so CredentialManager can
exchange the service account) and a service account credential. After
exchange the credential is an HTTP bearer token.

Args:
config: A ServiceAccount object containing the Google Service Account
Expand All @@ -196,11 +226,10 @@ def service_account_scheme_credential(
Returns:
Tuple: (AuthScheme, AuthCredential)
"""
auth_scheme = HTTPBearer(bearerFormat="JWT")
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, service_account=config
)
return auth_scheme, auth_credential
return _service_account_auth_scheme(), auth_credential


def openid_dict_to_scheme_credential(
Expand Down
56 changes: 52 additions & 4 deletions tests/unittests/tools/openapi_tool/auth/test_auth_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from unittest.mock import AsyncMock
from unittest.mock import Mock
from unittest.mock import patch

from fastapi.openapi.models import APIKey
Expand All @@ -28,6 +30,8 @@
from google.adk.auth.auth_credential import ServiceAccountCredential
from google.adk.auth.auth_schemes import AuthSchemeType
from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
from google.adk.auth.auth_tool import AuthConfig
from google.adk.auth.credential_manager import CredentialManager
from google.adk.tools.openapi_tool.auth.auth_helpers import credential_to_param
from google.adk.tools.openapi_tool.auth.auth_helpers import dict_to_auth_scheme
from google.adk.tools.openapi_tool.auth.auth_helpers import INTERNAL_AUTH_PREFIX
Expand Down Expand Up @@ -133,8 +137,10 @@ def test_service_account_dict_to_scheme_credential():

scheme, credential = service_account_dict_to_scheme_credential(config, scopes)

assert isinstance(scheme, HTTPBearer)
assert scheme.bearerFormat == "JWT"
assert isinstance(scheme, OAuth2)
assert scheme.flows is not None
assert scheme.flows.clientCredentials is not None
assert scheme.flows.clientCredentials.tokenUrl
assert credential.auth_type == AuthCredentialTypes.SERVICE_ACCOUNT
assert credential.service_account.scopes == scopes
assert (
Expand Down Expand Up @@ -163,12 +169,54 @@ def test_service_account_scheme_credential():

scheme, credential = service_account_scheme_credential(config)

assert isinstance(scheme, HTTPBearer)
assert scheme.bearerFormat == "JWT"
assert isinstance(scheme, OAuth2)
assert scheme.flows is not None
assert scheme.flows.clientCredentials is not None
assert scheme.flows.clientCredentials.tokenUrl
assert credential.auth_type == AuthCredentialTypes.SERVICE_ACCOUNT
assert credential.service_account == config


@pytest.mark.asyncio
async def test_service_account_helper_scheme_allows_credential_manager_exchange():
"""SA helpers must yield a client-credentials scheme (#6656).

With HTTPBearer, CredentialManager treated the SA as needing interactive
auth and returned None (adk_request_credential) instead of exchanging it.
"""
scheme, credential = service_account_scheme_credential(
ServiceAccount(
use_default_credential=True,
scopes=["https://www.googleapis.com/auth/cloud-platform"],
)
)
manager = CredentialManager(
AuthConfig(auth_scheme=scheme, raw_auth_credential=credential)
)
assert manager._is_client_credentials_flow() # pylint: disable=protected-access

exchanged = AuthCredential(
auth_type=AuthCredentialTypes.HTTP,
http=HttpAuth(
scheme="bearer",
credentials=HttpCredentials(token="sa-access-token"),
),
)
manager._load_existing_credential = AsyncMock(return_value=None) # pylint: disable=protected-access
manager._exchange_credential = AsyncMock(return_value=(exchanged, True)) # pylint: disable=protected-access
manager._refresh_credential = AsyncMock(return_value=(exchanged, False)) # pylint: disable=protected-access
manager._save_credential = AsyncMock() # pylint: disable=protected-access

ctx = Mock()
ctx.get_auth_response = Mock(return_value=None)
result = await manager.get_auth_credential(ctx)

assert result is not None
assert result.auth_type == AuthCredentialTypes.HTTP
assert result.http.credentials.token == "sa-access-token"
manager._exchange_credential.assert_awaited_once() # pylint: disable=protected-access


def test_openid_dict_to_scheme_credential():
config_dict = {
"authorization_endpoint": "auth_url",
Expand Down