From 874c79239b711563fafee62831a99ff1fcaab1c0 Mon Sep 17 00:00:00 2001 From: prasanna8585 Date: Wed, 16 Sep 2026 10:13:07 +0530 Subject: [PATCH] fix(bigtable): stop resolving view parameters from caller-writable state BigtableParameterizedViewTool wraps Bigtable's execute_sql to automatically inject values (like user_id) into a parameterized view's VIEW_PARAMETERS(), so the view itself can restrict which rows a query is allowed to see. The class's own docstring states this "securely restricts query execution to the logged-in user's data" -- but the resolution logic had an undocumented fallback: when a configured view_parameter_name did not match a strongly-typed attribute of tool_context (which user_id always does, via ReadonlyContext), it fell back to tool_context.state[param_name]. tool_context.state is writable by the caller -- the same fact this same PR's sibling fix (BigQuery/AgentEngineSandboxComputer, in query_tool.py and sandbox_computer.py) already established for a different pair of tools in this same file tree. A view parameter taken from state lets the caller choose which user's (or tenant's) rows a "securely restricted" parameterized view returns, defeating the row- level restriction the feature exists to provide. This is realistic for any application-defined scoping parameter beyond the documented user_id example -- tenant_id, org_id, customer_id, or any other custom name a multi-tenant deployment configures via view_parameter_names, none of which are built-in ToolContext/ReadonlyContext attributes. Verified this was never previously addressed: the feature was introduced in 14a24f2b (PR #6128, "Support parameterized views with secure parameter injection") and has had no subsequent fix; the only later commit touching this file is an unrelated mypy-typing refactor. Fix: remove the tool_context.state fallback entirely. Only names that resolve to a real tool_context attribute are honored now; anything else is silently omitted from view_parameters, exactly as an already- missing attribute was previously handled. Verified: existing test suite (14 tests) run and pass, including the two pre-existing tests that only exercise the safe, strongly-typed path (user_id), which are unaffected. The test that asserted the vulnerable fallback behavior (test_bigtable_parameterized_view_tool_execution_session_state_fallback) is replaced with test_bigtable_parameterized_view_tool_ignores_state_for_custom_parameters, asserting the opposite: a caller-controlled tenant_id sitting in session state is not used. Confirmed this new test actually catches the bug by temporarily reverting the fix locally and re-running it -- it fails against the vulnerable code and passes against the fix. test_bigtable_parameterized_view_tool_execution_multiple_parameters is updated to match: tenant_id/agent_id (state-only) are now correctly omitted from the result, while user_id (a real attribute) still resolves. --- .../adk/tools/bigtable/bigtable_toolset.py | 21 ++++++++++---- .../tools/bigtable/test_bigtable_toolset.py | 28 +++++++++++++------ 2 files changed, 35 insertions(+), 14 deletions(-) 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", }, }