codekingpro/portable-devtools

1# Copyright (C) Dnspython Contributors, see LICENSE for text of ISC license2 3import copy4import functools5import socket6import struct7import time8from typing import Any, Optional9 10import aioquic.quic.configuration  # type: ignore11import aioquic.quic.connection  # type: ignore12 13import dns.inet14 15QUIC_MAX_DATAGRAM = 204816MAX_SESSION_TICKETS = 817# If we hit the max sessions limit we will delete this many of the oldest connections.18# The value must be a integer > 0 and <= MAX_SESSION_TICKETS.19SESSIONS_TO_DELETE = MAX_SESSION_TICKETS // 420 21 22class UnexpectedEOF(Exception):23    pass24 25 26class Buffer:27    def __init__(self):28        self._buffer = b""29        self._seen_end = False30 31    def put(self, data, is_end):32        if self._seen_end:33            return34        self._buffer += data35        if is_end:36            self._seen_end = True37 38    def have(self, amount):39        if len(self._buffer) >= amount:40            return True41        if self._seen_end:42            raise UnexpectedEOF43        return False44 45    def seen_end(self):46        return self._seen_end47 48    def get(self, amount):49        assert self.have(amount)50        data = self._buffer[:amount]51        self._buffer = self._buffer[amount:]52        return data53 54 55class BaseQuicStream:56    def __init__(self, connection, stream_id):57        self._connection = connection58        self._stream_id = stream_id59        self._buffer = Buffer()60        self._expecting = 061 62    def id(self):63        return self._stream_id64 65    def _expiration_from_timeout(self, timeout):66        if timeout is not None:67            expiration = time.time() + timeout68        else:69            expiration = None70        return expiration71 72    def _timeout_from_expiration(self, expiration):73        if expiration is not None:74            timeout = max(expiration - time.time(), 0.0)75        else:76            timeout = None77        return timeout78 79    # Subclass must implement receive() as sync / async and which returns a message80    # or raises UnexpectedEOF.81 82    def _encapsulate(self, datagram):83        l = len(datagram)84        return struct.pack("!H", l) + datagram85 86    def _common_add_input(self, data, is_end):87        self._buffer.put(data, is_end)88        try:89            return self._expecting > 0 and self._buffer.have(self._expecting)90        except UnexpectedEOF:91            return True92 93    def _close(self):94        self._connection.close_stream(self._stream_id)95        self._buffer.put(b"", True)  # send EOF in case we haven't seen it.96 97 98class BaseQuicConnection:99    def __init__(100        self, connection, address, port, source=None, source_port=0, manager=None101    ):102        self._done = False103        self._connection = connection104        self._address = address105        self._port = port106        self._closed = False107        self._manager = manager108        self._streams = {}109        self._af = dns.inet.af_for_address(address)110        self._peer = dns.inet.low_level_address_tuple((address, port))111        if source is None and source_port != 0:112            if self._af == socket.AF_INET:113                source = "0.0.0.0"114            elif self._af == socket.AF_INET6:115                source = "::"116            else:117                raise NotImplementedError118        if source:119            self._source = (source, source_port)120        else:121            self._source = None122 123    def close_stream(self, stream_id):124        del self._streams[stream_id]125 126    def _get_timer_values(self, closed_is_special=True):127        now = time.time()128        expiration = self._connection.get_timer()129        if expiration is None:130            expiration = now + 3600  # arbitrary "big" value131        interval = max(expiration - now, 0)132        if self._closed and closed_is_special:133            # lower sleep interval to avoid a race in the closing process134            # which can lead to higher latency closing due to sleeping when135            # we have events.136            interval = min(interval, 0.05)137        return (expiration, interval)138 139    def _handle_timer(self, expiration):140        now = time.time()141        if expiration <= now:142            self._connection.handle_timer(now)143 144 145class AsyncQuicConnection(BaseQuicConnection):146    async def make_stream(self, timeout: Optional[float] = None) -> Any:147        pass148 149 150class BaseQuicManager:151    def __init__(self, conf, verify_mode, connection_factory, server_name=None):152        self._connections = {}153        self._connection_factory = connection_factory154        self._session_tickets = {}155        if conf is None:156            verify_path = None157            if isinstance(verify_mode, str):158                verify_path = verify_mode159                verify_mode = True160            conf = aioquic.quic.configuration.QuicConfiguration(161                alpn_protocols=["doq", "doq-i03"],162                verify_mode=verify_mode,163                server_name=server_name,164            )165            if verify_path is not None:166                conf.load_verify_locations(verify_path)167        self._conf = conf168 169    def _connect(170        self, address, port=853, source=None, source_port=0, want_session_ticket=True171    ):172        connection = self._connections.get((address, port))173        if connection is not None:174            return (connection, False)175        conf = self._conf176        if want_session_ticket:177            try:178                session_ticket = self._session_tickets.pop((address, port))179                # We found a session ticket, so make a configuration that uses it.180                conf = copy.copy(conf)181                conf.session_ticket = session_ticket182            except KeyError:183                # No session ticket.184                pass185            # Whether or not we found a session ticket, we want a handler to save186            # one.187            session_ticket_handler = functools.partial(188                self.save_session_ticket, address, port189            )190        else:191            session_ticket_handler = None192        qconn = aioquic.quic.connection.QuicConnection(193            configuration=conf,194            session_ticket_handler=session_ticket_handler,195        )196        lladdress = dns.inet.low_level_address_tuple((address, port))197        qconn.connect(lladdress, time.time())198        connection = self._connection_factory(199            qconn, address, port, source, source_port, self200        )201        self._connections[(address, port)] = connection202        return (connection, True)203 204    def closed(self, address, port):205        try:206            del self._connections[(address, port)]207        except KeyError:208            pass209 210    def save_session_ticket(self, address, port, ticket):211        # We rely on dictionaries keys() being in insertion order here.  We212        # can't just popitem() as that would be LIFO which is the opposite of213        # what we want.214        l = len(self._session_tickets)215        if l >= MAX_SESSION_TICKETS:216            keys_to_delete = list(self._session_tickets.keys())[0:SESSIONS_TO_DELETE]217            for key in keys_to_delete:218                del self._session_tickets[key]219        self._session_tickets[(address, port)] = ticket220 221 222class AsyncQuicManager(BaseQuicManager):223    def connect(self, address, port=853, source=None, source_port=0):224        raise NotImplementedError225