231 lines
7.3 KiB
Python
231 lines
7.3 KiB
Python
|
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,
|
||
|
|
version_string="HTTP/2",
|
||
|
|
# 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"]
|