From f8631500c7e46b3ce417fdc47a483cfcf4032c7a Mon Sep 17 00:00:00 2001 From: petrmarinec <54589756+petrmarinec@users.noreply.github.com> Date: Mon, 6 Jul 2026 08:48:27 -0700 Subject: [PATCH] fix: Constrain artifact references to the caller scope InMemoryArtifactService and GcsArtifactService currently accept caller-supplied artifact://... references and later dereference the embedded scope during load_artifact() without checking that the embedded app_name, user_id, and session_id still match the caller. This patch keeps same-scope references working while blocking references that escape the caller scope. - Validate parsed artifact://... references against the caller app/user/session scope before storing them. - Validate the scope again before dereferencing an artifact reference on load. - Allow same-session references. - Allow same-user user: references across sessions. Closes #6124 Closes #6125 Co-authored-by: Jason Zhang PiperOrigin-RevId: 943314213 --- src/google/adk/artifacts/artifact_util.py | 21 ++ .../adk/artifacts/gcs_artifact_service.py | 15 +- .../artifacts/in_memory_artifact_service.py | 17 +- .../artifacts/test_artifact_service.py | 250 ++++++++++++++++++ tests/unittests/cli/test_fast_api.py | 52 +++- 5 files changed, 352 insertions(+), 3 deletions(-) diff --git a/src/google/adk/artifacts/artifact_util.py b/src/google/adk/artifacts/artifact_util.py index c0e4d418dcd..bd70052f879 100644 --- a/src/google/adk/artifacts/artifact_util.py +++ b/src/google/adk/artifacts/artifact_util.py @@ -20,6 +20,8 @@ from google.genai import types +from ..errors import input_validation_error + class ParsedArtifactUri(NamedTuple): """The result of parsing an artifact URI.""" @@ -113,3 +115,22 @@ def is_artifact_ref(artifact: types.Part) -> bool: and artifact.file_data.file_uri and artifact.file_data.file_uri.startswith("artifact://") ) + + +def validate_artifact_reference_scope( + *, + app_name: str, + user_id: str, + session_id: str | None, + parsed_uri: ParsedArtifactUri, +) -> None: + """Ensures artifact references cannot escape the caller's scope.""" + if parsed_uri.app_name != app_name or parsed_uri.user_id != user_id: + raise input_validation_error.InputValidationError( + "Artifact references must stay within the same app and user scope." + ) + if parsed_uri.session_id is not None and parsed_uri.session_id != session_id: + raise input_validation_error.InputValidationError( + "Session-scoped artifact references must stay within the same" + " session scope." + ) diff --git a/src/google/adk/artifacts/gcs_artifact_service.py b/src/google/adk/artifacts/gcs_artifact_service.py index aaeebe4f3e9..2db8fa8ef8f 100644 --- a/src/google/adk/artifacts/gcs_artifact_service.py +++ b/src/google/adk/artifacts/gcs_artifact_service.py @@ -252,10 +252,17 @@ def _save_artifact( if not file_uri: raise InputValidationError("Artifact file_data must have a file_uri.") if artifact_util.is_artifact_ref(artifact): - if not artifact_util.parse_artifact_uri(file_uri): + parsed_uri = artifact_util.parse_artifact_uri(file_uri) + if not parsed_uri: raise InputValidationError( f"Invalid artifact reference URI: {file_uri}" ) + artifact_util.validate_artifact_reference_scope( + app_name=app_name, + user_id=user_id, + session_id=session_id, + parsed_uri=parsed_uri, + ) # Store the URI and mime_type (if any) as blob metadata; no content to upload. metadata = { **(blob.metadata or {}), @@ -315,6 +322,12 @@ def _load_artifact( raise InputValidationError( f"Invalid artifact reference URI: {file_uri}" ) + artifact_util.validate_artifact_reference_scope( + app_name=app_name, + user_id=user_id, + session_id=session_id, + parsed_uri=parsed_uri, + ) return self._load_artifact( app_name=parsed_uri.app_name, user_id=parsed_uri.user_id, diff --git a/src/google/adk/artifacts/in_memory_artifact_service.py b/src/google/adk/artifacts/in_memory_artifact_service.py index 48e7afca9ac..2782333f976 100644 --- a/src/google/adk/artifacts/in_memory_artifact_service.py +++ b/src/google/adk/artifacts/in_memory_artifact_service.py @@ -128,10 +128,19 @@ async def save_artifact( artifact_version.mime_type = "text/plain" elif artifact.file_data is not None: if artifact_util.is_artifact_ref(artifact): - if not artifact_util.parse_artifact_uri(artifact.file_data.file_uri): + parsed_uri = artifact_util.parse_artifact_uri( + artifact.file_data.file_uri + ) + if not parsed_uri: raise InputValidationError( f"Invalid artifact reference URI: {artifact.file_data.file_uri}" ) + artifact_util.validate_artifact_reference_scope( + app_name=app_name, + user_id=user_id, + session_id=session_id, + parsed_uri=parsed_uri, + ) # If it's a valid artifact URI, we store the artifact part as-is. # And we don't know the mime type until we load it. else: @@ -180,6 +189,12 @@ async def load_artifact( "Invalid artifact reference URI:" f" {artifact_data.file_data.file_uri}" ) + artifact_util.validate_artifact_reference_scope( + app_name=app_name, + user_id=user_id, + session_id=session_id, + parsed_uri=parsed_uri, + ) return await self.load_artifact( app_name=parsed_uri.app_name, user_id=parsed_uri.user_id, diff --git a/tests/unittests/artifacts/test_artifact_service.py b/tests/unittests/artifacts/test_artifact_service.py index a9802e66945..8dfa29aeab1 100644 --- a/tests/unittests/artifacts/test_artifact_service.py +++ b/tests/unittests/artifacts/test_artifact_service.py @@ -905,6 +905,256 @@ async def test_file_save_artifact_rejects_absolute_path_within_scope(tmp_path): ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ], +) +async def test_artifact_reference_allows_same_session_scope( + service_type, artifact_service_factory +): + """ArtifactService allows references inside the same session scope.""" + artifact_service = artifact_service_factory(service_type) + + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="sess0", + filename="source.txt", + artifact=types.Part(text="hello"), + ) + + ref = types.Part( + file_data=types.FileData( + file_uri=( + "artifact://apps/app0/users/user0/sessions/sess0/" + "artifacts/source.txt/versions/0" + ), + mime_type="text/plain", + ) + ) + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="sess0", + filename="ref.txt", + artifact=ref, + ) + + loaded = await artifact_service.load_artifact( + app_name="app0", + user_id="user0", + session_id="sess0", + filename="ref.txt", + ) + assert loaded == types.Part(text="hello") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ], +) +async def test_artifact_reference_allows_same_user_user_scope( + service_type, artifact_service_factory +): + """ArtifactService allows references to user-scoped files from same user.""" + artifact_service = artifact_service_factory(service_type) + + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="sess0", + filename="user:profile.txt", + artifact=types.Part(text="profile"), + ) + + ref = types.Part( + file_data=types.FileData( + file_uri=( + "artifact://apps/app0/users/user0/artifacts/" + "user:profile.txt/versions/0" + ), + mime_type="text/plain", + ) + ) + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="sess1", + filename="ref.txt", + artifact=ref, + ) + + loaded = await artifact_service.load_artifact( + app_name="app0", + user_id="user0", + session_id="sess1", + filename="ref.txt", + ) + assert loaded == types.Part(text="profile") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ], +) +async def test_artifact_reference_rejects_cross_user_on_save( + service_type, artifact_service_factory +): + """ArtifactService rejects references to different users on save.""" + artifact_service = artifact_service_factory(service_type) + + await artifact_service.save_artifact( + app_name="app0", + user_id="victim", + session_id="victim-sess", + filename="user:secret.txt", + artifact=types.Part(text="secret"), + ) + + ref = types.Part( + file_data=types.FileData( + file_uri=( + "artifact://apps/app0/users/victim/artifacts/" + "user:secret.txt/versions/0" + ), + mime_type="text/plain", + ) + ) + with pytest.raises(InputValidationError, match="same app and user scope"): + await artifact_service.save_artifact( + app_name="app0", + user_id="attacker", + session_id="attacker-sess", + filename="ref.txt", + artifact=ref, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ], +) +async def test_artifact_reference_rejects_cross_app_on_save( + service_type, artifact_service_factory +): + """ArtifactService rejects references to different apps on save.""" + artifact_service = artifact_service_factory(service_type) + + await artifact_service.save_artifact( + app_name="victim-app", + user_id="user0", + session_id="sess0", + filename="user:secret.txt", + artifact=types.Part(text="secret"), + ) + + ref = types.Part( + file_data=types.FileData( + file_uri=( + "artifact://apps/victim-app/users/user0/artifacts/" + "user:secret.txt/versions/0" + ), + mime_type="text/plain", + ) + ) + with pytest.raises(InputValidationError, match="same app and user scope"): + await artifact_service.save_artifact( + app_name="attacker-app", + user_id="user0", + session_id="sess0", + filename="ref.txt", + artifact=ref, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "service_type", + [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, + ], +) +async def test_artifact_reference_rejects_cross_session_on_load( + service_type, artifact_service_factory +): + """ArtifactService rejects modified references to different sessions on load.""" + artifact_service = artifact_service_factory(service_type) + + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="sess0", + filename="source.txt", + artifact=types.Part(text="source"), + ) + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="sess1", + filename="source.txt", + artifact=types.Part(text="other-session"), + ) + + ref = types.Part( + file_data=types.FileData( + file_uri=( + "artifact://apps/app0/users/user0/sessions/sess0/" + "artifacts/source.txt/versions/0" + ), + mime_type="text/plain", + ) + ) + await artifact_service.save_artifact( + app_name="app0", + user_id="user0", + session_id="sess0", + filename="ref.txt", + artifact=ref, + ) + + new_uri = ( + "artifact://apps/app0/users/user0/sessions/sess1/" + "artifacts/source.txt/versions/0" + ) + # Manually modify the stored reference URI to point to a different session. + if service_type == ArtifactServiceType.GCS: + blob_name = artifact_service._get_blob_name( + "app0", "user0", "ref.txt", 0, "sess0" + ) + blob = artifact_service.bucket.get_blob(blob_name) + blob.metadata["adkFileUri"] = new_uri + elif service_type == ArtifactServiceType.IN_MEMORY: + ref_path = artifact_service._artifact_path( + "app0", "user0", "ref.txt", "sess0" + ) + artifact_service.artifacts[ref_path][0].data.file_data.file_uri = new_uri + + with pytest.raises(InputValidationError, match="same session scope"): + await artifact_service.load_artifact( + app_name="app0", + user_id="user0", + session_id="sess0", + filename="ref.txt", + ) + + class TestEnsurePart: """Tests for the ensure_part normalization helper.""" diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index 8873af3bb06..d3cdd1ea791 100755 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -1746,6 +1746,53 @@ def test_save_artifact(test_app, create_test_session, mock_artifact_service): assert stored["artifact"].text == "hello world" +def test_save_artifact_reference( + test_app, create_test_session, mock_artifact_service +): + """Test saving an artifact reference through the FastAPI endpoint.""" + info = create_test_session + url = ( + f"/apps/{info['app_name']}/users/{info['user_id']}/sessions/" + f"{info['session_id']}/artifacts" + ) + payload = { + "filename": "reference.txt", + "artifact": { + "fileData": { + "fileUri": ( + f"artifact://apps/{info['app_name']}/users/{info['user_id']}/" + f"sessions/{info['session_id']}/artifacts/target_file/versions/0" + ), + "mimeType": "text/plain", + } + }, + } + + response = test_app.post(url, json=payload) + assert response.status_code == 200 + data = response.json() + assert data["version"] == 0 + assert data["customMetadata"] == {} + assert data["mimeType"] in (None, "text/plain") + assert data["canonicalUri"].endswith( + f"/sessions/{info['session_id']}/artifacts/" + f"{payload['filename']}/versions/0" + ) + assert isinstance(data["createTime"], float) + + key = ( + f"{info['app_name']}:{info['user_id']}:{info['session_id']}:" + f"{payload['filename']}" + ) + stored = mock_artifact_service._artifacts[key][0] + assert stored["artifact"].file_data is not None + assert ( + stored["artifact"].file_data.file_uri + == payload["artifact"]["fileData"]["fileUri"] + ) + assert stored["artifact"].file_data.mime_type == "text/plain" + + def test_artifact_endpoints_support_nested_names( test_app, create_test_session, mock_artifact_service ): @@ -2579,7 +2626,10 @@ def test_builder_save_rejects_nested_args_key(builder_test_client, tmp_path): def test_builder_get_rejects_non_yaml_file_paths(builder_test_client, tmp_path): - """GET /dev/apps/{app_name}/builder?file_path=... rejects non-YAML extensions.""" + """GET /dev/apps/{app_name}/builder?file_path=... + + rejects non-YAML extensions. + """ app_root = tmp_path / "app" app_root.mkdir(parents=True, exist_ok=True) (app_root / ".env").write_text("SECRET=supersecret\n")