scrapy/scrapy/core/http2/protocol.py

import logging

from h2.config import H2Configuration
from h2.connection import H2Connection
from h2.events import (
    ConnectionTerminated, DataReceived, ResponseReceived, RemoteSettingsChanged,
    StreamEnded, StreamReset, TrailersReceived, WindowUpdated
)
from twisted.internet.protocol import connectionDone, Protocol
from zope.interface import implementer, Interface

from scrapy.core.http2.stream import Stream
from scrapy.http import Request

LOGGER = logging.getLogger(__name__)


class IH2EventsHandler(Interface):
    def connection_terminated(event: ConnectionTerminated):
        pass

    def data_received(event: DataReceived):
        pass

    def response_received(event: ResponseReceived):
        pass

    def remote_settings_changed(event: RemoteSettingsChanged):
        pass

    def stream_ended(event: StreamEnded):
        pass

    def stream_reset(event: StreamReset):
        pass

    def trailers_received(event: TrailersReceived):
        pass

    def window_updated(event: WindowUpdated):
        pass


@implementer(IH2EventsHandler)
class H2ClientProtocol(Protocol):
    # TODO: Check for user-agent while testing
    # TODO: Add support for cookies
    # TODO: Handle priority updates

    def __init__(self):
        config = H2Configuration(client_side=True, header_encoding='utf-8')
        self.conn = H2Connection(config=config)

        # ID of the next request stream
        # Assuming each request stream creates a new response stream
        # we increment by 2 for each new request stream created
        self.next_stream_id = 1

        # Streams are stored in a dictionary keyed off their stream IDs
        self.streams = {}

        # Boolean to keep track the connection is made
        # If requests are received before connection is made
        # we keep all requests in a pool and send them as the connection
        # is made
        self.is_connection_made = False
        self._pending_request_stream_pool = []

    def _new_stream(self, request: Request):
        """Instantiates a new Stream object
        """
        stream = Stream(self.next_stream_id, request, self)
        self.next_stream_id += 2

        self.streams[stream.stream_id] = stream
        return stream

    def _write_to_transport(self):
        """ Write data to the underlying transport connection
        from the HTTP2 connection instance if any
        """
        data = self.conn.data_to_send()
        if data:
            self.transport.write(data)

    def request(self, _request: Request):
        stream = self._new_stream(_request)
        d = stream.get_response()

        # If connection is not yet established then add the
        # stream to pool or initiate request
        if self.is_connection_made:
            stream.initiate_request()
        else:
            self._pending_request_stream_pool.append(stream)

        return d

    def connectionMade(self):
        """Called by Twisted when the connection is established. We can start
        sending some data now: we should open with the connection preamble.
        """
        LOGGER.info("Connection made to {}".format(self.transport))
        self.conn.initiate_connection()
        self._write_to_transport()

        self.is_connection_made = True

        # Initiate all pending requests
        for stream in self._pending_request_stream_pool:
            assert isinstance(stream, Stream)
            stream.initiate_request()

        self._pending_request_stream_pool.clear()

    def dataReceived(self, data):
        events = self.conn.receive_data(data)
        self._handle_events(events)
        self._write_to_transport()

    def connectionLost(self, reason=connectionDone):

        """Called by Twisted when the transport connection is lost.
        """
        for stream_id in self.streams.keys():
            # TODO: Close each Stream instance in a clean manner
            self.conn.end_stream(stream_id)

    def _handle_events(self, events):
        """Private method which acts as a bridge between the events
        received from the HTTP/2 data and IH2EventsHandler

        Arguments:
            events {list} -- A list of events that the remote peer
                triggered by sending data
        """
        for event in events:
            if isinstance(event, ConnectionTerminated):
                self.connection_terminated(event)
            elif isinstance(event, DataReceived):
                self.data_received(event)
            elif isinstance(event, ResponseReceived):
                self.response_received(event)
            elif isinstance(event, StreamEnded):
                self.stream_ended(event)
            elif isinstance(event, StreamReset):
                self.stream_reset(event)
            elif isinstance(event, TrailersReceived):
                self.trailers_received(event)
            elif isinstance(event, WindowUpdated):
                self.window_updated(event)
            elif isinstance(event, RemoteSettingsChanged):
                self.remote_settings_changed(event)

    def send_headers(self, stream_id, headers):
        """ Send the headers for a given stream to the resource
        Initiates a new connection hence.

        Arguments:
            stream_id {int} -- Valid stream id
            headers {List[Tuple[str, str]]} -- Headers of the request
        """
        if stream_id in self.streams:
            self.conn.send_headers(stream_id, headers, end_stream=True)
            self._write_to_transport()
        else:
            pass

    def connection_terminated(self, event: ConnectionTerminated):
        pass

    def data_received(self, event: DataReceived):
        stream_id = event.stream_id
        # TODO: Stream do not exist in self.streams dict
        self.streams[stream_id].receive_data(event.data)

    def response_received(self, event: ResponseReceived):
        stream_id = event.stream_id
        # TODO: Stream do not exist in self.streams dict
        self.streams[stream_id].receive_headers(event.headers)

    def remote_settings_changed(self, event: RemoteSettingsChanged):
        pass

    def stream_ended(self, event: StreamEnded):
        stream_id = event.stream_id
        # TODO: Stream do not exist in self.streams dict
        self.streams[stream_id].end_stream()

    def stream_reset(self, event: StreamReset):
        pass

    def trailers_received(self, event: TrailersReceived):
        pass

    def window_updated(self, event: WindowUpdated):
        pass