# Copyright 2019 by Kurt Griffiths
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""ASGI BoundedStream class."""
from falcon.errors import OperationNotAllowed
__all__ = ['BoundedStream']
[docs]
class BoundedStream:
"""File-like input object for reading the body of the request, if any.
This class implements coroutine functions for asynchronous reading or
iteration, but otherwise provides an interface similar to that defined by
:class:`io.IOBase`.
If the request includes a Content-Length header, the number of bytes in the
stream will be truncated to the length specified by the header. Otherwise,
the stream will yield data until the ASGI server indicates that no more
bytes are available.
For large request bodies, the preferred method of using the stream object is
as an asynchronous iterator. In this mode, each body chunk is simply yielded
in its entirety, as it is received from the ASGI server. Because no data is
buffered by the framework, this is the most memory-efficient way of reading
the request body::
# If the request body is empty or has already be consumed, the iteration
# will immediately stop without yielding any data chunks. Otherwise, a
# series of byte # strings will be yielded until the entire request
# body has been yielded or the client disconnects.
async for data_chunk in req.stream
pass
The stream object also supports asynchronous ``read()`` and
``readall()`` methods::
# Read all of the data at once; use only when you are confident
# that the request body is small enough to not eat up all of
# your memory. For small bodies, this is the most performant
# option.
data = await req.stream.readall()
# ...or call read() without arguments
data = await req.stream.read()
# ...or read the data in chunks. You may choose to read more
# or less than 32 KiB as shown in this example. But note that
# this approach will generally be less efficient as compared
# to async iteration, resulting in more usage and
# copying of memory.
while True:
data_chunk = await req.stream.read(32 * 1024)
if not data_chunk:
break
Warning:
Apps may not use both ``read()`` and the asynchronous iterator
interface to consume the same request body; the only time that
it is safe to do so is when one or the other method is used to
completely read the entire body *before* the other method is
even attempted. Therefore, it is important to always call
:meth:`~.exhaust` or :meth:`~.close` if a body has only been
partially read and the remaining data is to be ignored.
Note:
The stream object provides a convenient abstraction over the series of
body chunks contained in any ASGI "http.request" events received by the
app. As such, some request body data may be temporarily buffered in
memory during and between calls to read from the stream. The framework
has been designed to minimize the amount of data that must be buffered
in this manner.
Args:
receive (awaitable): ASGI awaitable callable that will yield a new
request event dictionary when one is available.
Keyword Args:
first_event (dict): First ASGI event received from the client,
if one was preloaded (default ``None``).
content_length (int): Expected content length of the stream, derived
from the Content-Length header in the request (if available).
"""
__slots__ = [
'_buffer',
'_bytes_remaining',
'_closed',
'_iteration_started',
'_pos',
'_receive',
]
def __init__(self, receive, first_event=None, content_length=None):
self._closed = False
self._iteration_started = False
self._receive = receive
# NOTE(kgriffs): Outside of testing, first_event will always be set
# and we also assume a body is expected, otherwise why bother
# creating a stream object to read it? But just in case this
# object is created in other cases, use "in" here rather than
# EAFP.
if first_event and 'body' in first_event:
first_chunk = first_event['body']
else:
first_chunk = b''
if content_length is None:
self._buffer = first_chunk
# NOTE(kgriffs): If length is unknown we just set remaining bytes
# to a ridiculously high number so that we will keep reading
# until we get an event with more_body == False. We do not
# use sys.maxsize because 2**31 on 32-bit systems is not
# a large enough number (someone may have an API that accepts
# multi-GB payloads).
self._bytes_remaining = 2**63
else:
if len(first_chunk) > content_length:
self._buffer = first_chunk[:content_length]
else:
self._buffer = first_chunk
self._bytes_remaining = content_length - len(self._buffer)
self._pos = len(self._buffer)
if first_event and self._bytes_remaining:
# NOTE(kgriffs): Override if the event says there's no more data
if not ('more_body' in first_event and first_event['more_body']):
self._bytes_remaining = 0
def __aiter__(self):
# NOTE(kgriffs): This returns an async generator, but that's OK because
# it also implements the iterator protocol defined in PEP 492, albeit
# in a more efficient way than a regular async iterator.
return self._iter_content()
# -------------------------------------------------------------------------
# These methods are included to improve compatibility with Python's
# standard "file-like" IO interface.
# -------------------------------------------------------------------------
# NOTE(kgriffs): According to the Python docs, NotImplementedError is not
# meant to be used to mean "not supported"; rather, the method should
# just be left undefined; hence we do not implement readline(),
# readlines(), __iter__(), __next__(), flush(), seek(),
# truncate(), __del__().
[docs]
def fileno(self):
"""Raise an instance of OSError since a file descriptor is not used."""
raise OSError('This IO object does not use a file descriptor')
[docs]
def isatty(self):
"""Return ``False`` always."""
return False
[docs]
def readable(self):
"""Return ``True`` always."""
return True
[docs]
def seekable(self):
"""Return ``False`` always."""
return False
[docs]
def writable(self):
"""Return ``False`` always."""
return False
[docs]
def tell(self):
"""Return the number of bytes read from the stream so far."""
return self._pos
@property
def closed(self):
return self._closed
# -------------------------------------------------------------------------
@property
def eof(self):
return not self._buffer and self._bytes_remaining == 0
[docs]
def close(self):
"""Clear any buffered data and close this stream.
Once the stream is closed, any operation on it will
raise an instance of :class:`ValueError`.
As a convenience, it is allowed to call this method more than
once; only the first call, however, will have an effect.
"""
if not self._closed:
self._buffer = b''
self._bytes_remaining = 0
self._closed = True
[docs]
async def exhaust(self):
"""Consume and immediately discard any remaining data in the stream."""
if self._closed:
raise ValueError(
'This stream is closed; no further operations on it are permitted.'
)
self._buffer = b''
while self._bytes_remaining > 0:
event = await self._receive()
if event['type'] == 'http.disconnect':
self._bytes_remaining = 0
else:
try:
num_bytes = len(event['body'])
except KeyError:
# NOTE(kgriffs): The ASGI spec states that 'body' is optional.
num_bytes = 0
self._bytes_remaining -= num_bytes
self._pos += num_bytes
if not ('more_body' in event and event['more_body']):
self._bytes_remaining = 0
# Immediately dereference the data so it can be discarded ASAP
event = None
# NOTE(kgriffs): Ensure that if we read more than expected, this
# value is normalized to zero.
self._bytes_remaining = 0
[docs]
async def readall(self):
"""Read and return all remaining data in the request body.
Warning:
Only use this method when you can be certain that you have
enough free memory for the entire request body, and that you
have configured your web server to limit request bodies to a
reasonable size (to guard against malicious requests).
Returns:
bytes: The request body data, or ``b''`` if the body is empty or
has already been consumed.
"""
if self._closed:
raise OperationNotAllowed(
'This stream is closed; no further operations on it are permitted.'
)
if self.eof:
return b''
if self._buffer:
next_chunk = self._buffer
self._buffer = b''
chunks = [next_chunk]
else:
chunks = []
while self._bytes_remaining > 0:
event = await self._receive()
# PERF(kgriffs): Use try..except because we normally expect the
# 'body' key to be present.
try:
next_chunk = event['body']
except KeyError:
pass
else:
next_chunk_len = len(next_chunk)
if next_chunk_len <= self._bytes_remaining:
chunks.append(next_chunk)
self._bytes_remaining -= next_chunk_len
else:
# NOTE(kgriffs): Do not read more data than we are
# expecting. This *should* never happen if the
# server enforces the content-length header, but
# it is better to be safe than sorry.
chunks.append(next_chunk[: self._bytes_remaining])
self._bytes_remaining = 0
# NOTE(kgriffs): This also handles the case of receiving
# the event: {'type': 'http.disconnect'}
if not ('more_body' in event and event['more_body']):
self._bytes_remaining = 0
data = chunks[0] if len(chunks) == 1 else b''.join(chunks)
self._pos += len(data)
return data
[docs]
async def read(self, size=None):
"""Read some or all of the remaining bytes in the request body.
Warning:
A size should always be specified, unless you can be certain that
you have enough free memory for the entire request body, and that
you have configured your web server to limit request bodies to a
reasonable size (to guard against malicious requests).
Warning:
Apps may not use both ``read()`` and the asynchronous iterator
interface to consume the same request body; the only time that
it is safe to do so is when one or the other method is used to
completely read the entire body *before* the other method is
even attempted. Therefore, it is important to always call
:meth:`~.exhaust` or :meth:`~.close` if a body has only been
partially read and the remaining data is to be ignored.
Keyword Args:
size (int): The maximum number of bytes to read. The actual
amount of data that can be read will depend on how much is
available, and may be smaller than the amount requested. If the
size is -1 or not specified, all remaining data is read and
returned.
Returns:
bytes: The request body data, or ``b''`` if the body is empty or
has already been consumed.
"""
if self._closed:
raise OperationNotAllowed(
'This stream is closed; no further operations on it are permitted.'
)
if self.eof:
return b''
if size is None or size == -1:
return await self.readall()
if size <= 0:
return b''
if self._buffer:
num_bytes_available = len(self._buffer)
chunks = [self._buffer]
else:
num_bytes_available = 0
chunks = []
while self._bytes_remaining > 0 and num_bytes_available < size:
event = await self._receive()
# PERF(kgriffs): Use try..except because we normally expect the
# 'body' key to be present.
try:
next_chunk = event['body']
except KeyError:
pass
else:
next_chunk_len = len(next_chunk)
if next_chunk_len <= self._bytes_remaining:
chunks.append(next_chunk)
self._bytes_remaining -= next_chunk_len
num_bytes_available += next_chunk_len
else:
# NOTE(kgriffs): Do not read more data than we are
# expecting. This *should* never happen, but better
# safe than sorry.
chunks.append(next_chunk[: self._bytes_remaining])
self._bytes_remaining = 0
num_bytes_available += self._bytes_remaining
# NOTE(kgriffs): This also handles the case of receiving
# the event: {'type': 'http.disconnect'}
if not ('more_body' in event and event['more_body']):
self._bytes_remaining = 0
self._buffer = chunks[0] if len(chunks) == 1 else b''.join(chunks)
if num_bytes_available <= size:
data = self._buffer
self._buffer = b''
else:
data = self._buffer[:size]
self._buffer = self._buffer[size:]
self._pos += len(data)
return data
async def _iter_content(self):
if self._closed:
raise OperationNotAllowed(
'This stream is closed; no further operations on it are permitted.'
)
if self.eof:
return
if self._iteration_started:
raise OperationNotAllowed('This stream is already being iterated over.')
self._iteration_started = True
if self._buffer:
next_chunk = self._buffer
self._buffer = b''
self._pos += len(next_chunk)
yield next_chunk
while self._bytes_remaining > 0:
event = await self._receive()
# PERF(kgriffs): Use try...except because we normally expect the
# 'body' key to be present.
try:
next_chunk = event['body']
except KeyError:
pass
else:
# NOTE(kgriffs): No need to yield empty body chunks.
if next_chunk:
next_chunk_len = len(next_chunk)
if next_chunk_len <= self._bytes_remaining:
self._bytes_remaining -= next_chunk_len
self._pos += next_chunk_len
else:
# NOTE(kgriffs): We received more data than expected,
# so truncate to the expected length.
next_chunk = next_chunk[: self._bytes_remaining]
self._pos += self._bytes_remaining
self._bytes_remaining = 0
yield next_chunk
# NOTE(kgriffs): Per the ASGI spec, more_body is optional
# and should be considered False if not present.
# NOTE(kgriffs): This also handles the case of receiving
# the event: {'type': 'http.disconnect'}
# PERF(kgriffs): event.get() is more elegant, but uses a
# few more CPU cycles.
if not ('more_body' in event and event['more_body']):
self._bytes_remaining = 0