Spaces:
Sleeping
Sleeping
File size: 7,517 Bytes
392d02b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 |
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"]
|