diff --git a/python/_brotli.c b/python/_brotli.c index 88a4a9aa8..84dadeca6 100644 --- a/python/_brotli.c +++ b/python/_brotli.c @@ -72,6 +72,8 @@ static const char kDecompressSinkError[] = "brotli: decoder process called with data when 'can_accept_more_data()' is " "False"; static const char kDecompressError[] = "brotli: decoder failed"; +static const char kDecompressOutputLimitError[] = + "brotli: decompressed output exceeds output_buffer_limit"; /* clang-format off */ PyDoc_STRVAR(brotli_error_doc, @@ -207,17 +209,26 @@ PyDoc_STRVAR(brotli_decompress__doc__, "Decompress a compressed byte string.\n" "\n" "Signature:\n" -" decompress(data) -> bytes\n" +" decompress(data, output_buffer_limit=0) -> bytes\n" "\n" "Args:\n" " data (bytes): The input data;\n" " MUST provide C-contiguous one-dimensional bytes view.\n" +" output_buffer_limit (int): Maximum bytes of decompressed output to\n" +" allow. If the decompressed data would be larger than\n" +" this value, brotli.error is raised. Set this when\n" +" decompressing untrusted input to prevent decompression\n" +" bombs. Default is 0, which means no limit.\n" +" Note: memory is allocated in power-of-2 blocks, so\n" +" the process may briefly use up to 2x this value in\n" +" memory before the limit is enforced.\n" "\n" "Returns:\n" " The decompressed byte string.\n" "\n" "Raises:\n" -" brotli.error: If decompressor fails.\n"); +" brotli.error: If decompressor fails or output exceeds\n" +" output_buffer_limit.\n"); PyDoc_STRVAR(brotli_doc, "Implementation module for the Brotli library."); /* clang-format on */ @@ -860,7 +871,7 @@ static PyObject* brotli_Decompressor_can_accept_more_data( static PyObject* brotli_decompress(PyObject* m, PyObject* args, PyObject* keywds) { - static const char* kwlist[] = {"string", NULL}; + static const char* kwlist[] = {"string", "output_buffer_limit", NULL}; BrotliDecoderState* state = NULL; BrotliDecoderResult result = BROTLI_DECODER_RESULT_ERROR; @@ -871,9 +882,12 @@ static PyObject* brotli_decompress(PyObject* m, PyObject* args, PyObject* input_object = NULL; Py_buffer input; int oom = 0; + int limit_exceeded = 0; + Py_ssize_t output_buffer_limit = 0; /* 0 = unlimited */ - if (!PyArg_ParseTupleAndKeywords(args, keywds, "O|:decompress", - (char**)kwlist, &input_object)) { + if (!PyArg_ParseTupleAndKeywords(args, keywds, "O|n:decompress", + (char**)kwlist, &input_object, + &output_buffer_limit)) { return NULL; } if (!get_data_view(input_object, &input)) { @@ -902,6 +916,11 @@ static PyObject* brotli_decompress(PyObject* m, PyObject* args, state, &available_in, &next_in, &buffer.avail_out, &buffer.next_out, 0); if (result == BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT) { assert(buffer.avail_out == 0); + if (output_buffer_limit > 0 && + buffer.total_allocated >= (uint64_t)output_buffer_limit) { + limit_exceeded = 1; + break; + } if (Buffer_Grow(&buffer) < 0) { oom = 1; break; @@ -914,6 +933,9 @@ static PyObject* brotli_decompress(PyObject* m, PyObject* args, if (oom) { goto finally; + } else if (limit_exceeded) { + set_brotli_exception_from_module(m, kDecompressOutputLimitError); + goto finally; } else if (result != BROTLI_DECODER_RESULT_SUCCESS || available_in > 0) { set_brotli_exception_from_module(m, kDecompressError); goto finally; diff --git a/python/tests/decompress_test.py b/python/tests/decompress_test.py index e59c96d8b..66ff5228a 100644 --- a/python/tests/decompress_test.py +++ b/python/tests/decompress_test.py @@ -22,3 +22,97 @@ def test_decompress(compressed_name, original_name): def test_garbage_appended(): with pytest.raises(brotli.error): brotli.decompress(brotli.compress(b'a') + b'a') + + +# --- output_buffer_limit tests --- + +def test_decompress_with_limit_succeeds_within_limit(): + """Decompression succeeds when output fits within the limit.""" + original = b'Hello, brotli!' * 100 + compressed = brotli.compress(original) + # Set limit well above actual output size + result = brotli.decompress(compressed, output_buffer_limit=len(original) * 2) + assert result == original + + +def test_decompress_with_limit_raises_when_exceeded(): + """Decompression raises brotli.error when output exceeds the limit.""" + original = b'\x00' * (1024 * 1024) # 1 MB of zeros + compressed = brotli.compress(original, quality=1) + with pytest.raises(brotli.error): + brotli.decompress(compressed, output_buffer_limit=1024) + + +def test_decompress_with_zero_limit_is_unlimited(): + """output_buffer_limit=0 means no limit (backward compatibility).""" + original = b'A' * (256 * 1024) + compressed = brotli.compress(original) + result = brotli.decompress(compressed, output_buffer_limit=0) + assert result == original + + +def test_decompress_without_limit_is_unlimited(): + """Omitting output_buffer_limit is equivalent to 0 (unlimited).""" + original = b'B' * (256 * 1024) + compressed = brotli.compress(original) + result = brotli.decompress(compressed) + assert result == original + + +def test_decompress_with_limit_at_exact_boundary(): + """Decompression succeeds when limit equals the decompressed size.""" + original = b'C' * (64 * 1024) + compressed = brotli.compress(original) + # The internal buffer allocates in power-of-2 blocks, so the total allocated + # will be >= len(original). Setting limit to len(original) should succeed + # since total_allocated is checked against the limit only when the buffer + # needs to grow beyond the current allocation. + result = brotli.decompress(compressed, output_buffer_limit=len(original)) + assert result == original + + +def test_decompress_with_limit_just_below_output(): + """Decompression raises when limit is below the required output size.""" + original = b'\x00' * (512 * 1024) # 512 KB + compressed = brotli.compress(original, quality=1) + # Use half the output size as limit; the buffer grows in power-of-2 blocks, + # so the limit must be below a block boundary to trigger. + with pytest.raises(brotli.error): + brotli.decompress(compressed, output_buffer_limit=len(original) // 2) + + +@pytest.mark.parametrize( + 'compressed_name, original_name', _test_utils.gather_compressed_inputs() +) +def test_decompress_with_large_limit(compressed_name, original_name): + """All test vectors decompress successfully with a generous limit.""" + compressed = _test_utils.take_input(compressed_name) + original = _test_utils.take_input(original_name) + # Use a limit of at least 1 MB or 2x the original size, whichever is larger. + generous_limit = max(len(original) * 2, 1024 * 1024) + result = brotli.decompress(compressed, output_buffer_limit=generous_limit) + assert result == original + + +def test_decompress_limit_error_message(): + """The error message for limit exceeded is specific and informative.""" + original = b'\x00' * (1024 * 1024) + compressed = brotli.compress(original, quality=1) + with pytest.raises(brotli.error, match='output_buffer_limit'): + brotli.decompress(compressed, output_buffer_limit=1024) + + +def test_decompress_limit_keyword_only(): + """output_buffer_limit can be passed as a keyword argument.""" + original = b'test data for keyword arg' + compressed = brotli.compress(original) + result = brotli.decompress(compressed, output_buffer_limit=1024) + assert result == original + + +def test_decompress_with_negative_limit_is_unlimited(): + """Negative output_buffer_limit is treated as unlimited (no limit).""" + original = b'D' * (256 * 1024) + compressed = brotli.compress(original) + result = brotli.decompress(compressed, output_buffer_limit=-1) + assert result == original