Add support for lifespan state

This allows ASGI apps to store state during startup that is then
passed in every scope.

src/hypercorn/asyncio/lifespan.py

@@ -5,7 +5,7 @@
 from typing import Any, Callable
 
 from ..config import Config
-from ..typing import AppWrapper, ASGIReceiveEvent, ASGISendEvent, LifespanScope
+from ..typing import AppWrapper, ASGIReceiveEvent, ASGISendEvent, LifespanScope, LifespanState
 from ..utils import LifespanFailureError, LifespanTimeoutError
 
 
@@ -14,14 +14,21 @@ class UnexpectedMessageError(Exception):
 
 
 class Lifespan:
-    def __init__(self, app: AppWrapper, config: Config, loop: asyncio.AbstractEventLoop) -> None:
+    def __init__(
+        self,
+        app: AppWrapper,
+        config: Config,
+        loop: asyncio.AbstractEventLoop,
+        lifespan_state: LifespanState,
+    ) -> None:
         self.app = app
         self.config = config
         self.startup = asyncio.Event()
         self.shutdown = asyncio.Event()
         self.app_queue: asyncio.Queue = asyncio.Queue(config.max_app_queue_size)
         self.supported = True
         self.loop = loop
+        self.state = lifespan_state
 
         # This mimics the Trio nursery.start task_status and is
         # required to ensure the support has been checked before
@@ -33,6 +40,7 @@ async def handle_lifespan(self) -> None:
         scope: LifespanScope = {
             "type": "lifespan",
             "asgi": {"spec_version": "2.0", "version": "3.0"},
+            "state": self.state,
         }
 
         def _call_soon(func: Callable, *args: Any) -> Any:

src/hypercorn/asyncio/run.py

@@ -18,7 +18,7 @@
 from .udp_server import UDPServer
 from .worker_context import WorkerContext
 from ..config import Config, Sockets
-from ..typing import AppWrapper
+from ..typing import AppWrapper, LifespanState
 from ..utils import (
     check_multiprocess_shutdown_event,
     load_application,
@@ -77,7 +77,8 @@ def _signal_handler(*_: Any) -> None:  # noqa: N803
 
         shutdown_trigger = signal_event.wait
 
-    lifespan = Lifespan(app, config, loop)
+    lifespan_state: LifespanState = {}
+    lifespan = Lifespan(app, config, loop, lifespan_state)
 
     lifespan_task = loop.create_task(lifespan.handle_lifespan())
     await lifespan.wait_for_startup()
@@ -106,7 +107,7 @@ async def _server_callback(reader: asyncio.StreamReader, writer: asyncio.StreamW
         task = asyncio.current_task(loop)
         server_tasks.add(task)
         task.add_done_callback(server_tasks.discard)
-        await TCPServer(app, loop, config, context, reader, writer)
+        await TCPServer(app, loop, config, context, lifespan_state, reader, writer)
 
     servers = []
     for sock in sockets.secure_sockets:
@@ -140,7 +141,7 @@ async def _server_callback(reader: asyncio.StreamReader, writer: asyncio.StreamW
             sock = _share_socket(sock)
 
         _, protocol = await loop.create_datagram_endpoint(
-            lambda: UDPServer(app, loop, config, context), sock=sock
+            lambda: UDPServer(app, loop, config, context, lifespan_state), sock=sock
         )
         task = loop.create_task(protocol.run())
         server_tasks.add(task)

src/hypercorn/asyncio/tcp_server.py

@@ -9,7 +9,7 @@
 from ..config import Config
 from ..events import Closed, Event, RawData, Updated
 from ..protocol import ProtocolWrapper
-from ..typing import AppWrapper
+from ..typing import AppWrapper, ConnectionState, LifespanState
 from ..utils import parse_socket_addr
 
 MAX_RECV = 2**16
@@ -22,6 +22,7 @@ def __init__(
         loop: asyncio.AbstractEventLoop,
         config: Config,
         context: WorkerContext,
+        state: LifespanState,
         reader: asyncio.StreamReader,
         writer: asyncio.StreamWriter,
     ) -> None:
@@ -33,6 +34,7 @@ def __init__(
         self.reader = reader
         self.writer = writer
         self.send_lock = asyncio.Lock()
+        self.state = state
         self.idle_task = AsyncioSingleTask()
 
     def __await__(self) -> Generator[Any, None, None]:
@@ -58,6 +60,7 @@ async def run(self) -> None:
                     self.config,
                     self.context,
                     task_group,
+                    ConnectionState(self.state.copy()),
                     ssl,
                     client,
                     server,

src/hypercorn/asyncio/udp_server.py

@@ -7,7 +7,7 @@
 from .worker_context import WorkerContext
 from ..config import Config
 from ..events import Event, RawData
-from ..typing import AppWrapper
+from ..typing import AppWrapper, ConnectionState, LifespanState
 from ..utils import parse_socket_addr
 
 if TYPE_CHECKING:
@@ -22,6 +22,7 @@ def __init__(
         loop: asyncio.AbstractEventLoop,
         config: Config,
         context: WorkerContext,
+        state: LifespanState,
     ) -> None:
         self.app = app
         self.config = config
@@ -30,6 +31,7 @@ def __init__(
         self.protocol: "QuicProtocol"
         self.protocol_queue: asyncio.Queue = asyncio.Queue(10)
         self.transport: Optional[asyncio.DatagramTransport] = None
+        self.state = state
 
     def connection_made(self, transport: asyncio.DatagramTransport) -> None:  # type: ignore
         self.transport = transport
@@ -48,7 +50,13 @@ async def run(self) -> None:
         server = parse_socket_addr(socket.family, socket.getsockname())
         async with TaskGroup(self.loop) as task_group:
             self.protocol = QuicProtocol(
-                self.app, self.config, self.context, task_group, server, self.protocol_send
+                self.app,
+                self.config,
+                self.context,
+                task_group,
+                ConnectionState(self.state.copy()),
+                server,
+                self.protocol_send,
             )
 
             while not self.context.terminated.is_set() or not self.protocol.idle:

src/hypercorn/protocol/__init__.py

@@ -6,7 +6,7 @@
 from .h11 import H2CProtocolRequiredError, H2ProtocolAssumedError, H11Protocol
 from ..config import Config
 from ..events import Event, RawData
-from ..typing import AppWrapper, TaskGroup, WorkerContext
+from ..typing import AppWrapper, ConnectionState, TaskGroup, WorkerContext
 
 
 class ProtocolWrapper:
@@ -16,6 +16,7 @@ def __init__(
         config: Config,
         context: WorkerContext,
         task_group: TaskGroup,
+        state: ConnectionState,
         ssl: bool,
         client: Optional[Tuple[str, int]],
         server: Optional[Tuple[str, int]],
@@ -30,13 +31,15 @@ def __init__(
         self.client = client
         self.server = server
         self.send = send
+        self.state = state
         self.protocol: Union[H11Protocol, H2Protocol]
         if alpn_protocol == "h2":
             self.protocol = H2Protocol(
                 self.app,
                 self.config,
                 self.context,
                 self.task_group,
+                self.state,
                 self.ssl,
                 self.client,
                 self.server,
@@ -48,6 +51,7 @@ def __init__(
                 self.config,
                 self.context,
                 self.task_group,
+                self.state,
                 self.ssl,
                 self.client,
                 self.server,
@@ -66,6 +70,7 @@ async def handle(self, event: Event) -> None:
                 self.config,
                 self.context,
                 self.task_group,
+                self.state,
                 self.ssl,
                 self.client,
                 self.server,
@@ -80,6 +85,7 @@ async def handle(self, event: Event) -> None:
                 self.config,
                 self.context,
                 self.task_group,
+                self.state,
                 self.ssl,
                 self.client,
                 self.server,

src/hypercorn/protocol/events.py

@@ -3,6 +3,8 @@
 from dataclasses import dataclass
 from typing import List, Tuple
 
+from hypercorn.typing import ConnectionState
+
 
 @dataclass(frozen=True)
 class Event:
@@ -15,6 +17,7 @@ class Request(Event):
     http_version: str
     method: str
     raw_path: bytes
+    state: ConnectionState
 
 
 @dataclass(frozen=True)

src/hypercorn/protocol/h11.py

@@ -20,7 +20,7 @@
 from .ws_stream import WSStream
 from ..config import Config
 from ..events import Closed, Event, RawData, Updated
-from ..typing import AppWrapper, H11SendableEvent, TaskGroup, WorkerContext
+from ..typing import AppWrapper, ConnectionState, H11SendableEvent, TaskGroup, WorkerContext
 
 STREAM_ID = 1
 
@@ -84,6 +84,7 @@ def __init__(
         config: Config,
         context: WorkerContext,
         task_group: TaskGroup,
+        connection_state: ConnectionState,
         ssl: bool,
         client: Optional[Tuple[str, int]],
         server: Optional[Tuple[str, int]],
@@ -103,6 +104,7 @@ def __init__(
         self.ssl = ssl
         self.stream: Optional[Union[HTTPStream, WSStream]] = None
         self.task_group = task_group
+        self.connection_state = connection_state
 
     async def initiate(self) -> None:
         pass
@@ -236,6 +238,7 @@ async def _create_stream(self, request: h11.Request) -> None:
                 http_version=request.http_version.decode(),
                 method=request.method.decode("ascii").upper(),
                 raw_path=request.target,
+                state=self.connection_state,
             )
         )
         self.keep_alive_requests += 1

src/hypercorn/protocol/h2.py

@@ -24,7 +24,7 @@
 from .ws_stream import WSStream
 from ..config import Config
 from ..events import Closed, Event, RawData, Updated
-from ..typing import AppWrapper, Event as IOEvent, TaskGroup, WorkerContext
+from ..typing import AppWrapper, ConnectionState, Event as IOEvent, TaskGroup, WorkerContext
 from ..utils import filter_pseudo_headers
 
 BUFFER_HIGH_WATER = 2 * 2**14  # Twice the default max frame size (two frames worth)
@@ -85,6 +85,7 @@ def __init__(
         config: Config,
         context: WorkerContext,
         task_group: TaskGroup,
+        connection_state: ConnectionState,
         ssl: bool,
         client: Optional[Tuple[str, int]],
         server: Optional[Tuple[str, int]],
@@ -96,6 +97,7 @@ def __init__(
         self.config = config
         self.context = context
         self.task_group = task_group
+        self.connection_state = connection_state
 
         self.connection = h2.connection.H2Connection(
             config=h2.config.H2Configuration(client_side=False, header_encoding=None)
@@ -360,6 +362,7 @@ async def _create_stream(self, request: h2.events.RequestReceived) -> None:
                 http_version="2",
                 method=method,
                 raw_path=raw_path,
+                state=self.connection_state,
             )
         )
         self.keep_alive_requests += 1

src/hypercorn/protocol/h3.py

@@ -23,7 +23,7 @@
 from .http_stream import HTTPStream
 from .ws_stream import WSStream
 from ..config import Config
-from ..typing import AppWrapper, TaskGroup, WorkerContext
+from ..typing import AppWrapper, ConnectionState, TaskGroup, WorkerContext
 from ..utils import filter_pseudo_headers
 
 
@@ -34,6 +34,7 @@ def __init__(
         config: Config,
         context: WorkerContext,
         task_group: TaskGroup,
+        state: ConnectionState,
         client: Optional[Tuple[str, int]],
         server: Optional[Tuple[str, int]],
         quic: QuicConnection,
@@ -48,6 +49,7 @@ def __init__(
         self.server = server
         self.streams: Dict[int, Union[HTTPStream, WSStream]] = {}
         self.task_group = task_group
+        self.state = state
 
     async def handle(self, quic_event: QuicEvent) -> None:
         for event in self.connection.handle_event(quic_event):
@@ -127,6 +129,7 @@ async def _create_stream(self, request: HeadersReceived) -> None:
                 http_version="3",
                 method=method,
                 raw_path=raw_path,
+                state=self.state,
             )
         )
         await self.context.mark_request()

src/hypercorn/protocol/http_stream.py

@@ -97,6 +97,7 @@ async def handle(self, event: Event) -> None:
                 "headers": event.headers,
                 "client": self.client,
                 "server": self.server,
+                "state": event.state,
                 "extensions": {},
             }
 
@@ -158,6 +159,7 @@ async def app_send(self, message: Optional[ASGISendEvent]) -> None:
                         http_version=self.scope["http_version"],
                         method="GET",
                         raw_path=message["path"].encode(),
+                        state=self.scope["state"],
                     )
                 )
             elif (

src/hypercorn/protocol/quic.py

@@ -23,7 +23,7 @@
 from .h3 import H3Protocol
 from ..config import Config
 from ..events import Closed, Event, RawData
-from ..typing import AppWrapper, SingleTask, TaskGroup, WorkerContext
+from ..typing import AppWrapper, ConnectionState, SingleTask, TaskGroup, WorkerContext
 
 
 @dataclass
@@ -41,6 +41,7 @@ def __init__(
         config: Config,
         context: WorkerContext,
         task_group: TaskGroup,
+        state: ConnectionState,
         server: Optional[Tuple[str, int]],
         send: Callable[[Event], Awaitable[None]],
     ) -> None:
@@ -51,6 +52,7 @@ def __init__(
         self.send = send
         self.server = server
         self.task_group = task_group
+        self.state = state
 
         self.quic_config = QuicConfiguration(alpn_protocols=H3_ALPN, is_client=False)
         self.quic_config.load_cert_chain(certfile=config.certfile, keyfile=config.keyfile)
@@ -128,6 +130,7 @@ async def _handle_events(
                     self.config,
                     self.context,
                     self.task_group,
+                    self.state,
                     client,
                     self.server,
                     connection,

src/hypercorn/protocol/ws_stream.py

@@ -219,6 +219,7 @@ async def handle(self, event: Event) -> None:
                 "headers": event.headers,
                 "client": self.client,
                 "server": self.server,
+                "state": event.state,
                 "subprotocols": self.handshake.subprotocols or [],
                 "extensions": {"websocket.http.response": {}},
             }