Source code
Revision control
Copy as Markdown
Other Tools
# mypy: allow-untyped-defs
"""WebSocket over HTTP/3 server for WPT.
This is the main entry point for the WebSocket H3 server. It owns the QUIC
server lifecycle, creates the HTTP/3 connection object, dispatches incoming H3
events, handles Extended CONNECT requests for WebSocket streams, and serves
normal WPT resources over the same H3 port through wptserve adapters.
"""
import asyncio
import contextlib
import logging
import os
import ssl
import sys
import threading
from io import BytesIO
from pathlib import Path
from typing import Any, Dict, List, Optional, Set, Tuple, cast
from aioquic.asyncio import QuicConnectionProtocol, serve
from aioquic.asyncio.client import connect
from aioquic.asyncio.protocol import QuicStreamAdapter
from aioquic.h3.connection import H3_ALPN, H3Connection
from aioquic.h3.events import DataReceived, H3Event, HeadersReceived
from aioquic.quic.configuration import QuicConfiguration
from aioquic.quic.connection import logger as quic_connection_logger
from aioquic.quic.connection import stream_is_client_initiated
from aioquic.quic.connection import stream_is_unidirectional
from aioquic.quic.events import (
ConnectionTerminated,
ProtocolNegotiated,
QuicEvent,
StreamReset,
)
from aioquic.tls import SessionTicket
from pywebsocket3 import dispatch
from pywebsocket3.handshake.base import AbortedByUserException
from pywebsocket3.handshake.base import HandshakeException
from tools import localpaths # noqa: F401
from .h3_wptserve_adapter import (
_H3_FORBIDDEN_RESPONSE_HEADERS,
_H3RequestHandlerAdapter,
_H3ServerAdapter,
create_h3_request_response_classes,
)
from .headers import H3Headers
from .websocket_h3_session import (
DEFAULT_MAX_UNREAD_WEBSOCKET_BYTES,
_WebSocketH3Request,
_WebSocketH3Session,
)
from .ws_h3_handshake import WsH3Handshaker
SERVER_NAME = "websocket-h3-server"
DEFAULT_MAX_PENDING_HANDSHAKES = 16
DEFAULT_MAX_ACTIVE_REQUESTS = 16
_logger: logging.Logger = logging.getLogger(__name__)
quic_connection_logger.setLevel(logging.WARNING)
class H3ConnectionForWebSocket(H3Connection):
"""H3Connection subclass that enables Extended CONNECT for WebSocket."""
ENABLE_CONNECT_PROTOCOL = 0x08
def _get_local_settings(self) -> Dict[int, int]:
settings = super()._get_local_settings()
settings[self.ENABLE_CONNECT_PROTOCOL] = 1
return settings
class WebSocketH3Protocol(QuicConnectionProtocol):
"""Handles one client connection to the WebSocket-over-H3 server.
One QUIC connection can carry many HTTP/3 streams. WebSocket CONNECT
streams are passed to pywebsocket handlers. Normal GET and HEAD streams are
served through the H3 wptserve adapter so test pages and resources can load
from the same origin.
"""
def __init__(
self,
*args: Any,
ws_doc_root: str = "",
router: Any = None,
h3_server_adapter: Optional[_H3ServerAdapter] = None,
h3_request_cls: Any = None,
h3_response_cls: Any = None,
logger: Optional[logging.Logger] = None,
max_pending_handshakes: int = DEFAULT_MAX_PENDING_HANDSHAKES,
max_active_requests: int = DEFAULT_MAX_ACTIVE_REQUESTS,
max_unread_websocket_bytes: int = DEFAULT_MAX_UNREAD_WEBSOCKET_BYTES,
**kwargs: Any,
) -> None:
"""Stores the state needed while handling one QUIC connection."""
super().__init__(*args, **kwargs)
self._http: Optional[H3ConnectionForWebSocket] = None
self._sessions: Dict[int, _WebSocketH3Session] = {}
self._active_requests = 0
self._pending_handshakes = 0
self._workers: Set[asyncio.Task[Any]] = set()
self._queued_workers: Dict[int, asyncio.Task[Any]] = {}
self._handshake_slots = asyncio.Semaphore(max_pending_handshakes)
self._request_slots = asyncio.Semaphore(max_active_requests)
self._max_unread_websocket_bytes = max_unread_websocket_bytes
self._ws_doc_root = ws_doc_root
self._router = router
self._h3_server_adapter = h3_server_adapter
self._h3_request_cls = h3_request_cls
self._h3_response_cls = h3_response_cls
self._logger = logger if logger is not None else _logger
def quic_event_received(self, event: QuicEvent) -> None:
"""Receives raw QUIC events from aioquic and handles related sessions."""
if isinstance(event, ProtocolNegotiated):
self._http = H3ConnectionForWebSocket(self._quic)
if self._http is not None:
for h3_event in self._http.handle_event(event):
self._h3_event_received(h3_event)
if isinstance(event, ConnectionTerminated):
self._close_all_sessions()
elif isinstance(event, StreamReset):
queued = self._queued_workers.pop(event.stream_id, None)
if queued is not None:
queued.cancel()
session = self._sessions.pop(event.stream_id, None)
if session is not None:
session.abort()
def _h3_event_received(self, event: H3Event) -> None:
"""Routes HTTP/3 events to the matching headers or data handler."""
if isinstance(event, HeadersReceived):
self._handle_headers(event)
elif isinstance(event, DataReceived):
self._handle_data(event)
def _handle_headers(self, event: HeadersReceived) -> None:
"""Chooses how to handle a stream from its H3 request headers."""
headers = dict(event.headers)
method = headers.get(b":method")
protocol = headers.get(b":protocol")
if method == b"CONNECT" and protocol == b"websocket":
self._handle_websocket_connect(event)
if getattr(event, "stream_ended", False):
self._handle_data(DataReceived(
data=b"", stream_id=event.stream_id, stream_ended=True))
elif method == b"CONNECT":
self._send_error(event.stream_id, 501)
elif method in (b"GET", b"HEAD"):
self._handle_request(event)
else:
self._send_error(event.stream_id, 405)
def _handle_websocket_connect(self, event: HeadersReceived) -> None:
"""Handles a WebSocket CONNECT request on an HTTP/3 stream.
This finds the matching pywebsocket handler, validates the WebSocket
handshake, sends the H3 response headers, stores the session for later
DATA frames, and starts the handler thread for the stream.
"""
assert self._http is not None
stream_id = event.stream_id
h3_headers = H3Headers(event.headers)
path = h3_headers.get("path", "/")
ws_dispatcher = dispatch.Dispatcher( # type: ignore
self._ws_doc_root, None, False)
loop = asyncio.get_running_loop()
session = _WebSocketH3Session(
self, stream_id, h3_headers, ws_dispatcher, loop, self._logger,
max_unread_websocket_bytes=self._max_unread_websocket_bytes)
self._sessions[stream_id] = session
self._track_worker(self._finish_websocket_connect(
session, path, ws_dispatcher), stream_id)
def _track_worker(self, work: Any, stream_id: int) -> None:
task = asyncio.create_task(work)
self._workers.add(task)
self._queued_workers[stream_id] = task
task.add_done_callback(self._workers.discard)
task.add_done_callback(
lambda finished: self._queued_workers.pop(stream_id, None))
async def _finish_websocket_connect(
self, session: _WebSocketH3Session, path: str,
ws_dispatcher: dispatch.Dispatcher,
) -> None:
try:
async with self._handshake_slots:
self._queued_workers.pop(session.stream_id, None)
if self._sessions.get(session.stream_id) is not session:
return
self._pending_handshakes += 1
try:
await self._do_websocket_connect(session, path, ws_dispatcher)
finally:
self._pending_handshakes -= 1
except asyncio.CancelledError:
session.abort()
if self._sessions.get(session.stream_id) is session:
self._sessions.pop(session.stream_id)
raise
async def _do_websocket_connect(
self, session: _WebSocketH3Session, path: str,
ws_dispatcher: dispatch.Dispatcher,
) -> None:
stream_id = session.stream_id
request = session.request
try:
def handshake() -> bool:
if not ws_dispatcher.get_handler_suite(path): # type: ignore
return False
WsH3Handshaker(request, ws_dispatcher).do_handshake()
return True
if not await asyncio.to_thread(handshake):
if self._sessions.get(stream_id) is not session:
session.abort()
return
self._logger.warning("No WebSocket handler found for %s", path)
session.abort()
self._sessions.pop(stream_id, None)
self._send_error(stream_id, 404)
return
except asyncio.CancelledError:
session.abort()
self._sessions.pop(stream_id, None)
raise
except HandshakeException as error:
status = error.status if error.status else 400
self._logger.info(
"WebSocket/H3 handshake failed for stream %d: %s",
stream_id, error)
session.abort()
self._sessions.pop(stream_id, None)
self._send_error(stream_id, status)
return
except AbortedByUserException:
self._logger.info("WebSocket/H3 handshake aborted for stream %d",
stream_id)
self._send_websocket_response(stream_id, request, end_stream=True)
session.abort()
self._sessions.pop(stream_id, None)
return
except Exception:
self._logger.exception("WebSocket/H3 handshake failed for stream %d",
stream_id)
session.abort()
self._sessions.pop(stream_id, None)
self._send_error(stream_id, 500)
return
if self._sessions.get(stream_id) is not session:
session.abort()
return
request._dispatcher = ws_dispatcher
self._send_websocket_response(stream_id, request, end_stream=False)
request.connection.finish_handshake()
session.start()
def _handle_request(self, event: HeadersReceived) -> None:
"""Serves a regular HTTP/3 request through wptserve.
This handles non-WebSocket streams such as HTML, JavaScript, CSS,
`.sub.` resources, and Python handlers. It adapts the H3 request to the
request and response objects expected by wptserve, then lets the WPT
router run the matching handler.
"""
assert self._http is not None
if self._router is None or self._h3_server_adapter is None:
self._send_error(event.stream_id, 404)
return
self._track_worker(self._serve_request(event), event.stream_id)
async def _serve_request(self, event: HeadersReceived) -> None:
async with self._request_slots:
self._queued_workers.pop(event.stream_id, None)
self._active_requests += 1
try:
await asyncio.to_thread(self._serve_request_sync, event)
finally:
self._active_requests -= 1
def _serve_request_sync(self, event: HeadersReceived) -> None:
assert self._h3_server_adapter is not None
h3_headers = H3Headers(event.headers)
authority = h3_headers.get("authority")
if authority and "Host" not in h3_headers:
h3_headers["Host"] = authority
request_handler = _H3RequestHandlerAdapter(
server=self._h3_server_adapter,
headers=h3_headers,
rfile=BytesIO(b""),
stream_id=event.stream_id,
protocol=self,
)
try:
request_handler.server.rewriter.rewrite(request_handler)
with self._h3_request_cls(request_handler) as request:
response = self._h3_response_cls(request_handler, request)
handler = self._router.get_handler(request)
if handler is None:
response.set_error(404)
else:
if hasattr(handler, "base_path") and handler.base_path:
request.doc_root = handler.base_path
if hasattr(handler, "url_base") and handler.url_base != "/":
request.url_base = handler.url_base
try:
handler(request, response)
except Exception as error:
from wptserve.utils import HTTPException, get_error_cause
if isinstance(error, HTTPException):
exc = (get_error_cause(error)
if 500 <= error.code < 600 else error)
response.set_error(error.code, exc)
else:
response.set_error(500, error)
if not response.writer.content_written:
response.write()
if not response.writer.stream_ended:
response.writer.end_stream()
except Exception:
self._logger.exception("Failed to serve HTTP/3 stream %d",
event.stream_id)
if self._http is not None:
try:
self._call_on_loop(
self._send_error, event.stream_id, 500)
except Exception:
self._logger.exception(
"Failed to send HTTP/3 error response for stream %d",
event.stream_id)
def _call_on_loop(self, callback: Any, *args: Any) -> Any:
async def call() -> Any:
return callback(*args)
action = call()
try:
future = asyncio.run_coroutine_threadsafe(action, self._loop)
except RuntimeError:
action.close()
raise
return future.result(timeout=30)
def _handle_data(self, event: DataReceived) -> None:
"""Forwards H3 DATA frames to the matching WebSocket session."""
session = self._sessions.get(event.stream_id)
if session is None:
return
if event.data:
if session.feed_data(event.data) is False:
self._sessions.pop(event.stream_id, None)
session.abort()
self._quic.stop_stream(event.stream_id, 0x101)
self._quic.reset_stream(event.stream_id, 0x101)
self.transmit()
return
if event.stream_ended:
session.close()
def _send_error(self, stream_id: int, status_code: int) -> None:
"""Sends a simple H3 error response and closes the stream."""
assert self._http is not None
self._http.send_headers(
stream_id=stream_id,
headers=[
(b":status", str(status_code).encode()),
(b"server", SERVER_NAME.encode()),
],
end_stream=True,
)
self.transmit()
def _send_websocket_response(
self,
stream_id: int,
request: _WebSocketH3Request,
end_stream: bool,
) -> None:
"""Sends the H3 response headers prepared by the WebSocket handshaker."""
assert self._http is not None
status = request.status if request.status else 403
response_headers: List[Tuple[bytes, bytes]] = [
(b":status", str(status).encode()),
(b"server", SERVER_NAME.encode()),
]
headers_out = list(request.headers_out.items())
headers_out.extend(request.extra_headers)
for name, value in headers_out:
header_name = (name.decode("ascii")
if isinstance(name, bytes) else str(name)).lower()
if header_name in _H3_FORBIDDEN_RESPONSE_HEADERS:
continue
header_value = (value if isinstance(value, bytes)
else str(value).encode("latin-1"))
response_headers.append((header_name.encode(), header_value))
self._http.send_headers(
stream_id=stream_id,
headers=response_headers,
end_stream=end_stream)
self.transmit()
def _close_all_sessions(self) -> None:
"""Closes all WebSocket sessions active on this QUIC connection."""
for task in self._queued_workers.values():
task.cancel()
self._queued_workers.clear()
for session in list(self._sessions.values()):
session.abort()
self._sessions.clear()
class SessionTicketStore:
"""Simple in-memory store for TLS session tickets."""
def __init__(self) -> None:
self.tickets: Dict[bytes, SessionTicket] = {}
def add(self, ticket: SessionTicket) -> None:
self.tickets[ticket.ticket] = ticket
def pop(self, label: bytes) -> Optional[SessionTicket]:
return self.tickets.pop(label, None)
class WebSocketH3Server:
"""A WebSocket over HTTP/3 server for testing.
The server handles WebSocket Extended CONNECT streams and also serves the
test page/resources from the same H3 origin. The route-related parameters
are passed to wptserve so regular WPT requests use the same routing,
rewrite, and template-substitution behavior as the HTTP/HTTPS/H2 servers.
:param host: Host to bind.
:param port: UDP port to bind for QUIC.
:param ws_doc_root: Directory containing WebSocket *_wsh.py handlers.
:param cert_path: Path to the TLS certificate file.
:param key_path: Path to the TLS private key file.
:param doc_root: Directory containing WPT test files and resources.
:param routes: wptserve route table for regular WPT resources.
:param rewrites: wptserve rewrite rules applied before routing.
:param config: wptserve config used by requests and template substitution.
:param logger: Logger object for this server.
:param max_pending_handshakes: Maximum concurrent handshakes per QUIC
connection; additional CONNECT streams wait for a slot.
:param max_active_requests: Maximum concurrent resource handlers per QUIC
connection; additional requests wait for a slot.
:param max_unread_websocket_bytes: Maximum bytes waiting to be read by
the WebSocket handler per stream; exceeding this resets that stream.
"""
def __init__(
self,
host: str,
port: int,
ws_doc_root: str,
cert_path: str,
key_path: str,
doc_root: str = "",
routes: Any = None,
rewrites: Any = None,
config: Any = None,
logger: Optional[logging.Logger] = None,
max_pending_handshakes: int = DEFAULT_MAX_PENDING_HANDSHAKES,
max_active_requests: int = DEFAULT_MAX_ACTIVE_REQUESTS,
max_unread_websocket_bytes: int = DEFAULT_MAX_UNREAD_WEBSOCKET_BYTES,
) -> None:
for name, value in (
("max_pending_handshakes", max_pending_handshakes),
("max_active_requests", max_active_requests),
("max_unread_websocket_bytes", max_unread_websocket_bytes),
):
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
raise ValueError(f"{name} must be a positive integer")
self.host = host
self.port = port
self.ws_doc_root = ws_doc_root
self.cert_path = cert_path
self.key_path = key_path
self.doc_root = doc_root
self.routes = routes
self.rewrites = rewrites
self.config = config
self.max_pending_handshakes = max_pending_handshakes
self.max_active_requests = max_active_requests
self.max_unread_websocket_bytes = max_unread_websocket_bytes
self._logger = logger if logger is not None else _logger
self._router: Any = None
self._h3_server_adapter: Optional[_H3ServerAdapter] = None
self._h3_request_cls: Any = None
self._h3_response_cls: Any = None
self.started = False
self.loop: Optional[asyncio.AbstractEventLoop] = None
self.server_thread: Optional[threading.Thread] = None
self._server: Any = None
self._startup_complete = threading.Event()
self._startup_error: Optional[BaseException] = None
if self.doc_root:
from wptserve.router import Router
route_list = self.routes
if route_list is None:
from wptserve import routes as default_routes
route_list = default_routes.routes
rewriter = None
if self.rewrites is not None:
from wptserve.server import RequestRewriter
request_rewriter_class: Any = RequestRewriter
rewriter = request_rewriter_class(self.rewrites)
router_class: Any = Router
self._router = router_class(self.doc_root, route_list)
self._h3_server_adapter = _H3ServerAdapter(
self._router, self.host, self.port, rewriter, self.config)
(self._h3_request_cls,
self._h3_response_cls) = create_h3_request_response_classes()
def start(self) -> None:
"""Start the server."""
self._startup_complete.clear()
self._startup_error = None
self.server_thread = threading.Thread(
target=self._start_on_server_thread, daemon=True)
self.server_thread.start()
self._startup_complete.wait()
startup_error: Optional[BaseException] = self._startup_error
if startup_error is not None:
self.server_thread.join()
raise startup_error
self.started = True
def _start_on_server_thread(self) -> None:
secrets_log_file = None
try:
if "SSLKEYLOGFILE" in os.environ:
try:
secrets_log_file = open(os.environ["SSLKEYLOGFILE"], "a")
except OSError as error:
self._logger.warning(
"Failed to open SSLKEYLOGFILE: %s", error)
if secrets_log_file is not None:
configuration = QuicConfiguration(
alpn_protocols=H3_ALPN,
is_client=False,
max_datagram_frame_size=65536,
secrets_log_file=secrets_log_file,
)
else:
configuration = QuicConfiguration(
alpn_protocols=H3_ALPN,
is_client=False,
max_datagram_frame_size=65536,
)
configuration.load_cert_chain(
Path(self.cert_path), Path(self.key_path))
ticket_store = SessionTicketStore()
if sys.platform == "win32":
asyncio.set_event_loop_policy(
asyncio.WindowsSelectorEventLoopPolicy())
self.loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.loop)
self._server = self.loop.run_until_complete(
serve(
self.host,
self.port,
configuration=configuration,
create_protocol=self._create_protocol,
session_ticket_fetcher=ticket_store.pop,
session_ticket_handler=ticket_store.add,
))
self._startup_complete.set()
self.loop.run_forever()
except BaseException as error:
self._startup_error = error
self._startup_complete.set()
self._logger.exception("WebSocket/H3 server thread failed")
finally:
if self.loop is not None:
if self._server is not None:
for protocol in set(getattr(
self._server, "_protocols", {}).values()):
protocol._close_all_sessions()
self._server.close()
self.loop.run_until_complete(asyncio.sleep(0))
self._server = None
pending = asyncio.all_tasks(self.loop)
for task in pending:
task.cancel()
if pending:
self.loop.run_until_complete(
asyncio.gather(*pending, return_exceptions=True))
self.loop.close()
if secrets_log_file is not None:
secrets_log_file.close()
def _create_protocol(self, *args: Any, **kwargs: Any) -> WebSocketH3Protocol:
return WebSocketH3Protocol(
*args,
ws_doc_root=self.ws_doc_root,
router=self._router,
h3_server_adapter=self._h3_server_adapter,
h3_request_cls=self._h3_request_cls,
h3_response_cls=self._h3_response_cls,
logger=self._logger,
max_pending_handshakes=self.max_pending_handshakes,
max_active_requests=self.max_active_requests,
max_unread_websocket_bytes=self.max_unread_websocket_bytes,
**kwargs)
def stop(self) -> None:
"""Stop the server."""
if self.started and self.loop is not None:
self.loop.call_soon_threadsafe(
lambda: asyncio.create_task(self._stop_on_server_thread()))
if self.server_thread is not None:
self.server_thread.join()
self._logger.info("Stopped %s on %s:%d",
SERVER_NAME, self.host, self.port)
self.started = False
async def _stop_on_server_thread(self) -> None:
assert self.loop is not None
try:
if self._server is not None:
# Stop accepting peers before draining work on existing connections.
self._server._transport.close()
protocols = set(getattr(self._server, "_protocols", {}).values())
workers = [task for protocol in protocols
for task in protocol._workers]
if workers:
_, pending = await asyncio.wait(workers, timeout=10)
if pending:
self._logger.warning(
"Timed out waiting for %d WebSocket/H3 requests",
len(pending))
self._server.close()
for protocol in protocols:
protocol._close_all_sessions()
finally:
self.loop.call_soon(self.loop.stop)
def server_is_running(host: str, port: int, timeout: float) -> bool:
"""Open a temporary QUIC connection to check if the server is reachable."""
try:
loop = asyncio.get_event_loop()
except RuntimeError:
loop = asyncio.new_event_loop()
return loop.run_until_complete(
_server_is_running_with_timeout(host, port, timeout))
async def _server_is_running_with_timeout(
host: str, port: int, timeout: float
) -> bool:
"""Run the QUIC readiness probe with the caller-provided timeout."""
try:
await asyncio.wait_for(_connect_to_server(host, port), timeout=timeout)
except asyncio.TimeoutError:
_logger.warning(
"Failed to connect to WebSocket over HTTP/3 server at %s:%d",
host, port)
return False
return True
def _close_unusable_writer(reader: asyncio.StreamReader,
writer: asyncio.StreamWriter) -> None:
"""Close peer-initiated unidirectional streams opened during the probe."""
stream_id = cast(QuicStreamAdapter, writer.transport).stream_id
if (stream_is_unidirectional(stream_id) and
not stream_is_client_initiated(stream_id)):
with contextlib.suppress(ValueError):
writer.close()
async def _connect_to_server(host: str, port: int) -> None:
"""Connect to the H3 server and ping it to verify it is accepting QUIC."""
configuration = QuicConfiguration(
alpn_protocols=H3_ALPN,
is_client=True,
verify_mode=ssl.CERT_NONE,
)
async with connect(host, port, configuration=configuration,
stream_handler=_close_unusable_writer) as protocol:
await protocol.ping()