diff --git a/docs/release.md b/docs/release.md index ee770933..5c5e6aa1 100644 --- a/docs/release.md +++ b/docs/release.md @@ -14,6 +14,16 @@ ## Unreleased +### Fixes + +* Improve error handling in the Zstandard codec. Empty frames can now be decoded, so + `Zstd().decode(Zstd().encode(b""))` round-trips. Fix a double free and a free of an + uninitialized pointer when stream decompression failed during initialization, check for + failed allocations, no longer leak the compression context when setting parameters fails, + and report zstd error names as text rather than as the `repr` of a `bytes` object. + Streaming decompression now grows its output buffer geometrically. `Zstd.__repr__` now + includes `checksum` when it is enabled. + (release_0.17.1)= ## 0.17.1 diff --git a/src/numcodecs/zstd.pyx b/src/numcodecs/zstd.pyx index 7b6ff3c0..e57f11de 100644 --- a/src/numcodecs/zstd.pyx +++ b/src/numcodecs/zstd.pyx @@ -103,6 +103,11 @@ DEFAULT_CLEVEL = 0 MAX_CLEVEL = ZSTD_maxCLevel() +cdef str error_name(size_t code): + """Return the human readable name of a zstd error code as a str.""" + return (ZSTD_getErrorName(code)).decode('utf-8', 'replace') + + def compress(source, int level=DEFAULT_CLEVEL, bint checksum=False): """Compress data. @@ -126,9 +131,10 @@ def compress(source, int level=DEFAULT_CLEVEL, bint checksum=False): memoryview source_mv const Py_buffer* source_pb const char* source_ptr - size_t source_size, dest_size, compressed_size + size_t source_size, dest_size, compressed_size, status bytes dest char* dest_ptr + ZSTD_CCtx* cctx # check level if level > MAX_CLEVEL: @@ -143,22 +149,23 @@ def compress(source, int level=DEFAULT_CLEVEL, bint checksum=False): source_size = source_pb.len cctx = ZSTD_createCCtx() - param_set_result = ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, level) + if cctx == NULL: + raise MemoryError('Zstd compression error: could not create compression context') - if ZSTD_isError(param_set_result): - error = ZSTD_getErrorName(param_set_result) - raise RuntimeError('Could not set zstd compression level: %s' % error) + try: - param_set_result = ZSTD_CCtx_setParameter(cctx, ZSTD_c_checksumFlag, 1 if checksum else 0) + status = ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, level) + if ZSTD_isError(status): + raise RuntimeError('Could not set zstd compression level: %s' % error_name(status)) - if ZSTD_isError(param_set_result): - error = ZSTD_getErrorName(param_set_result) - raise RuntimeError('Could not set zstd checksum flag: %s' % error) - - try: + status = ZSTD_CCtx_setParameter(cctx, ZSTD_c_checksumFlag, 1 if checksum else 0) + if ZSTD_isError(status): + raise RuntimeError('Could not set zstd checksum flag: %s' % error_name(status)) # setup destination dest_size = ZSTD_compressBound(source_size) + if ZSTD_isError(dest_size): + raise RuntimeError('Zstd compression error: source too large to compress') dest = PyBytes_FromStringAndSize(NULL, dest_size) dest_ptr = PyBytes_AS_STRING(dest) @@ -167,13 +174,11 @@ def compress(source, int level=DEFAULT_CLEVEL, bint checksum=False): compressed_size = ZSTD_compress2(cctx, dest_ptr, dest_size, source_ptr, source_size) finally: - if cctx: - ZSTD_freeCCtx(cctx) + ZSTD_freeCCtx(cctx) # check compression was successful if ZSTD_isError(compressed_size): - error = ZSTD_getErrorName(compressed_size) - raise RuntimeError('Zstd compression error: %s' % error) + raise RuntimeError('Zstd compression error: %s' % error_name(compressed_size)) # resize after compression PyBytes_RESIZE(dest, compressed_size) @@ -207,7 +212,7 @@ def decompress(source, dest=None): memoryview dest_mv Py_buffer* dest_pb char* dest_ptr - size_t source_size, dest_size, decompressed_size + size_t source_size, dest_size, decompressed_size = 0 size_t dest_nbytes unsigned long long content_size @@ -219,65 +224,64 @@ def decompress(source, dest=None): source_ptr = source_pb.buf source_size = source_pb.len + if source_size == 0: + raise RuntimeError('Zstd decompression error: invalid input data') + + # determine uncompressed size using unsigned long long for full range try: - # determine uncompressed size using unsigned long long for full range - try: - content_size = findTotalContentSize(source_ptr, source_size) - except RuntimeError: - raise RuntimeError('Zstd decompression error: invalid input data') + content_size = findTotalContentSize(source_ptr, source_size) + except RuntimeError as e: + raise RuntimeError('Zstd decompression error: invalid input data: %s' % e) from None - if content_size == ZSTD_CONTENTSIZE_UNKNOWN and dest is None: + if content_size == ZSTD_CONTENTSIZE_ERROR: + raise RuntimeError('Zstd decompression error: invalid input data') + elif content_size == ZSTD_CONTENTSIZE_UNKNOWN: + if dest is None: return stream_decompress(source_pb) - elif content_size == ZSTD_CONTENTSIZE_UNKNOWN: - # dest is not None - # set dest_size based on dest - pass - elif content_size == ZSTD_CONTENTSIZE_ERROR or content_size == 0: - raise RuntimeError('Zstd decompression error: invalid input data') - elif content_size > (SIZE_MAX): - raise RuntimeError('Zstd decompression error: content size too large for platform') - + # dest_size is set based on dest below + dest_size = 0 + elif content_size > ((SIZE_MAX) >> 1): + # larger than any Python buffer (PY_SSIZE_T_MAX) can be + raise RuntimeError('Zstd decompression error: content size too large for platform') + else: + # note that a content size of 0 is valid: it is an empty frame dest_size = content_size - # setup destination buffer - if dest is None: - # allocate memory - dest_1d = dest = PyBytes_FromStringAndSize(NULL, dest_size) - else: - dest_1d = ensure_contiguous_ndarray(dest) - - # obtain dest memoryview - dest_mv = memoryview(dest_1d) - dest_pb = PyMemoryView_GET_BUFFER(dest_mv) - dest_ptr = dest_pb.buf - dest_nbytes = dest_pb.len - - if content_size == ZSTD_CONTENTSIZE_UNKNOWN: - dest_size = dest_nbytes - - # validate output buffer - if dest_nbytes < dest_size: - raise ValueError('destination buffer too small; expected at least %s, ' - 'got %s' % (dest_size, dest_nbytes)) - - # perform decompression - with nogil: - decompressed_size = ZSTD_decompress(dest_ptr, dest_size, source_ptr, source_size) + # setup destination buffer + if dest is None: + # allocate memory + dest_1d = dest = PyBytes_FromStringAndSize(NULL, dest_size) + else: + dest_1d = ensure_contiguous_ndarray(dest) - finally: - pass + # obtain dest memoryview + dest_mv = memoryview(dest_1d) + dest_pb = PyMemoryView_GET_BUFFER(dest_mv) + dest_ptr = dest_pb.buf + dest_nbytes = dest_pb.len + + if content_size == ZSTD_CONTENTSIZE_UNKNOWN: + dest_size = dest_nbytes + + # validate output buffer + if dest_nbytes < dest_size: + raise ValueError('destination buffer too small; expected at least %s, ' + 'got %s' % (dest_size, dest_nbytes)) + + # perform decompression + with nogil: + decompressed_size = ZSTD_decompress(dest_ptr, dest_size, source_ptr, source_size) # check decompression was successful if ZSTD_isError(decompressed_size): - error = ZSTD_getErrorName(decompressed_size) - raise RuntimeError('Zstd decompression error: %s' % error) + raise RuntimeError('Zstd decompression error: %s' % error_name(decompressed_size)) elif decompressed_size != dest_size: raise RuntimeError('Zstd decompression error: expected to decompress %s, got %s' % (dest_size, decompressed_size)) return dest -cdef stream_decompress(const Py_buffer* source_pb): +cdef bytes stream_decompress(const Py_buffer* source_pb): """Decompress data of unknown size Parameters @@ -293,47 +297,58 @@ cdef stream_decompress(const Py_buffer* source_pb): cdef: const char *source_ptr - void *dest_ptr void *new_dst - size_t source_size, dest_size, decompressed_size - size_t DEST_GROWTH_SIZE, status + size_t source_size, dest_size, new_size, growth + size_t DEST_GROWTH_SIZE, status = 0 ZSTD_inBuffer input ZSTD_outBuffer output ZSTD_DStream *zds + # 0: success, 1: zstd error, 2: size overflow, 3: out of memory, 4: truncated input + int failure = 0 + bytes dest # Recommended size for output buffer, guaranteed to flush at least # one completely block in all circumstances - DEST_GROWTH_SIZE = ZSTD_DStreamOutSize(); + DEST_GROWTH_SIZE = ZSTD_DStreamOutSize() source_ptr = source_pb.buf source_size = source_pb.len # unknown content size, guess it is twice the size as the source dest_size = source_size * 2 + if dest_size < source_size: + # multiplication overflowed + dest_size = source_size if dest_size < DEST_GROWTH_SIZE: # minimum dest_size is DEST_GROWTH_SIZE dest_size = DEST_GROWTH_SIZE - dest_ptr = malloc(dest_size) - zds = ZSTD_createDStream() + # output.dst must be valid for free() before anything can fail + output = ZSTD_outBuffer(NULL, 0, 0) + zds = NULL try: + output.dst = malloc(dest_size) + if output.dst == NULL: + raise MemoryError('Zstd stream decompression error: could not allocate output buffer') + output.size = dest_size - with nogil: + zds = ZSTD_createDStream() + if zds == NULL: + raise MemoryError('Zstd stream decompression error: could not create decompression stream') - status = ZSTD_initDStream(zds) - if ZSTD_isError(status): - error = ZSTD_getErrorName(status) - ZSTD_freeDStream(zds); - raise RuntimeError('Zstd stream decompression error on ZSTD_initDStream: %s' % error) + status = ZSTD_initDStream(zds) + if ZSTD_isError(status): + raise RuntimeError('Zstd stream decompression error on ZSTD_initDStream: %s' % error_name(status)) - input = ZSTD_inBuffer(source_ptr, source_size, 0) - output = ZSTD_outBuffer(dest_ptr, dest_size, 0) + input = ZSTD_inBuffer(source_ptr, source_size, 0) + + with nogil: # Initialize to 1 to force a loop iteration status = 1 - while(status > 0 or input.pos < input.size): + while status > 0 or input.pos < input.size: # Possible returned values of ZSTD_decompressStream: # 0: frame is completely decoded and fully flushed # error (<0) @@ -341,35 +356,55 @@ cdef stream_decompress(const Py_buffer* source_pb): status = ZSTD_decompressStream(zds, &output, &input) if ZSTD_isError(status): - error = ZSTD_getErrorName(status) - raise RuntimeError('Zstd stream decompression error on ZSTD_decompressStream: %s' % error) - - # There is more to decompress, grow the buffer - if status > 0 and output.pos == output.size: - new_size = output.size + DEST_GROWTH_SIZE - - if new_size < output.size or new_size < DEST_GROWTH_SIZE: - raise RuntimeError('Zstd stream decompression error: output buffer overflow') - - new_dst = realloc(output.dst, new_size) - - if new_dst == NULL: - # output.dst freed in finally block - raise RuntimeError('Zstd stream decompression error on realloc: could not expand output buffer') - - output.dst = new_dst - output.size = new_size + failure = 1 + break + + if output.pos == output.size: + # The output buffer is full. If a frame is incomplete or + # another frame follows, grow the buffer geometrically. + if status > 0 or input.pos < input.size: + growth = output.size if output.size > DEST_GROWTH_SIZE else DEST_GROWTH_SIZE + if growth > SIZE_MAX - output.size: + growth = DEST_GROWTH_SIZE + if growth > SIZE_MAX - output.size: + failure = 2 + break + new_size = output.size + growth + + new_dst = realloc(output.dst, new_size) + if new_dst == NULL: + # output.dst freed in finally block + failure = 3 + break + + output.dst = new_dst + output.size = new_size + elif status > 0 and input.pos == input.size: + # All of the input has been consumed, there is room for more + # output, but the frame is not complete. + failure = 4 + break + + if failure == 1: + raise RuntimeError('Zstd stream decompression error on ZSTD_decompressStream: %s' % error_name(status)) + elif failure == 2: + raise RuntimeError('Zstd stream decompression error: output buffer overflow') + elif failure == 3: + raise MemoryError('Zstd stream decompression error on realloc: could not expand output buffer') + elif failure == 4: + raise RuntimeError('Zstd stream decompression error: truncated input data') # Copy the output to a bytes object dest = PyBytes_FromStringAndSize(output.dst, output.pos) finally: - ZSTD_freeDStream(zds) + if zds != NULL: + ZSTD_freeDStream(zds) free(output.dst) return dest -cdef unsigned long long findTotalContentSize(const char* source_ptr, size_t source_size): +cdef unsigned long long findTotalContentSize(const char* source_ptr, size_t source_size) except? 0: """Find the total uncompressed content size of all frames in the source buffer Parameters @@ -382,6 +417,13 @@ cdef unsigned long long findTotalContentSize(const char* source_ptr, size_t sour total_content_size: Sum of the content size of all frames within the source buffer If any of the frame sizes is unknown, return ZSTD_CONTENTSIZE_UNKNOWN. If any of the frames causes ZSTD_getFrameContentSize to error, return ZSTD_CONTENTSIZE_ERROR. + + Raises + ------ + RuntimeError + If the size of a frame cannot be determined, e.g. the input is truncated. + OverflowError + If the total content size does not fit in an unsigned long long. """ cdef: unsigned long long frame_content_size = 0 @@ -390,13 +432,12 @@ cdef unsigned long long findTotalContentSize(const char* source_ptr, size_t sour size_t offset = 0 while offset < source_size: - frame_compressed_size = ZSTD_findFrameCompressedSize(source_ptr + offset, source_size - offset); + frame_compressed_size = ZSTD_findFrameCompressedSize(source_ptr + offset, source_size - offset) if ZSTD_isError(frame_compressed_size): - error = ZSTD_getErrorName(frame_compressed_size) - raise RuntimeError('Could not set determine zstd frame size: %s' % error) + raise RuntimeError('Could not determine zstd frame size: %s' % error_name(frame_compressed_size)) - frame_content_size = ZSTD_getFrameContentSize(source_ptr + offset, frame_compressed_size); + frame_content_size = ZSTD_getFrameContentSize(source_ptr + offset, frame_compressed_size) if frame_content_size == ZSTD_CONTENTSIZE_ERROR: return ZSTD_CONTENTSIZE_ERROR @@ -404,6 +445,10 @@ cdef unsigned long long findTotalContentSize(const char* source_ptr, size_t sour if frame_content_size == ZSTD_CONTENTSIZE_UNKNOWN: return ZSTD_CONTENTSIZE_UNKNOWN + if frame_content_size > ZSTD_CONTENTSIZE_ERROR - 1 - total_content_size: + # Keep the sum below the ZSTD_CONTENTSIZE_* sentinel values + raise OverflowError('Zstd decompression error: total content size too large') + total_content_size += frame_content_size offset += frame_compressed_size @@ -444,10 +489,10 @@ class Zstd(Codec): return decompress(buf, out) def __repr__(self): - r = '%s(level=%r)' % \ - (type(self).__name__, - self.level) - return r + r = '%s(level=%r' % (type(self).__name__, self.level) + if self.checksum: + r += ', checksum=%r' % self.checksum + return r + ')' @classmethod def default_level(cls): diff --git a/tests/test_zstd.py b/tests/test_zstd.py index 09342088..e04a299f 100644 --- a/tests/test_zstd.py +++ b/tests/test_zstd.py @@ -1,10 +1,12 @@ import itertools +import struct import subprocess import numpy as np import pytest from numcodecs.zstd import Zstd +from numcodecs import zstd from tests.common import ( check_backwards_compatibility, check_config, @@ -187,3 +189,155 @@ def test_multi_frame(): assert codec.decode(bytes_val + hola + mundo) == b"Hello World!Hola Mundo!" # unknown size frame in the middle assert codec.decode(hola + bytes_val + mundo) == b"Hola Hello World!Mundo!" + + +def test_empty_frame(): + codec = Zstd() + encoded = codec.encode(b"") + assert codec.decode(encoded) == b"" + # decode into an empty destination buffer + out = bytearray() + codec.decode(encoded, out=out) + assert out == b"" + # an empty frame can be concatenated with other frames + assert codec.decode(encoded + codec.encode(b"abc") + encoded) == b"abc" + + +def test_decode_empty_input(): + with pytest.raises(RuntimeError, match='Zstd decompression error: invalid input data'): + Zstd().decode(b"") + + +def test_truncated_input(): + codec = Zstd() + known = codec.encode(b"Hello world!" * 100) + unknown = b'(\xb5/\xfd\x00Xa\x00\x00Hello World!' + for data in (known[:-3], unknown[:-3], known + known[:-3]): + with pytest.raises(RuntimeError, match='invalid input data'): + codec.decode(data) + + +def test_error_messages_are_text(): + # error names from zstd must be formatted as text, not as the repr of bytes + codec = Zstd() + encoded = bytearray(Zstd(checksum=True).encode(b"abc" * 10)) + encoded[-1] ^= 0xFF + with pytest.raises(RuntimeError) as excinfo: + codec.decode(bytes(encoded)) + assert "b'" not in str(excinfo.value) + assert "b\"" not in str(excinfo.value) + assert "checksum" in str(excinfo.value) + + +def test_decode_unknown_size_into_out(): + codec = Zstd() + unknown = b'(\xb5/\xfd\x00Xa\x00\x00Hello World!' + out = bytearray(12) + codec.decode(unknown, out=out) + assert out == b"Hello World!" + + +def test_repr_checksum(): + check_repr("Zstd(level=3)") + check_repr("Zstd(level=3, checksum=True)") + + +def test_stream_decompress_growth(): + # Highly compressible data with unknown content size requires the output + # buffer to grow many times. + codec = Zstd() + frame = b'(\xb5/\xfd\x00X$\x02\x00\xa4\x03ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz\x01\x00:\xfc\xdfs\x05\x05L\x00\x00\x08s\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08k\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08c\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08[\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08S\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08K\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08C\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08u\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08m\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08e\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08]\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08U\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08M\x01\x00\xfc\xff9\x10\x02M\x00\x00\x08E\x01\x00\xfc\x7f\x1d\x08\x01' + expected = b'ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz' * 1024 * 32 + assert codec.decode(frame * 3) == expected * 3 + + +def _frame_with_content_size(n): + """A minimal single-segment frame whose header claims ``n`` bytes of content.""" + return b'\x28\xb5\x2f\xfd' + bytes([0xE0]) + struct.pack('