diff --git a/.cargo/config b/.cargo/config index e5007178..af951327 100644 --- a/.cargo/config +++ b/.cargo/config @@ -1,2 +1,5 @@ [target.x86_64-apple-darwin] rustflags = ["-C", "link-arg=-undefined", "-C", "link-arg=dynamic_lookup"] + +[target.aarch64-apple-darwin] +rustflags = ["-C", "link-arg=-undefined", "-C", "link-arg=dynamic_lookup"] diff --git a/c-ext/decompressor.c b/c-ext/decompressor.c index b8e7f82c..d0b9591a 100644 --- a/c-ext/decompressor.c +++ b/c-ext/decompressor.c @@ -299,15 +299,15 @@ PyObject *Decompressor_decompress(ZstdDecompressor *self, PyObject *args, goto finally; } - decompressedSize = ZSTD_getFrameContentSize(source.buf, source.len); - - if (ZSTD_CONTENTSIZE_ERROR == decompressedSize) { + ZSTD_frameHeader frameHeader; + if (ZSTD_getFrameHeader_advanced(&frameHeader, source.buf, source.len, self->format) != 0) { PyErr_SetString(ZstdError, "error determining content size from frame header"); goto finally; } + decompressedSize=frameHeader.frameContentSize; /* Special case of empty frame. */ - else if (0 == decompressedSize) { + if (0 == decompressedSize) { result = PyBytes_FromStringAndSize("", 0); goto finally; } diff --git a/c-ext/frameparams.c b/c-ext/frameparams.c index 951a91d6..9d37b398 100644 --- a/c-ext/frameparams.c +++ b/c-ext/frameparams.c @@ -12,19 +12,21 @@ extern PyObject *ZstdError; FrameParametersObject *get_frame_parameters(PyObject *self, PyObject *args, PyObject *kwargs) { - static char *kwlist[] = {"data", NULL}; + static char *kwlist[] = {"data", "format", NULL}; Py_buffer source; + ZSTD_frameHeader header; + ZSTD_format_e format = ZSTD_f_zstd1; FrameParametersObject *result = NULL; size_t zresult; - if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:get_frame_parameters", - kwlist, &source)) { + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*|I:get_frame_parameters", + kwlist, &source, &format)) { return NULL; } - zresult = ZSTD_getFrameHeader(&header, source.buf, source.len); + zresult = ZSTD_getFrameHeader_advanced(&header, source.buf, source.len, format); if (ZSTD_isError(zresult)) { PyErr_Format(ZstdError, "cannot get frame parameters: %s", diff --git a/rust-ext/src/decompressor.rs b/rust-ext/src/decompressor.rs index 2964d9a3..07f0573f 100644 --- a/rust-ext/src/decompressor.rs +++ b/rust-ext/src/decompressor.rs @@ -176,17 +176,31 @@ impl ZstdDecompressor { self.setup_dctx(py, true)?; - let output_size = - unsafe { zstd_sys::ZSTD_getFrameContentSize(buffer.buf_ptr(), buffer.len_bytes()) }; + let mut header = zstd_sys::ZSTD_frameHeader { + frameContentSize: 0, + windowSize: 0, + blockSizeMax: 0, + frameType: zstd_sys::ZSTD_frameType_e::ZSTD_frame, + headerSize: 0, + dictID: 0, + checksumFlag: 0, + _reserved1: 0, + _reserved2: 0, + }; + let zresult = unsafe { + zstd_sys::ZSTD_getFrameHeader_advanced(&mut header, buffer.buf_ptr(), buffer.len_bytes(), self.format) + }; + + if zresult != 0 { + return Err(ZstdError::new_err( + "error determining content size from frame header" + )) + } let (output_buffer_size, output_size) = - if output_size == zstd_sys::ZSTD_CONTENTSIZE_ERROR as _ { - return Err(ZstdError::new_err( - "error determining content size from frame header", - )); - } else if output_size == 0 { + if header.frameContentSize == 0 { return Ok(PyBytes::new(py, &[])); - } else if output_size == zstd_sys::ZSTD_CONTENTSIZE_UNKNOWN as _ { + } else if header.frameContentSize == zstd_sys::ZSTD_CONTENTSIZE_UNKNOWN as _ { if max_output_size == 0 { return Err(ZstdError::new_err( "could not determine content size in frame header", @@ -195,7 +209,7 @@ impl ZstdDecompressor { (max_output_size, 0) } else { - (output_size as _, output_size) + (header.frameContentSize as _, header.frameContentSize) }; let mut dest_buffer: Vec = Vec::new(); diff --git a/rust-ext/src/frame_parameters.rs b/rust-ext/src/frame_parameters.rs index 13d7fd3f..7628dad5 100644 --- a/rust-ext/src/frame_parameters.rs +++ b/rust-ext/src/frame_parameters.rs @@ -6,7 +6,7 @@ use { crate::ZstdError, - pyo3::{buffer::PyBuffer, prelude::*, wrap_pyfunction}, + pyo3::{buffer::PyBuffer, prelude::*, wrap_pyfunction, exceptions::{PyValueError}}, }; #[pyclass(module = "zstandard.backend_rust")] @@ -67,10 +67,19 @@ fn frame_header_size(data: PyBuffer) -> PyResult { } #[pyfunction] -fn get_frame_parameters(py: Python, buffer: PyBuffer) -> PyResult> { +#[pyo3(signature = (buffer, format=zstd_sys::ZSTD_format_e::ZSTD_f_zstd1 as u32))] +fn get_frame_parameters(py: Python, buffer: PyBuffer, format: u32) -> PyResult> { let raw_data = unsafe { std::slice::from_raw_parts::(buffer.buf_ptr() as *const _, buffer.len_bytes()) }; + let format = if format == zstd_sys::ZSTD_format_e::ZSTD_f_zstd1 as _ { + zstd_sys::ZSTD_format_e::ZSTD_f_zstd1 + } else if format == zstd_sys::ZSTD_format_e::ZSTD_f_zstd1_magicless as _ { + zstd_sys::ZSTD_format_e::ZSTD_f_zstd1_magicless + } else { + return Err(PyValueError::new_err(format!("invalid format value"))); + }; + let mut header = zstd_sys::ZSTD_frameHeader { frameContentSize: 0, @@ -84,7 +93,7 @@ fn get_frame_parameters(py: Python, buffer: PyBuffer) -> PyResult int: ... def frame_content_size(data: ByteString) -> int: ... def frame_header_size(data: ByteString) -> int: ... -def get_frame_parameters(data: ByteString) -> FrameParameters: ... +def get_frame_parameters(data: ByteString, format: Optional[int]) -> FrameParameters: ... def train_dictionary( dict_size: int, samples: list[ByteString], diff --git a/zstandard/backend_cffi.py b/zstandard/backend_cffi.py index 7137542f..9f0e2ea9 100644 --- a/zstandard/backend_cffi.py +++ b/zstandard/backend_cffi.py @@ -2558,7 +2558,7 @@ def frame_header_size(data): return zresult -def get_frame_parameters(data): +def get_frame_parameters(data, format=FORMAT_ZSTD1): """ Parse a zstd frame header into frame parameters. @@ -2569,13 +2569,15 @@ def get_frame_parameters(data): :param data: Data from which to read frame parameters. + :param format: + Set the format of data for the decoder. :return: :py:class:`FrameParameters` """ params = ffi.new("ZSTD_frameHeader *") data_buffer = ffi.from_buffer(data) - zresult = lib.ZSTD_getFrameHeader(params, data_buffer, len(data_buffer)) + zresult = lib.ZSTD_getFrameHeader_advanced(params, data_buffer, len(data_buffer), format) if lib.ZSTD_isError(zresult): raise ZstdError( "cannot get frame parameters: %s" % _zstd_error(zresult) @@ -3820,13 +3822,12 @@ def decompress( data_buffer = ffi.from_buffer(data) - output_size = lib.ZSTD_getFrameContentSize( - data_buffer, len(data_buffer) - ) - - if output_size == lib.ZSTD_CONTENTSIZE_ERROR: + params = ffi.new("ZSTD_frameHeader *") + zresult = lib.ZSTD_getFrameHeader_advanced(params, data_buffer, len(data_buffer), self._format) + if zresult != 0: raise ZstdError("error determining content size from frame header") - elif output_size == 0: + output_size = params.frameContentSize + if output_size == 0: return b"" elif output_size == lib.ZSTD_CONTENTSIZE_UNKNOWN: if not max_output_size: