Python asyncio 模块，IncompleteReadError() 实例源码

def test_readexactly_eof(self):
        # Read exact number of bytes (eof).
        stream = asyncio.StreamReader(loop=self.loop)
        n = 2 * len(self.DATA)
        read_task = asyncio.Task(stream.readexactly(n), loop=self.loop)

        def cb():
            stream.feed_data(self.DATA)
            stream.feed_eof()
        self.loop.call_soon(cb)

        with self.assertRaises(asyncio.IncompleteReadError) as cm:
            self.loop.run_until_complete(read_task)
        self.assertEqual(cm.exception.partial, self.DATA)
        self.assertEqual(cm.exception.expected, n)
        self.assertEqual(str(cm.exception),
                         '18 bytes read on a total of 36 expected bytes')
        self.assertEqual(b'', stream._buffer)

def read(self):
        try:
            data = await self.reader.readexactly(HEADER_LENGTH)
        except IncompleteReadError:
            self.closed = True
            return None

        box = SecretBox(self.key)

        header = box.decrypt(data, self.nonce)

        if header == TERMINATION_HEADER:
            self.closed = True
            return None

        length = struct.unpack('>H', header[:2])[0]
        mac = header[2:]

        data = await self.reader.readexactly(length)

        body = box.decrypt(mac + data, inc_nonce(self.nonce))

        self.nonce = inc_nonce(inc_nonce(self.nonce))
        return body

def run(self):
        try:
            await self.on_run()
        except asyncio.CancelledError as ex:
            logger.debug("Connection was closed.")
            self.set_error(ex)
        except ProtocolError as ex:
            logger.debug("Protocol error: %s", ex)
            self.set_error(ex)
        except (asyncio.IncompleteReadError, ConnectionError) as ex:
            logger.debug("Remote end was closed. Terminating connection.")
            self.set_error(ex)
        except Exception as ex:
            logger.exception("Unexpected error. Terminating connection.")
            self.set_error(ex)
        finally:
            self.close()

def test_tcp_server_auto_close_connection(event_loop, unused_tcp_port):
    """Connections are closed automatically when sessions finish."""

    host = '127.0.0.1'
    port = unused_tcp_port

    # Intentionally do not close the writer.
    async def noop(reader, writer):
        pass

    async with TCPServer(host, port, noop, loop=event_loop) as server:
        reader, writer = await asyncio.open_connection(
            host=host,
            port=port,
            loop=event_loop,
        )
        try:
            with pytest.raises(asyncio.IncompleteReadError):
                await asyncio.wait_for(
                    reader.readexactly(1), timeout=0.1,
                )
            server.close()
            await server.wait_closed()
        finally:
            writer.close()

def test_unix_server_auto_close_connection(event_loop, tempcwd):
    """Connections are closed automatically when sessions finish."""

    # Intentionally do not close the writer.
    async def noop(reader, writer):
        pass

    path = './echo.sock'
    async with UnixSocketServer(path, noop, loop=event_loop) as server:
        reader, writer = await asyncio.open_unix_connection(
            path=path,
            loop=event_loop,
        )
        try:
            with pytest.raises(asyncio.IncompleteReadError):
                await asyncio.wait_for(
                    reader.readexactly(1), timeout=0.1,
                )
            server.close()
            await server.wait_closed()
        finally:
            writer.close()

def test_readexactly_eof(self):
        # Read exact number of bytes (eof).
        stream = asyncio.StreamReader(loop=self.loop)
        n = 2 * len(self.DATA)
        read_task = asyncio.Task(stream.readexactly(n), loop=self.loop)

        def cb():
            stream.feed_data(self.DATA)
            stream.feed_eof()
        self.loop.call_soon(cb)

        with self.assertRaises(asyncio.IncompleteReadError) as cm:
            self.loop.run_until_complete(read_task)
        self.assertEqual(cm.exception.partial, self.DATA)
        self.assertEqual(cm.exception.expected, n)
        self.assertEqual(str(cm.exception),
                         '18 bytes read on a total of 36 expected bytes')
        self.assertEqual(b'', stream._buffer)

def connect(self):
        loop = asyncio.get_event_loop()
        while True:
            try:
                self.reader, self.writer = await asyncio.open_connection(self.server.config['masterserver_ip'],
                                                                         self.server.config['masterserver_port'],
                                                                         loop=loop)
                await self.handle_connection()
            except (ConnectionRefusedError, TimeoutError):
                pass
            except (ConnectionResetError, asyncio.IncompleteReadError):
                self.writer = None
                self.reader = None
            finally:
                logger.log_debug("Couldn't connect to the master server, retrying in 30 seconds.")
                await asyncio.sleep(30)

def __call__(self, reader, writer):
        iproto = self.protocol_cls(reader)
        oproto = self.protocol_cls(writer)
        while not reader.at_eof():
            try:
                with async_timeout.timeout(self.timeout):
                    yield from self.processor.process(iproto, oproto)
            except ConnectionError:
                logger.debug('client has closed the connection')
                writer.close()
            except asyncio.TimeoutError:
                logger.debug('timeout when processing the client request')
                writer.close()
            except asyncio.IncompleteReadError:
                logger.debug('client has closed the connection')
                writer.close()
            except Exception:
                # app exception
                logger.exception('unhandled app exception')
                writer.close()
        writer.close()

def test_readexactly_eof(self):
        # Read exact number of bytes (eof).
        stream = asyncio.StreamReader(loop=self.loop)
        n = 2 * len(self.DATA)
        read_task = asyncio.Task(stream.readexactly(n), loop=self.loop)

        def cb():
            stream.feed_data(self.DATA)
            stream.feed_eof()
        self.loop.call_soon(cb)

        with self.assertRaises(asyncio.IncompleteReadError) as cm:
            self.loop.run_until_complete(read_task)
        self.assertEqual(cm.exception.partial, self.DATA)
        self.assertEqual(cm.exception.expected, n)
        self.assertEqual(str(cm.exception),
                         '18 bytes read on a total of 36 expected bytes')
        self.assertEqual(b'', stream._buffer)

def readuntil(self, separator):
        seplen = len(separator)
        offset = 0
        while True:
            buflen = len(self._buffer)
            if buflen - offset >= seplen:
                isep = self._buffer.find(separator, offset)
                if isep != -1:
                    break
                offset = buflen + 1 - seplen
            if self._eof:
                chunk = bytes(self._buffer)
                self._buffer.clear()
                raise asyncio.IncompleteReadError(chunk, None)
            yield from self._wait_for_data('readuntil')
        chunk = self._buffer[:isep + seplen]
        del self._buffer[:isep + seplen]
        self._maybe_resume_transport()
        return bytes(chunk)

def handle_jobs(job_handler, host, port, *, loop):
    """
    Connects to the remote master and continuously receives calls, executes
    them, then returns a response until interrupted.
    """

    try:

        try:
            reader, writer = await asyncio.open_connection(host, port, loop=loop)
        except OSError:
            logging.error("worker could not connect to server")
            return

        while True:

            try:
                call_encoded = await reader.readuntil(b"\n")
            except (asyncio.IncompleteReadError, ConnectionResetError):
                break
            logging.debug("worker got call")
            call_json = call_encoded.decode("utf-8")
            call = json.loads(call_json)

            response = job_handler(call)

            response_json = json.dumps(response) + "\n"
            response_encoded = response_json.encode("utf-8")
            writer.write(response_encoded)
            logging.debug("worker returned response")

    except KeyboardInterrupt:

        pass

def read_request_data(request, chunk_size):
    '''
    cachable request data reader to help with conflict error requests
    '''
    if getattr(request, '_retry_attempt', 0) > 0:
        # we are on a retry request, see if we have read cached data yet...
        if request._retry_attempt > getattr(request, '_last_cache_data_retry_count', 0):
            if getattr(request, '_cache_data', None) is None:
                # request payload was too large to fit into request cache.
                # so retrying this request is not supported and we need to throw
                # another error
                raise UnRetryableRequestError()
            data = request._cache_data[request._last_read_pos:request._last_read_pos + chunk_size]
            request._last_read_pos += len(data)
            if request._last_read_pos >= len(request._cache_data):
                # done reading cache data
                request._last_cache_data_retry_count = request._retry_attempt
            return data

    if not hasattr(request, '_cache_data'):
        request._cache_data = b''

    try:
        data = await request.content.readexactly(chunk_size)
    except asyncio.IncompleteReadError as e:
        data = e.partial

    if request._cache_data is not None:
        if len(request._cache_data) + len(data) > MAX_REQUEST_CACHE_SIZE:
            # we only allow caching up to chunk size, otherwise, no cache data..
            request._cache_data = None
        else:
            request._cache_data += data

    request._last_read_pos += len(data)
    return data

def connection_handler(execute_rpc, reader, writer):
    '''
    Catch fatal errors, log them, and close the connection
    '''
    logger = logging.getLogger('evm.rpc.ipc')
    try:
        await connection_loop(execute_rpc, reader, writer, logger)
    except (ConnectionResetError, asyncio.IncompleteReadError):
        logger.debug("Client closed connection")
    except Exception:
        logger.exception("Unrecognized exception while handling requests")
    finally:
        writer.close()

def read(self, n: int) -> bytes:
        self.logger.debug("Waiting for %s bytes from %s", n, self.remote)
        try:
            data = await asyncio.wait_for(self.reader.readexactly(n), self.conn_idle_timeout)
        except (asyncio.IncompleteReadError, ConnectionResetError):
            raise PeerConnectionLost("EOF reading from stream")
        return data

def test_base_incomplete_error(loop):
    proto = BaseSocksProtocol(None, None, ('python.org', 80),
                              None, None, reader_limit=10, loop=loop)
    proto._stream_reader.readexactly = make_mocked_coro(
        raise_exception=asyncio.IncompleteReadError(b'part', 5))
    with pytest.raises(aiosocks.InvalidServerReply):
        await proto.read_response(4)

def read_response(self, n):
        try:
            return (await self._stream_reader.readexactly(n))
        except asyncio.IncompleteReadError as e:
            raise InvalidServerReply(
                'Server sent fewer bytes than required (%s)' % str(e))

def test_readuntil_eof(self):
        stream = asyncio.StreamReader(loop=self.loop)
        stream.feed_data(b'some dataAA')
        stream.feed_eof()

        with self.assertRaises(asyncio.IncompleteReadError) as cm:
            self.loop.run_until_complete(stream.readuntil(b'AAA'))
        self.assertEqual(cm.exception.partial, b'some dataAA')
        self.assertIsNone(cm.exception.expected)
        self.assertEqual(b'', stream._buffer)

def readuntil(self, delim):
        if not self.current_event:
            if not self.events:
                raise asyncio.IncompleteReadError(b'', delim)
            yield from asyncio.sleep(self.delay)
            self.current_event = self.events.pop(0)
        data, rest = self.current_event.split(delim, 1)
        self.current_event = rest
        return data + delim

def _client_connected(self, reader, writer):
        connection_id = self.next_connection_id
        self.next_connection_id += 1

        logger.debug('%d connected', connection_id)
        while True:
            if reader.at_eof():
                logger.debug('%d connection closed', connection_id)
                break

            try:
                request = await self.parser.read_request(reader)
                request.connection_id = connection_id
                resp = await self.process_request(request)
                await self.parser.write_response(writer, resp)
            except asyncio.IncompleteReadError:
                writer.close()
            except Exception as e:
                logger.exception(e)

                if isinstance(e, ParseError):
                    status_code = 400
                else:
                    status_code = 500

                resp = Response(
                    body=str(e) + '\n',
                    status_code=status_code,
                )
                await self.parser.write_response(writer, resp)
                writer.close()

def connect(self):
        loop = asyncio.get_event_loop()
        while True:
            try:
                self.reader, self.writer = await asyncio.open_connection(self.server.config['district_ip'],
                                                                         self.server.config['district_port'], loop=loop)
                await self.handle_connection()
            except (ConnectionRefusedError, TimeoutError):
                pass
            except (ConnectionResetError, asyncio.IncompleteReadError):
                self.writer = None
                self.reader = None
            finally:
                logger.log_debug("Couldn't connect to the district, retrying in 30 seconds.")
                await asyncio.sleep(30)

def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        addr_repr = ':'.join(map(str, writer.get_extra_info('peername')))
        logger.info('accepted connection from %s', addr_repr)

        try:
            writer.write(ControlServer.HANDSHAKE_MESSAGE)

            while True:
                # FIXME: maybe do not allow to execute arbitrary object
                action = cast(Callable[[ControlManager], Any], await ControlServer.receive_object(reader))

                try:
                    result = action(self._control)
                    if asyncio.iscoroutine(result):
                        result = await result
                except asyncio.CancelledError:
                    raise
                except Exception as e:
                    result = e

                ControlServer.send_object(result, writer)

                if isinstance(result, DaemonExit):
                    logger.info('stop command received')
                    if self._daemon_stop_handler is not None:
                        self._daemon_stop_handler(self)
                    return
        except asyncio.IncompleteReadError:
            pass
        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.warning('%s disconnected because of %r', addr_repr, e)
        finally:
            writer.close()

def _response_reader(self, reader: asyncio.StreamReader):
        try:
            while True:
                try:
                    response_header = await reader.readexactly(4 + 4 + 8)
                except asyncio.IncompleteReadError as err:
                    if err.partial:
                        raise RuntimeError('Partial bytes was read.')
                    log.debug('Clean end of stream.')
                    break
                (magic, errno, handle) = struct.unpack(">LLQ", response_header)
                if magic != NBD_REPLY_MAGIC:
                    raise RuntimeError('Protocol error')
                # log.debug('Found response')
                (fut, expected_length) = self._tasks.pop(handle)
                if errno:
                    # TODO: some errno MUST broke connection (!) like request format error (!)
                    fut.set_exception(RemoteNBDError('Remote NBD Error', errno))
                    continue

                if not expected_length:
                    fut.set_result(None)
                    continue

                try:
                    fut.set_result(await reader.readexactly(expected_length))
                except Exception as e:
                    # TODO: set_exception(RemoteRBDError() from e)
                    fut.set_exception(RemoteNBDError('Remote NBD Error: Can\'t read associated data: %r' % e))
                    raise
        finally:
            if self._tasks:
                log.error('Incomplete tasks in the queue. Aborting them.')
                for (fut, expected_length) in self._tasks.values():
                    try:
                        fut.set_exception(RemoteNBDError('Protocol error', None))
                    except Exception:
                        # may be already cancelled.
                        pass
                self._tasks.clear()

def test_tcp_server_wait_closed_timeout(event_loop, unused_tcp_port):
    """Not finished closing until all sessions complete."""

    host = '127.0.0.1'
    port = unused_tcp_port

    # This session will (intentionally) never receive enough data.  After a
    # while, it will get cancelled and (intentionally) ignore the cancellation
    # request to simulate a badly designed session handler.  Once our test
    # verifies this behavior, it will close the client socket, which will
    # finally trigger the shutdown sequence.
    async def noop(reader, writer):
        try:
            await reader.readexactly(1024)
        except asyncio.CancelledError:
            try:
                await reader.readexactly(1024)
            except asyncio.IncompleteReadError:
                pass
        finally:
            writer.close()

    server = TCPServer(host, port, noop, loop=event_loop)
    server.start()
    await server.wait_started()
    try:
        reader, writer = await asyncio.open_connection(
            host=host,
            port=port,
            loop=event_loop,
        )
        try:
            writer.write(b'REQ')
            with pytest.raises(asyncio.TimeoutError):
                await asyncio.wait_for(
                    reader.readexactly(1), timeout=0.1,
                )
            server.close()
            with pytest.raises(asyncio.TimeoutError):
                await server.wait_closed(timeout=0.1)
        finally:
            writer.close()
    finally:
        server.close()
        await server.wait_closed()

def test_unix_server_wait_closed_timeout(event_loop, tempcwd):
    """Not finished closing until all sessions complete."""

    # This session will (intentionally) never receive enough data.  After a
    # while, it will get cancelled and (intentionally) ignore the cancellation
    # request to simulate a badly designed session handler.  Once our test
    # verifies this behavior, it will close the client socket, which will
    # finally trigger the shutdown sequence.
    async def noop(reader, writer):
        try:
            await reader.readexactly(1024)
        except asyncio.CancelledError:
            try:
                await reader.readexactly(1024)
            except asyncio.IncompleteReadError:
                pass
        finally:
            writer.close()

    path = './echo.sock'
    server = UnixSocketServer(path, noop, loop=event_loop)
    server.start()
    await server.wait_started()
    try:
        reader, writer = await asyncio.open_unix_connection(
            path=path,
            loop=event_loop,
        )
        try:
            writer.write(b'REQ')
            with pytest.raises(asyncio.TimeoutError):
                await asyncio.wait_for(
                    reader.readexactly(1), timeout=0.1,
                )
            server.close()
            with pytest.raises(asyncio.TimeoutError):
                await server.wait_closed(timeout=0.1)
        finally:
            writer.close()
    finally:
        server.close()
        await server.wait_closed()

def _listen_for_events(self, vm=None):
        '''
        Listen for events and call appropriate handlers.
        This function do not exit until manually terminated.

        This is coroutine.

        :param vm: Listen for events only for this VM, use None to listen for
        events about all VMs and not related to any particular VM.
        :return: True if any event was received, otherwise False
        :rtype: bool
        '''

        reader, cleanup_func = yield from self._get_events_reader(vm)
        try:
            some_event_received = False
            while not reader.at_eof():
                try:
                    event_header = yield from reader.readuntil(b'\0')
                    if event_header != b'1\0':
                        raise qubesadmin.exc.QubesDaemonCommunicationError(
                            'Non-event received on events connection: '
                            + repr(event_header))
                    subject = (yield from reader.readuntil(b'\0'))[:-1].decode(
                        'utf-8')
                    event = (yield from reader.readuntil(b'\0'))[:-1].decode(
                        'utf-8')
                    kwargs = {}
                    while True:
                        key = (yield from reader.readuntil(b'\0'))[:-1].decode(
                            'utf-8')
                        if not key:
                            break
                        value = (yield from reader.readuntil(b'\0'))[:-1].\
                            decode('utf-8')
                        kwargs[key] = value
                except asyncio.IncompleteReadError as err:
                    if err.partial == b'':
                        break
                    else:
                        raise

                if not subject:
                    subject = None
                self.handle(subject, event, **kwargs)

                some_event_received = True
        finally:
            cleanup_func()
        return some_event_received

def execute(self, api, *args, **kwargs):
        """
        Execute a rpc call by api name. This is function is a :ref:`coroutine <coroutine>`.

        :param api: api name defined in thrift file
        :param args: positional arguments passed to api function
        :param kwargs:  keyword arguments passed to api function
        :return: result of this rpc call
        :raises: :class:`~asyncio.TimeoutError` if this task has exceeded the `timeout`
        :raises: :class:`ThriftAppError` if thrift response is an exception defined in thrift.
        :raises: :class:`ConnectionClosedError`: if server has closed this connection.
        """
        if self.closed:
            raise ConnectionClosedError('Connection closed')

        try:
            with async_timeout.timeout(self.timeout):
                kw = args2kwargs(getattr(self.service, api + "_args").thrift_spec, *args)
                kwargs.update(kw)
                result_cls = getattr(self.service, api + "_result")

                self._seqid += 1
                self._oprot.write_message_begin(api, TMessageType.CALL, self._seqid)
                args = getattr(self.service, api + '_args')()
                for k, v in kwargs.items():
                    setattr(args, k, v)
                args.write(self._oprot)
                self._oprot.write_message_end()
                yield from self._oprot.trans.drain()
                if not getattr(result_cls, "oneway"):
                    result = yield from self._recv(api)
                    return result
        except asyncio.TimeoutError:
            self.close()
            raise
        except ConnectionError as e:
            self.close()
            logger.debug('connection error {}'.format(str(e)))
            raise ConnectionClosedError('the server has closed this connection') from e
        except asyncio.IncompleteReadError as e:
            self.close()
            raise ConnectionClosedError('Server connection has closed') from e