1 files changed, 308 insertions, 0 deletions
diff --git a/src/waitress/buffers.py b/src/waitress/buffers.py
new file mode 100644
index 0000000..04f6b42
--- /dev/null
+++ b/src/waitress/buffers.py
@@ -0,0 +1,308 @@
+##############################################################################
+#
+# Copyright (c) 2001-2004 Zope Foundation and Contributors.
+# All Rights Reserved.
+#
+# This software is subject to the provisions of the Zope Public License,
+# Version 2.1 (ZPL).  A copy of the ZPL should accompany this distribution.
+# THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED
+# WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS
+# FOR A PARTICULAR PURPOSE.
+#
+##############################################################################
+"""Buffers
+"""
+from io import BytesIO
+
+# copy_bytes controls the size of temp. strings for shuffling data around.
+COPY_BYTES = 1 << 18  # 256K
+
+# The maximum number of bytes to buffer in a simple string.
+STRBUF_LIMIT = 8192
+
+
+class FileBasedBuffer(object):
+
+    remain = 0
+
+    def __init__(self, file, from_buffer=None):
+        self.file = file
+        if from_buffer is not None:
+            from_file = from_buffer.getfile()
+            read_pos = from_file.tell()
+            from_file.seek(0)
+            while True:
+                data = from_file.read(COPY_BYTES)
+                if not data:
+                    break
+                file.write(data)
+            self.remain = int(file.tell() - read_pos)
+            from_file.seek(read_pos)
+            file.seek(read_pos)
+
+    def __len__(self):
+        return self.remain
+
+    def __nonzero__(self):
+        return True
+
+    __bool__ = __nonzero__  # py3
+
+    def append(self, s):
+        file = self.file
+        read_pos = file.tell()
+        file.seek(0, 2)
+        file.write(s)
+        file.seek(read_pos)
+        self.remain = self.remain + len(s)
+
+    def get(self, numbytes=-1, skip=False):
+        file = self.file
+        if not skip:
+            read_pos = file.tell()
+        if numbytes < 0:
+            # Read all
+            res = file.read()
+        else:
+            res = file.read(numbytes)
+        if skip:
+            self.remain -= len(res)
+        else:
+            file.seek(read_pos)
+        return res
+
+    def skip(self, numbytes, allow_prune=0):
+        if self.remain < numbytes:
+            raise ValueError(
+                "Can't skip %d bytes in buffer of %d bytes" % (numbytes, self.remain)
+            )
+        self.file.seek(numbytes, 1)
+        self.remain = self.remain - numbytes
+
+    def newfile(self):
+        raise NotImplementedError()
+
+    def prune(self):
+        file = self.file
+        if self.remain == 0:
+            read_pos = file.tell()
+            file.seek(0, 2)
+            sz = file.tell()
+            file.seek(read_pos)
+            if sz == 0:
+                # Nothing to prune.
+                return
+        nf = self.newfile()
+        while True:
+            data = file.read(COPY_BYTES)
+            if not data:
+                break
+            nf.write(data)
+        self.file = nf
+
+    def getfile(self):
+        return self.file
+
+    def close(self):
+        if hasattr(self.file, "close"):
+            self.file.close()
+        self.remain = 0
+
+
+class TempfileBasedBuffer(FileBasedBuffer):
+    def __init__(self, from_buffer=None):
+        FileBasedBuffer.__init__(self, self.newfile(), from_buffer)
+
+    def newfile(self):
+        from tempfile import TemporaryFile
+
+        return TemporaryFile("w+b")
+
+
+class BytesIOBasedBuffer(FileBasedBuffer):
+    def __init__(self, from_buffer=None):
+        if from_buffer is not None:
+            FileBasedBuffer.__init__(self, BytesIO(), from_buffer)
+        else:
+            # Shortcut. :-)
+            self.file = BytesIO()
+
+    def newfile(self):
+        return BytesIO()
+
+
+def _is_seekable(fp):
+    if hasattr(fp, "seekable"):
+        return fp.seekable()
+    return hasattr(fp, "seek") and hasattr(fp, "tell")
+
+
+class ReadOnlyFileBasedBuffer(FileBasedBuffer):
+    # used as wsgi.file_wrapper
+
+    def __init__(self, file, block_size=32768):
+        self.file = file
+        self.block_size = block_size  # for __iter__
+
+    def prepare(self, size=None):
+        if _is_seekable(self.file):
+            start_pos = self.file.tell()
+            self.file.seek(0, 2)
+            end_pos = self.file.tell()
+            self.file.seek(start_pos)
+            fsize = end_pos - start_pos
+            if size is None:
+                self.remain = fsize
+            else:
+                self.remain = min(fsize, size)
+        return self.remain
+
+    def get(self, numbytes=-1, skip=False):
+        # never read more than self.remain (it can be user-specified)
+        if numbytes == -1 or numbytes > self.remain:
+            numbytes = self.remain
+        file = self.file
+        if not skip:
+            read_pos = file.tell()
+        res = file.read(numbytes)
+        if skip:
+            self.remain -= len(res)
+        else:
+            file.seek(read_pos)
+        return res
+
+    def __iter__(self):  # called by task if self.filelike has no seek/tell
+        return self
+
+    def next(self):
+        val = self.file.read(self.block_size)
+        if not val:
+            raise StopIteration
+        return val
+
+    __next__ = next  # py3
+
+    def append(self, s):
+        raise NotImplementedError
+
+
+class OverflowableBuffer(object):
+    """
+    This buffer implementation has four stages:
+    - No data
+    - Bytes-based buffer
+    - BytesIO-based buffer
+    - Temporary file storage
+    The first two stages are fastest for simple transfers.
+    """
+
+    overflowed = False
+    buf = None
+    strbuf = b""  # Bytes-based buffer.
+
+    def __init__(self, overflow):
+        # overflow is the maximum to be stored in a StringIO buffer.
+        self.overflow = overflow
+
+    def __len__(self):
+        buf = self.buf
+        if buf is not None:
+            # use buf.__len__ rather than len(buf) FBO of not getting
+            # OverflowError on Python 2
+            return buf.__len__()
+        else:
+            return self.strbuf.__len__()
+
+    def __nonzero__(self):
+        # use self.__len__ rather than len(self) FBO of not getting
+        # OverflowError on Python 2
+        return self.__len__() > 0
+
+    __bool__ = __nonzero__  # py3
+
+    def _create_buffer(self):
+        strbuf = self.strbuf
+        if len(strbuf) >= self.overflow:
+            self._set_large_buffer()
+        else:
+            self._set_small_buffer()
+        buf = self.buf
+        if strbuf:
+            buf.append(self.strbuf)
+            self.strbuf = b""
+        return buf
+
+    def _set_small_buffer(self):
+        self.buf = BytesIOBasedBuffer(self.buf)
+        self.overflowed = False
+
+    def _set_large_buffer(self):
+        self.buf = TempfileBasedBuffer(self.buf)
+        self.overflowed = True
+
+    def append(self, s):
+        buf = self.buf
+        if buf is None:
+            strbuf = self.strbuf
+            if len(strbuf) + len(s) < STRBUF_LIMIT:
+                self.strbuf = strbuf + s
+                return
+            buf = self._create_buffer()
+        buf.append(s)
+        # use buf.__len__ rather than len(buf) FBO of not getting
+        # OverflowError on Python 2
+        sz = buf.__len__()
+        if not self.overflowed:
+            if sz >= self.overflow:
+                self._set_large_buffer()
+
+    def get(self, numbytes=-1, skip=False):
+        buf = self.buf
+        if buf is None:
+            strbuf = self.strbuf
+            if not skip:
+                return strbuf
+            buf = self._create_buffer()
+        return buf.get(numbytes, skip)
+
+    def skip(self, numbytes, allow_prune=False):
+        buf = self.buf
+        if buf is None:
+            if allow_prune and numbytes == len(self.strbuf):
+                # We could slice instead of converting to
+                # a buffer, but that would eat up memory in
+                # large transfers.
+                self.strbuf = b""
+                return
+            buf = self._create_buffer()
+        buf.skip(numbytes, allow_prune)
+
+    def prune(self):
+        """
+        A potentially expensive operation that removes all data
+        already retrieved from the buffer.
+        """
+        buf = self.buf
+        if buf is None:
+            self.strbuf = b""
+            return
+        buf.prune()
+        if self.overflowed:
+            # use buf.__len__ rather than len(buf) FBO of not getting
+            # OverflowError on Python 2
+            sz = buf.__len__()
+            if sz < self.overflow:
+                # Revert to a faster buffer.
+                self._set_small_buffer()
+
+    def getfile(self):
+        buf = self.buf
+        if buf is None:
+            buf = self._create_buffer()
+        return buf.getfile()
+
+    def close(self):
+        buf = self.buf
+        if buf is not None:
+            buf.close()