[pull] main from google:main

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."
+    )

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,

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,

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."""