Merge pull request #156 from pycompression/zstandardwindow

rhpvorderman · web-flow · commit 63a10ffcd1c0 · 2024-03-27T16:32:46.000+01:00
Always set the maximum allowed window size for zstandard decompression
diff --git a/src/xopen/__init__.py b/src/xopen/__init__.py
@@ -508,15 +508,26 @@ def _open_zst(
     assert compresslevel != 0
     if compresslevel is None:
         compresslevel = XOPEN_DEFAULT_ZST_COMPRESSION
+    if zstandard:
+        max_window_bits = zstandard.WINDOWLOG_MAX
+    else:
+        max_window_bits = 31
     if threads != 0:
         try:
             # zstd can compress using multiple cores
+            program_args: Tuple[str, ...] = ("zstd",)
+            if "r" in mode:
+                # Only use --long=31 for decompression. Using it for
+                # compression overrides level settings for window size and
+                # forces other zstd users to use `--long=31` to decompress any
+                # archive that has been compressed by xopen.
+                program_args += (f"--long={max_window_bits}",)
             return _PipedCompressionProgram(
                 filename,
                 mode,
                 compresslevel,
                 threads,
-                _PROGRAM_SETTINGS["zstd"],
+                _ProgramSettings(program_args, tuple(range(1, 20)), "-T"),
             )
         except OSError:
             if zstandard is None:
@@ -525,11 +536,9 @@ def _open_zst(
 
     if zstandard is None:
         raise ImportError("zstandard module (python-zstandard) not available")
-    if compresslevel is not None and "r" not in mode:
-        cctx = zstandard.ZstdCompressor(level=compresslevel)
-    else:
-        cctx = None
-    f = zstandard.open(filename, mode, cctx=cctx)  # type: ignore
+    dctx = zstandard.ZstdDecompressor(max_window_size=2**max_window_bits)
+    cctx = zstandard.ZstdCompressor(level=compresslevel)
+    f = zstandard.open(filename, mode, cctx=cctx, dctx=dctx)  # type: ignore
     if mode == "rb":
         return io.BufferedReader(f)
     return io.BufferedWriter(f)  # mode "ab" and "wb"
diff --git a/tests/only_zeroes.zst b/tests/only_zeroes.zst
diff --git a/tests/test_xopen.py b/tests/test_xopen.py
@@ -576,6 +576,22 @@ def test_xopen_zst_fails_when_zstandard_not_available(monkeypatch):
             f.read()
 
 
+@pytest.mark.parametrize("threads", (0, 1))
+def test_xopen_zst_long_window_size(threads):
+    if threads == 0 and zstandard is None:
+        return
+    elif threads == 1 and not shutil.which("zstd"):
+        return
+    # File created with:
+    # cat /dev/zero | head -c 2147483648 > only_zeroes
+    # Then compressed with
+    # zstd --long=31 -19 only_zeroes
+    test_zst = Path(__file__).parent / "only_zeroes.zst"
+    with xopen(test_zst, "rb", threads=threads) as f:
+        data = f.read(1024)
+    assert data == bytes(1024)
+
+
 @pytest.mark.parametrize("threads", (0, 1))
 @pytest.mark.parametrize("ext", extensions)
 def test_pass_file_object_for_reading(ext, threads):
