99 lines
3.7 KiB
Python
99 lines
3.7 KiB
Python
import asyncio
|
|
import socket
|
|
from contextlib import asynccontextmanager
|
|
from typing import AsyncGenerator, Callable, Optional, cast
|
|
|
|
from ..quic.configuration import QuicConfiguration
|
|
from ..quic.connection import QuicConnection, QuicTokenHandler
|
|
from ..tls import SessionTicketHandler
|
|
from .protocol import QuicConnectionProtocol, QuicStreamHandler
|
|
|
|
__all__ = ["connect"]
|
|
|
|
|
|
@asynccontextmanager
|
|
async def connect(
|
|
host: str,
|
|
port: int,
|
|
*,
|
|
configuration: Optional[QuicConfiguration] = None,
|
|
create_protocol: Optional[Callable] = QuicConnectionProtocol,
|
|
session_ticket_handler: Optional[SessionTicketHandler] = None,
|
|
stream_handler: Optional[QuicStreamHandler] = None,
|
|
token_handler: Optional[QuicTokenHandler] = None,
|
|
wait_connected: bool = True,
|
|
local_port: int = 0,
|
|
) -> AsyncGenerator[QuicConnectionProtocol, None]:
|
|
"""
|
|
Connect to a QUIC server at the given `host` and `port`.
|
|
|
|
:meth:`connect()` returns an awaitable. Awaiting it yields a
|
|
:class:`~aioquic.asyncio.QuicConnectionProtocol` which can be used to
|
|
create streams.
|
|
|
|
:func:`connect` also accepts the following optional arguments:
|
|
|
|
* ``configuration`` is a :class:`~aioquic.quic.configuration.QuicConfiguration`
|
|
configuration object.
|
|
* ``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_handler`` is a callback which is invoked by the TLS
|
|
engine when a new session ticket is received.
|
|
* ``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`.
|
|
* ``wait_connected`` indicates whether the context manager should wait for the
|
|
connection to be established before yielding the
|
|
:class:`~aioquic.asyncio.QuicConnectionProtocol`. By default this is `True` but
|
|
you can set it to `False` if you want to immediately start sending data using
|
|
0-RTT.
|
|
* ``local_port`` is the UDP port number that this client wants to bind.
|
|
"""
|
|
loop = asyncio.get_event_loop()
|
|
local_host = "::"
|
|
|
|
# lookup remote address
|
|
infos = await loop.getaddrinfo(host, port, type=socket.SOCK_DGRAM)
|
|
addr = infos[0][4]
|
|
if len(addr) == 2:
|
|
addr = ("::ffff:" + addr[0], addr[1], 0, 0)
|
|
|
|
# prepare QUIC connection
|
|
if configuration is None:
|
|
configuration = QuicConfiguration(is_client=True)
|
|
if configuration.server_name is None:
|
|
configuration.server_name = host
|
|
connection = QuicConnection(
|
|
configuration=configuration,
|
|
session_ticket_handler=session_ticket_handler,
|
|
token_handler=token_handler,
|
|
)
|
|
|
|
# explicitly enable IPv4/IPv6 dual stack
|
|
sock = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
|
|
completed = False
|
|
try:
|
|
sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
|
|
sock.bind((local_host, local_port, 0, 0))
|
|
completed = True
|
|
finally:
|
|
if not completed:
|
|
sock.close()
|
|
# connect
|
|
transport, protocol = await loop.create_datagram_endpoint(
|
|
lambda: create_protocol(connection, stream_handler=stream_handler),
|
|
sock=sock,
|
|
)
|
|
protocol = cast(QuicConnectionProtocol, protocol)
|
|
try:
|
|
protocol.connect(addr, transmit=wait_connected)
|
|
if wait_connected:
|
|
await protocol.wait_connected()
|
|
yield protocol
|
|
finally:
|
|
protocol.close()
|
|
await protocol.wait_closed()
|
|
transport.close()
|