342 lines
12 KiB
Python
342 lines
12 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
"""
|
||
|
|
hyper/http20/stream
|
||
|
|
~~~~~~~~~~~~~~~~~~~
|
||
|
|
|
||
|
|
Objects that make up the stream-level abstraction of hyper's HTTP/2 support.
|
||
|
|
|
||
|
|
These objects are not expected to be part of the public HTTP/2 API: they're
|
||
|
|
intended purely for use inside hyper's HTTP/2 abstraction.
|
||
|
|
|
||
|
|
Conceptually, a single HTTP/2 connection is made up of many streams: each
|
||
|
|
stream is an independent, bi-directional sequence of HTTP headers and data.
|
||
|
|
Each stream is identified by a monotonically increasing integer, assigned to
|
||
|
|
the stream by the endpoint that initiated the stream.
|
||
|
|
"""
|
||
|
|
import vinetrimmer.vendor.h2.exceptions
|
||
|
|
|
||
|
|
from ..common.headers import HTTPHeaderMap
|
||
|
|
from .util import h2_safe_headers
|
||
|
|
import logging
|
||
|
|
|
||
|
|
log = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
# Define the largest chunk of data we'll send in one go. Realistically, we
|
||
|
|
# should take the MSS into account but that's pretty dull, so let's just say
|
||
|
|
# 1kB and call it a day.
|
||
|
|
MAX_CHUNK = 1024
|
||
|
|
|
||
|
|
|
||
|
|
class Stream(object):
|
||
|
|
"""
|
||
|
|
A single HTTP/2 stream.
|
||
|
|
|
||
|
|
A stream is an independent, bi-directional sequence of HTTP headers and
|
||
|
|
data. Each stream is identified by a single integer. From a HTTP
|
||
|
|
perspective, a stream _approximately_ matches a single request-response
|
||
|
|
pair.
|
||
|
|
"""
|
||
|
|
def __init__(self,
|
||
|
|
stream_id,
|
||
|
|
window_manager,
|
||
|
|
connection,
|
||
|
|
send_outstanding_data,
|
||
|
|
recv_cb,
|
||
|
|
close_cb):
|
||
|
|
self.stream_id = stream_id
|
||
|
|
self.headers = HTTPHeaderMap()
|
||
|
|
|
||
|
|
# Set to a key-value set of the response headers once their
|
||
|
|
# HEADERS..CONTINUATION frame sequence finishes.
|
||
|
|
self.response_headers = None
|
||
|
|
|
||
|
|
# Set to a key-value set of the response trailers once their
|
||
|
|
# HEADERS..CONTINUATION frame sequence finishes.
|
||
|
|
self.response_trailers = None
|
||
|
|
|
||
|
|
# A dict mapping the promised stream ID of a pushed resource to a
|
||
|
|
# key-value set of its request headers. Entries are added once their
|
||
|
|
# PUSH_PROMISE..CONTINUATION frame sequence finishes.
|
||
|
|
self.promised_headers = {}
|
||
|
|
|
||
|
|
# Unconsumed response data chunks. Empties after every call to _read().
|
||
|
|
self.data = []
|
||
|
|
|
||
|
|
# Whether the remote side has completed the stream.
|
||
|
|
self.remote_closed = False
|
||
|
|
|
||
|
|
# Whether we have closed the stream.
|
||
|
|
self.local_closed = False
|
||
|
|
|
||
|
|
# There are two flow control windows: one for data we're sending,
|
||
|
|
# one for data being sent to us.
|
||
|
|
self._in_window_manager = window_manager
|
||
|
|
|
||
|
|
# Save off a reference to the state machine wrapped with lock.
|
||
|
|
self._conn = connection
|
||
|
|
|
||
|
|
# Save off a data callback.
|
||
|
|
self._send_outstanding_data = send_outstanding_data
|
||
|
|
self._recv_cb = recv_cb
|
||
|
|
self._close_cb = close_cb
|
||
|
|
|
||
|
|
def add_header(self, name, value, replace=False):
|
||
|
|
"""
|
||
|
|
Adds a single HTTP header to the headers to be sent on the request.
|
||
|
|
"""
|
||
|
|
if not replace:
|
||
|
|
self.headers[name] = value
|
||
|
|
else:
|
||
|
|
self.headers.replace(name, value)
|
||
|
|
|
||
|
|
def send_headers(self, end_stream=False):
|
||
|
|
"""
|
||
|
|
Sends the complete saved header block on the stream.
|
||
|
|
"""
|
||
|
|
headers = self.get_headers()
|
||
|
|
with self._conn as conn:
|
||
|
|
conn.send_headers(self.stream_id, headers, end_stream)
|
||
|
|
self._send_outstanding_data()
|
||
|
|
|
||
|
|
if end_stream:
|
||
|
|
self.local_closed = True
|
||
|
|
|
||
|
|
def send_data(self, data, final):
|
||
|
|
"""
|
||
|
|
Send some data on the stream. If this is the end of the data to be
|
||
|
|
sent, the ``final`` flag _must_ be set to True. If no data is to be
|
||
|
|
sent, set ``data`` to ``None``.
|
||
|
|
"""
|
||
|
|
# Define a utility iterator for file objects.
|
||
|
|
def file_iterator(fobj):
|
||
|
|
while True:
|
||
|
|
data = fobj.read(MAX_CHUNK)
|
||
|
|
yield data
|
||
|
|
if len(data) < MAX_CHUNK:
|
||
|
|
break
|
||
|
|
|
||
|
|
# Build the appropriate iterator for the data, in chunks of CHUNK_SIZE.
|
||
|
|
if hasattr(data, 'read'):
|
||
|
|
chunks = file_iterator(data)
|
||
|
|
else:
|
||
|
|
chunks = (data[i:i+MAX_CHUNK]
|
||
|
|
for i in range(0, len(data), MAX_CHUNK))
|
||
|
|
|
||
|
|
for chunk in chunks:
|
||
|
|
self._send_chunk(chunk, final)
|
||
|
|
|
||
|
|
def _read(self, amt=None):
|
||
|
|
"""
|
||
|
|
Read data from the stream. Unlike a normal read behaviour, this
|
||
|
|
function returns _at least_ ``amt`` data, but may return more.
|
||
|
|
"""
|
||
|
|
def listlen(list):
|
||
|
|
return sum(map(len, list))
|
||
|
|
|
||
|
|
# Keep reading until the stream is closed or we get enough data.
|
||
|
|
while (not self.remote_closed and
|
||
|
|
(amt is None or listlen(self.data) < amt)):
|
||
|
|
self._recv_cb(stream_id=self.stream_id)
|
||
|
|
|
||
|
|
result = b''.join(self.data)
|
||
|
|
self.data = []
|
||
|
|
return result
|
||
|
|
|
||
|
|
def _read_one_frame(self):
|
||
|
|
"""
|
||
|
|
Reads a single data frame from the stream and returns it.
|
||
|
|
"""
|
||
|
|
# Keep reading until the stream is closed or we have a data frame.
|
||
|
|
while not self.remote_closed and not self.data:
|
||
|
|
self._recv_cb(stream_id=self.stream_id)
|
||
|
|
|
||
|
|
try:
|
||
|
|
return self.data.pop(0)
|
||
|
|
except IndexError:
|
||
|
|
return None
|
||
|
|
|
||
|
|
def receive_response(self, event):
|
||
|
|
"""
|
||
|
|
Receive response headers.
|
||
|
|
"""
|
||
|
|
# TODO: If this is called while we're still sending data, we may want
|
||
|
|
# to stop sending that data and check the response. Early responses to
|
||
|
|
# big uploads are almost always a problem.
|
||
|
|
self.response_headers = HTTPHeaderMap(event.headers)
|
||
|
|
|
||
|
|
def receive_trailers(self, event):
|
||
|
|
"""
|
||
|
|
Receive response trailers.
|
||
|
|
"""
|
||
|
|
self.response_trailers = HTTPHeaderMap(event.headers)
|
||
|
|
|
||
|
|
def receive_push(self, event):
|
||
|
|
"""
|
||
|
|
Receive the request headers for a pushed stream.
|
||
|
|
"""
|
||
|
|
self.promised_headers[event.pushed_stream_id] = event.headers
|
||
|
|
|
||
|
|
def receive_data(self, event):
|
||
|
|
"""
|
||
|
|
Receive a chunk of data.
|
||
|
|
"""
|
||
|
|
size = event.flow_controlled_length
|
||
|
|
increment = self._in_window_manager._handle_frame(size)
|
||
|
|
|
||
|
|
# Append the data to the buffer.
|
||
|
|
self.data.append(event.data)
|
||
|
|
|
||
|
|
if increment:
|
||
|
|
try:
|
||
|
|
with self._conn as conn:
|
||
|
|
conn.increment_flow_control_window(
|
||
|
|
increment, stream_id=self.stream_id
|
||
|
|
)
|
||
|
|
except vinetrimmer.vendor.h2.exceptions.StreamClosedError:
|
||
|
|
# We haven't got to it yet, but the stream is already
|
||
|
|
# closed. We don't need to increment the window in this
|
||
|
|
# case!
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
self._send_outstanding_data()
|
||
|
|
|
||
|
|
def receive_end_stream(self, event):
|
||
|
|
"""
|
||
|
|
All of the data is returned now.
|
||
|
|
"""
|
||
|
|
self.remote_closed = True
|
||
|
|
|
||
|
|
def receive_reset(self, event):
|
||
|
|
"""
|
||
|
|
Stream forcefully reset.
|
||
|
|
"""
|
||
|
|
self.remote_closed = True
|
||
|
|
self._close_cb(self.stream_id)
|
||
|
|
|
||
|
|
def get_headers(self):
|
||
|
|
"""
|
||
|
|
Provides the headers to the connection object.
|
||
|
|
"""
|
||
|
|
# Strip any headers invalid in H2.
|
||
|
|
return h2_safe_headers(self.headers)
|
||
|
|
|
||
|
|
def getheaders(self):
|
||
|
|
"""
|
||
|
|
Once all data has been sent on this connection, returns a key-value set
|
||
|
|
of the headers of the response to the original request.
|
||
|
|
"""
|
||
|
|
# Keep reading until all headers are received.
|
||
|
|
while self.response_headers is None:
|
||
|
|
self._recv_cb(stream_id=self.stream_id)
|
||
|
|
|
||
|
|
# Find the Content-Length header if present.
|
||
|
|
self._in_window_manager.document_size = (
|
||
|
|
int(self.response_headers.get(b'content-length', [0])[0])
|
||
|
|
)
|
||
|
|
|
||
|
|
return self.response_headers
|
||
|
|
|
||
|
|
def gettrailers(self):
|
||
|
|
"""
|
||
|
|
Once all data has been sent on this connection, returns a key-value set
|
||
|
|
of the trailers of the response to the original request.
|
||
|
|
|
||
|
|
.. warning:: Note that this method requires that the stream is
|
||
|
|
totally exhausted. This means that, if you have not
|
||
|
|
completely read from the stream, all stream data will be
|
||
|
|
read into memory.
|
||
|
|
|
||
|
|
:returns: The key-value set of the trailers, or ``None`` if no trailers
|
||
|
|
were sent.
|
||
|
|
"""
|
||
|
|
# Keep reading until the stream is done.
|
||
|
|
while not self.remote_closed:
|
||
|
|
self._recv_cb(stream_id=self.stream_id)
|
||
|
|
|
||
|
|
return self.response_trailers
|
||
|
|
|
||
|
|
def get_pushes(self, capture_all=False):
|
||
|
|
"""
|
||
|
|
Returns a generator that yields push promises from the server. Note
|
||
|
|
that this method is not idempotent; promises returned in one call will
|
||
|
|
not be returned in subsequent calls. Iterating through generators
|
||
|
|
returned by multiple calls to this method simultaneously results in
|
||
|
|
undefined behavior.
|
||
|
|
|
||
|
|
:param capture_all: If ``False``, the generator will yield all buffered
|
||
|
|
push promises without blocking. If ``True``, the generator will
|
||
|
|
first yield all buffered push promises, then yield additional ones
|
||
|
|
as they arrive, and terminate when the original stream closes.
|
||
|
|
"""
|
||
|
|
while True:
|
||
|
|
for pair in self.promised_headers.items():
|
||
|
|
yield pair
|
||
|
|
self.promised_headers = {}
|
||
|
|
if not capture_all or self.remote_closed:
|
||
|
|
break
|
||
|
|
self._recv_cb(stream_id=self.stream_id)
|
||
|
|
|
||
|
|
def close(self, error_code=None):
|
||
|
|
"""
|
||
|
|
Closes the stream. If the stream is currently open, attempts to close
|
||
|
|
it as gracefully as possible.
|
||
|
|
|
||
|
|
:param error_code: (optional) The error code to reset the stream with.
|
||
|
|
:returns: Nothing.
|
||
|
|
"""
|
||
|
|
# FIXME: I think this is overbroad, but for now it's probably ok.
|
||
|
|
if not (self.remote_closed and self.local_closed):
|
||
|
|
try:
|
||
|
|
with self._conn as conn:
|
||
|
|
conn.reset_stream(self.stream_id, error_code or 0)
|
||
|
|
except h2.exceptions.ProtocolError:
|
||
|
|
# If for any reason we can't reset the stream, just
|
||
|
|
# tolerate it.
|
||
|
|
pass
|
||
|
|
else:
|
||
|
|
self._send_outstanding_data(tolerate_peer_gone=True)
|
||
|
|
self.remote_closed = True
|
||
|
|
self.local_closed = True
|
||
|
|
|
||
|
|
self._close_cb(self.stream_id)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def _out_flow_control_window(self):
|
||
|
|
"""
|
||
|
|
The size of our outbound flow control window.
|
||
|
|
"""
|
||
|
|
|
||
|
|
with self._conn as conn:
|
||
|
|
return conn.local_flow_control_window(self.stream_id)
|
||
|
|
|
||
|
|
def _send_chunk(self, data, final):
|
||
|
|
"""
|
||
|
|
Implements most of the sending logic.
|
||
|
|
|
||
|
|
Takes a single chunk of size at most MAX_CHUNK, wraps it in a frame and
|
||
|
|
sends it. Optionally sets the END_STREAM flag if this is the last chunk
|
||
|
|
(determined by being of size less than MAX_CHUNK) and no more data is
|
||
|
|
to be sent.
|
||
|
|
"""
|
||
|
|
# If we don't fit in the connection window, try popping frames off the
|
||
|
|
# connection in hope that one might be a window update frame.
|
||
|
|
while len(data) > self._out_flow_control_window:
|
||
|
|
self._recv_cb()
|
||
|
|
|
||
|
|
# If the length of the data is less than MAX_CHUNK, we're probably
|
||
|
|
# at the end of the file. If this is the end of the data, mark it
|
||
|
|
# as END_STREAM.
|
||
|
|
end_stream = False
|
||
|
|
if len(data) < MAX_CHUNK and final:
|
||
|
|
end_stream = True
|
||
|
|
|
||
|
|
# Send the frame and decrement the flow control window.
|
||
|
|
with self._conn as conn:
|
||
|
|
conn.send_data(
|
||
|
|
stream_id=self.stream_id, data=data, end_stream=end_stream
|
||
|
|
)
|
||
|
|
self._send_outstanding_data()
|
||
|
|
|
||
|
|
if end_stream:
|
||
|
|
self.local_closed = True
|