codekingpro/portable-devtools

1import asyncio2import sys3import time4 5from . import base_socket6from . import exceptions7from . import packet8from . import payload9 10 11class AsyncSocket(base_socket.BaseSocket):12    async def poll(self):13        """Wait for packets to send to the client."""14        try:15            packets = [await asyncio.wait_for(16                self.queue.get(),17                self.server.ping_interval + self.server.ping_timeout)]18            self.queue.task_done()19        except (asyncio.TimeoutError, asyncio.CancelledError):20            raise exceptions.QueueEmpty()21        if packets == [None]:22            return []23        while True:24            try:25                pkt = self.queue.get_nowait()26                self.queue.task_done()27                if pkt is None:28                    self.queue.put_nowait(None)29                    break30                packets.append(pkt)31            except asyncio.QueueEmpty:32                break33        return packets34 35    async def receive(self, pkt):36        """Receive packet from the client."""37        self.server.logger.info('%s: Received packet %s data %s',38                                self.sid, packet.packet_names[pkt.packet_type],39                                pkt.data if not isinstance(pkt.data, bytes)40                                else '<binary>')41        if pkt.packet_type == packet.PONG:42            self.schedule_ping()43        elif pkt.packet_type == packet.MESSAGE:44            await self.server._trigger_event(45                'message', self.sid, pkt.data,46                run_async=self.server.async_handlers)47        elif pkt.packet_type == packet.UPGRADE:48            await self.send(packet.Packet(packet.NOOP))49        elif pkt.packet_type == packet.CLOSE:50            await self.close(wait=False, abort=True)51        else:52            raise exceptions.UnknownPacketError()53 54    async def check_ping_timeout(self):55        """Make sure the client is still sending pings."""56        if self.closed:57            raise exceptions.SocketIsClosedError()58        if self.last_ping and \59                time.time() - self.last_ping > self.server.ping_timeout:60            self.server.logger.info('%s: Client is gone, closing socket',61                                    self.sid)62            # Passing abort=False here will cause close() to write a63            # CLOSE packet. This has the effect of updating half-open sockets64            # to their correct state of disconnected65            await self.close(wait=False, abort=False)66            return False67        return True68 69    async def send(self, pkt):70        """Send a packet to the client."""71        if not await self.check_ping_timeout():72            return73        else:74            await self.queue.put(pkt)75        self.server.logger.info('%s: Sending packet %s data %s',76                                self.sid, packet.packet_names[pkt.packet_type],77                                pkt.data if not isinstance(pkt.data, bytes)78                                else '<binary>')79 80    async def handle_get_request(self, environ):81        """Handle a long-polling GET request from the client."""82        connections = [83            s.strip()84            for s in environ.get('HTTP_CONNECTION', '').lower().split(',')]85        transport = environ.get('HTTP_UPGRADE', '').lower()86        if 'upgrade' in connections and transport in self.upgrade_protocols:87            self.server.logger.info('%s: Received request to upgrade to %s',88                                    self.sid, transport)89            return await getattr(self, '_upgrade_' + transport)(environ)90        if self.upgrading or self.upgraded:91            # we are upgrading to WebSocket, do not return any more packets92            # through the polling endpoint93            return [packet.Packet(packet.NOOP)]94        try:95            packets = await self.poll()96        except exceptions.QueueEmpty:97            exc = sys.exc_info()98            await self.close(wait=False)99            raise exc[1].with_traceback(exc[2])100        return packets101 102    async def handle_post_request(self, environ):103        """Handle a long-polling POST request from the client."""104        length = int(environ.get('CONTENT_LENGTH', '0'))105        if length > self.server.max_http_buffer_size:106            raise exceptions.ContentTooLongError()107        else:108            body = (await environ['wsgi.input'].read(length)).decode('utf-8')109            p = payload.Payload(encoded_payload=body)110            for pkt in p.packets:111                await self.receive(pkt)112 113    async def close(self, wait=True, abort=False):114        """Close the socket connection."""115        if not self.closed and not self.closing:116            self.closing = True117            await self.server._trigger_event('disconnect', self.sid)118            if not abort:119                await self.send(packet.Packet(packet.CLOSE))120            self.closed = True121            if wait:122                await self.queue.join()123 124    def schedule_ping(self):125        self.server.start_background_task(self._send_ping)126 127    async def _send_ping(self):128        self.last_ping = None129        await asyncio.sleep(self.server.ping_interval)130        if not self.closing and not self.closed:131            self.last_ping = time.time()132            await self.send(packet.Packet(packet.PING))133 134    async def _upgrade_websocket(self, environ):135        """Upgrade the connection from polling to websocket."""136        if self.upgraded:137            raise IOError('Socket has been upgraded already')138        if self.server._async['websocket'] is None:139            # the selected async mode does not support websocket140            return self.server._bad_request()141        ws = self.server._async['websocket'](142            self._websocket_handler, self.server)143        return await ws(environ)144 145    async def _websocket_handler(self, ws):146        """Engine.IO handler for websocket transport."""147        async def websocket_wait():148            data = await ws.wait()149            if data and len(data) > self.server.max_http_buffer_size:150                raise ValueError('packet is too large')151            return data152 153        if self.connected:154            # the socket was already connected, so this is an upgrade155            self.upgrading = True  # hold packet sends during the upgrade156 157            try:158                pkt = await websocket_wait()159            except IOError:  # pragma: no cover160                return161            decoded_pkt = packet.Packet(encoded_packet=pkt)162            if decoded_pkt.packet_type != packet.PING or \163                    decoded_pkt.data != 'probe':164                self.server.logger.info(165                    '%s: Failed websocket upgrade, no PING packet', self.sid)166                self.upgrading = False167                return168            await ws.send(packet.Packet(packet.PONG, data='probe').encode())169            await self.queue.put(packet.Packet(packet.NOOP))  # end poll170 171            try:172                pkt = await websocket_wait()173            except IOError:  # pragma: no cover174                self.upgrading = False175                return176            decoded_pkt = packet.Packet(encoded_packet=pkt)177            if decoded_pkt.packet_type != packet.UPGRADE:178                self.upgraded = False179                self.server.logger.info(180                    ('%s: Failed websocket upgrade, expected UPGRADE packet, '181                     'received %s instead.'),182                    self.sid, pkt)183                self.upgrading = False184                return185            self.upgraded = True186            self.upgrading = False187        else:188            self.connected = True189            self.upgraded = True190 191        # start separate writer thread192        async def writer():193            while True:194                packets = None195                try:196                    packets = await self.poll()197                except exceptions.QueueEmpty:198                    break199                if not packets:200                    # empty packet list returned -> connection closed201                    break202                try:203                    for pkt in packets:204                        await ws.send(pkt.encode())205                except:206                    break207            await ws.close()208 209        writer_task = asyncio.ensure_future(writer())210 211        self.server.logger.info(212            '%s: Upgrade to websocket successful', self.sid)213 214        while True:215            p = None216            wait_task = asyncio.ensure_future(websocket_wait())217            try:218                p = await asyncio.wait_for(219                    wait_task,220                    self.server.ping_interval + self.server.ping_timeout)221            except asyncio.CancelledError:  # pragma: no cover222                # there is a bug (https://bugs.python.org/issue30508) in223                # asyncio that causes a "Task exception never retrieved" error224                # to appear when wait_task raises an exception before it gets225                # cancelled. Calling wait_task.exception() prevents the error226                # from being issued in Python 3.6, but causes other errors in227                # other versions, so we run it with all errors suppressed and228                # hope for the best.229                try:230                    wait_task.exception()231                except:232                    pass233                break234            except:235                break236            if p is None:237                # connection closed by client238                break239            pkt = packet.Packet(encoded_packet=p)240            try:241                await self.receive(pkt)242            except exceptions.UnknownPacketError:  # pragma: no cover243                pass244            except exceptions.SocketIsClosedError:  # pragma: no cover245                self.server.logger.info('Receive error -- socket is closed')246                break247            except:  # pragma: no cover248                # if we get an unexpected exception we log the error and exit249                # the connection properly250                self.server.logger.exception('Unknown receive error')251 252        await self.queue.put(None)  # unlock the writer task so it can exit253        await asyncio.wait_for(writer_task, timeout=None)254        await self.close(wait=False, abort=True)255