[Python] Add output_buffer_limit to brotli.decompress()

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;

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