Patched code sample

import zlib
import struct
from io import BytesIO

from thrift.transport.TTransport import TTransportBase, TTransportException

MAX_UNCOMPRESSED_DATA_SIZE = 100 * 1024 * 1024  # 100MB


class TZlibTransport(TTransportBase):
    """Transport that wraps another, compressing/decompressing with zlib."""

    def __init__(self, trans):
        self._trans = trans
        self._rbuf = BytesIO()

    def _read_frame(self):
        """Reads a single zlib-compressed frame from the underlying transport."""
        header = self._trans.readAll(4)
        (sz,) = struct.unpack("!i", header)
        if sz < 0:
            raise TTransportException(TTransportException.NEGATIVE_SIZE)

        compressed = self._trans.readAll(sz)

        # FIX: Enforce a maximum size on decompressed data using max_length to
        # prevent a decompression bomb from exhausting system memory.
        try:
            decompressor = zlib.decompressobj()
            uncompressed = decompressor.decompress(compressed, MAX_UNCOMPRESSED_DATA_SIZE)
            if decompressor.unconsumed_tail or decompressor.flush():
                raise zlib.error("Compressed data exceeds size limit")
        except zlib.error as e:
            raise TTransportException(TTransportException.SIZE_LIMIT, f"Decompression error: {e}")

        self._rbuf = BytesIO(uncompressed)

    # other methods omitted for brevity