Upload 1579 files

8070315 verified over 2 years ago

7.48 kB

from __future__ import annotations

import threading
import types
import typing

import h2.config  # type: ignore[import-untyped]
import h2.connection  # type: ignore[import-untyped]
import h2.events  # type: ignore[import-untyped]

import urllib3.connection
import urllib3.util.ssl_
from urllib3.response import BaseHTTPResponse

from ._collections import HTTPHeaderDict
from .connection import HTTPSConnection
from .connectionpool import HTTPSConnectionPool

orig_HTTPSConnection = HTTPSConnection

T = typing.TypeVar("T")


class _LockedObject(typing.Generic[T]):
    """
    A wrapper class that hides a specific object behind a lock.
    The goal here is to provide a simple way to protect access to an object
    that cannot safely be simultaneously accessed from multiple threads. The
    intended use of this class is simple: take hold of it with a context
    manager, which returns the protected object.
    """

    def __init__(self, obj: T):
        self.lock = threading.RLock()
        self._obj = obj

    def __enter__(self) -> T:
        self.lock.acquire()
        return self._obj

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: types.TracebackType | None,
    ) -> None:
        self.lock.release()


class HTTP2Connection(HTTPSConnection):
    def __init__(
        self, host: str, port: int | None = None, **kwargs: typing.Any
    ) -> None:
        self._h2_conn = self._new_h2_conn()
        self._h2_stream: int | None = None
        self._h2_headers: list[tuple[bytes, bytes]] = []

        if "proxy" in kwargs or "proxy_config" in kwargs:  # Defensive:
            raise NotImplementedError("Proxies aren't supported with HTTP/2")

        super().__init__(host, port, **kwargs)

    def _new_h2_conn(self) -> _LockedObject[h2.connection.H2Connection]:
        config = h2.config.H2Configuration(client_side=True)
        return _LockedObject(h2.connection.H2Connection(config=config))

    def connect(self) -> None:
        super().connect()

        with self._h2_conn as h2_conn:
            h2_conn.initiate_connection()
            self.sock.sendall(h2_conn.data_to_send())

    def putrequest(
        self,
        method: str,
        url: str,
        skip_host: bool = False,
        skip_accept_encoding: bool = False,
    ) -> None:
        with self._h2_conn as h2_conn:
            self._request_url = url
            self._h2_stream = h2_conn.get_next_available_stream_id()

            if ":" in self.host:
                authority = f"[{self.host}]:{self.port or 443}"
            else:
                authority = f"{self.host}:{self.port or 443}"

            self._h2_headers.extend(
                (
                    (b":scheme", b"https"),
                    (b":method", method.encode()),
                    (b":authority", authority.encode()),
                    (b":path", url.encode()),
                )
            )

    def putheader(self, header: str, *values: str) -> None:  # type: ignore[override]
        for value in values:
            self._h2_headers.append(
                (header.encode("utf-8").lower(), value.encode("utf-8"))
            )

    def endheaders(self) -> None:  # type: ignore[override]
        with self._h2_conn as h2_conn:
            h2_conn.send_headers(
                stream_id=self._h2_stream,
                headers=self._h2_headers,
                end_stream=True,
            )
            if data_to_send := h2_conn.data_to_send():
                self.sock.sendall(data_to_send)

    def send(self, data: bytes) -> None:  # type: ignore[override]  # Defensive:
        if not data:
            return
        raise NotImplementedError("Sending data isn't supported yet")

    def getresponse(  # type: ignore[override]
        self,
    ) -> HTTP2Response:
        status = None
        data = bytearray()
        with self._h2_conn as h2_conn:
            end_stream = False
            while not end_stream:
                # TODO: Arbitrary read value.
                if received_data := self.sock.recv(65535):
                    events = h2_conn.receive_data(received_data)
                    for event in events:
                        if isinstance(event, h2.events.ResponseReceived):
                            headers = HTTPHeaderDict()
                            for header, value in event.headers:
                                if header == b":status":
                                    status = int(value.decode())
                                else:
                                    headers.add(
                                        header.decode("ascii"), value.decode("ascii")
                                    )

                        elif isinstance(event, h2.events.DataReceived):
                            data += event.data
                            h2_conn.acknowledge_received_data(
                                event.flow_controlled_length, event.stream_id
                            )

                        elif isinstance(event, h2.events.StreamEnded):
                            end_stream = True

                if data_to_send := h2_conn.data_to_send():
                    self.sock.sendall(data_to_send)

        # We always close to not have to handle connection management.
        self.close()

        assert status is not None
        return HTTP2Response(
            status=status,
            headers=headers,
            request_url=self._request_url,
            data=bytes(data),
        )

    def close(self) -> None:
        with self._h2_conn as h2_conn:
            try:
                h2_conn.close_connection()
                if data := h2_conn.data_to_send():
                    self.sock.sendall(data)
            except Exception:
                pass

        # Reset all our HTTP/2 connection state.
        self._h2_conn = self._new_h2_conn()
        self._h2_stream = None
        self._h2_headers = []

        super().close()


class HTTP2Response(BaseHTTPResponse):
    # TODO: This is a woefully incomplete response object, but works for non-streaming.
    def __init__(
        self,
        status: int,
        headers: HTTPHeaderDict,
        request_url: str,
        data: bytes,
        decode_content: bool = False,  # TODO: support decoding
    ) -> None:
        super().__init__(
            status=status,
            headers=headers,
            # Following CPython, we map HTTP versions to major * 10 + minor integers
            version=20,
            # No reason phrase in HTTP/2
            reason=None,
            decode_content=decode_content,
            request_url=request_url,
        )
        self._data = data
        self.length_remaining = 0

    @property
    def data(self) -> bytes:
        return self._data

    def get_redirect_location(self) -> None:
        return None

    def close(self) -> None:
        pass


def inject_into_urllib3() -> None:
    HTTPSConnectionPool.ConnectionCls = HTTP2Connection
    urllib3.connection.HTTPSConnection = HTTP2Connection  # type: ignore[misc]

    # TODO: Offer 'http/1.1' as well, but for testing purposes this is handy.
    urllib3.util.ssl_.ALPN_PROTOCOLS = ["h2"]


def extract_from_urllib3() -> None:
    HTTPSConnectionPool.ConnectionCls = orig_HTTPSConnection
    urllib3.connection.HTTPSConnection = orig_HTTPSConnection  # type: ignore[misc]

    urllib3.util.ssl_.ALPN_PROTOCOLS = ["http/1.1"]

	from __future__ import annotations

	import threading
	import types
	import typing

	import h2.config # type: ignore[import-untyped]
	import h2.connection # type: ignore[import-untyped]
	import h2.events # type: ignore[import-untyped]

	import urllib3.connection
	import urllib3.util.ssl_
	from urllib3.response import BaseHTTPResponse

	from ._collections import HTTPHeaderDict
	from .connection import HTTPSConnection
	from .connectionpool import HTTPSConnectionPool

	orig_HTTPSConnection = HTTPSConnection

	T = typing.TypeVar("T")


	class _LockedObject(typing.Generic[T]):
	"""
	A wrapper class that hides a specific object behind a lock.

	The goal here is to provide a simple way to protect access to an object
	that cannot safely be simultaneously accessed from multiple threads. The
	intended use of this class is simple: take hold of it with a context
	manager, which returns the protected object.
	"""

	def __init__(self, obj: T):
	self.lock = threading.RLock()
	self._obj = obj

	def __enter__(self) -> T:
	self.lock.acquire()
	return self._obj

	def __exit__(
	self,
	exc_type: type[BaseException] \| None,
	exc_val: BaseException \| None,
	exc_tb: types.TracebackType \| None,
	) -> None:
	self.lock.release()


	class HTTP2Connection(HTTPSConnection):
	def __init__(
	self, host: str, port: int \| None = None, **kwargs: typing.Any
	) -> None:
	self._h2_conn = self._new_h2_conn()
	self._h2_stream: int \| None = None
	self._h2_headers: list[tuple[bytes, bytes]] = []

	if "proxy" in kwargs or "proxy_config" in kwargs: # Defensive:
	raise NotImplementedError("Proxies aren't supported with HTTP/2")

	super().__init__(host, port, **kwargs)

	def _new_h2_conn(self) -> _LockedObject[h2.connection.H2Connection]:
	config = h2.config.H2Configuration(client_side=True)
	return _LockedObject(h2.connection.H2Connection(config=config))

	def connect(self) -> None:
	super().connect()

	with self._h2_conn as h2_conn:
	h2_conn.initiate_connection()
	self.sock.sendall(h2_conn.data_to_send())

	def putrequest(
	self,
	method: str,
	url: str,
	skip_host: bool = False,
	skip_accept_encoding: bool = False,
	) -> None:
	with self._h2_conn as h2_conn:
	self._request_url = url
	self._h2_stream = h2_conn.get_next_available_stream_id()

	if ":" in self.host:
	authority = f"[{self.host}]:{self.port or 443}"
	else:
	authority = f"{self.host}:{self.port or 443}"

	self._h2_headers.extend(
	(
	(b":scheme", b"https"),
	(b":method", method.encode()),
	(b":authority", authority.encode()),
	(b":path", url.encode()),
	)
	)

	def putheader(self, header: str, *values: str) -> None: # type: ignore[override]
	for value in values:
	self._h2_headers.append(
	(header.encode("utf-8").lower(), value.encode("utf-8"))
	)

	def endheaders(self) -> None: # type: ignore[override]
	with self._h2_conn as h2_conn:
	h2_conn.send_headers(
	stream_id=self._h2_stream,
	headers=self._h2_headers,
	end_stream=True,
	)
	if data_to_send := h2_conn.data_to_send():
	self.sock.sendall(data_to_send)

	def send(self, data: bytes) -> None: # type: ignore[override] # Defensive:
	if not data:
	return
	raise NotImplementedError("Sending data isn't supported yet")

	def getresponse( # type: ignore[override]
	self,
	) -> HTTP2Response:
	status = None
	data = bytearray()
	with self._h2_conn as h2_conn:
	end_stream = False
	while not end_stream:
	# TODO: Arbitrary read value.
	if received_data := self.sock.recv(65535):
	events = h2_conn.receive_data(received_data)
	for event in events:
	if isinstance(event, h2.events.ResponseReceived):
	headers = HTTPHeaderDict()
	for header, value in event.headers:
	if header == b":status":
	status = int(value.decode())
	else:
	headers.add(
	header.decode("ascii"), value.decode("ascii")
	)

	elif isinstance(event, h2.events.DataReceived):
	data += event.data
	h2_conn.acknowledge_received_data(
	event.flow_controlled_length, event.stream_id
	)

	elif isinstance(event, h2.events.StreamEnded):
	end_stream = True

	if data_to_send := h2_conn.data_to_send():
	self.sock.sendall(data_to_send)

	# We always close to not have to handle connection management.
	self.close()

	assert status is not None
	return HTTP2Response(
	status=status,
	headers=headers,
	request_url=self._request_url,
	data=bytes(data),
	)

	def close(self) -> None:
	with self._h2_conn as h2_conn:
	try:
	h2_conn.close_connection()
	if data := h2_conn.data_to_send():
	self.sock.sendall(data)
	except Exception:
	pass

	# Reset all our HTTP/2 connection state.
	self._h2_conn = self._new_h2_conn()
	self._h2_stream = None
	self._h2_headers = []

	super().close()


	class HTTP2Response(BaseHTTPResponse):
	# TODO: This is a woefully incomplete response object, but works for non-streaming.
	def __init__(
	self,
	status: int,
	headers: HTTPHeaderDict,
	request_url: str,
	data: bytes,
	decode_content: bool = False, # TODO: support decoding
	) -> None:
	super().__init__(
	status=status,
	headers=headers,
	# Following CPython, we map HTTP versions to major * 10 + minor integers
	version=20,
	# No reason phrase in HTTP/2
	reason=None,
	decode_content=decode_content,
	request_url=request_url,
	)
	self._data = data
	self.length_remaining = 0

	@property
	def data(self) -> bytes:
	return self._data

	def get_redirect_location(self) -> None:
	return None

	def close(self) -> None:
	pass


	def inject_into_urllib3() -> None:
	HTTPSConnectionPool.ConnectionCls = HTTP2Connection
	urllib3.connection.HTTPSConnection = HTTP2Connection # type: ignore[misc]

	# TODO: Offer 'http/1.1' as well, but for testing purposes this is handy.
	urllib3.util.ssl_.ALPN_PROTOCOLS = ["h2"]


	def extract_from_urllib3() -> None:
	HTTPSConnectionPool.ConnectionCls = orig_HTTPSConnection
	urllib3.connection.HTTPSConnection = orig_HTTPSConnection # type: ignore[misc]

	urllib3.util.ssl_.ALPN_PROTOCOLS = ["http/1.1"]