Skip to content

Commit 8d99154

Browse files
fix: add cumulative and fallback-path tar bomb tests for decompress_to_cache
Added test cases to `tests/test_model_management.py` to cover: 1. Cumulative decompression bomb detection (multiple small files exceeding 20GB). 2. Decompression bomb detection in the fallback path (when `tarfile.data_filter` is unavailable). These tests ensure full coverage of the security logic in `ModelManagement.decompress_to_cache`. Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
1 parent c6eb038 commit 8d99154

1 file changed

Lines changed: 67 additions & 0 deletions

File tree

tests/test_model_management.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -626,6 +626,73 @@ def fake_extractall(*args, **kwargs):
626626

627627
assert not cache_dir.exists()
628628

629+
def test_decompress_tar_bomb_cumulative(self, tmp_path):
630+
"""Tar bomb prevention: multiple small files exceeding 20GB limit."""
631+
cache_dir = tmp_path / "tmp_cache_dir_bomb_cumulative"
632+
cache_dir.mkdir()
633+
fake_tar_gz = tmp_path / "fake.tar.gz"
634+
fake_tar_gz.write_text("fake tar")
635+
636+
m1 = MagicMock(spec=tarfile.TarInfo)
637+
m1.name = "f1.txt"
638+
m1.size = 15 * 1024 * 1024 * 1024
639+
m1.issym.return_value = False
640+
m1.islnk.return_value = False
641+
m1.isdir.return_value = False
642+
m1.isreg.return_value = True
643+
644+
m2 = MagicMock(spec=tarfile.TarInfo)
645+
m2.name = "f2.txt"
646+
m2.size = 6 * 1024 * 1024 * 1024
647+
m2.issym.return_value = False
648+
m2.islnk.return_value = False
649+
m2.isdir.return_value = False
650+
m2.isreg.return_value = True
651+
652+
with patch("tarfile.open") as mock_tar_open:
653+
mock_tar = MagicMock()
654+
mock_tar.__iter__.return_value = iter([m1, m2])
655+
# Simulate extractall consuming the members generator
656+
mock_tar.extractall.side_effect = lambda path, members, filter=None: list(members)
657+
mock_tar_open.return_value.__enter__.return_value = mock_tar
658+
659+
with pytest.raises(tarfile.TarError, match="Decompression bomb detected"):
660+
ModelManagement.decompress_to_cache(str(fake_tar_gz), str(cache_dir))
661+
662+
assert not cache_dir.exists()
663+
664+
def test_decompress_tar_bomb_fallback_path(self, tmp_path):
665+
"""Tar bomb prevention in fallback path (no data_filter)."""
666+
cache_dir = tmp_path / "tmp_cache_dir_bomb_fallback"
667+
cache_dir.mkdir()
668+
fake_tar_gz = tmp_path / "fake.tar.gz"
669+
fake_tar_gz.write_text("fake tar")
670+
671+
mock_member = MagicMock(spec=tarfile.TarInfo)
672+
mock_member.name = "huge_file.txt"
673+
mock_member.size = 21 * 1024 * 1024 * 1024
674+
mock_member.issym.return_value = False
675+
mock_member.islnk.return_value = False
676+
mock_member.isdir.return_value = False
677+
mock_member.isreg.return_value = True
678+
679+
with patch("qwen3_embed.common.model_management.tarfile") as mock_tarfile_mod:
680+
# Mock tarfile.TarError since it is used in "except" and "raise"
681+
mock_tarfile_mod.TarError = tarfile.TarError
682+
# Mock tarfile.open to return a mock_tar
683+
mock_tar = MagicMock()
684+
mock_tar.__iter__.return_value = iter([mock_member])
685+
mock_tarfile_mod.open.return_value.__enter__.return_value = mock_tar
686+
687+
# Ensure it DOES NOT have data_filter
688+
if hasattr(mock_tarfile_mod, "data_filter"):
689+
del mock_tarfile_mod.data_filter
690+
691+
with pytest.raises(tarfile.TarError, match="Decompression bomb detected"):
692+
ModelManagement.decompress_to_cache(str(fake_tar_gz), str(cache_dir))
693+
694+
assert not cache_dir.exists()
695+
629696

630697
# ---------------------------------------------------------------------------
631698
# TestDownloadFilesFromHuggingFace

0 commit comments

Comments
 (0)