diff --git a/src/google/adk/artifacts/artifact_util.py b/src/google/adk/artifacts/artifact_util.py index 7eea062a99e..dfbadbb4d0a 100644 --- a/src/google/adk/artifacts/artifact_util.py +++ b/src/google/adk/artifacts/artifact_util.py @@ -21,6 +21,8 @@ from google.genai import types +from ..errors import input_validation_error + class ParsedArtifactUri(NamedTuple): """The result of parsing an artifact URI.""" @@ -32,6 +34,8 @@ class ParsedArtifactUri(NamedTuple): version: int +_WINDOWS_DRIVE_RE = re.compile(r"[A-Za-z]:") + _SESSION_SCOPED_ARTIFACT_URI_RE = re.compile( r"artifact://apps/([^/]+)/users/([^/]+)/sessions/([^/]+)/artifacts/([^/]+)/versions/(\d+)" ) @@ -114,3 +118,63 @@ 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: Optional[str], + 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." + ) + + +def _is_drive_qualified(value: str) -> bool: + """Checks whether a value starts with a Windows drive letter such as ``C:``.""" + return _WINDOWS_DRIVE_RE.match(value) is not None + + +def validate_path_segment(value: str, field_name: str) -> None: + """Rejects values that could alter the constructed path. + + Args: + value: The caller-supplied identifier (e.g. user_id or session_id). + field_name: Human-readable name used in the error message. + + Raises: + InputValidationError: If the value contains traversal segments, null bytes, + is an absolute path / starts with a slash, or is drive-qualified. + """ + if not value: + raise input_validation_error.InputValidationError( + f"{field_name} must not be empty." + ) + if "\x00" in value: + raise input_validation_error.InputValidationError( + f"{field_name} must not contain null bytes." + ) + if isinstance(value, str) and ( + value.startswith("/") or value.startswith("\\") + ): + raise input_validation_error.InputValidationError( + f"{field_name} {value!r} must not be an absolute path or start with a" + " slash." + ) + if isinstance(value, str) and _is_drive_qualified(value): + raise input_validation_error.InputValidationError( + f"{field_name} {value!r} must not be drive-qualified." + ) + if value in (".", "..") or ".." in value.replace("\\", "/").split("/"): + raise input_validation_error.InputValidationError( + f"{field_name} {value!r} must not contain traversal segments." + ) diff --git a/src/google/adk/artifacts/file_artifact_service.py b/src/google/adk/artifacts/file_artifact_service.py index 9c3870b6e3c..e4a07be7e54 100644 --- a/src/google/adk/artifacts/file_artifact_service.py +++ b/src/google/adk/artifacts/file_artifact_service.py @@ -33,6 +33,7 @@ from pydantic import ValidationError from typing_extensions import override +from . import artifact_util from ..errors.input_validation_error import InputValidationError from .base_artifact_service import ArtifactVersion from .base_artifact_service import BaseArtifactService @@ -85,14 +86,32 @@ def _to_posix_path(path_value: str) -> PurePosixPath: return PurePosixPath(path_value) +def _is_rooted_or_drive_qualified(path_value: str) -> bool: + """Checks POSIX and Windows rooted or drive-qualified path forms.""" + # A Windows root covers POSIX absolute paths, UNC and device prefixes alike; + # only the drive-relative form (`C:name`) has no root of its own. + if artifact_util._is_drive_qualified(path_value): + return True + return bool(PureWindowsPath(path_value).root) + + +def _has_parent_reference(path_value: str) -> bool: + """Checks parent traversal using either platform's separators.""" + return ( + ".." in PurePosixPath(path_value).parts + or ".." in PureWindowsPath(path_value).parts + ) + + def _resolve_scoped_artifact_path( scope_root: Path, filename: str ) -> tuple[Path, Path]: """Returns the absolute artifact directory and its relative path. The caller is expected to pass the scope root directory (user or session). - This helper joins the filename under that root, resolves traversal segments, - and guards against paths that escape the scope root. + Filenames that are rooted, drive-qualified, or contain a parent reference are + rejected outright, including parent references that would resolve back inside + the scope root. Whatever remains is joined under the scope root. Args: scope_root: Directory that defines the storage scope. @@ -103,17 +122,23 @@ def _resolve_scoped_artifact_path( to `scope_root`. Raises: - InputValidationError: If `filename` resolves outside of `scope_root`. + InputValidationError: If `filename` is rooted, drive-qualified, contains a + parent reference, or otherwise resolves outside of `scope_root`. """ stripped = _strip_user_namespace(filename).strip() - pure_path = _to_posix_path(stripped) - scope_root_resolved = scope_root.resolve(strict=False) - if pure_path.is_absolute(): + if _is_rooted_or_drive_qualified(stripped): raise InputValidationError( - f"Absolute artifact filename {filename!r} is not permitted; " - "provide a path relative to the storage scope." + f"Rooted or drive-qualified artifact filename {filename!r} is not " + "permitted; provide a path relative to the storage scope." ) + if _has_parent_reference(stripped): + raise InputValidationError( + f"Artifact filename {filename!r} must not contain parent traversal." + ) + + scope_root_resolved = scope_root.resolve(strict=False) + pure_path = _to_posix_path(stripped) candidate = scope_root_resolved / Path(pure_path) candidate = candidate.resolve(strict=False) @@ -147,7 +172,7 @@ def _validate_path_segment(value: str, field_name: str) -> None: Raises: InputValidationError: If the value contains path separators, traversal - segments, or null bytes. + segments, null bytes, or is drive-qualified. """ if not value: raise InputValidationError(f"{field_name} must not be empty.") @@ -157,6 +182,10 @@ def _validate_path_segment(value: str, field_name: str) -> None: raise InputValidationError( f"{field_name} {value!r} must not contain path separators." ) + if artifact_util._is_drive_qualified(value): + raise InputValidationError( + f"{field_name} {value!r} must not be drive-qualified." + ) if value in (".", "..") or ".." in value.split("/"): raise InputValidationError( f"{field_name} {value!r} must not contain traversal segments." @@ -259,7 +288,7 @@ def _scope_root( base = self._base_root(user_id) if _is_user_scoped(session_id, filename): return _user_artifacts_dir(base) - if not session_id: + if session_id is None: raise InputValidationError( "Session ID must be provided for session-scoped artifacts." ) @@ -524,7 +553,7 @@ def _list_artifact_keys_sync( base_root = self._base_root(user_id) - if session_id: + if session_id is not None: session_root = _session_artifacts_dir(base_root, session_id) for artifact_dir in _iter_artifact_dirs(session_root): metadata = self._latest_metadata(artifact_dir) diff --git a/src/google/adk/artifacts/gcs_artifact_service.py b/src/google/adk/artifacts/gcs_artifact_service.py index f8706dedbd2..ee4a6b0fa3e 100644 --- a/src/google/adk/artifacts/gcs_artifact_service.py +++ b/src/google/adk/artifacts/gcs_artifact_service.py @@ -32,6 +32,7 @@ from google.genai import types from typing_extensions import override +from . import artifact_util from ..errors.input_validation_error import InputValidationError from .base_artifact_service import ArtifactVersion from .base_artifact_service import BaseArtifactService @@ -161,6 +162,8 @@ def _get_blob_prefix( session_id: Optional[str] = None, ) -> str: """Constructs the blob name prefix in GCS for a given artifact.""" + artifact_util.validate_path_segment(app_name, "app_name") + artifact_util.validate_path_segment(user_id, "user_id") if self._file_has_user_namespace(filename): return f"{app_name}/{user_id}/user/{filename}" @@ -168,6 +171,7 @@ def _get_blob_prefix( raise InputValidationError( "Session ID must be provided for session-scoped artifacts." ) + artifact_util.validate_path_segment(session_id, "session_id") return f"{app_name}/{user_id}/{session_id}/{filename}" def _get_blob_name( @@ -276,6 +280,10 @@ def _load_artifact( def _list_artifact_keys( self, app_name: str, user_id: str, session_id: Optional[str] ) -> list[str]: + artifact_util.validate_path_segment(app_name, "app_name") + artifact_util.validate_path_segment(user_id, "user_id") + if session_id is not None: + artifact_util.validate_path_segment(session_id, "session_id") filenames = set() if session_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..f1ddc9e564e 100644 --- a/src/google/adk/artifacts/in_memory_artifact_service.py +++ b/src/google/adk/artifacts/in_memory_artifact_service.py @@ -85,6 +85,8 @@ def _artifact_path( Returns: The constructed artifact path. """ + artifact_util.validate_path_segment(app_name, "app_name") + artifact_util.validate_path_segment(user_id, "user_id") if self._file_has_user_namespace(filename): return f"{app_name}/{user_id}/user/{filename}" @@ -92,6 +94,7 @@ def _artifact_path( raise InputValidationError( "Session ID must be provided for session-scoped artifacts." ) + artifact_util.validate_path_segment(session_id, "session_id") return f"{app_name}/{user_id}/{session_id}/{filename}" @override @@ -128,10 +131,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 +192,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, @@ -200,6 +218,10 @@ async def load_artifact( async def list_artifact_keys( self, *, app_name: str, user_id: str, session_id: Optional[str] = None ) -> list[str]: + artifact_util.validate_path_segment(app_name, "app_name") + artifact_util.validate_path_segment(user_id, "user_id") + if session_id is not None: + artifact_util.validate_path_segment(session_id, "session_id") usernamespace_prefix = f"{app_name}/{user_id}/user/" session_prefix = ( f"{app_name}/{user_id}/{session_id}/" if session_id else None diff --git a/src/google/adk/errors/input_validation_error.py b/src/google/adk/errors/input_validation_error.py index 0bcd55f8471..080114c40a9 100644 --- a/src/google/adk/errors/input_validation_error.py +++ b/src/google/adk/errors/input_validation_error.py @@ -18,7 +18,7 @@ class InputValidationError(ValueError): """Represents an error raised when user input fails validation.""" - def __init__(self, message="Invalid input."): + def __init__(self, message: str = "Invalid input.") -> None: """Initializes the InputValidationError exception. Args: diff --git a/tests/unittests/artifacts/test_artifact_service.py b/tests/unittests/artifacts/test_artifact_service.py index 8b82397097a..379d2ab8eaf 100644 --- a/tests/unittests/artifacts/test_artifact_service.py +++ b/tests/unittests/artifacts/test_artifact_service.py @@ -723,9 +723,21 @@ async def test_file_list_artifact_versions(tmp_path, artifact_service_factory): ("filename", "session_id"), [ ("../escape.txt", "sess123"), + (r"..\escape.txt", "sess123"), + ("folder/../alias.txt", "sess123"), + (r"folder\..\alias.txt", "sess123"), + (r"folder/..\alias.txt", "sess123"), ("user:../escape.txt", "sess123"), + (r"user:..\escape.txt", "sess123"), + (r"user:folder\..\alias.txt", "sess123"), ("/absolute/path.txt", "sess123"), ("user:/absolute/path.txt", None), + (r"C:\absolute\path.txt", "sess123"), + ("C:/absolute/path.txt", "sess123"), + ("C:drive-relative.txt", "sess123"), + (r"\\server\share\file.txt", "sess123"), + ("//server/share/file.txt", "sess123"), + (r"\rooted\file.txt", "sess123"), ], ) async def test_file_save_artifact_rejects_out_of_scope_paths( @@ -756,6 +768,8 @@ async def test_file_save_artifact_rejects_out_of_scope_paths( ".", "has/slash", "back\\slash", + "C:evil", + "C:", "null\x00byte", "", ], @@ -787,6 +801,8 @@ async def test_file_save_artifact_rejects_traversal_in_user_id( ".", "has/slash", "back\\slash", + "C:evil", + "C:", "null\x00byte", "", ], @@ -832,6 +848,483 @@ async def test_file_save_artifact_rejects_absolute_path_within_scope(tmp_path): ) +@pytest.mark.asyncio +async def test_file_empty_session_id_is_rejected(tmp_path): + """An empty session_id is an error rather than meaning "no session".""" + artifact_service = FileArtifactService(root_dir=tmp_path / "artifacts") + with pytest.raises(InputValidationError, match="must not be empty"): + await artifact_service.save_artifact( + app_name="myapp", + user_id="user123", + session_id="", + filename="safe.txt", + artifact=types.Part(text="content"), + ) + with pytest.raises(InputValidationError, match="must not be empty"): + await artifact_service.list_artifact_keys( + app_name="myapp", + user_id="user123", + session_id="", + ) + + +INVALID_PATH_SEGMENT_CASES = ( + ("../escape", "must not contain traversal segments"), + ("../../etc", "must not contain traversal segments"), + ("foo/../../bar", "must not contain traversal segments"), + ("..", "must not contain traversal segments"), + (".", "must not contain traversal segments"), + ("null\x00byte", "must not contain null bytes"), + ("", "must not be empty"), + ("/etc/passwd", "must not be an absolute path or start with a slash"), + ("/leading/slash", "must not be an absolute path or start with a slash"), + ( + "\\leading\\backslash", + "must not be an absolute path or start with a slash", + ), + (r"C:\absolute", "must not be drive-qualified"), + ("C:/absolute", "must not be drive-qualified"), + ("C:drive-relative", "must not be drive-qualified"), +) + +# FileArtifactService keeps its own stricter validator, which additionally +# rejects embedded path separators, so it is not covered by these cases. +SHARED_VALIDATOR_SERVICE_TYPES = [ + ArtifactServiceType.IN_MEMORY, + ArtifactServiceType.GCS, +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +async def test_save_and_load_namespaced_user_id_succeeds( + service_type, artifact_service_factory +): + """The in-memory and GCS services permit namespaced user IDs.""" + service = artifact_service_factory(service_type) + artifact = types.Part.from_bytes(data=b"data", mime_type="text/plain") + await service.save_artifact( + app_name="myapp", + user_id="group/user123", + session_id="sess123", + filename="safe.txt", + artifact=artifact, + ) + loaded = await service.load_artifact( + app_name="myapp", + user_id="group/user123", + session_id="sess123", + filename="safe.txt", + ) + assert loaded.inline_data.data == b"data" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("app_name,match", INVALID_PATH_SEGMENT_CASES) +async def test_save_artifact_rejects_traversal_in_app_name( + service_type, app_name, match, artifact_service_factory +): + """Saving rejects app_name values that could alter the storage key.""" + service = artifact_service_factory(service_type) + artifact = types.Part.from_bytes(data=b"data", mime_type="text/plain") + with pytest.raises(InputValidationError, match=match): + await service.save_artifact( + app_name=app_name, + user_id="user123", + filename="user:safe.txt", + artifact=artifact, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("user_id,match", INVALID_PATH_SEGMENT_CASES) +async def test_save_artifact_rejects_traversal_in_user_id( + service_type, user_id, match, artifact_service_factory +): + """Saving rejects user_id values that could alter the storage key.""" + service = artifact_service_factory(service_type) + artifact = types.Part.from_bytes(data=b"data", mime_type="text/plain") + with pytest.raises(InputValidationError, match=match): + await service.save_artifact( + app_name="myapp", + user_id=user_id, + filename="user:safe.txt", + artifact=artifact, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("session_id,match", INVALID_PATH_SEGMENT_CASES) +async def test_save_artifact_rejects_traversal_in_session_id( + service_type, session_id, match, artifact_service_factory +): + """Saving rejects session_id values that could alter the storage key.""" + service = artifact_service_factory(service_type) + artifact = types.Part.from_bytes(data=b"data", mime_type="text/plain") + with pytest.raises(InputValidationError, match=match): + await service.save_artifact( + app_name="myapp", + user_id="user123", + session_id=session_id, + filename="safe.txt", + artifact=artifact, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("app_name,match", INVALID_PATH_SEGMENT_CASES) +async def test_load_artifact_rejects_traversal_in_app_name( + service_type, app_name, match, artifact_service_factory +): + """Loading rejects app_name values that could alter the storage key.""" + service = artifact_service_factory(service_type) + with pytest.raises(InputValidationError, match=match): + await service.load_artifact( + app_name=app_name, + user_id="user123", + filename="user:safe.txt", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("user_id,match", INVALID_PATH_SEGMENT_CASES) +async def test_load_artifact_rejects_traversal_in_user_id( + service_type, user_id, match, artifact_service_factory +): + """Loading rejects user_id values that could alter the storage key.""" + service = artifact_service_factory(service_type) + with pytest.raises(InputValidationError, match=match): + await service.load_artifact( + app_name="myapp", + user_id=user_id, + filename="user:safe.txt", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("session_id,match", INVALID_PATH_SEGMENT_CASES) +async def test_load_artifact_rejects_traversal_in_session_id( + service_type, session_id, match, artifact_service_factory +): + """Loading rejects session_id values that could alter the storage key.""" + service = artifact_service_factory(service_type) + with pytest.raises(InputValidationError, match=match): + await service.load_artifact( + app_name="myapp", + user_id="user123", + session_id=session_id, + filename="safe.txt", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("app_name,match", INVALID_PATH_SEGMENT_CASES) +async def test_delete_artifact_rejects_traversal_in_app_name( + service_type, app_name, match, artifact_service_factory +): + """Deleting rejects app_name values that could alter the storage key.""" + service = artifact_service_factory(service_type) + with pytest.raises(InputValidationError, match=match): + await service.delete_artifact( + app_name=app_name, + user_id="user123", + filename="user:safe.txt", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("user_id,match", INVALID_PATH_SEGMENT_CASES) +async def test_delete_artifact_rejects_traversal_in_user_id( + service_type, user_id, match, artifact_service_factory +): + """Deleting rejects user_id values that could alter the storage key.""" + service = artifact_service_factory(service_type) + with pytest.raises(InputValidationError, match=match): + await service.delete_artifact( + app_name="myapp", + user_id=user_id, + filename="user:safe.txt", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("session_id,match", INVALID_PATH_SEGMENT_CASES) +async def test_delete_artifact_rejects_traversal_in_session_id( + service_type, session_id, match, artifact_service_factory +): + """Deleting rejects session_id values that could alter the storage key.""" + service = artifact_service_factory(service_type) + with pytest.raises(InputValidationError, match=match): + await service.delete_artifact( + app_name="myapp", + user_id="user123", + session_id=session_id, + filename="safe.txt", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("app_name,match", INVALID_PATH_SEGMENT_CASES) +async def test_list_artifact_keys_rejects_traversal_in_app_name( + service_type, app_name, match, artifact_service_factory +): + """Listing rejects app_name values that could alter the storage key.""" + service = artifact_service_factory(service_type) + with pytest.raises(InputValidationError, match=match): + await service.list_artifact_keys( + app_name=app_name, + user_id="user123", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("user_id,match", INVALID_PATH_SEGMENT_CASES) +async def test_list_artifact_keys_rejects_traversal_in_user_id( + service_type, user_id, match, artifact_service_factory +): + """Listing rejects user_id values that could alter the storage key.""" + service = artifact_service_factory(service_type) + with pytest.raises(InputValidationError, match=match): + await service.list_artifact_keys( + app_name="myapp", + user_id=user_id, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_type", SHARED_VALIDATOR_SERVICE_TYPES) +@pytest.mark.parametrize("session_id,match", INVALID_PATH_SEGMENT_CASES) +async def test_list_artifact_keys_rejects_traversal_in_session_id( + service_type, session_id, match, artifact_service_factory +): + """Listing rejects session_id values that could alter the storage key.""" + service = artifact_service_factory(service_type) + with pytest.raises(InputValidationError, match=match): + await service.list_artifact_keys( + app_name="myapp", + user_id="user123", + session_id=session_id, + ) + + +@pytest.mark.asyncio +async def test_artifact_reference_allows_same_session_scope( + artifact_service_factory, +): + """InMemoryArtifactService allows references inside the same session scope.""" + artifact_service = artifact_service_factory(ArtifactServiceType.IN_MEMORY) + + 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 +async def test_artifact_reference_allows_same_user_user_scope( + artifact_service_factory, +): + """InMemoryArtifactService allows user-scoped references from same user.""" + artifact_service = artifact_service_factory(ArtifactServiceType.IN_MEMORY) + + 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 +async def test_artifact_reference_rejects_cross_user_on_save( + artifact_service_factory, +): + """InMemoryArtifactService rejects references to different users on save.""" + artifact_service = artifact_service_factory(ArtifactServiceType.IN_MEMORY) + + 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 +async def test_artifact_reference_rejects_cross_app_on_save( + artifact_service_factory, +): + """InMemoryArtifactService rejects references to different apps on save.""" + artifact_service = artifact_service_factory(ArtifactServiceType.IN_MEMORY) + + 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 +async def test_artifact_reference_rejects_cross_session_on_load( + artifact_service_factory, +): + """A stored reference retargeted at another session is rejected on load.""" + artifact_service = artifact_service_factory(ArtifactServiceType.IN_MEMORY) + + 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, + ) + + # Manually retarget the stored reference URI at a different session. + ref_path = artifact_service._artifact_path( + "app0", "user0", "ref.txt", "sess0" + ) + artifact_service.artifacts[ref_path][0].data.file_data.file_uri = ( + "artifact://apps/app0/users/user0/sessions/sess1/" + "artifacts/source.txt/versions/0" + ) + + 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/artifacts/test_artifact_util.py b/tests/unittests/artifacts/test_artifact_util.py index 1c4f411f14f..a2931b04888 100644 --- a/tests/unittests/artifacts/test_artifact_util.py +++ b/tests/unittests/artifacts/test_artifact_util.py @@ -14,7 +14,10 @@ """Tests for artifact_util.""" +from unittest import mock + from google.adk.artifacts import artifact_util +from google.adk.errors.input_validation_error import InputValidationError from google.genai import types import pytest @@ -107,3 +110,178 @@ def test_is_artifact_ref_true(): def test_is_artifact_ref_false(part): """Tests is_artifact_ref with non-reference parts.""" assert artifact_util.is_artifact_ref(part) is False + + +@pytest.mark.parametrize( + "field_name", + ["user_id", "app_name", "session_id"], +) +@pytest.mark.parametrize( + "value", + [ + "user123", + "myapp", + "sess123", + "group/user123", + "has/slash", + "back\\slash", + mock.MagicMock(), + ], +) +def test_validate_path_segment_valid(value, field_name): + """Normal and namespaced segments should pass validation.""" + artifact_util.validate_path_segment(value, field_name) + + +@pytest.mark.parametrize( + "field_name", + ["user_id", "app_name", "session_id"], +) +@pytest.mark.parametrize( + "value", + [ + "../escape", + "../../etc", + "foo/../../bar", + "mixed/..\\separators", + "./..\\", + ".\\../", + "..", + ".", + "null\x00byte", + "", + "/etc/passwd", + "/leading/slash", + "\\leading\\backslash", + "C:\\absolute", + "C:/absolute", + "C:drive-relative", + ], +) +def test_validate_path_segment_invalid(value, field_name): + """Traversal segments, null bytes, and absolute paths should raise InputValidationError.""" + with pytest.raises(InputValidationError): + artifact_util.validate_path_segment(value, field_name) + + +@pytest.mark.parametrize( + "caller_session_id, uri_session_id", + [ + # Session-scoped reference read from the session that owns it. + ("session1", "session1"), + # User-scoped reference (no session in the URI) is readable from any + # session of the same user, including outside of a session. + ("session1", None), + (None, None), + ], +) +def test_validate_artifact_reference_scope_within_scope_is_allowed( + caller_session_id, uri_session_id +): + """References that stay inside the caller's app/user/session scope pass.""" + parsed = artifact_util.ParsedArtifactUri( + app_name="app1", + user_id="user1", + session_id=uri_session_id, + filename="file1", + version=1, + ) + + artifact_util.validate_artifact_reference_scope( + app_name="app1", + user_id="user1", + session_id=caller_session_id, + parsed_uri=parsed, + ) + + +@pytest.mark.parametrize( + "uri_app_name, uri_user_id", + [ + ("other_app", "user1"), + ("app1", "other_user"), + ("other_app", "other_user"), + ], +) +def test_validate_artifact_reference_scope_other_app_or_user_raises( + uri_app_name, uri_user_id +): + """A reference owned by another app or user must be rejected.""" + parsed = artifact_util.ParsedArtifactUri( + app_name=uri_app_name, + user_id=uri_user_id, + session_id="session1", + filename="file1", + version=1, + ) + + with pytest.raises(InputValidationError) as exc_info: + artifact_util.validate_artifact_reference_scope( + app_name="app1", + user_id="user1", + session_id="session1", + parsed_uri=parsed, + ) + + assert "same app and user scope" in str(exc_info.value) + + +def test_validate_artifact_reference_scope_other_session_raises(): + """A session-scoped reference from another session must be rejected.""" + parsed = artifact_util.ParsedArtifactUri( + app_name="app1", + user_id="user1", + session_id="other_session", + filename="file1", + version=1, + ) + + with pytest.raises(InputValidationError) as exc_info: + artifact_util.validate_artifact_reference_scope( + app_name="app1", + user_id="user1", + session_id="session1", + parsed_uri=parsed, + ) + + assert "same session scope" in str(exc_info.value) + + +def test_validate_artifact_reference_scope_session_uri_without_caller_session_raises(): + """A session-scoped reference cannot be used outside of any session.""" + parsed = artifact_util.ParsedArtifactUri( + app_name="app1", + user_id="user1", + session_id="session1", + filename="file1", + version=1, + ) + + with pytest.raises(InputValidationError) as exc_info: + artifact_util.validate_artifact_reference_scope( + app_name="app1", + user_id="user1", + session_id=None, + parsed_uri=parsed, + ) + + assert "same session scope" in str(exc_info.value) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("C:", True), + ("c:/data", True), + ("Z:relative", True), + ("1:x", False), + ("_:x", False), + ("é:x", False), + (":x", False), + ("user:profile.txt", False), + ("plain", False), + ], +) +def test_is_drive_qualified_matches_only_drive_letters(value, expected): + """Only a single ASCII letter followed by a colon counts as a drive.""" + assert artifact_util._is_drive_qualified(value) is expected