hyper/http20/stream.py

# -*- 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.
"""
from ..h2 import exceptions as h2Exceptions

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 h2Exceptions.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 h2Exceptions.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