Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
connection.py357 linesDownload Raw Back to http2
1from __future__ import annotations2 3import logging4import re5import threading6import types7import typing8 9import h2.config10import h2.connection11import h2.events12 13from .._base_connection import _TYPE_BODY14from .._collections import HTTPHeaderDict15from ..connection import HTTPSConnection, _get_default_user_agent16from ..exceptions import ConnectionError17from ..response import BaseHTTPResponse18 19orig_HTTPSConnection = HTTPSConnection20 21T = typing.TypeVar("T")22 23log = logging.getLogger(__name__)24 25RE_IS_LEGAL_HEADER_NAME = re.compile(rb"^[!#$%&'*+\-.^_`|~0-9a-z]+$")26RE_IS_ILLEGAL_HEADER_VALUE = re.compile(rb"[\0\x00\x0a\x0d\r\n]|^[ \r\n\t]|[ \r\n\t]$")27 28 29def _is_legal_header_name(name: bytes) -> bool:30    """31    "An implementation that validates fields according to the definitions in Sections32    5.1 and 5.5 of [HTTP] only needs an additional check that field names do not33    include uppercase characters." (https://httpwg.org/specs/rfc9113.html#n-field-validity)34 35    `http.client._is_legal_header_name` does not validate the field name according to the36    HTTP 1.1 spec, so we do that here, in addition to checking for uppercase characters.37 38    This does not allow for the `:` character in the header name, so should not39    be used to validate pseudo-headers.40    """41    return bool(RE_IS_LEGAL_HEADER_NAME.match(name))42 43 44def _is_illegal_header_value(value: bytes) -> bool:45    """46    "A field value MUST NOT contain the zero value (ASCII NUL, 0x00), line feed47    (ASCII LF, 0x0a), or carriage return (ASCII CR, 0x0d) at any position. A field48    value MUST NOT start or end with an ASCII whitespace character (ASCII SP or HTAB,49    0x20 or 0x09)." (https://httpwg.org/specs/rfc9113.html#n-field-validity)50    """51    return bool(RE_IS_ILLEGAL_HEADER_VALUE.search(value))52 53 54class _LockedObject(typing.Generic[T]):55    """56    A wrapper class that hides a specific object behind a lock.57    The goal here is to provide a simple way to protect access to an object58    that cannot safely be simultaneously accessed from multiple threads. The59    intended use of this class is simple: take hold of it with a context60    manager, which returns the protected object.61    """62 63    __slots__ = (64        "lock",65        "_obj",66    )67 68    def __init__(self, obj: T):69        self.lock = threading.RLock()70        self._obj = obj71 72    def __enter__(self) -> T:73        self.lock.acquire()74        return self._obj75 76    def __exit__(77        self,78        exc_type: type[BaseException] | None,79        exc_val: BaseException | None,80        exc_tb: types.TracebackType | None,81    ) -> None:82        self.lock.release()83 84 85class HTTP2Connection(HTTPSConnection):86    def __init__(87        self, host: str, port: int | None = None, **kwargs: typing.Any88    ) -> None:89        self._h2_conn = self._new_h2_conn()90        self._h2_stream: int | None = None91        self._headers: list[tuple[bytes, bytes]] = []92 93        if "proxy" in kwargs or "proxy_config" in kwargs:  # Defensive:94            raise NotImplementedError("Proxies aren't supported with HTTP/2")95 96        super().__init__(host, port, **kwargs)97 98        if self._tunnel_host is not None:99            raise NotImplementedError("Tunneling isn't supported with HTTP/2")100 101    def _new_h2_conn(self) -> _LockedObject[h2.connection.H2Connection]:102        config = h2.config.H2Configuration(client_side=True)103        return _LockedObject(h2.connection.H2Connection(config=config))104 105    def connect(self) -> None:106        super().connect()107        with self._h2_conn as conn:108            conn.initiate_connection()109            if data_to_send := conn.data_to_send():110                self.sock.sendall(data_to_send)111 112    def putrequest(  # type: ignore[override]113        self,114        method: str,115        url: str,116        **kwargs: typing.Any,117    ) -> None:118        """putrequest119        This deviates from the HTTPConnection method signature since we never need to override120        sending accept-encoding headers or the host header.121        """122        if "skip_host" in kwargs:123            raise NotImplementedError("`skip_host` isn't supported")124        if "skip_accept_encoding" in kwargs:125            raise NotImplementedError("`skip_accept_encoding` isn't supported")126 127        self._request_url = url or "/"128        self._validate_path(url)  # type: ignore[attr-defined]129 130        if ":" in self.host:131            authority = f"[{self.host}]:{self.port or 443}"132        else:133            authority = f"{self.host}:{self.port or 443}"134 135        self._headers.append((b":scheme", b"https"))136        self._headers.append((b":method", method.encode()))137        self._headers.append((b":authority", authority.encode()))138        self._headers.append((b":path", url.encode()))139 140        with self._h2_conn as conn:141            self._h2_stream = conn.get_next_available_stream_id()142 143    def putheader(self, header: str | bytes, *values: str | bytes) -> None:  # type: ignore[override]144        # TODO SKIPPABLE_HEADERS from urllib3 are ignored.145        header = header.encode() if isinstance(header, str) else header146        header = header.lower()  # A lot of upstream code uses capitalized headers.147        if not _is_legal_header_name(header):148            raise ValueError(f"Illegal header name {str(header)}")149 150        for value in values:151            value = value.encode() if isinstance(value, str) else value152            if _is_illegal_header_value(value):153                raise ValueError(f"Illegal header value {str(value)}")154            self._headers.append((header, value))155 156    def endheaders(self, message_body: typing.Any = None) -> None:  # type: ignore[override]157        if self._h2_stream is None:158            raise ConnectionError("Must call `putrequest` first.")159 160        with self._h2_conn as conn:161            conn.send_headers(162                stream_id=self._h2_stream,163                headers=self._headers,164                end_stream=(message_body is None),165            )166            if data_to_send := conn.data_to_send():167                self.sock.sendall(data_to_send)168        self._headers = []  # Reset headers for the next request.169 170    def send(self, data: typing.Any) -> None:171        """Send data to the server.172        `data` can be: `str`, `bytes`, an iterable, or file-like objects173        that support a .read() method.174        """175        if self._h2_stream is None:176            raise ConnectionError("Must call `putrequest` first.")177 178        with self._h2_conn as conn:179            if data_to_send := conn.data_to_send():180                self.sock.sendall(data_to_send)181 182            if hasattr(data, "read"):  # file-like objects183                while True:184                    chunk = data.read(self.blocksize)185                    if not chunk:186                        break187                    if isinstance(chunk, str):188                        chunk = chunk.encode()189                    conn.send_data(self._h2_stream, chunk, end_stream=False)190                    if data_to_send := conn.data_to_send():191                        self.sock.sendall(data_to_send)192                conn.end_stream(self._h2_stream)193                return194 195            if isinstance(data, str):  # str -> bytes196                data = data.encode()197 198            try:199                if isinstance(data, bytes):200                    conn.send_data(self._h2_stream, data, end_stream=True)201                    if data_to_send := conn.data_to_send():202                        self.sock.sendall(data_to_send)203                else:204                    for chunk in data:205                        conn.send_data(self._h2_stream, chunk, end_stream=False)206                        if data_to_send := conn.data_to_send():207                            self.sock.sendall(data_to_send)208                    conn.end_stream(self._h2_stream)209            except TypeError:210                raise TypeError(211                    "`data` should be str, bytes, iterable, or file. got %r"212                    % type(data)213                )214 215    def set_tunnel(216        self,217        host: str,218        port: int | None = None,219        headers: typing.Mapping[str, str] | None = None,220        scheme: str = "http",221    ) -> None:222        raise NotImplementedError(223            "HTTP/2 does not support setting up a tunnel through a proxy"224        )225 226    def getresponse(  # type: ignore[override]227        self,228    ) -> HTTP2Response:229        status = None230        data = bytearray()231        with self._h2_conn as conn:232            end_stream = False233            while not end_stream:234                # TODO: Arbitrary read value.235                if received_data := self.sock.recv(65535):236                    events = conn.receive_data(received_data)237                    for event in events:238                        if isinstance(event, h2.events.ResponseReceived):239                            headers = HTTPHeaderDict()240                            for header, value in event.headers:241                                if header == b":status":242                                    status = int(value.decode())243                                else:244                                    headers.add(245                                        header.decode("ascii"), value.decode("ascii")246                                    )247 248                        elif isinstance(event, h2.events.DataReceived):249                            data += event.data250                            conn.acknowledge_received_data(251                                event.flow_controlled_length, event.stream_id252                            )253 254                        elif isinstance(event, h2.events.StreamEnded):255                            end_stream = True256 257                if data_to_send := conn.data_to_send():258                    self.sock.sendall(data_to_send)259 260        assert status is not None261        return HTTP2Response(262            status=status,263            headers=headers,264            request_url=self._request_url,265            data=bytes(data),266        )267 268    def request(  # type: ignore[override]269        self,270        method: str,271        url: str,272        body: _TYPE_BODY | None = None,273        headers: typing.Mapping[str, str] | None = None,274        *,275        preload_content: bool = True,276        decode_content: bool = True,277        enforce_content_length: bool = True,278        **kwargs: typing.Any,279    ) -> None:280        """Send an HTTP/2 request"""281        if "chunked" in kwargs:282            # TODO this is often present from upstream.283            # raise NotImplementedError("`chunked` isn't supported with HTTP/2")284            pass285 286        if self.sock is not None:287            self.sock.settimeout(self.timeout)288 289        self.putrequest(method, url)290 291        headers = headers or {}292        for k, v in headers.items():293            if k.lower() == "transfer-encoding" and v == "chunked":294                continue295            else:296                self.putheader(k, v)297 298        if b"user-agent" not in dict(self._headers):299            self.putheader(b"user-agent", _get_default_user_agent())300 301        if body:302            self.endheaders(message_body=body)303            self.send(body)304        else:305            self.endheaders()306 307    def close(self) -> None:308        with self._h2_conn as conn:309            try:310                conn.close_connection()311                if data := conn.data_to_send():312                    self.sock.sendall(data)313            except Exception:314                pass315 316        # Reset all our HTTP/2 connection state.317        self._h2_conn = self._new_h2_conn()318        self._h2_stream = None319        self._headers = []320 321        super().close()322 323 324class HTTP2Response(BaseHTTPResponse):325    # TODO: This is a woefully incomplete response object, but works for non-streaming.326    def __init__(327        self,328        status: int,329        headers: HTTPHeaderDict,330        request_url: str,331        data: bytes,332        decode_content: bool = False,  # TODO: support decoding333    ) -> None:334        super().__init__(335            status=status,336            headers=headers,337            # Following CPython, we map HTTP versions to major * 10 + minor integers338            version=20,339            version_string="HTTP/2",340            # No reason phrase in HTTP/2341            reason=None,342            decode_content=decode_content,343            request_url=request_url,344        )345        self._data = data346        self.length_remaining = 0347 348    @property349    def data(self) -> bytes:350        return self._data351 352    def get_redirect_location(self) -> None:353        return None354 355    def close(self) -> None:356        pass357 
codekingpro/portable-devtools · Team Ai