diff --git a/src/google/adk/code_executors/container_code_executor.py b/src/google/adk/code_executors/container_code_executor.py index 6dfc445d55c..8d6c39b4f61 100644 --- a/src/google/adk/code_executors/container_code_executor.py +++ b/src/google/adk/code_executors/container_code_executor.py @@ -449,13 +449,34 @@ def _get_output_files(self) -> list[File]: return [] tar_buffer = io.BytesIO(tar_bytes) output_files = [] + total_output_bytes = 0 with tarfile.open(fileobj=tar_buffer, mode='r') as tar: + # Only regular files and hard links ultimately pointing to regular files + # are eligible. Track the latest entry for each path in archive order. + regular_targets: set[str] = set() for member in tar.getmembers(): + member_name = os.path.normpath(member.name) if member.isfile(): + regular_targets.add(member_name) + elif member.islnk(): + if os.path.normpath(member.linkname) in regular_targets: + regular_targets.add(member_name) + else: + regular_targets.discard(member_name) + else: + regular_targets.discard(member_name) + if member_name in regular_targets: file_obj = tar.extractfile(member) if file_obj: content = file_obj.read() + # Hard links repeat file contents without growing the TAR archive. + total_output_bytes += len(content) + if total_output_bytes > _MAX_OUTPUT_ARCHIVE_BYTES: + raise RuntimeError( + f'Output files in {self.output_dir} exceed' + f' {_MAX_OUTPUT_ARCHIVE_BYTES} bytes after resolving links.' + ) # The archive's top-level entry is output_dir itself. file_name = member.name.partition('/')[2] if file_name.startswith('user:'): diff --git a/tests/unittests/code_executors/test_container_code_executor.py b/tests/unittests/code_executors/test_container_code_executor.py index 4305924d379..19c757790b5 100644 --- a/tests/unittests/code_executors/test_container_code_executor.py +++ b/tests/unittests/code_executors/test_container_code_executor.py @@ -696,6 +696,142 @@ def test_execute_code_with_output_files(mock_docker): ] +@mock.patch('google.adk.code_executors.container_code_executor.docker') +def test_execute_code_with_hardlinked_output_files(mock_docker): + """Hard-linked output files retain both names and contents.""" + client = _mock_docker_client() + container = client.containers.run.return_value + mock_docker.from_env.return_value = client + executor = ContainerCodeExecutor(image='test-image:latest') + + content = b'shared file content' + tar_buffer = io.BytesIO() + + with tarfile.open(fileobj=tar_buffer, mode='w') as tar: + original = tarfile.TarInfo(name='outputs/original.txt') + original.size = len(content) + tar.addfile(original, io.BytesIO(content)) + + linked = tarfile.TarInfo(name='outputs/linked.txt') + linked.type = tarfile.LNKTYPE + linked.linkname = 'outputs/original.txt' + tar.addfile(linked) + + symlink = tarfile.TarInfo(name='outputs/symlink.txt') + symlink.type = tarfile.SYMTYPE + symlink.linkname = 'original.txt' + tar.addfile(symlink) + + container.get_archive.return_value = ( + [tar_buffer.getvalue()], + mock.MagicMock(), + ) + + context = mock.MagicMock(spec=InvocationContext) + result = executor.execute_code( + context, + CodeExecutionInput(code='print("test")'), + ) + + assert result.output_files == [ + File( + name='original.txt', + content=content, + mime_type='text/plain', + ), + File( + name='linked.txt', + content=content, + mime_type='text/plain', + ), + ] + + +@mock.patch('google.adk.code_executors.container_code_executor.docker') +def test_execute_code_skips_hardlinks_to_symlinks(mock_docker): + """Hard links to symlinks are excluded without losing unrelated outputs.""" + client = _mock_docker_client() + container = client.containers.run.return_value + mock_docker.from_env.return_value = client + executor = ContainerCodeExecutor(image='test-image:latest') + + content = b'ordinary output' + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode='w') as tar: + ordinary = tarfile.TarInfo('outputs/ordinary.txt') + ordinary.size = len(content) + tar.addfile(ordinary, io.BytesIO(content)) + + linked = tarfile.TarInfo('outputs/linked.txt') + linked.type = tarfile.LNKTYPE + linked.linkname = ordinary.name + tar.addfile(linked) + + chained = tarfile.TarInfo('outputs/chained.txt') + chained.type = tarfile.LNKTYPE + chained.linkname = linked.name + tar.addfile(chained) + + for symlink_name, target in ( + ('broken-symlink.txt', 'missing-target'), + ('valid-symlink.txt', 'ordinary.txt'), + ): + symlink = tarfile.TarInfo(f'outputs/{symlink_name}') + symlink.type = tarfile.SYMTYPE + symlink.linkname = target + tar.addfile(symlink) + + alias = tarfile.TarInfo(f'outputs/hard-{symlink_name}') + alias.type = tarfile.LNKTYPE + alias.linkname = symlink.name + tar.addfile(alias) + + container.get_archive.return_value = ( + [tar_buffer.getvalue()], + mock.MagicMock(), + ) + context = mock.MagicMock(spec=InvocationContext) + result = executor.execute_code( + context, CodeExecutionInput(code='print("test")') + ) + + assert result.output_files == [ + File(name='ordinary.txt', content=content, mime_type='text/plain'), + File(name='linked.txt', content=content, mime_type='text/plain'), + File(name='chained.txt', content=content, mime_type='text/plain'), + ] + + +@mock.patch.object(container_code_executor, '_MAX_OUTPUT_ARCHIVE_BYTES', 16_384) +@mock.patch('google.adk.code_executors.container_code_executor.docker') +def test_execute_code_limits_hardlink_expansion(mock_docker): + """Hard links cannot bypass the output size budget with duplicate data.""" + client = _mock_docker_client() + container = client.containers.run.return_value + mock_docker.from_env.return_value = client + executor = ContainerCodeExecutor(image='test-image:latest') + + content = b'x' * 4096 + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode='w') as tar: + original = tarfile.TarInfo(name='outputs/original.bin') + original.size = len(content) + tar.addfile(original, io.BytesIO(content)) + for index in range(5): + linked = tarfile.TarInfo(name=f'outputs/linked-{index}.bin') + linked.type = tarfile.LNKTYPE + linked.linkname = 'outputs/original.bin' + tar.addfile(linked) + + archive_bytes = tar_buffer.getvalue() + assert len(archive_bytes) < 16_384 + container.get_archive.return_value = ([archive_bytes], mock.MagicMock()) + context = mock.MagicMock(spec=InvocationContext) + + with pytest.raises(RuntimeError, match='after resolving links'): + executor.execute_code(context, CodeExecutionInput(code='print("test")')) + + @mock.patch('google.adk.code_executors.container_code_executor.docker') def test_execute_code_reraises_non_404_archive_error(mock_docker): """Non-404 APIErrors from get_archive, even without a response, propagate."""