Skip to content

Commit 5c614e4

Browse files
petrmarinecpetrmarinec
committed
fix: constrain in-memory artifact references to caller scope
1 parent 71b936b commit 5c614e4

2 files changed

Lines changed: 229 additions & 1 deletion

File tree

‎src/google/adk/artifacts/in_memory_artifact_service.py‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,28 @@ def _artifact_path(
9494
)
9595
return f"{app_name}/{user_id}/{session_id}/{filename}"
9696

97+
def _validate_artifact_reference_scope(
98+
self,
99+
*,
100+
app_name: str,
101+
user_id: str,
102+
session_id: Optional[str],
103+
parsed_uri: artifact_util.ParsedArtifactUri,
104+
) -> None:
105+
"""Ensures artifact references cannot escape the caller's scope."""
106+
if parsed_uri.app_name != app_name or parsed_uri.user_id != user_id:
107+
raise InputValidationError(
108+
"Artifact references must stay within the same app and user scope."
109+
)
110+
if (
111+
parsed_uri.session_id is not None
112+
and parsed_uri.session_id != session_id
113+
):
114+
raise InputValidationError(
115+
"Session-scoped artifact references must stay within the same"
116+
" session scope."
117+
)
118+
97119
@override
98120
async def save_artifact(
99121
self,
@@ -128,10 +150,19 @@ async def save_artifact(
128150
artifact_version.mime_type = "text/plain"
129151
elif artifact.file_data is not None:
130152
if artifact_util.is_artifact_ref(artifact):
131-
if not artifact_util.parse_artifact_uri(artifact.file_data.file_uri):
153+
parsed_uri = artifact_util.parse_artifact_uri(
154+
artifact.file_data.file_uri
155+
)
156+
if not parsed_uri:
132157
raise InputValidationError(
133158
f"Invalid artifact reference URI: {artifact.file_data.file_uri}"
134159
)
160+
self._validate_artifact_reference_scope(
161+
app_name=app_name,
162+
user_id=user_id,
163+
session_id=session_id,
164+
parsed_uri=parsed_uri,
165+
)
135166
# If it's a valid artifact URI, we store the artifact part as-is.
136167
# And we don't know the mime type until we load it.
137168
else:
@@ -180,6 +211,12 @@ async def load_artifact(
180211
"Invalid artifact reference URI:"
181212
f" {artifact_data.file_data.file_uri}"
182213
)
214+
self._validate_artifact_reference_scope(
215+
app_name=app_name,
216+
user_id=user_id,
217+
session_id=session_id,
218+
parsed_uri=parsed_uri,
219+
)
183220
return await self.load_artifact(
184221
app_name=parsed_uri.app_name,
185222
user_id=parsed_uri.user_id,

‎tests/unittests/artifacts/test_artifact_service.py‎

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -878,6 +878,197 @@ async def test_file_save_artifact_rejects_absolute_path_within_scope(tmp_path):
878878
)
879879

880880

881+
@pytest.mark.asyncio
882+
async def test_in_memory_artifact_reference_allows_same_session_scope():
883+
artifact_service = InMemoryArtifactService()
884+
885+
await artifact_service.save_artifact(
886+
app_name="app0",
887+
user_id="user0",
888+
session_id="sess0",
889+
filename="source.txt",
890+
artifact=types.Part(text="hello"),
891+
)
892+
893+
ref = types.Part(
894+
file_data=types.FileData(
895+
file_uri=(
896+
"artifact://apps/app0/users/user0/sessions/sess0/"
897+
"artifacts/source.txt/versions/0"
898+
),
899+
mime_type="text/plain",
900+
)
901+
)
902+
await artifact_service.save_artifact(
903+
app_name="app0",
904+
user_id="user0",
905+
session_id="sess0",
906+
filename="ref.txt",
907+
artifact=ref,
908+
)
909+
910+
loaded = await artifact_service.load_artifact(
911+
app_name="app0",
912+
user_id="user0",
913+
session_id="sess0",
914+
filename="ref.txt",
915+
)
916+
assert loaded == types.Part(text="hello")
917+
918+
919+
@pytest.mark.asyncio
920+
async def test_in_memory_artifact_reference_allows_same_user_user_scope():
921+
artifact_service = InMemoryArtifactService()
922+
923+
await artifact_service.save_artifact(
924+
app_name="app0",
925+
user_id="user0",
926+
session_id="sess0",
927+
filename="user:profile.txt",
928+
artifact=types.Part(text="profile"),
929+
)
930+
931+
ref = types.Part(
932+
file_data=types.FileData(
933+
file_uri=(
934+
"artifact://apps/app0/users/user0/artifacts/"
935+
"user:profile.txt/versions/0"
936+
),
937+
mime_type="text/plain",
938+
)
939+
)
940+
await artifact_service.save_artifact(
941+
app_name="app0",
942+
user_id="user0",
943+
session_id="sess1",
944+
filename="ref.txt",
945+
artifact=ref,
946+
)
947+
948+
loaded = await artifact_service.load_artifact(
949+
app_name="app0",
950+
user_id="user0",
951+
session_id="sess1",
952+
filename="ref.txt",
953+
)
954+
assert loaded == types.Part(text="profile")
955+
956+
957+
@pytest.mark.asyncio
958+
async def test_in_memory_artifact_reference_rejects_cross_user_on_save():
959+
artifact_service = InMemoryArtifactService()
960+
961+
await artifact_service.save_artifact(
962+
app_name="app0",
963+
user_id="victim",
964+
session_id="victim-sess",
965+
filename="user:secret.txt",
966+
artifact=types.Part(text="secret"),
967+
)
968+
969+
ref = types.Part(
970+
file_data=types.FileData(
971+
file_uri=(
972+
"artifact://apps/app0/users/victim/artifacts/"
973+
"user:secret.txt/versions/0"
974+
),
975+
mime_type="text/plain",
976+
)
977+
)
978+
with pytest.raises(InputValidationError, match="same app and user scope"):
979+
await artifact_service.save_artifact(
980+
app_name="app0",
981+
user_id="attacker",
982+
session_id="attacker-sess",
983+
filename="ref.txt",
984+
artifact=ref,
985+
)
986+
987+
988+
@pytest.mark.asyncio
989+
async def test_in_memory_artifact_reference_rejects_cross_app_on_save():
990+
artifact_service = InMemoryArtifactService()
991+
992+
await artifact_service.save_artifact(
993+
app_name="victim-app",
994+
user_id="user0",
995+
session_id="sess0",
996+
filename="user:secret.txt",
997+
artifact=types.Part(text="secret"),
998+
)
999+
1000+
ref = types.Part(
1001+
file_data=types.FileData(
1002+
file_uri=(
1003+
"artifact://apps/victim-app/users/user0/artifacts/"
1004+
"user:secret.txt/versions/0"
1005+
),
1006+
mime_type="text/plain",
1007+
)
1008+
)
1009+
with pytest.raises(InputValidationError, match="same app and user scope"):
1010+
await artifact_service.save_artifact(
1011+
app_name="attacker-app",
1012+
user_id="user0",
1013+
session_id="sess0",
1014+
filename="ref.txt",
1015+
artifact=ref,
1016+
)
1017+
1018+
1019+
@pytest.mark.asyncio
1020+
async def test_in_memory_artifact_reference_rejects_cross_session_on_load():
1021+
artifact_service = InMemoryArtifactService()
1022+
1023+
await artifact_service.save_artifact(
1024+
app_name="app0",
1025+
user_id="user0",
1026+
session_id="sess0",
1027+
filename="source.txt",
1028+
artifact=types.Part(text="source"),
1029+
)
1030+
await artifact_service.save_artifact(
1031+
app_name="app0",
1032+
user_id="user0",
1033+
session_id="sess1",
1034+
filename="source.txt",
1035+
artifact=types.Part(text="other-session"),
1036+
)
1037+
1038+
ref = types.Part(
1039+
file_data=types.FileData(
1040+
file_uri=(
1041+
"artifact://apps/app0/users/user0/sessions/sess0/"
1042+
"artifacts/source.txt/versions/0"
1043+
),
1044+
mime_type="text/plain",
1045+
)
1046+
)
1047+
await artifact_service.save_artifact(
1048+
app_name="app0",
1049+
user_id="user0",
1050+
session_id="sess0",
1051+
filename="ref.txt",
1052+
artifact=ref,
1053+
)
1054+
1055+
ref_path = artifact_service._artifact_path(
1056+
"app0", "user0", "ref.txt", "sess0"
1057+
)
1058+
artifact_service.artifacts[ref_path][0].data.file_data.file_uri = (
1059+
"artifact://apps/app0/users/user0/sessions/sess1/"
1060+
"artifacts/source.txt/versions/0"
1061+
)
1062+
1063+
with pytest.raises(InputValidationError, match="same session scope"):
1064+
await artifact_service.load_artifact(
1065+
app_name="app0",
1066+
user_id="user0",
1067+
session_id="sess0",
1068+
filename="ref.txt",
1069+
)
1070+
1071+
8811072
class TestEnsurePart:
8821073
"""Tests for the ensure_part normalization helper."""
8831074

0 commit comments

Comments
 (0)