Code
This commit is contained in:
mofixx
215
Code/venv/lib/python3.13/site-packages/aioquic/asyncio/server.py
Normal file
215
Code/venv/lib/python3.13/site-packages/aioquic/asyncio/server.py
Normal file
@ -0,0 +1,215 @@
|
||||
import asyncio
|
||||
import os
|
||||
from functools import partial
|
||||
from typing import Callable, Dict, Optional, Text, Union, cast
|
||||
|
||||
from ..buffer import Buffer
|
||||
from ..quic.configuration import SMALLEST_MAX_DATAGRAM_SIZE, QuicConfiguration
|
||||
from ..quic.connection import NetworkAddress, QuicConnection
|
||||
from ..quic.packet import (
|
||||
QuicPacketType,
|
||||
encode_quic_retry,
|
||||
encode_quic_version_negotiation,
|
||||
pull_quic_header,
|
||||
)
|
||||
from ..quic.retry import QuicRetryTokenHandler
|
||||
from ..tls import SessionTicketFetcher, SessionTicketHandler
|
||||
from .protocol import QuicConnectionProtocol, QuicStreamHandler
|
||||
|
||||
__all__ = ["serve"]
|
||||
|
||||
|
||||
class QuicServer(asyncio.DatagramProtocol):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
configuration: QuicConfiguration,
|
||||
create_protocol: Callable = QuicConnectionProtocol,
|
||||
session_ticket_fetcher: Optional[SessionTicketFetcher] = None,
|
||||
session_ticket_handler: Optional[SessionTicketHandler] = None,
|
||||
retry: bool = False,
|
||||
stream_handler: Optional[QuicStreamHandler] = None,
|
||||
) -> None:
|
||||
self._configuration = configuration
|
||||
self._create_protocol = create_protocol
|
||||
self._loop = asyncio.get_event_loop()
|
||||
self._protocols: Dict[bytes, QuicConnectionProtocol] = {}
|
||||
self._session_ticket_fetcher = session_ticket_fetcher
|
||||
self._session_ticket_handler = session_ticket_handler
|
||||
self._transport: Optional[asyncio.DatagramTransport] = None
|
||||
|
||||
self._stream_handler = stream_handler
|
||||
|
||||
if retry:
|
||||
self._retry = QuicRetryTokenHandler()
|
||||
else:
|
||||
self._retry = None
|
||||
|
||||
def close(self):
|
||||
for protocol in set(self._protocols.values()):
|
||||
protocol.close()
|
||||
self._protocols.clear()
|
||||
self._transport.close()
|
||||
|
||||
def connection_made(self, transport: asyncio.BaseTransport) -> None:
|
||||
self._transport = cast(asyncio.DatagramTransport, transport)
|
||||
|
||||
def datagram_received(self, data: Union[bytes, Text], addr: NetworkAddress) -> None:
|
||||
data = cast(bytes, data)
|
||||
buf = Buffer(data=data)
|
||||
|
||||
try:
|
||||
header = pull_quic_header(
|
||||
buf, host_cid_length=self._configuration.connection_id_length
|
||||
)
|
||||
except ValueError:
|
||||
return
|
||||
|
||||
# version negotiation
|
||||
if (
|
||||
header.version is not None
|
||||
and header.version not in self._configuration.supported_versions
|
||||
):
|
||||
self._transport.sendto(
|
||||
encode_quic_version_negotiation(
|
||||
source_cid=header.destination_cid,
|
||||
destination_cid=header.source_cid,
|
||||
supported_versions=self._configuration.supported_versions,
|
||||
),
|
||||
addr,
|
||||
)
|
||||
return
|
||||
|
||||
protocol = self._protocols.get(header.destination_cid, None)
|
||||
original_destination_connection_id: Optional[bytes] = None
|
||||
retry_source_connection_id: Optional[bytes] = None
|
||||
if (
|
||||
protocol is None
|
||||
and len(data) >= SMALLEST_MAX_DATAGRAM_SIZE
|
||||
and header.packet_type == QuicPacketType.INITIAL
|
||||
):
|
||||
# retry
|
||||
if self._retry is not None:
|
||||
if not header.token:
|
||||
# create a retry token
|
||||
source_cid = os.urandom(8)
|
||||
self._transport.sendto(
|
||||
encode_quic_retry(
|
||||
version=header.version,
|
||||
source_cid=source_cid,
|
||||
destination_cid=header.source_cid,
|
||||
original_destination_cid=header.destination_cid,
|
||||
retry_token=self._retry.create_token(
|
||||
addr, header.destination_cid, source_cid
|
||||
),
|
||||
),
|
||||
addr,
|
||||
)
|
||||
return
|
||||
else:
|
||||
# validate retry token
|
||||
try:
|
||||
(
|
||||
original_destination_connection_id,
|
||||
retry_source_connection_id,
|
||||
) = self._retry.validate_token(addr, header.token)
|
||||
except ValueError:
|
||||
return
|
||||
else:
|
||||
original_destination_connection_id = header.destination_cid
|
||||
|
||||
# create new connection
|
||||
connection = QuicConnection(
|
||||
configuration=self._configuration,
|
||||
original_destination_connection_id=original_destination_connection_id,
|
||||
retry_source_connection_id=retry_source_connection_id,
|
||||
session_ticket_fetcher=self._session_ticket_fetcher,
|
||||
session_ticket_handler=self._session_ticket_handler,
|
||||
)
|
||||
protocol = self._create_protocol(
|
||||
connection, stream_handler=self._stream_handler
|
||||
)
|
||||
protocol.connection_made(self._transport)
|
||||
|
||||
# register callbacks
|
||||
protocol._connection_id_issued_handler = partial(
|
||||
self._connection_id_issued, protocol=protocol
|
||||
)
|
||||
protocol._connection_id_retired_handler = partial(
|
||||
self._connection_id_retired, protocol=protocol
|
||||
)
|
||||
protocol._connection_terminated_handler = partial(
|
||||
self._connection_terminated, protocol=protocol
|
||||
)
|
||||
|
||||
self._protocols[header.destination_cid] = protocol
|
||||
self._protocols[connection.host_cid] = protocol
|
||||
|
||||
if protocol is not None:
|
||||
protocol.datagram_received(data, addr)
|
||||
|
||||
def _connection_id_issued(self, cid: bytes, protocol: QuicConnectionProtocol):
|
||||
self._protocols[cid] = protocol
|
||||
|
||||
def _connection_id_retired(
|
||||
self, cid: bytes, protocol: QuicConnectionProtocol
|
||||
) -> None:
|
||||
assert self._protocols[cid] == protocol
|
||||
del self._protocols[cid]
|
||||
|
||||
def _connection_terminated(self, protocol: QuicConnectionProtocol):
|
||||
for cid, proto in list(self._protocols.items()):
|
||||
if proto == protocol:
|
||||
del self._protocols[cid]
|
||||
|
||||
|
||||
async def serve(
|
||||
host: str,
|
||||
port: int,
|
||||
*,
|
||||
configuration: QuicConfiguration,
|
||||
create_protocol: Callable = QuicConnectionProtocol,
|
||||
session_ticket_fetcher: Optional[SessionTicketFetcher] = None,
|
||||
session_ticket_handler: Optional[SessionTicketHandler] = None,
|
||||
retry: bool = False,
|
||||
stream_handler: QuicStreamHandler = None,
|
||||
) -> QuicServer:
|
||||
"""
|
||||
Start a QUIC server at the given `host` and `port`.
|
||||
|
||||
:func:`serve` requires a :class:`~aioquic.quic.configuration.QuicConfiguration`
|
||||
containing TLS certificate and private key as the ``configuration`` argument.
|
||||
|
||||
:func:`serve` also accepts the following optional arguments:
|
||||
|
||||
* ``create_protocol`` allows customizing the :class:`~asyncio.Protocol` that
|
||||
manages the connection. It should be a callable or class accepting the same
|
||||
arguments as :class:`~aioquic.asyncio.QuicConnectionProtocol` and returning
|
||||
an instance of :class:`~aioquic.asyncio.QuicConnectionProtocol` or a subclass.
|
||||
* ``session_ticket_fetcher`` is a callback which is invoked by the TLS
|
||||
engine when a session ticket is presented by the peer. It should return
|
||||
the session ticket with the specified ID or `None` if it is not found.
|
||||
* ``session_ticket_handler`` is a callback which is invoked by the TLS
|
||||
engine when a new session ticket is issued. It should store the session
|
||||
ticket for future lookup.
|
||||
* ``retry`` specifies whether client addresses should be validated prior to
|
||||
the cryptographic handshake using a retry packet.
|
||||
* ``stream_handler`` is a callback which is invoked whenever a stream is
|
||||
created. It must accept two arguments: a :class:`asyncio.StreamReader`
|
||||
and a :class:`asyncio.StreamWriter`.
|
||||
"""
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
_, protocol = await loop.create_datagram_endpoint(
|
||||
lambda: QuicServer(
|
||||
configuration=configuration,
|
||||
create_protocol=create_protocol,
|
||||
session_ticket_fetcher=session_ticket_fetcher,
|
||||
session_ticket_handler=session_ticket_handler,
|
||||
retry=retry,
|
||||
stream_handler=stream_handler,
|
||||
),
|
||||
local_addr=(host, port),
|
||||
)
|
||||
return protocol
|
||||
Reference in New Issue
Block a user