multiplexing media through cfproxy

proxy/h2_transport.py

@@ -0,0 +1,332 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from collections import deque
+
+import httpx
+from h2.config import H2Configuration
+from h2.connection import H2Connection
+from h2.errors import ErrorCodes
+from h2.events import (
+    ConnectionTerminated, DataReceived, RemoteSettingsChanged,
+    ResponseReceived, StreamEnded, StreamReset, WindowUpdated,
+)
+from h2.exceptions import H2Error, StreamClosedError
+from h2.settings import SettingCodes, Settings
+
+STREAM_RECEIVE_WINDOW = 256 * 1024
+CONNECTION_RECEIVE_WINDOW = 4 * 1024 * 1024
+
+
+class _ResponseStream(httpx.AsyncByteStream):
+    def __init__(self, connection, stream_id, read_timeout):
+        self.connection = connection
+        self.stream_id = stream_id
+        self.read_timeout = read_timeout
+        self.changed = asyncio.Event()
+        self.chunks = deque()
+        self.headers = None
+        self.response_started = None
+        self.error = None
+        self.ended = False
+        self.closed = False
+        self.extensions = {'http_version': b'HTTP/2', 'stream_id': stream_id}
+
+    async def wait(self):
+        try:
+            await asyncio.wait_for(self.changed.wait(), self.read_timeout)
+        except asyncio.TimeoutError as exc:
+            raise httpx.ReadTimeout('Timed out waiting for H2 stream %d' % self.stream_id) from exc
+
+    def check_error(self):
+        if self.error is not None:
+            raise self.error
+
+    async def __aiter__(self):
+        while not self.closed:
+            self.changed.clear()
+            self.check_error()
+            if self.chunks:
+                data, credit = self.chunks.popleft()
+                yield data
+                self.connection.consumed(self, credit)
+            elif self.ended:
+                return
+            else:
+                await self.wait()
+
+    async def aclose(self):
+        if not self.closed:
+            self.closed = True
+            self.chunks.clear()
+            self.connection.release(self)
+            self.changed.set()
+
+
+class _Connection:
+    def __init__(self, reader, writer, max_streams, write_timeout):
+        self.reader, self.writer = reader, writer
+        self.max_streams = max_streams
+        self.write_timeout = write_timeout
+        self.h2 = H2Connection(config=H2Configuration(client_side=True))
+        self.h2.local_settings = Settings(
+            client=True, initial_values={SettingCodes.ENABLE_PUSH: 0,
+                                         SettingCodes.INITIAL_WINDOW_SIZE: STREAM_RECEIVE_WINDOW})
+        self.streams = {}
+        self.changed = asyncio.Event()
+        self.write_lock = asyncio.Lock()
+        self.settings_received = False
+        self.error = None
+        self.h2.initiate_connection()
+        self.h2.increment_flow_control_window(CONNECTION_RECEIVE_WINDOW - 65535)
+        self.flush()
+        self.reader_task = asyncio.create_task(self.read_loop())
+
+    def get_extra_info(self, name):
+        return self.writer.get_extra_info(name)
+
+    def flush(self):
+        data = self.h2.data_to_send()
+        if data:
+            self.writer.write(data)
+
+    async def drain(self):
+        try:
+            async with self.write_lock:
+                await asyncio.wait_for(self.writer.drain(), self.write_timeout)
+        except asyncio.TimeoutError as exc:
+            error = httpx.WriteTimeout('H2 connection write timed out')
+            self.fail(error)
+            raise error from exc
+        except OSError as exc:
+            error = httpx.WriteError(str(exc))
+            self.fail(error)
+            raise error from exc
+
+    def fail(self, error):
+        if self.error is None:
+            self.error = error
+            for stream in self.streams.values():
+                if not stream.ended:
+                    stream.error = error
+                stream.changed.set()
+            self.changed.set()
+            self.writer.close()
+
+    async def read_loop(self):
+        try:
+            while self.error is None:
+                data = await self.reader.read(64 * 1024)
+                if not data:
+                    raise httpx.ReadError('H2 peer closed the connection')
+                events = self.h2.receive_data(data)
+                terminated = any(isinstance(e, ConnectionTerminated) for e in events)
+                for event in events:
+                    if isinstance(event, RemoteSettingsChanged):
+                        self.settings_received = True
+                        self.changed.set()
+                    elif isinstance(event, WindowUpdated):
+                        self.changed.set()
+                    elif isinstance(event, ConnectionTerminated):
+                        raise httpx.RemoteProtocolError('H2 GOAWAY error=%s last_stream=%s' % (
+                            event.error_code, event.last_stream_id))
+                    else:
+                        stream = self.streams.get(getattr(event, 'stream_id', None))
+                        if isinstance(event, DataReceived):
+                            if event.flow_controlled_length and not terminated:
+                                self.h2.increment_flow_control_window(event.flow_controlled_length)
+                            if stream is not None:
+                                if event.data:
+                                    stream.chunks.append((event.data, event.flow_controlled_length))
+                                elif not terminated:
+                                    self.consumed(stream, event.flow_controlled_length)
+                        elif stream is not None:
+                            if isinstance(event, ResponseReceived):
+                                stream.headers = event.headers
+                                if stream.response_started is not None:
+                                    stream.response_started()
+                            elif isinstance(event, StreamEnded):
+                                stream.ended = True
+                            elif isinstance(event, StreamReset):
+                                stream.error = httpx.RemoteProtocolError(
+                                    'H2 stream %d reset: %s' % (event.stream_id, event.error_code))
+                                self.changed.set()
+                        if stream is not None:
+                            stream.changed.set()
+                self.flush()
+                await self.drain()
+        except asyncio.CancelledError:
+            self.fail(httpx.ReadError('H2 reader stopped'))
+            raise
+        except (OSError, H2Error, httpx.HTTPError) as exc:
+            error = exc if isinstance(exc, httpx.HTTPError) else httpx.RemoteProtocolError(str(exc))
+            self.fail(error)
+
+    def check_error(self):
+        if self.error is not None:
+            raise self.error
+
+    async def reserve(self, timeouts):
+        timeout = timeouts.get('pool')
+        deadline = None if timeout is None else time.monotonic() + timeout
+        while True:
+            self.changed.clear()
+            self.check_error()
+            if self.settings_received and len(self.streams) < min(
+                    self.max_streams, self.h2.remote_settings.max_concurrent_streams):
+                stream_id = self.h2.get_next_available_stream_id()
+                stream = _ResponseStream(self, stream_id, timeouts.get('read'))
+                self.streams[stream_id] = stream
+                return stream
+            remaining = None if deadline is None else max(0, deadline - time.monotonic())
+            try:
+                await asyncio.wait_for(self.changed.wait(), remaining)
+            except asyncio.TimeoutError as exc:
+                raise httpx.PoolTimeout('H2 stream capacity exhausted') from exc
+
+    async def send_body(self, request, stream):
+        async for data in request.stream:
+            offset = 0
+            while offset < len(data):
+                self.changed.clear()
+                self.check_error()
+                stream.check_error()
+                count = min(len(data) - offset, self.h2.max_outbound_frame_size,
+                            self.h2.local_flow_control_window(stream.stream_id))
+                if count <= 0:
+                    try:
+                        await asyncio.wait_for(self.changed.wait(), self.write_timeout)
+                    except asyncio.TimeoutError as exc:
+                        raise httpx.WriteTimeout('H2 upload flow control timed out') from exc
+                    continue
+                self.h2.send_data(stream.stream_id, data[offset:offset + count])
+                offset += count
+                self.flush()
+                await self.drain()
+        self.h2.end_stream(stream.stream_id)
+        self.flush()
+        await self.drain()
+
+    async def request(self, request):
+        timeouts = request.extensions.get('timeout', {})
+        trace = request.extensions.get('trace')
+        if trace:
+            await trace('http2.wait_for_stream.started', {})
+        stream = await self.reserve(timeouts)
+        stream.response_started = request.extensions.get('h2_response_started')
+        stream_id = stream.stream_id
+        try:
+            headers = [(b':method', request.method.encode('ascii')),
+                       (b':scheme', request.url.raw_scheme),
+                       (b':authority', request.headers['host'].encode('ascii')),
+                       (b':path', request.url.raw_path)]
+            headers.extend((k.lower(), v) for k, v in request.headers.raw if k.lower() not in (
+                b'host', b'connection', b'keep-alive', b'proxy-connection',
+                b'upgrade', b'transfer-encoding'))
+            self.h2.send_headers(stream_id, headers)
+            self.flush()
+            if trace:
+                await trace('http2.send_request_headers.started', {'stream_id': stream_id})
+            await self.send_body(request, stream)
+            if trace:
+                await trace('http2.send_request_body.complete', {'stream_id': stream_id})
+            while stream.headers is None:
+                stream.changed.clear()
+                stream.check_error()
+                await stream.wait()
+            stream.check_error()
+            status = next(int(v) for k, v in stream.headers if k == b':status')
+            return httpx.Response(status, headers=[(k, v) for k, v in stream.headers
+                                                  if not k.startswith(b':')],
+                                  stream=stream, extensions=stream.extensions)
+        except H2Error as exc:
+            await stream.aclose()
+            raise httpx.LocalProtocolError(str(exc)) from exc
+        except BaseException:
+            await stream.aclose()
+            raise
+
+    def consumed(self, stream, credit):
+        if credit and not stream.ended and not stream.closed and self.error is None:
+            try:
+                self.h2.increment_flow_control_window(credit, stream.stream_id)
+                self.flush()
+            except StreamClosedError:
+                pass
+
+    def release(self, stream):
+        self.streams.pop(stream.stream_id, None)
+        if self.error is None:
+            try:
+                self.h2.reset_stream(stream.stream_id, ErrorCodes.CANCEL)
+                self.flush()
+            except StreamClosedError:
+                pass
+            except OSError as exc:
+                self.fail(httpx.WriteError(str(exc)))
+        self.changed.set()
+
+    async def aclose(self):
+        self.fail(httpx.ReadError('H2 connection closed'))
+        self.reader_task.cancel()
+        await asyncio.gather(self.reader_task, return_exceptions=True)
+        try:
+            await asyncio.wait_for(self.writer.wait_closed(), 2)
+        except (OSError, asyncio.TimeoutError):
+            pass
+
+
+class H2Transport(httpx.AsyncBaseTransport):
+    def __init__(self, ssl_context, max_streams=64):
+        self.ssl_context = ssl_context
+        self.ssl_context.set_alpn_protocols(['h2'])
+        self.max_streams = max_streams
+        self.connection = None
+        self.origin = None
+        self.connect_lock = asyncio.Lock()
+        self.closed = False
+
+    async def handle_async_request(self, request):
+        origin = (request.url.host, request.url.port or 443)
+        if request.url.scheme != 'https' or self.origin not in (None, origin):
+            raise httpx.UnsupportedProtocol('CF H2 transport requires a single HTTPS origin')
+        async with self.connect_lock:
+            if self.closed:
+                raise httpx.ConnectError('H2 transport closed')
+            self.origin = origin
+            if self.connection is None or self.connection.error is not None:
+                if self.connection is not None:
+                    await self.connection.aclose()
+                timeouts = request.extensions.get('timeout', {})
+                trace = request.extensions.get('trace')
+                if trace:
+                    await trace('connection.connect_tcp.started', {})
+                writer = None
+                try:
+                    reader, writer = await asyncio.wait_for(asyncio.open_connection(
+                        *origin, ssl=self.ssl_context, server_hostname=origin[0]),
+                        timeouts.get('connect'))
+                    if writer.get_extra_info('ssl_object').selected_alpn_protocol() != 'h2':
+                        raise httpx.RemoteProtocolError('CF origin did not negotiate h2')
+                    self.connection = _Connection(reader, writer, self.max_streams, timeouts.get('write'))
+                    if trace:
+                        await trace('connection.start_tls.complete', {'return_value': self.connection})
+                except BaseException as exc:
+                    if self.connection is not None:
+                        self.connection.fail(httpx.ConnectError('H2 connection setup interrupted'))
+                    if writer is not None:
+                        writer.close()
+                    if isinstance(exc, asyncio.TimeoutError):
+                        raise httpx.ConnectTimeout('H2 TLS connect timed out') from exc
+                    if isinstance(exc, OSError):
+                        raise httpx.ConnectError(str(exc)) from exc
+                    raise
+            connection = self.connection
+        return await connection.request(request)
+
+    async def aclose(self):
+        self.closed = True
+        async with self.connect_lock:
+            if self.connection is not None:
+                await self.connection.aclose()