diff --git a/src/google/adk/tools/bigtable/bigtable_toolset.py b/src/google/adk/tools/bigtable/bigtable_toolset.py index eaa1dae61d8..0aef34d288f 100644 --- a/src/google/adk/tools/bigtable/bigtable_toolset.py +++ b/src/google/adk/tools/bigtable/bigtable_toolset.py @@ -54,6 +54,12 @@ class BigtableParameterizedViewTool(GoogleTool): pass it as `view_parameters={"user_id": user_id}`. This securely restricts query execution to the logged-in user's data without exposing the `user_id` parameter to the LLM. + + Only names that resolve to an attribute of tool_context itself (such as + user_id) are honored. tool_context.state is not consulted: state is + writable by the caller, so a name that fell back to state could be + overridden to run the query in a different user's scope, defeating the + view's own row-level restriction. """ def __init__( @@ -71,8 +77,9 @@ def __init__( credentials_config: The credentials configuration. tool_settings: The tool settings. view_parameter_names: A list of parameter names to resolve from - tool_context and pass into view_parameters. This is configured on the - toolset (BigtableToolset) and forwarded here. + tool_context's own attributes (not tool_context.state) and pass + into view_parameters. This is configured on the toolset + (BigtableToolset) and forwarded here. """ super().__init__( func=func, @@ -101,12 +108,14 @@ async def _run_async_with_credential( if "_view_parameters" in signature.parameters and self.view_parameter_names: view_params = {} for param_name in self.view_parameter_names: - # 1. Check if it's a strongly-typed top-level property (like 'user_id') + # Only resolve from strongly-typed, framework-set attributes of + # tool_context (like user_id). tool_context.state is deliberately not + # consulted here: it is writable by the caller, and a view parameter + # taken from there would let the caller pick which user's (or + # tenant's) rows the parameterized view returns, defeating the + # view's own row-level restriction. if (val := getattr(tool_context, param_name, None)) is not None: view_params[param_name] = val - # 2. Fallback to checking application-level session state - elif tool_context.state and param_name in tool_context.state: - view_params[param_name] = tool_context.state[param_name] args_to_call["_view_parameters"] = view_params return await super()._run_async_with_credential( diff --git a/tests/unittests/tools/bigtable/test_bigtable_toolset.py b/tests/unittests/tools/bigtable/test_bigtable_toolset.py index aa5fd440be2..ccdbc849b2d 100644 --- a/tests/unittests/tools/bigtable/test_bigtable_toolset.py +++ b/tests/unittests/tools/bigtable/test_bigtable_toolset.py @@ -253,18 +253,27 @@ def mock_execute_sql(_view_parameters=None): @pytest.mark.asyncio -async def test_bigtable_parameterized_view_tool_execution_session_state_fallback(): - """Test that BigtableParameterizedViewTool falls back to tool_context.state for custom parameters.""" +async def test_bigtable_parameterized_view_tool_ignores_state_for_custom_parameters(): + """Test that BigtableParameterizedViewTool never resolves view parameters + + from tool_context.state, even when a name matching view_parameter_names is + present there. tool_context.state is writable by the caller, so a view + parameter taken from it (e.g. a caller-chosen tenant_id) would let the + caller pick whose rows a parameterized view returns, defeating the view's + own row-level restriction. See the class docstring for the security + rationale. + """ def mock_execute_sql(_view_parameters=None): return {"status": "SUCCESS", "_view_parameters": _view_parameters} - # Create session with application-level state + # A caller-controlled tenant_id sitting in session state, as if a prior + # tool call (or the model itself) had written it there. session = Session( id="session-1", app_name="test-app", user_id="user-123", - state={"tenant_id": "tenant-xyz"}, + state={"tenant_id": "attacker-chosen-tenant"}, ) invocation_context = mock.create_autospec(InvocationContext, instance=True) @@ -273,7 +282,8 @@ def mock_execute_sql(_view_parameters=None): tool_context = Context(invocation_context=invocation_context) - # Ensure 'tenant_id' is NOT a top-level property or attribute on tool_context + # Confirm the test setup: 'tenant_id' is not a top-level attribute of + # tool_context, only a key in its (caller-writable) state. assert not hasattr(tool_context, "tenant_id") assert "tenant_id" in tool_context.state @@ -290,9 +300,10 @@ def mock_execute_sql(_view_parameters=None): tool_context=tool_context, ) + # tenant_id is silently omitted rather than taken from state. assert res == { "status": "SUCCESS", - "_view_parameters": {"tenant_id": "tenant-xyz"}, + "_view_parameters": {}, } @@ -330,11 +341,12 @@ def mock_execute_sql(_view_parameters=None): tool_context=tool_context, ) + # Only user_id resolves, since it is the only one of the three that is a + # real tool_context attribute; tenant_id and agent_id are only present in + # state, which is never consulted, so they are silently omitted. assert res == { "status": "SUCCESS", "_view_parameters": { "user_id": "user-123", - "tenant_id": "tenant-xyz", - "agent_id": "agent-123", }, }