Code

Code/venv/lib/python3.13/site-packages/aioquic/asyncio/client.py

@@ -0,0 +1,98 @@
+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()