codekingpro/portable-devtools

1import asyncio
2import socket
3from contextlib import asynccontextmanager
4from typing import AsyncGenerator, Callable, Optional, cast
5
6from ..quic.configuration import QuicConfiguration
7from ..quic.connection import QuicConnection, QuicTokenHandler
8from ..tls import SessionTicketHandler
9from .protocol import QuicConnectionProtocol, QuicStreamHandler
10
11__all__ = ["connect"]
12
13
14@asynccontextmanager
15async def connect(
16    host: str,
17    port: int,
18    *,
19    configuration: Optional[QuicConfiguration] = None,
20    create_protocol: Optional[Callable] = QuicConnectionProtocol,
21    session_ticket_handler: Optional[SessionTicketHandler] = None,
22    stream_handler: Optional[QuicStreamHandler] = None,
23    token_handler: Optional[QuicTokenHandler] = None,
24    wait_connected: bool = True,
25    local_port: int = 0,
26) -> AsyncGenerator[QuicConnectionProtocol, None]:
27    """
28    Connect to a QUIC server at the given `host` and `port`.
29
30    :meth:`connect()` returns an awaitable. Awaiting it yields a
31    :class:`~aioquic.asyncio.QuicConnectionProtocol` which can be used to
32    create streams.
33
34    :func:`connect` also accepts the following optional arguments:
35
36    * ``configuration`` is a :class:`~aioquic.quic.configuration.QuicConfiguration`
37      configuration object.
38    * ``create_protocol`` allows customizing the :class:`~asyncio.Protocol` that
39      manages the connection. It should be a callable or class accepting the same
40      arguments as :class:`~aioquic.asyncio.QuicConnectionProtocol` and returning
41      an instance of :class:`~aioquic.asyncio.QuicConnectionProtocol` or a subclass.
42    * ``session_ticket_handler`` is a callback which is invoked by the TLS
43      engine when a new session ticket is received.
44    * ``stream_handler`` is a callback which is invoked whenever a stream is
45      created. It must accept two arguments: a :class:`asyncio.StreamReader`
46      and a :class:`asyncio.StreamWriter`.
47    * ``wait_connected`` indicates whether the context manager should wait for the
48      connection to be established before yielding the
49      :class:`~aioquic.asyncio.QuicConnectionProtocol`. By default this is `True` but
50      you can set it to `False` if you want to immediately start sending data using
51      0-RTT.
52    * ``local_port`` is the UDP port number that this client wants to bind.
53    """
54    loop = asyncio.get_event_loop()
55    local_host = "::"
56
57    # lookup remote address
58    infos = await loop.getaddrinfo(host, port, type=socket.SOCK_DGRAM)
59    addr = infos[0][4]
60    if len(addr) == 2:
61        addr = ("::ffff:" + addr[0], addr[1], 0, 0)
62
63    # prepare QUIC connection
64    if configuration is None:
65        configuration = QuicConfiguration(is_client=True)
66    if configuration.server_name is None:
67        configuration.server_name = host
68    connection = QuicConnection(
69        configuration=configuration,
70        session_ticket_handler=session_ticket_handler,
71        token_handler=token_handler,
72    )
73
74    # explicitly enable IPv4/IPv6 dual stack
75    sock = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
76    completed = False
77    try:
78        sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
79        sock.bind((local_host, local_port, 0, 0))
80        completed = True
81    finally:
82        if not completed:
83            sock.close()
84    # connect
85    transport, protocol = await loop.create_datagram_endpoint(
86        lambda: create_protocol(connection, stream_handler=stream_handler),
87        sock=sock,
88    )
89    protocol = cast(QuicConnectionProtocol, protocol)
90    try:
91        protocol.connect(addr, transmit=wait_connected)
92        if wait_connected:
93            await protocol.wait_connected()
94        yield protocol
95    finally:
96        protocol.close()
97        await protocol.wait_closed()
98        transport.close()
99