bpo-41486: Fix initial buffer size can't > UINT32_MAX in zlib module (GH-25738)

* Fix initial buffer size can't > UINT32_MAX in zlib module

After commit f9bedb630e8a0b7d94e1c7e609b20dfaa2b22231, in 64-bit build,
if the initial buffer size > UINT32_MAX, ValueError will be raised.

These two functions are affected:
1. zlib.decompress(data, /, wbits=MAX_WBITS, bufsize=DEF_BUF_SIZE)
2. zlib.Decompress.flush([length])

This commit re-allows the size > UINT32_MAX.

* adds curly braces per PEP 7.

* Renames `Buffer_*` to `OutputBuffer_*` for clarity

3 files changed, 77 insertions, 71 deletions

diff --git a/Modules/_bz2module.c b/Modules/_bz2module.c
index 9893a63..d75bb32 100644
--- a/Modules/_bz2module.c
+++ b/Modules/_bz2module.c
@@ -18,8 +18,8 @@
 /* On success, return value >= 0
    On failure, return -1 */
 static inline Py_ssize_t
-Buffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
-                   char **next_out, uint32_t *avail_out)
+OutputBuffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
+                         char **next_out, uint32_t *avail_out)
 {
     Py_ssize_t allocated;
 
@@ -32,8 +32,8 @@ Buffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
 /* On success, return value >= 0
    On failure, return -1 */
 static inline Py_ssize_t
-Buffer_Grow(_BlocksOutputBuffer *buffer,
-            char **next_out, uint32_t *avail_out)
+OutputBuffer_Grow(_BlocksOutputBuffer *buffer,
+                  char **next_out, uint32_t *avail_out)
 {
     Py_ssize_t allocated;
 
@@ -44,19 +44,19 @@ Buffer_Grow(_BlocksOutputBuffer *buffer,
 }
 
 static inline Py_ssize_t
-Buffer_GetDataSize(_BlocksOutputBuffer *buffer, uint32_t avail_out)
+OutputBuffer_GetDataSize(_BlocksOutputBuffer *buffer, uint32_t avail_out)
 {
     return _BlocksOutputBuffer_GetDataSize(buffer, (Py_ssize_t) avail_out);
 }
 
 static inline PyObject *
-Buffer_Finish(_BlocksOutputBuffer *buffer, uint32_t avail_out)
+OutputBuffer_Finish(_BlocksOutputBuffer *buffer, uint32_t avail_out)
 {
     return _BlocksOutputBuffer_Finish(buffer, (Py_ssize_t) avail_out);
 }
 
 static inline void
-Buffer_OnError(_BlocksOutputBuffer *buffer)
+OutputBuffer_OnError(_BlocksOutputBuffer *buffer)
 {
     _BlocksOutputBuffer_OnError(buffer);
 }
@@ -177,7 +177,7 @@ compress(BZ2Compressor *c, char *data, size_t len, int action)
     PyObject *result;
     _BlocksOutputBuffer buffer = {.list = NULL};
 
-    if (Buffer_InitAndGrow(&buffer, -1, &c->bzs.next_out, &c->bzs.avail_out) < 0) {
+    if (OutputBuffer_InitAndGrow(&buffer, -1, &c->bzs.next_out, &c->bzs.avail_out) < 0) {
         goto error;
     }
     c->bzs.next_in = data;
@@ -198,7 +198,7 @@ compress(BZ2Compressor *c, char *data, size_t len, int action)
             break;
 
         if (c->bzs.avail_out == 0) {
-            if (Buffer_Grow(&buffer, &c->bzs.next_out, &c->bzs.avail_out) < 0) {
+            if (OutputBuffer_Grow(&buffer, &c->bzs.next_out, &c->bzs.avail_out) < 0) {
                 goto error;
             }
         }
@@ -215,13 +215,13 @@ compress(BZ2Compressor *c, char *data, size_t len, int action)
             break;
     }
 
-    result = Buffer_Finish(&buffer, c->bzs.avail_out);
+    result = OutputBuffer_Finish(&buffer, c->bzs.avail_out);
     if (result != NULL) {
         return result;
     }
 
 error:
-    Buffer_OnError(&buffer);
+    OutputBuffer_OnError(&buffer);
     return NULL;
 }
 
@@ -442,7 +442,7 @@ decompress_buf(BZ2Decompressor *d, Py_ssize_t max_length)
     _BlocksOutputBuffer buffer = {.list = NULL};
     bz_stream *bzs = &d->bzs;
 
-    if (Buffer_InitAndGrow(&buffer, max_length, &bzs->next_out, &bzs->avail_out) < 0) {
+    if (OutputBuffer_InitAndGrow(&buffer, max_length, &bzs->next_out, &bzs->avail_out) < 0) {
         goto error;
     }
 
@@ -469,21 +469,22 @@ decompress_buf(BZ2Decompressor *d, Py_ssize_t max_length)
         } else if (d->bzs_avail_in_real == 0) {
             break;
         } else if (bzs->avail_out == 0) {
-            if (Buffer_GetDataSize(&buffer, bzs->avail_out) == max_length)
+            if (OutputBuffer_GetDataSize(&buffer, bzs->avail_out) == max_length) {
                 break;
-            if (Buffer_Grow(&buffer, &bzs->next_out, &bzs->avail_out) < 0) {
+            }
+            if (OutputBuffer_Grow(&buffer, &bzs->next_out, &bzs->avail_out) < 0) {
                 goto error;
             }
         }
     }
 
-    result = Buffer_Finish(&buffer, bzs->avail_out);
+    result = OutputBuffer_Finish(&buffer, bzs->avail_out);
     if (result != NULL) {
         return result;
     }
 
 error:
-    Buffer_OnError(&buffer);
+    OutputBuffer_OnError(&buffer);
     return NULL;
 }
 
diff --git a/Modules/_lzmamodule.c b/Modules/_lzmamodule.c
index 0d62319..2f80bf0 100644
--- a/Modules/_lzmamodule.c
+++ b/Modules/_lzmamodule.c
@@ -25,8 +25,8 @@
 /* On success, return value >= 0
    On failure, return -1 */
 static inline Py_ssize_t
-Buffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
-                   uint8_t **next_out, size_t *avail_out)
+OutputBuffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
+                         uint8_t **next_out, size_t *avail_out)
 {
     Py_ssize_t allocated;
 
@@ -39,8 +39,8 @@ Buffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
 /* On success, return value >= 0
    On failure, return -1 */
 static inline Py_ssize_t
-Buffer_Grow(_BlocksOutputBuffer *buffer,
-            uint8_t **next_out, size_t *avail_out)
+OutputBuffer_Grow(_BlocksOutputBuffer *buffer,
+                  uint8_t **next_out, size_t *avail_out)
 {
     Py_ssize_t allocated;
 
@@ -51,19 +51,19 @@ Buffer_Grow(_BlocksOutputBuffer *buffer,
 }
 
 static inline Py_ssize_t
-Buffer_GetDataSize(_BlocksOutputBuffer *buffer, size_t avail_out)
+OutputBuffer_GetDataSize(_BlocksOutputBuffer *buffer, size_t avail_out)
 {
     return _BlocksOutputBuffer_GetDataSize(buffer, (Py_ssize_t) avail_out);
 }
 
 static inline PyObject *
-Buffer_Finish(_BlocksOutputBuffer *buffer, size_t avail_out)
+OutputBuffer_Finish(_BlocksOutputBuffer *buffer, size_t avail_out)
 {
     return _BlocksOutputBuffer_Finish(buffer, (Py_ssize_t) avail_out);
 }
 
 static inline void
-Buffer_OnError(_BlocksOutputBuffer *buffer)
+OutputBuffer_OnError(_BlocksOutputBuffer *buffer)
 {
     _BlocksOutputBuffer_OnError(buffer);
 }
@@ -550,7 +550,7 @@ compress(Compressor *c, uint8_t *data, size_t len, lzma_action action)
     _lzma_state *state = PyType_GetModuleState(Py_TYPE(c));
     assert(state != NULL);
 
-    if (Buffer_InitAndGrow(&buffer, -1, &c->lzs.next_out, &c->lzs.avail_out) < 0) {
+    if (OutputBuffer_InitAndGrow(&buffer, -1, &c->lzs.next_out, &c->lzs.avail_out) < 0) {
         goto error;
     }
     c->lzs.next_in = data;
@@ -573,19 +573,19 @@ compress(Compressor *c, uint8_t *data, size_t len, lzma_action action)
             (action == LZMA_FINISH && lzret == LZMA_STREAM_END)) {
             break;
         } else if (c->lzs.avail_out == 0) {
-            if (Buffer_Grow(&buffer, &c->lzs.next_out, &c->lzs.avail_out) < 0) {
+            if (OutputBuffer_Grow(&buffer, &c->lzs.next_out, &c->lzs.avail_out) < 0) {
                 goto error;
             }
         }
     }
 
-    result = Buffer_Finish(&buffer, c->lzs.avail_out);
+    result = OutputBuffer_Finish(&buffer, c->lzs.avail_out);
     if (result != NULL) {
         return result;
     }
 
 error:
-    Buffer_OnError(&buffer);
+    OutputBuffer_OnError(&buffer);
     return NULL;
 }
 
@@ -934,7 +934,7 @@ decompress_buf(Decompressor *d, Py_ssize_t max_length)
     _lzma_state *state = PyType_GetModuleState(Py_TYPE(d));
     assert(state != NULL);
 
-    if (Buffer_InitAndGrow(&buffer, max_length, &lzs->next_out, &lzs->avail_out) < 0) {
+    if (OutputBuffer_InitAndGrow(&buffer, max_length, &lzs->next_out, &lzs->avail_out) < 0) {
         goto error;
     }
 
@@ -962,10 +962,10 @@ decompress_buf(Decompressor *d, Py_ssize_t max_length)
                Maybe lzs's internal state still have a few bytes
                can be output, grow the output buffer and continue
                if max_lengh < 0. */
-            if (Buffer_GetDataSize(&buffer, lzs->avail_out) == max_length) {
+            if (OutputBuffer_GetDataSize(&buffer, lzs->avail_out) == max_length) {
                 break;
             }
-            if (Buffer_Grow(&buffer, &lzs->next_out, &lzs->avail_out) < 0) {
+            if (OutputBuffer_Grow(&buffer, &lzs->next_out, &lzs->avail_out) < 0) {
                 goto error;
             }
         } else if (lzs->avail_in == 0) {
@@ -973,13 +973,13 @@ decompress_buf(Decompressor *d, Py_ssize_t max_length)
         }
     }
 
-    result = Buffer_Finish(&buffer, lzs->avail_out);
+    result = OutputBuffer_Finish(&buffer, lzs->avail_out);
     if (result != NULL) {
         return result;
     }
 
 error:
-    Buffer_OnError(&buffer);
+    OutputBuffer_OnError(&buffer);
     return NULL;
 }
 
diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c
index fc63ca9..fe06094 100644
--- a/Modules/zlibmodule.c
+++ b/Modules/zlibmodule.c
@@ -19,8 +19,8 @@
 /* On success, return value >= 0
    On failure, return -1 */
 static inline Py_ssize_t
-Buffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
-                   Bytef **next_out, uint32_t *avail_out)
+OutputBuffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
+                         Bytef **next_out, uint32_t *avail_out)
 {
     Py_ssize_t allocated;
 
@@ -33,15 +33,17 @@ Buffer_InitAndGrow(_BlocksOutputBuffer *buffer, Py_ssize_t max_length,
 /* On success, return value >= 0
    On failure, return -1 */
 static inline Py_ssize_t
-Buffer_InitWithSize(_BlocksOutputBuffer *buffer, Py_ssize_t init_size,
-                    Bytef **next_out, uint32_t *avail_out)
+OutputBuffer_InitWithSize(_BlocksOutputBuffer *buffer, Py_ssize_t init_size,
+                          Bytef **next_out, uint32_t *avail_out)
 {
     Py_ssize_t allocated;
 
-    if (init_size < 0 || (size_t)init_size > UINT32_MAX) {
-        PyErr_SetString(PyExc_ValueError,
-                        "Initial buffer size should (0 <= size <= UINT32_MAX)");
-        return -1;
+    if (init_size >= 0 &&  // ensure (size_t) cast is safe
+        (size_t)init_size > UINT32_MAX)
+    {
+        /* In 32-bit build, never reach this conditional branch.
+           The maximum block size accepted by zlib is UINT32_MAX. */
+        init_size = UINT32_MAX;
     }
 
     allocated = _BlocksOutputBuffer_InitWithSize(
@@ -53,8 +55,8 @@ Buffer_InitWithSize(_BlocksOutputBuffer *buffer, Py_ssize_t init_size,
 /* On success, return value >= 0
    On failure, return -1 */
 static inline Py_ssize_t
-Buffer_Grow(_BlocksOutputBuffer *buffer,
-            Bytef **next_out, uint32_t *avail_out)
+OutputBuffer_Grow(_BlocksOutputBuffer *buffer,
+                  Bytef **next_out, uint32_t *avail_out)
 {
     Py_ssize_t allocated;
 
@@ -65,19 +67,19 @@ Buffer_Grow(_BlocksOutputBuffer *buffer,
 }
 
 static inline Py_ssize_t
-Buffer_GetDataSize(_BlocksOutputBuffer *buffer, uint32_t avail_out)
+OutputBuffer_GetDataSize(_BlocksOutputBuffer *buffer, uint32_t avail_out)
 {
     return _BlocksOutputBuffer_GetDataSize(buffer, (Py_ssize_t) avail_out);
 }
 
 static inline PyObject *
-Buffer_Finish(_BlocksOutputBuffer *buffer, uint32_t avail_out)
+OutputBuffer_Finish(_BlocksOutputBuffer *buffer, uint32_t avail_out)
 {
     return _BlocksOutputBuffer_Finish(buffer, (Py_ssize_t) avail_out);
 }
 
 static inline void
-Buffer_OnError(_BlocksOutputBuffer *buffer)
+OutputBuffer_OnError(_BlocksOutputBuffer *buffer)
 {
     _BlocksOutputBuffer_OnError(buffer);
 }
@@ -248,7 +250,7 @@ zlib_compress_impl(PyObject *module, Py_buffer *data, int level)
     Byte *ibuf = data->buf;
     Py_ssize_t ibuflen = data->len;
 
-    if (Buffer_InitAndGrow(&buffer, -1, &zst.next_out, &zst.avail_out) < 0) {
+    if (OutputBuffer_InitAndGrow(&buffer, -1, &zst.next_out, &zst.avail_out) < 0) {
         goto error;
     }
 
@@ -280,7 +282,7 @@ zlib_compress_impl(PyObject *module, Py_buffer *data, int level)
 
         do {
             if (zst.avail_out == 0) {
-                if (Buffer_Grow(&buffer, &zst.next_out, &zst.avail_out) < 0) {
+                if (OutputBuffer_Grow(&buffer, &zst.next_out, &zst.avail_out) < 0) {
                     deflateEnd(&zst);
                     goto error;
                 }
@@ -304,7 +306,7 @@ zlib_compress_impl(PyObject *module, Py_buffer *data, int level)
 
     err = deflateEnd(&zst);
     if (err == Z_OK) {
-        RetVal = Buffer_Finish(&buffer, zst.avail_out);
+        RetVal = OutputBuffer_Finish(&buffer, zst.avail_out);
         if (RetVal == NULL) {
             goto error;
         }
@@ -313,7 +315,7 @@ zlib_compress_impl(PyObject *module, Py_buffer *data, int level)
     else
         zlib_error(state, zst, err, "while finishing compression");
  error:
-    Buffer_OnError(&buffer);
+    OutputBuffer_OnError(&buffer);
     return NULL;
 }
 
@@ -352,7 +354,7 @@ zlib_decompress_impl(PyObject *module, Py_buffer *data, int wbits,
         bufsize = 1;
     }
 
-    if (Buffer_InitWithSize(&buffer, bufsize, &zst.next_out, &zst.avail_out) < 0) {
+    if (OutputBuffer_InitWithSize(&buffer, bufsize, &zst.next_out, &zst.avail_out) < 0) {
         goto error;
     }
 
@@ -385,7 +387,7 @@ zlib_decompress_impl(PyObject *module, Py_buffer *data, int wbits,
 
         do {
             if (zst.avail_out == 0) {
-                if (Buffer_Grow(&buffer, &zst.next_out, &zst.avail_out) < 0) {
+                if (OutputBuffer_Grow(&buffer, &zst.next_out, &zst.avail_out) < 0) {
                     inflateEnd(&zst);
                     goto error;
                 }
@@ -428,13 +430,13 @@ zlib_decompress_impl(PyObject *module, Py_buffer *data, int wbits,
         goto error;
     }
 
-    RetVal = Buffer_Finish(&buffer, zst.avail_out);
+    RetVal = OutputBuffer_Finish(&buffer, zst.avail_out);
     if (RetVal != NULL) {
         return RetVal;
     }
 
  error:
-    Buffer_OnError(&buffer);
+    OutputBuffer_OnError(&buffer);
     return NULL;
 }
 
@@ -677,7 +679,7 @@ zlib_Compress_compress_impl(compobject *self, PyTypeObject *cls,
     self->zst.next_in = data->buf;
     Py_ssize_t ibuflen = data->len;
 
-    if (Buffer_InitAndGrow(&buffer, -1, &self->zst.next_out, &self->zst.avail_out) < 0) {
+    if (OutputBuffer_InitAndGrow(&buffer, -1, &self->zst.next_out, &self->zst.avail_out) < 0) {
         goto error;
     }
 
@@ -686,8 +688,9 @@ zlib_Compress_compress_impl(compobject *self, PyTypeObject *cls,
 
         do {
             if (self->zst.avail_out == 0) {
-                if (Buffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0)
+                if (OutputBuffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0) {
                     goto error;
+                }
             }
 
             Py_BEGIN_ALLOW_THREADS
@@ -704,13 +707,13 @@ zlib_Compress_compress_impl(compobject *self, PyTypeObject *cls,
 
     } while (ibuflen != 0);
 
-    RetVal = Buffer_Finish(&buffer, self->zst.avail_out);
+    RetVal = OutputBuffer_Finish(&buffer, self->zst.avail_out);
     if (RetVal != NULL) {
         goto success;
     }
 
  error:
-    Buffer_OnError(&buffer);
+    OutputBuffer_OnError(&buffer);
     RetVal = NULL;
  success:
     LEAVE_ZLIB(self);
@@ -799,15 +802,16 @@ zlib_Decompress_decompress_impl(compobject *self, PyTypeObject *cls,
     if (max_length < 0) {
         PyErr_SetString(PyExc_ValueError, "max_length must be non-negative");
         return NULL;
-    } else if (max_length == 0)
+    } else if (max_length == 0) {
         max_length = -1;
+    }
 
     ENTER_ZLIB(self);
 
     self->zst.next_in = data->buf;
     ibuflen = data->len;
 
-    if (Buffer_InitAndGrow(&buffer, max_length, &self->zst.next_out, &self->zst.avail_out) < 0) {
+    if (OutputBuffer_InitAndGrow(&buffer, max_length, &self->zst.next_out, &self->zst.avail_out) < 0) {
         goto abort;
     }
 
@@ -816,10 +820,10 @@ zlib_Decompress_decompress_impl(compobject *self, PyTypeObject *cls,
 
         do {
             if (self->zst.avail_out == 0) {
-                if (Buffer_GetDataSize(&buffer, self->zst.avail_out) == max_length) {
+                if (OutputBuffer_GetDataSize(&buffer, self->zst.avail_out) == max_length) {
                     goto save;
                 }
-                if (Buffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0) {
+                if (OutputBuffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0) {
                     goto abort;
                 }
             }
@@ -865,13 +869,13 @@ zlib_Decompress_decompress_impl(compobject *self, PyTypeObject *cls,
         goto abort;
     }
 
-    RetVal = Buffer_Finish(&buffer, self->zst.avail_out);
+    RetVal = OutputBuffer_Finish(&buffer, self->zst.avail_out);
     if (RetVal != NULL) {
         goto success;
     }
 
  abort:
-    Buffer_OnError(&buffer);
+    OutputBuffer_OnError(&buffer);
     RetVal = NULL;
  success:
     LEAVE_ZLIB(self);
@@ -911,13 +915,13 @@ zlib_Compress_flush_impl(compobject *self, PyTypeObject *cls, int mode)
 
     self->zst.avail_in = 0;
 
-    if (Buffer_InitAndGrow(&buffer, -1, &self->zst.next_out, &self->zst.avail_out) < 0) {
+    if (OutputBuffer_InitAndGrow(&buffer, -1, &self->zst.next_out, &self->zst.avail_out) < 0) {
         goto error;
     }
 
     do {
         if (self->zst.avail_out == 0) {
-            if (Buffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0) {
+            if (OutputBuffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0) {
                 goto error;
             }
         }
@@ -954,13 +958,13 @@ zlib_Compress_flush_impl(compobject *self, PyTypeObject *cls, int mode)
         goto error;
     }
 
-    RetVal = Buffer_Finish(&buffer, self->zst.avail_out);
+    RetVal = OutputBuffer_Finish(&buffer, self->zst.avail_out);
     if (RetVal != NULL) {
         goto success;
     }
 
 error:
-    Buffer_OnError(&buffer);
+    OutputBuffer_OnError(&buffer);
     RetVal = NULL;
 success:
     LEAVE_ZLIB(self);
@@ -1189,7 +1193,7 @@ zlib_Decompress_flush_impl(compobject *self, PyTypeObject *cls,
     self->zst.next_in = data.buf;
     ibuflen = data.len;
 
-    if (Buffer_InitWithSize(&buffer, length, &self->zst.next_out, &self->zst.avail_out) < 0) {
+    if (OutputBuffer_InitWithSize(&buffer, length, &self->zst.next_out, &self->zst.avail_out) < 0) {
         goto abort;
     }
 
@@ -1199,8 +1203,9 @@ zlib_Decompress_flush_impl(compobject *self, PyTypeObject *cls,
 
         do {
             if (self->zst.avail_out == 0) {
-                if (Buffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0)
+                if (OutputBuffer_Grow(&buffer, &self->zst.next_out, &self->zst.avail_out) < 0) {
                     goto abort;
+                }
             }
 
             Py_BEGIN_ALLOW_THREADS
@@ -1243,13 +1248,13 @@ zlib_Decompress_flush_impl(compobject *self, PyTypeObject *cls,
         }
     }
 
-    RetVal = Buffer_Finish(&buffer, self->zst.avail_out);
+    RetVal = OutputBuffer_Finish(&buffer, self->zst.avail_out);
     if (RetVal != NULL) {
         goto success;
     }
 
  abort:
-    Buffer_OnError(&buffer);
+    OutputBuffer_OnError(&buffer);
     RetVal = NULL;
  success:
     PyBuffer_Release(&data);