diff --git a/util/__init__.py b/util/__init__.py index e823efea..037b0d8a 100644 --- a/util/__init__.py +++ b/util/__init__.py @@ -196,9 +196,11 @@ def unzstd(contents): """ unzstd contents in memory and return the data """ try: - if hasattr(zstd, 'decompress'): - return zstd.decompress(contents) - return zstd.ZstdDecompressor().stream_reader(contents).read() + if zstd.__name__ == 'zstandard': + # zstandard.decompress() fails without a content size in the frame + # header and stops after the first frame, so stream instead + return zstd.ZstdDecompressor().stream_reader(contents, read_across_frames=True).read() + return zstd.decompress(contents) except (zstd.ZstdError, Exception) as e: error_message(text=f'zstd: {e}') diff --git a/util/tests/test_utils.py b/util/tests/test_utils.py index 4205150f..23735348 100644 --- a/util/tests/test_utils.py +++ b/util/tests/test_utils.py @@ -17,6 +17,7 @@ import gzip import hashlib from io import BytesIO +from unittest import skipIf from unittest.mock import MagicMock from django.test import TestCase, override_settings @@ -24,9 +25,14 @@ from util import ( Checksum, bunzip2, extract, get_checksum, get_md5, get_sha1, get_sha256, get_sha512, gunzip, has_setting_of_type, is_epoch_time, response_is_valid, - sanitize_filter_params, tz_aware_datetime, + sanitize_filter_params, tz_aware_datetime, unzstd, ) +try: + import zstandard +except ImportError: + zstandard = None + @override_settings( CELERY_TASK_ALWAYS_EAGER=True, @@ -138,6 +144,20 @@ def test_extract_gzip(self): result = extract(compressed, 'gz') self.assertEqual(result, original) + @skipIf(zstandard is None, 'zstandard not installed') + def test_unzstd_without_content_size(self): + """Test unzstd with a frame that does not declare its content size.""" + original = b'x' * 100000 + compressed = zstandard.ZstdCompressor(write_content_size=False).compress(original) + self.assertEqual(unzstd(compressed), original) + + @skipIf(zstandard is None, 'zstandard not installed') + def test_extract_zstd_multiple_frames(self): + """Test extract with zstd data made of multiple frames.""" + first, second = b'a' * 50000, b'b' * 50000 + compressed = zstandard.ZstdCompressor().compress(first) + zstandard.ZstdCompressor().compress(second) + self.assertEqual(extract(compressed, 'zst'), first + second) + def test_extract_unknown_format(self): """Test extract with unknown format returns original.""" data = b'unchanged data'