codekingpro/portable-devtools

1from __future__ import annotations2 3import enum4import logging5import time6import types7import typing8 9import h2.config10import h2.connection11import h2.events12import h2.exceptions13import h2.settings14 15from .._backends.base import NetworkStream16from .._exceptions import (17    ConnectionNotAvailable,18    LocalProtocolError,19    RemoteProtocolError,20)21from .._models import Origin, Request, Response22from .._synchronization import Lock, Semaphore, ShieldCancellation23from .._trace import Trace24from .interfaces import ConnectionInterface25 26logger = logging.getLogger("httpcore.http2")27 28 29def has_body_headers(request: Request) -> bool:30    return any(31        k.lower() == b"content-length" or k.lower() == b"transfer-encoding"32        for k, v in request.headers33    )34 35 36class HTTPConnectionState(enum.IntEnum):37    ACTIVE = 138    IDLE = 239    CLOSED = 340 41 42class HTTP2Connection(ConnectionInterface):43    READ_NUM_BYTES = 64 * 102444    CONFIG = h2.config.H2Configuration(validate_inbound_headers=False)45 46    def __init__(47        self,48        origin: Origin,49        stream: NetworkStream,50        keepalive_expiry: float | None = None,51    ):52        self._origin = origin53        self._network_stream = stream54        self._keepalive_expiry: float | None = keepalive_expiry55        self._h2_state = h2.connection.H2Connection(config=self.CONFIG)56        self._state = HTTPConnectionState.IDLE57        self._expire_at: float | None = None58        self._request_count = 059        self._init_lock = Lock()60        self._state_lock = Lock()61        self._read_lock = Lock()62        self._write_lock = Lock()63        self._sent_connection_init = False64        self._used_all_stream_ids = False65        self._connection_error = False66 67        # Mapping from stream ID to response stream events.68        self._events: dict[69            int,70            list[71                h2.events.ResponseReceived72                | h2.events.DataReceived73                | h2.events.StreamEnded74                | h2.events.StreamReset,75            ],76        ] = {}77 78        # Connection terminated events are stored as state since79        # we need to handle them for all streams.80        self._connection_terminated: h2.events.ConnectionTerminated | None = None81 82        self._read_exception: Exception | None = None83        self._write_exception: Exception | None = None84 85    def handle_request(self, request: Request) -> Response:86        if not self.can_handle_request(request.url.origin):87            # This cannot occur in normal operation, since the connection pool88            # will only send requests on connections that handle them.89            # It's in place simply for resilience as a guard against incorrect90            # usage, for anyone working directly with httpcore connections.91            raise RuntimeError(92                f"Attempted to send request to {request.url.origin} on connection "93                f"to {self._origin}"94            )95 96        with self._state_lock:97            if self._state in (HTTPConnectionState.ACTIVE, HTTPConnectionState.IDLE):98                self._request_count += 199                self._expire_at = None100                self._state = HTTPConnectionState.ACTIVE101            else:102                raise ConnectionNotAvailable()103 104        with self._init_lock:105            if not self._sent_connection_init:106                try:107                    sci_kwargs = {"request": request}108                    with Trace(109                        "send_connection_init", logger, request, sci_kwargs110                    ):111                        self._send_connection_init(**sci_kwargs)112                except BaseException as exc:113                    with ShieldCancellation():114                        self.close()115                    raise exc116 117                self._sent_connection_init = True118 119                # Initially start with just 1 until the remote server provides120                # its max_concurrent_streams value121                self._max_streams = 1122 123                local_settings_max_streams = (124                    self._h2_state.local_settings.max_concurrent_streams125                )126                self._max_streams_semaphore = Semaphore(local_settings_max_streams)127 128                for _ in range(local_settings_max_streams - self._max_streams):129                    self._max_streams_semaphore.acquire()130 131        self._max_streams_semaphore.acquire()132 133        try:134            stream_id = self._h2_state.get_next_available_stream_id()135            self._events[stream_id] = []136        except h2.exceptions.NoAvailableStreamIDError:  # pragma: nocover137            self._used_all_stream_ids = True138            self._request_count -= 1139            raise ConnectionNotAvailable()140 141        try:142            kwargs = {"request": request, "stream_id": stream_id}143            with Trace("send_request_headers", logger, request, kwargs):144                self._send_request_headers(request=request, stream_id=stream_id)145            with Trace("send_request_body", logger, request, kwargs):146                self._send_request_body(request=request, stream_id=stream_id)147            with Trace(148                "receive_response_headers", logger, request, kwargs149            ) as trace:150                status, headers = self._receive_response(151                    request=request, stream_id=stream_id152                )153                trace.return_value = (status, headers)154 155            return Response(156                status=status,157                headers=headers,158                content=HTTP2ConnectionByteStream(self, request, stream_id=stream_id),159                extensions={160                    "http_version": b"HTTP/2",161                    "network_stream": self._network_stream,162                    "stream_id": stream_id,163                },164            )165        except BaseException as exc:  # noqa: PIE786166            with ShieldCancellation():167                kwargs = {"stream_id": stream_id}168                with Trace("response_closed", logger, request, kwargs):169                    self._response_closed(stream_id=stream_id)170 171            if isinstance(exc, h2.exceptions.ProtocolError):172                # One case where h2 can raise a protocol error is when a173                # closed frame has been seen by the state machine.174                #175                # This happens when one stream is reading, and encounters176                # a GOAWAY event. Other flows of control may then raise177                # a protocol error at any point they interact with the 'h2_state'.178                #179                # In this case we'll have stored the event, and should raise180                # it as a RemoteProtocolError.181                if self._connection_terminated:  # pragma: nocover182                    raise RemoteProtocolError(self._connection_terminated)183                # If h2 raises a protocol error in some other state then we184                # must somehow have made a protocol violation.185                raise LocalProtocolError(exc)  # pragma: nocover186 187            raise exc188 189    def _send_connection_init(self, request: Request) -> None:190        """191        The HTTP/2 connection requires some initial setup before we can start192        using individual request/response streams on it.193        """194        # Need to set these manually here instead of manipulating via195        # __setitem__() otherwise the H2Connection will emit SettingsUpdate196        # frames in addition to sending the undesired defaults.197        self._h2_state.local_settings = h2.settings.Settings(198            client=True,199            initial_values={200                # Disable PUSH_PROMISE frames from the server since we don't do anything201                # with them for now.  Maybe when we support caching?202                h2.settings.SettingCodes.ENABLE_PUSH: 0,203                # These two are taken from h2 for safe defaults204                h2.settings.SettingCodes.MAX_CONCURRENT_STREAMS: 100,205                h2.settings.SettingCodes.MAX_HEADER_LIST_SIZE: 65536,206            },207        )208 209        # Some websites (*cough* Yahoo *cough*) balk at this setting being210        # present in the initial handshake since it's not defined in the original211        # RFC despite the RFC mandating ignoring settings you don't know about.212        del self._h2_state.local_settings[213            h2.settings.SettingCodes.ENABLE_CONNECT_PROTOCOL214        ]215 216        self._h2_state.initiate_connection()217        self._h2_state.increment_flow_control_window(2**24)218        self._write_outgoing_data(request)219 220    # Sending the request...221 222    def _send_request_headers(self, request: Request, stream_id: int) -> None:223        """224        Send the request headers to a given stream ID.225        """226        end_stream = not has_body_headers(request)227 228        # In HTTP/2 the ':authority' pseudo-header is used instead of 'Host'.229        # In order to gracefully handle HTTP/1.1 and HTTP/2 we always require230        # HTTP/1.1 style headers, and map them appropriately if we end up on231        # an HTTP/2 connection.232        authority = [v for k, v in request.headers if k.lower() == b"host"][0]233 234        headers = [235            (b":method", request.method),236            (b":authority", authority),237            (b":scheme", request.url.scheme),238            (b":path", request.url.target),239        ] + [240            (k.lower(), v)241            for k, v in request.headers242            if k.lower()243            not in (244                b"host",245                b"transfer-encoding",246            )247        ]248 249        self._h2_state.send_headers(stream_id, headers, end_stream=end_stream)250        self._h2_state.increment_flow_control_window(2**24, stream_id=stream_id)251        self._write_outgoing_data(request)252 253    def _send_request_body(self, request: Request, stream_id: int) -> None:254        """255        Iterate over the request body sending it to a given stream ID.256        """257        if not has_body_headers(request):258            return259 260        assert isinstance(request.stream, typing.Iterable)261        for data in request.stream:262            self._send_stream_data(request, stream_id, data)263        self._send_end_stream(request, stream_id)264 265    def _send_stream_data(266        self, request: Request, stream_id: int, data: bytes267    ) -> None:268        """269        Send a single chunk of data in one or more data frames.270        """271        while data:272            max_flow = self._wait_for_outgoing_flow(request, stream_id)273            chunk_size = min(len(data), max_flow)274            chunk, data = data[:chunk_size], data[chunk_size:]275            self._h2_state.send_data(stream_id, chunk)276            self._write_outgoing_data(request)277 278    def _send_end_stream(self, request: Request, stream_id: int) -> None:279        """280        Send an empty data frame on on a given stream ID with the END_STREAM flag set.281        """282        self._h2_state.end_stream(stream_id)283        self._write_outgoing_data(request)284 285    # Receiving the response...286 287    def _receive_response(288        self, request: Request, stream_id: int289    ) -> tuple[int, list[tuple[bytes, bytes]]]:290        """291        Return the response status code and headers for a given stream ID.292        """293        while True:294            event = self._receive_stream_event(request, stream_id)295            if isinstance(event, h2.events.ResponseReceived):296                break297 298        status_code = 200299        headers = []300        assert event.headers is not None301        for k, v in event.headers:302            if k == b":status":303                status_code = int(v.decode("ascii", errors="ignore"))304            elif not k.startswith(b":"):305                headers.append((k, v))306 307        return (status_code, headers)308 309    def _receive_response_body(310        self, request: Request, stream_id: int311    ) -> typing.Iterator[bytes]:312        """313        Iterator that returns the bytes of the response body for a given stream ID.314        """315        while True:316            event = self._receive_stream_event(request, stream_id)317            if isinstance(event, h2.events.DataReceived):318                assert event.flow_controlled_length is not None319                assert event.data is not None320                amount = event.flow_controlled_length321                self._h2_state.acknowledge_received_data(amount, stream_id)322                self._write_outgoing_data(request)323                yield event.data324            elif isinstance(event, h2.events.StreamEnded):325                break326 327    def _receive_stream_event(328        self, request: Request, stream_id: int329    ) -> h2.events.ResponseReceived | h2.events.DataReceived | h2.events.StreamEnded:330        """331        Return the next available event for a given stream ID.332 333        Will read more data from the network if required.334        """335        while not self._events.get(stream_id):336            self._receive_events(request, stream_id)337        event = self._events[stream_id].pop(0)338        if isinstance(event, h2.events.StreamReset):339            raise RemoteProtocolError(event)340        return event341 342    def _receive_events(343        self, request: Request, stream_id: int | None = None344    ) -> None:345        """346        Read some data from the network until we see one or more events347        for a given stream ID.348        """349        with self._read_lock:350            if self._connection_terminated is not None:351                last_stream_id = self._connection_terminated.last_stream_id352                if stream_id and last_stream_id and stream_id > last_stream_id:353                    self._request_count -= 1354                    raise ConnectionNotAvailable()355                raise RemoteProtocolError(self._connection_terminated)356 357            # This conditional is a bit icky. We don't want to block reading if we've358            # actually got an event to return for a given stream. We need to do that359            # check *within* the atomic read lock. Though it also need to be optional,360            # because when we call it from `_wait_for_outgoing_flow` we *do* want to361            # block until we've available flow control, event when we have events362            # pending for the stream ID we're attempting to send on.363            if stream_id is None or not self._events.get(stream_id):364                events = self._read_incoming_data(request)365                for event in events:366                    if isinstance(event, h2.events.RemoteSettingsChanged):367                        with Trace(368                            "receive_remote_settings", logger, request369                        ) as trace:370                            self._receive_remote_settings_change(event)371                            trace.return_value = event372 373                    elif isinstance(374                        event,375                        (376                            h2.events.ResponseReceived,377                            h2.events.DataReceived,378                            h2.events.StreamEnded,379                            h2.events.StreamReset,380                        ),381                    ):382                        if event.stream_id in self._events:383                            self._events[event.stream_id].append(event)384 385                    elif isinstance(event, h2.events.ConnectionTerminated):386                        self._connection_terminated = event387 388        self._write_outgoing_data(request)389 390    def _receive_remote_settings_change(391        self, event: h2.events.RemoteSettingsChanged392    ) -> None:393        max_concurrent_streams = event.changed_settings.get(394            h2.settings.SettingCodes.MAX_CONCURRENT_STREAMS395        )396        if max_concurrent_streams:397            new_max_streams = min(398                max_concurrent_streams.new_value,399                self._h2_state.local_settings.max_concurrent_streams,400            )401            if new_max_streams and new_max_streams != self._max_streams:402                while new_max_streams > self._max_streams:403                    self._max_streams_semaphore.release()404                    self._max_streams += 1405                while new_max_streams < self._max_streams:406                    self._max_streams_semaphore.acquire()407                    self._max_streams -= 1408 409    def _response_closed(self, stream_id: int) -> None:410        self._max_streams_semaphore.release()411        del self._events[stream_id]412        with self._state_lock:413            if self._connection_terminated and not self._events:414                self.close()415 416            elif self._state == HTTPConnectionState.ACTIVE and not self._events:417                self._state = HTTPConnectionState.IDLE418                if self._keepalive_expiry is not None:419                    now = time.monotonic()420                    self._expire_at = now + self._keepalive_expiry421                if self._used_all_stream_ids:  # pragma: nocover422                    self.close()423 424    def close(self) -> None:425        # Note that this method unilaterally closes the connection, and does426        # not have any kind of locking in place around it.427        self._h2_state.close_connection()428        self._state = HTTPConnectionState.CLOSED429        self._network_stream.close()430 431    # Wrappers around network read/write operations...432 433    def _read_incoming_data(self, request: Request) -> list[h2.events.Event]:434        timeouts = request.extensions.get("timeout", {})435        timeout = timeouts.get("read", None)436 437        if self._read_exception is not None:438            raise self._read_exception  # pragma: nocover439 440        try:441            data = self._network_stream.read(self.READ_NUM_BYTES, timeout)442            if data == b"":443                raise RemoteProtocolError("Server disconnected")444        except Exception as exc:445            # If we get a network error we should:446            #447            # 1. Save the exception and just raise it immediately on any future reads.448            #    (For example, this means that a single read timeout or disconnect will449            #    immediately close all pending streams. Without requiring multiple450            #    sequential timeouts.)451            # 2. Mark the connection as errored, so that we don't accept any other452            #    incoming requests.453            self._read_exception = exc454            self._connection_error = True455            raise exc456 457        events: list[h2.events.Event] = self._h2_state.receive_data(data)458 459        return events460 461    def _write_outgoing_data(self, request: Request) -> None:462        timeouts = request.extensions.get("timeout", {})463        timeout = timeouts.get("write", None)464 465        with self._write_lock:466            data_to_send = self._h2_state.data_to_send()467 468            if self._write_exception is not None:469                raise self._write_exception  # pragma: nocover470 471            try:472                self._network_stream.write(data_to_send, timeout)473            except Exception as exc:  # pragma: nocover474                # If we get a network error we should:475                #476                # 1. Save the exception and just raise it immediately on any future write.477                #    (For example, this means that a single write timeout or disconnect will478                #    immediately close all pending streams. Without requiring multiple479                #    sequential timeouts.)480                # 2. Mark the connection as errored, so that we don't accept any other481                #    incoming requests.482                self._write_exception = exc483                self._connection_error = True484                raise exc485 486    # Flow control...487 488    def _wait_for_outgoing_flow(self, request: Request, stream_id: int) -> int:489        """490        Returns the maximum allowable outgoing flow for a given stream.491 492        If the allowable flow is zero, then waits on the network until493        WindowUpdated frames have increased the flow rate.494        https://tools.ietf.org/html/rfc7540#section-6.9495        """496        local_flow: int = self._h2_state.local_flow_control_window(stream_id)497        max_frame_size: int = self._h2_state.max_outbound_frame_size498        flow = min(local_flow, max_frame_size)499        while flow == 0:500            self._receive_events(request)501            local_flow = self._h2_state.local_flow_control_window(stream_id)502            max_frame_size = self._h2_state.max_outbound_frame_size503            flow = min(local_flow, max_frame_size)504        return flow505 506    # Interface for connection pooling...507 508    def can_handle_request(self, origin: Origin) -> bool:509        return origin == self._origin510 511    def is_available(self) -> bool:512        return (513            self._state != HTTPConnectionState.CLOSED514            and not self._connection_error515            and not self._used_all_stream_ids516            and not (517                self._h2_state.state_machine.state518                == h2.connection.ConnectionState.CLOSED519            )520        )521 522    def has_expired(self) -> bool:523        now = time.monotonic()524        return self._expire_at is not None and now > self._expire_at525 526    def is_idle(self) -> bool:527        return self._state == HTTPConnectionState.IDLE528 529    def is_closed(self) -> bool:530        return self._state == HTTPConnectionState.CLOSED531 532    def info(self) -> str:533        origin = str(self._origin)534        return (535            f"{origin!r}, HTTP/2, {self._state.name}, "536            f"Request Count: {self._request_count}"537        )538 539    def __repr__(self) -> str:540        class_name = self.__class__.__name__541        origin = str(self._origin)542        return (543            f"<{class_name} [{origin!r}, {self._state.name}, "544            f"Request Count: {self._request_count}]>"545        )546 547    # These context managers are not used in the standard flow, but are548    # useful for testing or working with connection instances directly.549 550    def __enter__(self) -> HTTP2Connection:551        return self552 553    def __exit__(554        self,555        exc_type: type[BaseException] | None = None,556        exc_value: BaseException | None = None,557        traceback: types.TracebackType | None = None,558    ) -> None:559        self.close()560 561 562class HTTP2ConnectionByteStream:563    def __init__(564        self, connection: HTTP2Connection, request: Request, stream_id: int565    ) -> None:566        self._connection = connection567        self._request = request568        self._stream_id = stream_id569        self._closed = False570 571    def __iter__(self) -> typing.Iterator[bytes]:572        kwargs = {"request": self._request, "stream_id": self._stream_id}573        try:574            with Trace("receive_response_body", logger, self._request, kwargs):575                for chunk in self._connection._receive_response_body(576                    request=self._request, stream_id=self._stream_id577                ):578                    yield chunk579        except BaseException as exc:580            # If we get an exception while streaming the response,581            # we want to close the response (and possibly the connection)582            # before raising that exception.583            with ShieldCancellation():584                self.close()585            raise exc586 587    def close(self) -> None:588        if not self._closed:589            self._closed = True590            kwargs = {"stream_id": self._stream_id}591            with Trace("response_closed", logger, self._request, kwargs):592                self._connection._response_closed(stream_id=self._stream_id)593