codekingpro/portable-devtools

1# Copyright (C) Dnspython Contributors, see LICENSE for text of ISC license2 3import asyncio4import socket5import ssl6import struct7import time8 9import aioquic.quic.configuration  # type: ignore10import aioquic.quic.connection  # type: ignore11import aioquic.quic.events  # type: ignore12 13import dns.asyncbackend14import dns.exception15import dns.inet16from dns.quic._common import (17    QUIC_MAX_DATAGRAM,18    AsyncQuicConnection,19    AsyncQuicManager,20    BaseQuicStream,21    UnexpectedEOF,22)23 24 25class AsyncioQuicStream(BaseQuicStream):26    def __init__(self, connection, stream_id):27        super().__init__(connection, stream_id)28        self._wake_up = asyncio.Condition()29 30    async def _wait_for_wake_up(self):31        async with self._wake_up:32            await self._wake_up.wait()33 34    async def wait_for(self, amount, expiration):35        while True:36            timeout = self._timeout_from_expiration(expiration)37            if self._buffer.have(amount):38                return39            self._expecting = amount40            try:41                await asyncio.wait_for(self._wait_for_wake_up(), timeout)42            except TimeoutError:43                raise dns.exception.Timeout44            self._expecting = 045 46    async def receive(self, timeout=None):47        expiration = self._expiration_from_timeout(timeout)48        await self.wait_for(2, expiration)49        (size,) = struct.unpack("!H", self._buffer.get(2))50        await self.wait_for(size, expiration)51        return self._buffer.get(size)52 53    async def send(self, datagram, is_end=False):54        data = self._encapsulate(datagram)55        await self._connection.write(self._stream_id, data, is_end)56 57    async def _add_input(self, data, is_end):58        if self._common_add_input(data, is_end):59            async with self._wake_up:60                self._wake_up.notify()61 62    async def close(self):63        self._close()64 65    # Streams are async context managers66 67    async def __aenter__(self):68        return self69 70    async def __aexit__(self, exc_type, exc_val, exc_tb):71        await self.close()72        async with self._wake_up:73            self._wake_up.notify()74        return False75 76 77class AsyncioQuicConnection(AsyncQuicConnection):78    def __init__(self, connection, address, port, source, source_port, manager=None):79        super().__init__(connection, address, port, source, source_port, manager)80        self._socket = None81        self._handshake_complete = asyncio.Event()82        self._socket_created = asyncio.Event()83        self._wake_timer = asyncio.Condition()84        self._receiver_task = None85        self._sender_task = None86 87    async def _receiver(self):88        try:89            af = dns.inet.af_for_address(self._address)90            backend = dns.asyncbackend.get_backend("asyncio")91            # Note that peer is a low-level address tuple, but make_socket() wants92            # a high-level address tuple, so we convert.93            self._socket = await backend.make_socket(94                af, socket.SOCK_DGRAM, 0, self._source, (self._peer[0], self._peer[1])95            )96            self._socket_created.set()97            async with self._socket:98                while not self._done:99                    (datagram, address) = await self._socket.recvfrom(100                        QUIC_MAX_DATAGRAM, None101                    )102                    if address[0] != self._peer[0] or address[1] != self._peer[1]:103                        continue104                    self._connection.receive_datagram(datagram, address, time.time())105                    # Wake up the timer in case the sender is sleeping, as there may be106                    # stuff to send now.107                    async with self._wake_timer:108                        self._wake_timer.notify_all()109        except Exception:110            pass111        finally:112            self._done = True113            async with self._wake_timer:114                self._wake_timer.notify_all()115            self._handshake_complete.set()116 117    async def _wait_for_wake_timer(self):118        async with self._wake_timer:119            await self._wake_timer.wait()120 121    async def _sender(self):122        await self._socket_created.wait()123        while not self._done:124            datagrams = self._connection.datagrams_to_send(time.time())125            for datagram, address in datagrams:126                assert address == self._peer127                await self._socket.sendto(datagram, self._peer, None)128            (expiration, interval) = self._get_timer_values()129            try:130                await asyncio.wait_for(self._wait_for_wake_timer(), interval)131            except Exception:132                pass133            self._handle_timer(expiration)134            await self._handle_events()135 136    async def _handle_events(self):137        count = 0138        while True:139            event = self._connection.next_event()140            if event is None:141                return142            if isinstance(event, aioquic.quic.events.StreamDataReceived):143                stream = self._streams.get(event.stream_id)144                if stream:145                    await stream._add_input(event.data, event.end_stream)146            elif isinstance(event, aioquic.quic.events.HandshakeCompleted):147                self._handshake_complete.set()148            elif isinstance(event, aioquic.quic.events.ConnectionTerminated):149                self._done = True150                self._receiver_task.cancel()151            elif isinstance(event, aioquic.quic.events.StreamReset):152                stream = self._streams.get(event.stream_id)153                if stream:154                    await stream._add_input(b"", True)155 156            count += 1157            if count > 10:158                # yield159                count = 0160                await asyncio.sleep(0)161 162    async def write(self, stream, data, is_end=False):163        self._connection.send_stream_data(stream, data, is_end)164        async with self._wake_timer:165            self._wake_timer.notify_all()166 167    def run(self):168        if self._closed:169            return170        self._receiver_task = asyncio.Task(self._receiver())171        self._sender_task = asyncio.Task(self._sender())172 173    async def make_stream(self, timeout=None):174        try:175            await asyncio.wait_for(self._handshake_complete.wait(), timeout)176        except TimeoutError:177            raise dns.exception.Timeout178        if self._done:179            raise UnexpectedEOF180        stream_id = self._connection.get_next_available_stream_id(False)181        stream = AsyncioQuicStream(self, stream_id)182        self._streams[stream_id] = stream183        return stream184 185    async def close(self):186        if not self._closed:187            self._manager.closed(self._peer[0], self._peer[1])188            self._closed = True189            self._connection.close()190            # sender might be blocked on this, so set it191            self._socket_created.set()192            async with self._wake_timer:193                self._wake_timer.notify_all()194            try:195                await self._receiver_task196            except asyncio.CancelledError:197                pass198            try:199                await self._sender_task200            except asyncio.CancelledError:201                pass202            await self._socket.close()203 204 205class AsyncioQuicManager(AsyncQuicManager):206    def __init__(self, conf=None, verify_mode=ssl.CERT_REQUIRED, server_name=None):207        super().__init__(conf, verify_mode, AsyncioQuicConnection, server_name)208 209    def connect(210        self, address, port=853, source=None, source_port=0, want_session_ticket=True211    ):212        (connection, start) = self._connect(213            address, port, source, source_port, want_session_ticket214        )215        if start:216            connection.run()217        return connection218 219    async def __aenter__(self):220        return self221 222    async def __aexit__(self, exc_type, exc_val, exc_tb):223        # Copy the iterator into a list as exiting things will mutate the connections224        # table.225        connections = list(self._connections.values())226        for connection in connections:227            await connection.close()228        return False229