Spaces:
Sleeping
Sleeping
| # Copyright (C) Dnspython Contributors, see LICENSE for text of ISC license | |
| import asyncio | |
| import socket | |
| import ssl | |
| import struct | |
| import time | |
| import aioquic.quic.configuration # type: ignore | |
| import aioquic.quic.connection # type: ignore | |
| import aioquic.quic.events # type: ignore | |
| import dns.asyncbackend | |
| import dns.exception | |
| import dns.inet | |
| from dns.quic._common import ( | |
| QUIC_MAX_DATAGRAM, | |
| AsyncQuicConnection, | |
| AsyncQuicManager, | |
| BaseQuicStream, | |
| UnexpectedEOF, | |
| ) | |
| class AsyncioQuicStream(BaseQuicStream): | |
| def __init__(self, connection, stream_id): | |
| super().__init__(connection, stream_id) | |
| self._wake_up = asyncio.Condition() | |
| async def _wait_for_wake_up(self): | |
| async with self._wake_up: | |
| await self._wake_up.wait() | |
| async def wait_for(self, amount, expiration): | |
| while True: | |
| timeout = self._timeout_from_expiration(expiration) | |
| if self._buffer.have(amount): | |
| return | |
| self._expecting = amount | |
| try: | |
| await asyncio.wait_for(self._wait_for_wake_up(), timeout) | |
| except TimeoutError: | |
| raise dns.exception.Timeout | |
| self._expecting = 0 | |
| async def receive(self, timeout=None): | |
| expiration = self._expiration_from_timeout(timeout) | |
| await self.wait_for(2, expiration) | |
| (size,) = struct.unpack("!H", self._buffer.get(2)) | |
| await self.wait_for(size, expiration) | |
| return self._buffer.get(size) | |
| async def send(self, datagram, is_end=False): | |
| data = self._encapsulate(datagram) | |
| await self._connection.write(self._stream_id, data, is_end) | |
| async def _add_input(self, data, is_end): | |
| if self._common_add_input(data, is_end): | |
| async with self._wake_up: | |
| self._wake_up.notify() | |
| async def close(self): | |
| self._close() | |
| # Streams are async context managers | |
| async def __aenter__(self): | |
| return self | |
| async def __aexit__(self, exc_type, exc_val, exc_tb): | |
| await self.close() | |
| async with self._wake_up: | |
| self._wake_up.notify() | |
| return False | |
| class AsyncioQuicConnection(AsyncQuicConnection): | |
| def __init__(self, connection, address, port, source, source_port, manager=None): | |
| super().__init__(connection, address, port, source, source_port, manager) | |
| self._socket = None | |
| self._handshake_complete = asyncio.Event() | |
| self._socket_created = asyncio.Event() | |
| self._wake_timer = asyncio.Condition() | |
| self._receiver_task = None | |
| self._sender_task = None | |
| async def _receiver(self): | |
| try: | |
| af = dns.inet.af_for_address(self._address) | |
| backend = dns.asyncbackend.get_backend("asyncio") | |
| # Note that peer is a low-level address tuple, but make_socket() wants | |
| # a high-level address tuple, so we convert. | |
| self._socket = await backend.make_socket( | |
| af, socket.SOCK_DGRAM, 0, self._source, (self._peer[0], self._peer[1]) | |
| ) | |
| self._socket_created.set() | |
| async with self._socket: | |
| while not self._done: | |
| (datagram, address) = await self._socket.recvfrom( | |
| QUIC_MAX_DATAGRAM, None | |
| ) | |
| if address[0] != self._peer[0] or address[1] != self._peer[1]: | |
| continue | |
| self._connection.receive_datagram( | |
| datagram, self._peer[0], time.time() | |
| ) | |
| # Wake up the timer in case the sender is sleeping, as there may be | |
| # stuff to send now. | |
| async with self._wake_timer: | |
| self._wake_timer.notify_all() | |
| except Exception: | |
| pass | |
| finally: | |
| self._done = True | |
| async with self._wake_timer: | |
| self._wake_timer.notify_all() | |
| self._handshake_complete.set() | |
| async def _wait_for_wake_timer(self): | |
| async with self._wake_timer: | |
| await self._wake_timer.wait() | |
| async def _sender(self): | |
| await self._socket_created.wait() | |
| while not self._done: | |
| datagrams = self._connection.datagrams_to_send(time.time()) | |
| for datagram, address in datagrams: | |
| assert address == self._peer[0] | |
| await self._socket.sendto(datagram, self._peer, None) | |
| (expiration, interval) = self._get_timer_values() | |
| try: | |
| await asyncio.wait_for(self._wait_for_wake_timer(), interval) | |
| except Exception: | |
| pass | |
| self._handle_timer(expiration) | |
| await self._handle_events() | |
| async def _handle_events(self): | |
| count = 0 | |
| while True: | |
| event = self._connection.next_event() | |
| if event is None: | |
| return | |
| if isinstance(event, aioquic.quic.events.StreamDataReceived): | |
| stream = self._streams.get(event.stream_id) | |
| if stream: | |
| await stream._add_input(event.data, event.end_stream) | |
| elif isinstance(event, aioquic.quic.events.HandshakeCompleted): | |
| self._handshake_complete.set() | |
| elif isinstance( | |
| event, aioquic.quic.events.ConnectionTerminated | |
| ) or isinstance(event, aioquic.quic.events.StreamReset): | |
| self._done = True | |
| self._receiver_task.cancel() | |
| count += 1 | |
| if count > 10: | |
| # yield | |
| count = 0 | |
| await asyncio.sleep(0) | |
| async def write(self, stream, data, is_end=False): | |
| self._connection.send_stream_data(stream, data, is_end) | |
| async with self._wake_timer: | |
| self._wake_timer.notify_all() | |
| def run(self): | |
| if self._closed: | |
| return | |
| self._receiver_task = asyncio.Task(self._receiver()) | |
| self._sender_task = asyncio.Task(self._sender()) | |
| async def make_stream(self, timeout=None): | |
| try: | |
| await asyncio.wait_for(self._handshake_complete.wait(), timeout) | |
| except TimeoutError: | |
| raise dns.exception.Timeout | |
| if self._done: | |
| raise UnexpectedEOF | |
| stream_id = self._connection.get_next_available_stream_id(False) | |
| stream = AsyncioQuicStream(self, stream_id) | |
| self._streams[stream_id] = stream | |
| return stream | |
| async def close(self): | |
| if not self._closed: | |
| self._manager.closed(self._peer[0], self._peer[1]) | |
| self._closed = True | |
| self._connection.close() | |
| # sender might be blocked on this, so set it | |
| self._socket_created.set() | |
| await self._socket.close() | |
| async with self._wake_timer: | |
| self._wake_timer.notify_all() | |
| try: | |
| await self._receiver_task | |
| except asyncio.CancelledError: | |
| pass | |
| try: | |
| await self._sender_task | |
| except asyncio.CancelledError: | |
| pass | |
| class AsyncioQuicManager(AsyncQuicManager): | |
| def __init__(self, conf=None, verify_mode=ssl.CERT_REQUIRED, server_name=None): | |
| super().__init__(conf, verify_mode, AsyncioQuicConnection, server_name) | |
| def connect(self, address, port=853, source=None, source_port=0): | |
| (connection, start) = self._connect(address, port, source, source_port) | |
| if start: | |
| connection.run() | |
| return connection | |
| async def __aenter__(self): | |
| return self | |
| async def __aexit__(self, exc_type, exc_val, exc_tb): | |
| # Copy the iterator into a list as exiting things will mutate the connections | |
| # table. | |
| connections = list(self._connections.values()) | |
| for connection in connections: | |
| await connection.close() | |
| return False | |