codekingpro/portable-devtools

1# Copyright (C) Dnspython Contributors, see LICENSE for text of ISC license2 3import socket4import ssl5import struct6import time7 8import aioquic.quic.configuration  # type: ignore9import aioquic.quic.connection  # type: ignore10import aioquic.quic.events  # type: ignore11import trio12 13import dns.exception14import dns.inet15from dns._asyncbackend import NullContext16from dns.quic._common import (17    QUIC_MAX_DATAGRAM,18    AsyncQuicConnection,19    AsyncQuicManager,20    BaseQuicStream,21    UnexpectedEOF,22)23 24 25class TrioQuicStream(BaseQuicStream):26    def __init__(self, connection, stream_id):27        super().__init__(connection, stream_id)28        self._wake_up = trio.Condition()29 30    async def wait_for(self, amount):31        while True:32            if self._buffer.have(amount):33                return34            self._expecting = amount35            async with self._wake_up:36                await self._wake_up.wait()37            self._expecting = 038 39    async def receive(self, timeout=None):40        if timeout is None:41            context = NullContext(None)42        else:43            context = trio.move_on_after(timeout)44        with context:45            await self.wait_for(2)46            (size,) = struct.unpack("!H", self._buffer.get(2))47            await self.wait_for(size)48            return self._buffer.get(size)49        raise dns.exception.Timeout50 51    async def send(self, datagram, is_end=False):52        data = self._encapsulate(datagram)53        await self._connection.write(self._stream_id, data, is_end)54 55    async def _add_input(self, data, is_end):56        if self._common_add_input(data, is_end):57            async with self._wake_up:58                self._wake_up.notify()59 60    async def close(self):61        self._close()62 63    # Streams are async context managers64 65    async def __aenter__(self):66        return self67 68    async def __aexit__(self, exc_type, exc_val, exc_tb):69        await self.close()70        async with self._wake_up:71            self._wake_up.notify()72        return False73 74 75class TrioQuicConnection(AsyncQuicConnection):76    def __init__(self, connection, address, port, source, source_port, manager=None):77        super().__init__(connection, address, port, source, source_port, manager)78        self._socket = trio.socket.socket(self._af, socket.SOCK_DGRAM, 0)79        self._handshake_complete = trio.Event()80        self._run_done = trio.Event()81        self._worker_scope = None82        self._send_pending = False83 84    async def _worker(self):85        try:86            if self._source:87                await self._socket.bind(88                    dns.inet.low_level_address_tuple(self._source, self._af)89                )90            await self._socket.connect(self._peer)91            while not self._done:92                (expiration, interval) = self._get_timer_values(False)93                if self._send_pending:94                    # Do not block forever if sends are pending.  Even though we95                    # have a wake-up mechanism if we've already started the blocking96                    # read, the possibility of context switching in send means that97                    # more writes can happen while we have no wake up context, so98                    # we need self._send_pending to avoid (effectively) a "lost wakeup"99                    # race.100                    interval = 0.0101                with trio.CancelScope(102                    deadline=trio.current_time() + interval103                ) as self._worker_scope:104                    datagram = await self._socket.recv(QUIC_MAX_DATAGRAM)105                    self._connection.receive_datagram(datagram, self._peer, time.time())106                self._worker_scope = None107                self._handle_timer(expiration)108                await self._handle_events()109                # We clear this now, before sending anything, as sending can cause110                # context switches that do more sends.  We want to know if that111                # happens so we don't block a long time on the recv() above.112                self._send_pending = False113                datagrams = self._connection.datagrams_to_send(time.time())114                for datagram, _ in datagrams:115                    await self._socket.send(datagram)116        finally:117            self._done = True118            self._handshake_complete.set()119 120    async def _handle_events(self):121        count = 0122        while True:123            event = self._connection.next_event()124            if event is None:125                return126            if isinstance(event, aioquic.quic.events.StreamDataReceived):127                stream = self._streams.get(event.stream_id)128                if stream:129                    await stream._add_input(event.data, event.end_stream)130            elif isinstance(event, aioquic.quic.events.HandshakeCompleted):131                self._handshake_complete.set()132            elif isinstance(event, aioquic.quic.events.ConnectionTerminated):133                self._done = True134                self._socket.close()135            elif isinstance(event, aioquic.quic.events.StreamReset):136                stream = self._streams.get(event.stream_id)137                if stream:138                    await stream._add_input(b"", True)139            count += 1140            if count > 10:141                # yield142                count = 0143                await trio.sleep(0)144 145    async def write(self, stream, data, is_end=False):146        self._connection.send_stream_data(stream, data, is_end)147        self._send_pending = True148        if self._worker_scope is not None:149            self._worker_scope.cancel()150 151    async def run(self):152        if self._closed:153            return154        async with trio.open_nursery() as nursery:155            nursery.start_soon(self._worker)156        self._run_done.set()157 158    async def make_stream(self, timeout=None):159        if timeout is None:160            context = NullContext(None)161        else:162            context = trio.move_on_after(timeout)163        with context:164            await self._handshake_complete.wait()165            if self._done:166                raise UnexpectedEOF167            stream_id = self._connection.get_next_available_stream_id(False)168            stream = TrioQuicStream(self, stream_id)169            self._streams[stream_id] = stream170            return stream171        raise dns.exception.Timeout172 173    async def close(self):174        if not self._closed:175            self._manager.closed(self._peer[0], self._peer[1])176            self._closed = True177            self._connection.close()178            self._send_pending = True179            if self._worker_scope is not None:180                self._worker_scope.cancel()181            await self._run_done.wait()182 183 184class TrioQuicManager(AsyncQuicManager):185    def __init__(186        self, nursery, conf=None, verify_mode=ssl.CERT_REQUIRED, server_name=None187    ):188        super().__init__(conf, verify_mode, TrioQuicConnection, server_name)189        self._nursery = nursery190 191    def connect(192        self, address, port=853, source=None, source_port=0, want_session_ticket=True193    ):194        (connection, start) = self._connect(195            address, port, source, source_port, want_session_ticket196        )197        if start:198            self._nursery.start_soon(connection.run)199        return connection200 201    async def __aenter__(self):202        return self203 204    async def __aexit__(self, exc_type, exc_val, exc_tb):205        # Copy the iterator into a list as exiting things will mutate the connections206        # table.207        connections = list(self._connections.values())208        for connection in connections:209            await connection.close()210        return False211