omserv 0.0.0.dev7__py3-none-any.whl

omserv/server/protocols/h11.py

@@ -0,0 +1,334 @@
+import itertools
+import typing as ta
+
+import h11
+
+from ..config import Config
+from ..events import Body
+from ..events import Closed
+from ..events import Data
+from ..events import EndBody
+from ..events import EndData
+from ..events import InformationalResponse
+from ..events import ProtocolEvent
+from ..events import RawData
+from ..events import Request
+from ..events import Response
+from ..events import ServerEvent
+from ..events import StreamClosed
+from ..events import Updated
+from ..headers import response_headers
+from ..streams.httpstream import HttpStream
+from ..streams.wsstream import WsStream
+from ..taskspawner import TaskSpawner
+from ..types import AppWrapper
+from ..workercontext import WorkerContext
+from .types import Protocol
+
+
+H11SendableEvent: ta.TypeAlias = ta.Union[  # noqa
+    h11.Data,
+    h11.EndOfMessage,
+    h11.InformationalResponse,
+    h11.Response,
+]
+
+
+STREAM_ID = 1
+
+
+class H2CProtocolRequiredError(Exception):
+    def __init__(self, data: bytes, request: h11.Request) -> None:
+        super().__init__()
+        settings = ''
+        headers = [(b':method', request.method), (b':path', request.target)]
+        for name, value in request.headers:
+            if name.lower() == b'http2-settings':
+                settings = value.decode()
+            elif name.lower() == b'host':
+                headers.append((b':authority', value))
+            headers.append((name, value))
+
+        self.data = data
+        self.headers = headers
+        self.settings = settings
+
+
+class H2ProtocolAssumedError(Exception):
+    def __init__(self, data: bytes) -> None:
+        super().__init__()
+        self.data = data
+
+
+class H11WsConnection:
+    # This class matches the h11 interface, and either passes data
+    # through without altering it (for Data, EndData) or sends h11
+    # events (Response, Body, EndBody).
+    our_state = None  # Prevents recycling the connection
+    they_are_waiting_for_100_continue = False
+    their_state = None
+    trailing_data = (b'', False)
+
+    def __init__(self, h11_connection: h11.Connection) -> None:
+        self.buffer = bytearray(h11_connection.trailing_data[0])
+        self.h11_connection = h11_connection
+
+    def receive_data(self, data: bytes) -> None:
+        self.buffer.extend(data)
+
+    def next_event(self) -> Data | type[h11.NEED_DATA]:
+        if self.buffer:
+            event = Data(stream_id=STREAM_ID, data=bytes(self.buffer))
+            self.buffer = bytearray()
+            return event
+        else:
+            return h11.NEED_DATA
+
+    def send(self, event: H11SendableEvent) -> bytes | None:
+        return self.h11_connection.send(event)
+
+    def start_next_cycle(self) -> None:
+        pass
+
+
+class H11Protocol(Protocol):
+    def __init__(
+            self,
+            app: AppWrapper,
+            config: Config,
+            context: WorkerContext,
+            task_spawner: TaskSpawner,
+            client: tuple[str, int] | None,
+            server: tuple[str, int] | None,
+            send: ta.Callable[[ServerEvent], ta.Awaitable[None]],
+    ) -> None:
+        self.app = app
+        self.can_read = context.event_class()
+        self.client = client
+        self.config = config
+        self.connection: h11.Connection | H11WsConnection = h11.Connection(
+            h11.SERVER, max_incomplete_event_size=self.config.h11_max_incomplete_size,
+        )
+        self.context = context
+        self.keep_alive_requests = 0
+        self.send = send
+        self.server = server
+        self.stream: HttpStream | WsStream | None = None
+        self.task_spawner = task_spawner
+
+    async def initiate(self) -> None:
+        pass
+
+    async def handle(self, event: ServerEvent) -> None:
+        if isinstance(event, RawData):
+            self.connection.receive_data(event.data)
+            await self._handle_events()
+        elif isinstance(event, Closed):
+            if self.stream is not None:
+                await self._close_stream()
+
+    async def stream_send(self, event: ProtocolEvent) -> None:
+        if isinstance(event, Response):
+            if event.status_code >= 200:
+                headers = list(itertools.chain(event.headers, response_headers(self.config, 'h11')))
+                if self.keep_alive_requests >= self.config.keep_alive_max_requests:
+                    headers.append((b'connection', b'close'))
+                await self._send_h11_event(h11.Response(
+                    headers=headers,
+                    status_code=event.status_code,
+                ))
+            else:
+                await self._send_h11_event(h11.InformationalResponse(
+                    headers=list(itertools.chain(event.headers, response_headers(self.config, 'h11'))),
+                    status_code=event.status_code,
+                ))
+
+        elif isinstance(event, InformationalResponse):
+            pass  # Ignore for HTTP/1
+
+        elif isinstance(event, Body):
+            await self._send_h11_event(h11.Data(data=event.data))
+
+        elif isinstance(event, EndBody):
+            await self._send_h11_event(h11.EndOfMessage())
+
+        elif isinstance(event, Data):
+            await self.send(RawData(data=event.data))
+
+        elif isinstance(event, EndData):
+            pass
+
+        elif isinstance(event, StreamClosed):
+            await self._maybe_recycle()
+
+    async def _handle_events(self) -> None:
+        while True:
+            if self.connection.they_are_waiting_for_100_continue:
+                await self._send_h11_event(h11.InformationalResponse(
+                    status_code=100, headers=response_headers(self.config, 'h11'),
+                ))
+
+            try:
+                event = self.connection.next_event()
+            except h11.RemoteProtocolError as error:
+                if self.connection.our_state in {h11.IDLE, h11.SEND_RESPONSE}:
+                    await self._send_error_response(error.error_status_hint)
+                await self.send(Closed())
+                break
+
+            else:
+                if isinstance(event, h11.Request):
+                    await self.send(Updated(idle=False))
+                    await self._check_protocol(event)
+                    await self._create_stream(event)
+
+                elif event is h11.PAUSED:
+                    await self.can_read.clear()
+                    await self.can_read.wait()
+
+                elif isinstance(event, h11.ConnectionClosed) or event is h11.NEED_DATA:
+                    break
+
+                elif self.stream is None:
+                    break
+
+                elif isinstance(event, h11.Data):
+                    await self.stream.handle(Body(stream_id=STREAM_ID, data=event.data))
+
+                elif isinstance(event, h11.EndOfMessage):
+                    await self.stream.handle(EndBody(stream_id=STREAM_ID))
+
+                elif isinstance(event, Data):
+                    # WebSocket pass through
+                    await self.stream.handle(event)
+
+    async def _create_stream(self, request: h11.Request) -> None:
+        upgrade_value = ''
+        connection_value = ''
+        for name, value in request.headers:
+            sanitised_name = name.decode('latin1').strip().lower()
+            if sanitised_name == 'upgrade':
+                upgrade_value = value.decode('latin1').strip()
+            elif sanitised_name == 'connection':
+                connection_value = value.decode('latin1').strip()
+
+        connection_tokens = connection_value.lower().split(',')
+        if (
+                any(token.strip() == 'upgrade' for token in connection_tokens)
+                and upgrade_value.lower() == 'websocket'
+                and request.method.decode('ascii').upper() == 'GET'
+        ):
+            self.stream = WsStream(
+                self.app,
+                self.config,
+                self.context,
+                self.task_spawner,
+                self.client,
+                self.server,
+                self.stream_send,
+                STREAM_ID,
+            )
+            self.connection = H11WsConnection(ta.cast(h11.Connection, self.connection))
+        else:
+            self.stream = HttpStream(
+                self.app,
+                self.config,
+                self.context,
+                self.task_spawner,
+                self.client,
+                self.server,
+                self.stream_send,
+                STREAM_ID,
+            )
+
+        if self.config.h11_pass_raw_headers:
+            headers = request.headers.raw_items()
+        else:
+            headers = list(request.headers)
+
+        await self.stream.handle(  # type: ignore
+            Request(
+                stream_id=STREAM_ID,
+                headers=headers,
+                http_version=request.http_version.decode(),
+                method=request.method.decode('ascii').upper(),
+                raw_path=request.target,
+            ),
+        )
+        self.keep_alive_requests += 1
+        await self.context.mark_request()
+
+    async def _send_h11_event(self, event: H11SendableEvent) -> None:
+        try:
+            data = self.connection.send(event)
+        except h11.LocalProtocolError:
+            if self.connection.their_state != h11.ERROR:
+                raise
+        else:
+            if data is None:
+                raise Exception('closed')  # FIXME
+            await self.send(RawData(data=data))
+
+    async def _send_error_response(self, status_code: int) -> None:
+        await self._send_h11_event(h11.Response(
+            status_code=status_code,
+            headers=list(
+                itertools.chain(
+                    [(b'content-length', b'0'), (b'connection', b'close')],
+                    response_headers(self.config, 'h11'),
+                ),
+            ),
+        ))
+        await self._send_h11_event(h11.EndOfMessage())
+
+    async def _maybe_recycle(self) -> None:
+        await self._close_stream()
+        if (
+                not self.context.terminated.is_set()
+                and self.connection.our_state is h11.DONE
+                and self.connection.their_state is h11.DONE
+        ):
+            try:
+                self.connection.start_next_cycle()
+            except h11.LocalProtocolError:
+                await self.send(Closed())
+            else:
+                self.response = None
+                self.scope = None
+                await self.can_read.set()
+                await self.send(Updated(idle=True))
+        else:
+            await self.can_read.set()
+            await self.send(Closed())
+
+    async def _close_stream(self) -> None:
+        if self.stream is not None:
+            await self.stream.handle(StreamClosed(stream_id=STREAM_ID))
+            self.stream = None
+
+    async def _check_protocol(self, event: h11.Request) -> None:
+        upgrade_value = ''
+        has_body = False
+        for name, value in event.headers:
+            sanitised_name = name.decode('latin1').strip().lower()
+            if sanitised_name == 'upgrade':
+                upgrade_value = value.decode('latin1').strip()
+            elif sanitised_name in {'content-length', 'transfer-encoding'}:
+                has_body = True
+
+        # h2c Upgrade requests with a body are a pain as the body must be fully recieved in HTTP/1.1 before the upgrade
+        # response and HTTP/2 takes over, so Hypercorn ignores the upgrade and responds in HTTP/1.1. Use a preflight
+        # OPTIONS request to initiate the upgrade if really required (or just use h2).
+        if upgrade_value.lower() == 'h2c' and not has_body:
+            await self._send_h11_event(h11.InformationalResponse(
+                status_code=101,
+                headers=[
+                    *response_headers(self.config, 'h11'),
+                    (b'connection', b'upgrade'),
+                    (b'upgrade', b'h2c'),
+                ],
+            ))
+            raise H2CProtocolRequiredError(self.connection.trailing_data[0], event)
+
+        elif event.method == b'PRI' and event.target == b'*' and event.http_version == b'2.0':
+            raise H2ProtocolAssumedError(b'PRI * HTTP/2.0\r\n\r\n' + self.connection.trailing_data[0])