encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: Make credential key generation stable and prevent cross-user credential leaks

Browse Source 
This change updates the credential key generation to use a stable hash (SHA256) instead of Python's built-in hash, which can vary based on PYTHONHASHSEED. It also makes sure that temporary or exchanged OAuth2 fields are excluded from the key calculation. I also added when saving credentials, a copy of the AuthConfig is used to avoid modifying the original shared AuthConfig instance with user-specific exchanged credentials.

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 864599326
This commit is contained in:
George Weale
2026-02-02 17:51:19 -08:00
committed by Copybara-Service
parent 666cebe369
commit 33012e6dda
11 changed files with 448 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/auth/auth_handler.py
+2
View File
@@ -115,6 +115,7 @@ class AuthHandler:
exchanged_auth_credential=self.auth_config.raw_auth_credential.model_copy(
deep=True
),
credential_key=self.auth_config.credential_key,
)
# Check for client_id and client_secret
@@ -133,6 +134,7 @@ class AuthHandler:
auth_scheme=self.auth_config.auth_scheme,
raw_auth_credential=self.auth_config.raw_auth_credential,
exchanged_auth_credential=exchanged_credential,
credential_key=self.auth_config.credential_key,
)
def generate_auth_uri(
src/google/adk/auth/auth_preprocessor.py
+32
View File
@@ -72,6 +72,7 @@ class _AuthLlmRequestProcessor(BaseLlmRequestProcessor):
if not responses:
return
requested_auth_config_by_request_id = {}
# look for auth response
for function_call_response in responses:
if function_call_response.name != REQUEST_EUC_FUNCTION_CALL_NAME:
@@ -79,7 +80,38 @@ class _AuthLlmRequestProcessor(BaseLlmRequestProcessor):
# found the function call response for the system long running request euc
# function call
request_euc_function_call_ids.add(function_call_response.id)
if request_euc_function_call_ids:
for event in events:
function_calls = event.get_function_calls()
if not function_calls:
continue
try:
for function_call in function_calls:
if (
function_call.id in request_euc_function_call_ids
and function_call.name == REQUEST_EUC_FUNCTION_CALL_NAME
):
args = AuthToolArguments.model_validate(function_call.args)
requested_auth_config_by_request_id[function_call.id] = (
args.auth_config
)
except TypeError:
continue
for function_call_response in responses:
if function_call_response.name != REQUEST_EUC_FUNCTION_CALL_NAME:
continue
auth_config = AuthConfig.model_validate(function_call_response.response)
requested_auth_config = requested_auth_config_by_request_id.get(
function_call_response.id
)
if (
requested_auth_config
and requested_auth_config.credential_key is not None
):
auth_config.credential_key = requested_auth_config.credential_key
await AuthHandler(auth_config=auth_config).parse_and_store_auth_response(
state=invocation_context.session.state
)
src/google/adk/auth/auth_tool.py
+49 -4
View File
@@ -14,8 +14,11 @@
from __future__ import annotations
import hashlib
import json
from typing import Optional
from pydantic import BaseModel
from typing_extensions import deprecated
from .auth_credential import AuthCredential
@@ -23,6 +26,28 @@ from .auth_credential import BaseModelWithConfig
from .auth_schemes import AuthScheme
def _stable_model_digest(model: BaseModel) -> str:
"""Returns a stable digest for a pydantic model.
The digest is stable across:
- Python hash seeds (does not use `hash()`).
- Dict insertion ordering differences (canonicalizes via `sort_keys=True`).
- Pydantic `model_extra` values (ignored).
"""
if getattr(model, "model_extra", None):
model = model.model_copy(deep=True)
model.model_extra.clear()
dumped = model.model_dump(by_alias=True, exclude_none=True, mode="json")
canonical_json = json.dumps(
dumped,
sort_keys=True,
ensure_ascii=False,
separators=(",", ":"),
)
return hashlib.sha256(canonical_json.encode("utf-8")).hexdigest()[:16]
class AuthConfig(BaseModelWithConfig):
"""The auth config sent by tool asking client to collect auth credentials and
@@ -58,12 +83,22 @@ class AuthConfig(BaseModelWithConfig):
super().__init__(**data)
if self.credential_key:
return
for obj in (self.raw_auth_credential, self.auth_scheme):
if not obj or not getattr(obj, "model_extra", None):
continue
for key in ("credential_key", "credentialKey"):
value = obj.model_extra.get(key)
if isinstance(value, str) and value:
self.credential_key = value
return
self.credential_key = self.get_credential_key()
@deprecated("This method is deprecated. Use credential_key instead.")
def get_credential_key(self):
"""Builds a hash key based on auth_scheme and raw_auth_credential used to
save / load this credential to / from a credentials service.
"""Builds a stable key based on auth_scheme and raw_auth_credential.
This is used to save/load credentials to/from a credential service when
`credential_key` is not explicitly provided.
"""
auth_scheme = self.auth_scheme
@@ -72,7 +107,7 @@ class AuthConfig(BaseModelWithConfig):
auth_scheme = auth_scheme.model_copy(deep=True)
auth_scheme.model_extra.clear()
scheme_name = (
f"{auth_scheme.type_.name}_{hash(auth_scheme.model_dump_json())}"
f"{auth_scheme.type_.name}_{_stable_model_digest(auth_scheme)}"
if auth_scheme
else ""
)
@@ -81,8 +116,18 @@ class AuthConfig(BaseModelWithConfig):
if auth_credential and auth_credential.model_extra:
auth_credential = auth_credential.model_copy(deep=True)
auth_credential.model_extra.clear()
if auth_credential and auth_credential.oauth2:
auth_credential = auth_credential.model_copy(deep=True)
auth_credential.oauth2.auth_uri = None
auth_credential.oauth2.state = None
auth_credential.oauth2.auth_response_uri = None
auth_credential.oauth2.auth_code = None
auth_credential.oauth2.access_token = None
auth_credential.oauth2.refresh_token = None
auth_credential.oauth2.expires_at = None
auth_credential.oauth2.expires_in = None
credential_name = (
f"{auth_credential.auth_type.value}_{hash(auth_credential.model_dump_json())}"
f"{auth_credential.auth_type.value}_{_stable_model_digest(auth_credential)}"
if auth_credential
else ""
)
src/google/adk/auth/credential_manager.py
+9 -6
View File
@@ -142,7 +142,9 @@ class CredentialManager:
# Step 2: Check if credential is already ready (no processing needed)
if self._is_credential_ready():
return self._auth_config.raw_auth_credential
# Return a copy to avoid leaking mutations across invocations/users when
# tools share a long-lived AuthConfig instance.
return self._auth_config.raw_auth_credential.model_copy(deep=True)
# Step 3: Try to load existing processed credential
credential = await self._load_existing_credential(context)
@@ -159,7 +161,9 @@ class CredentialManager:
if not credential:
# For client credentials flow, use raw credentials directly
if self._is_client_credentials_flow():
credential = self._auth_config.raw_auth_credential
# Exchange/refresh steps may mutate the credential object in-place, so
# do not operate on the shared tool config.
credential = self._auth_config.raw_auth_credential.model_copy(deep=True)
else:
# For authorization code flow, return None to trigger user authorization
return None
@@ -298,12 +302,11 @@ class CredentialManager:
self, context: CallbackContext, credential: AuthCredential
) -> None:
"""Save credential to credential service if available."""
# Update the exchanged credential in config
self._auth_config.exchanged_auth_credential = credential
credential_service = context._invocation_context.credential_service
if credential_service:
await context.save_credential(self._auth_config)
auth_config_to_save = self._auth_config.model_copy(deep=True)
auth_config_to_save.exchanged_auth_credential = credential
await context.save_credential(auth_config_to_save)
async def _populate_auth_scheme(self) -> bool:
"""Auto-discover server metadata and populate missing auth scheme info.
src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py
+16 -4
View File
@@ -70,6 +70,7 @@ class OpenAPIToolset(BaseToolset):
spec_str_type: Literal["json", "yaml"] = "json",
auth_scheme: Optional[AuthScheme] = None,
auth_credential: Optional[AuthCredential] = None,
credential_key: Optional[str] = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
tool_name_prefix: Optional[str] = None,
ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None,
@@ -108,6 +109,8 @@ class OpenAPIToolset(BaseToolset):
auth_credential: The auth credential to use for all tools. Use
AuthCredential or use helpers in
``google.adk.tools.openapi_tool.auth.auth_helpers``
credential_key: Optional stable key used for interactive auth and
credential caching across all tools in this toolset.
tool_filter: The filter used to filter the tools in the toolset. It can be
either a tool predicate or a list of tool names of the tools to expose.
tool_name_prefix: The prefix to prepend to the names of the tools returned
@@ -137,6 +140,7 @@ class OpenAPIToolset(BaseToolset):
AuthConfig(
auth_scheme=auth_scheme,
raw_auth_credential=auth_credential,
credential_key=credential_key,
)
if auth_scheme
else None
@@ -147,6 +151,8 @@ class OpenAPIToolset(BaseToolset):
self._tools: Final[List[RestApiTool]] = list(self._parse(spec_dict))
if auth_scheme or auth_credential:
self._configure_auth_all(auth_scheme, auth_credential)
if credential_key:
self._configure_credential_key_all(credential_key)
def _configure_auth_all(
self, auth_scheme: AuthScheme, auth_credential: AuthCredential
@@ -159,6 +165,11 @@ class OpenAPIToolset(BaseToolset):
if auth_credential:
tool.configure_auth_credential(auth_credential)
def _configure_credential_key_all(self, credential_key: str):
"""Configure credential key for all tools."""
for tool in self._tools:
tool.configure_credential_key(credential_key)
def configure_ssl_verify_all(
self, ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None
):
@@ -229,8 +240,9 @@ class OpenAPIToolset(BaseToolset):
def get_auth_config(self) -> Optional[AuthConfig]:
"""Returns the auth config for this toolset.
ADK will populate exchanged_auth_credential on this config before calling
get_tools(). The toolset can then access the ready-to-use credential via
self._auth_config.exchanged_auth_credential.
Note: This returns a copy so any exchanged credentials populated by the ADK
framework do not persist on the toolset instance across invocations.
"""
return self._auth_config
return (
self._auth_config.model_copy(deep=True) if self._auth_config else None
)
src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py
+13 -1
View File
@@ -100,6 +100,8 @@ class RestApiTool(BaseTool):
header_provider: Optional[
Callable[[ReadonlyContext], Dict[str, str]]
] = None,
*,
credential_key: Optional[str] = None,
):
"""Initializes the RestApiTool with the given parameters.
@@ -137,6 +139,8 @@ class RestApiTool(BaseTool):
an argument, allowing dynamic header generation based on the current
context. Useful for adding custom headers like correlation IDs,
authentication tokens, or other request metadata.
credential_key: Optional stable key used for interactive auth and
credential caching.
"""
# Gemini restrict the length of function name to be less than 64 characters
self.name = name[:60]
@@ -152,6 +156,7 @@ class RestApiTool(BaseTool):
else operation
)
self.auth_credential, self.auth_scheme = None, None
self.credential_key = credential_key
self.configure_auth_credential(auth_credential)
self.configure_auth_scheme(auth_scheme)
@@ -266,6 +271,10 @@ class RestApiTool(BaseTool):
auth_credential = AuthCredential.model_validate_json(auth_credential)
self.auth_credential = auth_credential
def configure_credential_key(self, credential_key: Optional[str] = None):
"""Configures the credential key for interactive auth / caching."""
self.credential_key = credential_key
def configure_ssl_verify(
self, ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None
):
@@ -449,7 +458,10 @@ class RestApiTool(BaseTool):
"""
# Prepare auth credentials for the API call
tool_auth_handler = ToolAuthHandler.from_tool_context(
tool_context, self.auth_scheme, self.auth_credential
tool_context,
self.auth_scheme,
self.auth_credential,
credential_key=self.credential_key,
)
auth_result = await tool_auth_handler.prepare_auth_credentials()
auth_state, auth_scheme, auth_credential = (
src/google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py
+89 -18
View File
@@ -14,6 +14,7 @@
from __future__ import annotations
import hashlib
import logging
from typing import Literal
from typing import Optional
@@ -24,6 +25,7 @@ from ....auth.auth_credential import AuthCredential
from ....auth.auth_credential import AuthCredentialTypes
from ....auth.auth_schemes import AuthScheme
from ....auth.auth_schemes import AuthSchemeType
from ....auth.auth_tool import _stable_model_digest
from ....auth.auth_tool import AuthConfig
from ....auth.refresher.oauth2_credential_refresher import OAuth2CredentialRefresher
from ...tool_context import ToolContext
@@ -50,19 +52,60 @@ class ToolContextCredentialStore:
def __init__(self, tool_context: ToolContext):
self.tool_context = tool_context
def _legacy_stable_digest(self, text: str) -> str:
return hashlib.sha256(text.encode("utf-8")).hexdigest()[:16]
def _get_legacy_credential_key(
self,
auth_scheme: Optional[AuthScheme],
auth_credential: Optional[AuthCredential],
) -> str:
if auth_credential and auth_credential.oauth2:
auth_credential = auth_credential.model_copy(deep=True)
auth_credential.oauth2.auth_uri = None
auth_credential.oauth2.state = None
auth_credential.oauth2.auth_response_uri = None
auth_credential.oauth2.auth_code = None
auth_credential.oauth2.access_token = None
auth_credential.oauth2.refresh_token = None
auth_credential.oauth2.expires_at = None
auth_credential.oauth2.expires_in = None
scheme_name = (
f"{auth_scheme.type_.name}_{self._legacy_stable_digest(auth_scheme.model_dump_json())}"
if auth_scheme
else ""
)
credential_name = (
f"{auth_credential.auth_type.value}_{self._legacy_stable_digest(auth_credential.model_dump_json())}"
if auth_credential
else ""
)
return f"{scheme_name}_{credential_name}_existing_exchanged_credential"
def get_credential_key(
self,
auth_scheme: Optional[AuthScheme],
auth_credential: Optional[AuthCredential],
) -> str:
"""Generates a unique key for the given auth scheme and credential."""
if auth_credential and auth_credential.oauth2:
auth_credential = auth_credential.model_copy(deep=True)
auth_credential.oauth2.auth_uri = None
auth_credential.oauth2.state = None
auth_credential.oauth2.auth_response_uri = None
auth_credential.oauth2.auth_code = None
auth_credential.oauth2.access_token = None
auth_credential.oauth2.refresh_token = None
auth_credential.oauth2.expires_at = None
auth_credential.oauth2.expires_in = None
scheme_name = (
f"{auth_scheme.type_.name}_{hash(auth_scheme.model_dump_json())}"
f"{auth_scheme.type_.name}_{_stable_model_digest(auth_scheme)}"
if auth_scheme
else ""
)
credential_name = (
f"{auth_credential.auth_type.value}_{hash(auth_credential.model_dump_json())}"
f"{auth_credential.auth_type.value}_{_stable_model_digest(auth_credential)}"
if auth_credential
else ""
)
@@ -86,9 +129,19 @@ class ToolContextCredentialStore:
# session implementation, we don't want session to persist the token,
# meanwhile we want the token shared across runs.
serialized_credential = self.tool_context.state.get(token_key)
if not serialized_credential:
if serialized_credential:
return AuthCredential.model_validate(serialized_credential)
legacy_key = self._get_legacy_credential_key(auth_scheme, auth_credential)
if legacy_key == token_key:
return None
return AuthCredential.model_validate(serialized_credential)
serialized_legacy_credential = self.tool_context.state.get(legacy_key)
if not serialized_legacy_credential:
return None
# Migrate to the current key for future lookups.
self.tool_context.state[token_key] = serialized_legacy_credential
return AuthCredential.model_validate(serialized_legacy_credential)
def store_credential(
self,
@@ -114,6 +167,8 @@ class ToolAuthHandler:
auth_credential: Optional[AuthCredential],
credential_exchanger: Optional[BaseAuthCredentialExchanger] = None,
credential_store: Optional["ToolContextCredentialStore"] = None,
*,
credential_key: Optional[str] = None,
):
self.tool_context = tool_context
self.auth_scheme = (
@@ -122,12 +177,35 @@ class ToolAuthHandler:
self.auth_credential = (
auth_credential.model_copy(deep=True) if auth_credential else None
)
self._credential_key = credential_key
self.credential_exchanger = (
credential_exchanger or AutoAuthCredentialExchanger()
)
self.credential_store = credential_store
self.should_store_credential = True
def _get_credential_key_override(self) -> Optional[str]:
"""Returns a user-provided credential_key if available."""
if self._credential_key:
return self._credential_key
for obj in (self.auth_credential, self.auth_scheme):
if not obj or not obj.model_extra:
continue
for key in ("credential_key", "credentialKey"):
value = obj.model_extra.get(key)
if isinstance(value, str) and value:
return value
return None
def _build_auth_config(self) -> AuthConfig:
return AuthConfig(
auth_scheme=self.auth_scheme,
raw_auth_credential=self.auth_credential,
credential_key=self._get_credential_key_override(),
)
@classmethod
def from_tool_context(
cls,
@@ -135,6 +213,8 @@ class ToolAuthHandler:
auth_scheme: Optional[AuthScheme],
auth_credential: Optional[AuthCredential],
credential_exchanger: Optional[BaseAuthCredentialExchanger] = None,
*,
credential_key: Optional[str] = None,
) -> "ToolAuthHandler":
"""Creates a ToolAuthHandler instance from a ToolContext."""
credential_store = ToolContextCredentialStore(tool_context)
@@ -142,8 +222,9 @@ class ToolAuthHandler:
tool_context,
auth_scheme,
auth_credential,
credential_exchanger,
credential_store,
credential_key=credential_key,
credential_exchanger=credential_exchanger,
credential_store=credential_store,
)
async def _get_existing_credential(
@@ -209,21 +290,11 @@ class ToolAuthHandler:
"OAuth2 credentials client_secret is missing."
)
self.tool_context.request_credential(
AuthConfig(
auth_scheme=self.auth_scheme,
raw_auth_credential=self.auth_credential,
)
)
self.tool_context.request_credential(self._build_auth_config())
return None
def _get_auth_response(self) -> AuthCredential:
return self.tool_context.get_auth_response(
AuthConfig(
auth_scheme=self.auth_scheme,
raw_auth_credential=self.auth_credential,
)
)
return self.tool_context.get_auth_response(self._build_auth_config())
def _external_exchange_required(self, credential) -> bool:
return (