Python asyncio.StreamReader() Examples

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 handshake(
    remote: kademlia.Node, privkey: datatypes.PrivateKey, token: CancelToken
) -> Tuple[
    bytes, bytes, BasePreImage, BasePreImage, asyncio.StreamReader, asyncio.StreamWriter
]:  # noqa: E501
    """
    Perform the auth handshake with given remote.

    Returns the established secrets and the StreamReader/StreamWriter pair already connected to
    the remote.
    """
    use_eip8 = False
    initiator = HandshakeInitiator(remote, privkey, use_eip8, token)
    reader, writer = await initiator.connect()
    opened_connections[remote.__repr__()] = (reader, writer)
    aes_secret, mac_secret, egress_mac, ingress_mac = await _handshake(
        initiator, reader, writer, token
    )
    return aes_secret, mac_secret, egress_mac, ingress_mac, reader, writer 

def test_readline(self):
        # Read one line. 'readline' will need to wait for the data
        # to come from 'cb'
        stream = asyncio.StreamReader(loop=self.loop)
        stream.feed_data(b'chunk1 ')
        read_task = asyncio.Task(stream.readline(), loop=self.loop)

        def cb():
            stream.feed_data(b'chunk2 ')
            stream.feed_data(b'chunk3 ')
            stream.feed_data(b'\n chunk4')
        self.loop.call_soon(cb)

        line = self.loop.run_until_complete(read_task)
        self.assertEqual(b'chunk1 chunk2 chunk3 \n', line)
        self.assertEqual(b' chunk4', stream._buffer) 

def test_readexactly(self):
        # Read exact number of bytes.
        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_data(self.DATA)
            stream.feed_data(self.DATA)
        self.loop.call_soon(cb)

        data = self.loop.run_until_complete(read_task)
        self.assertEqual(self.DATA + self.DATA, data)
        self.assertEqual(self.DATA, stream._buffer) 

def test_read_all_from_pipe_reader(self):
        # See asyncio issue 168.  This test is derived from the example
        # subprocess_attach_read_pipe.py, but we configure the
        # StreamReader's limit so that twice it is less than the size
        # of the data writter.  Also we must explicitly attach a child
        # watcher to the event loop.

        code = """\
import os, sys
fd = int(sys.argv[1])
os.write(fd, b'data')
os.close(fd)
"""
        rfd, wfd = os.pipe()
        args = [sys.executable, '-c', code, str(wfd)]

        pipe = open(rfd, 'rb', 0)
        reader = asyncio.StreamReader(loop=self.loop, limit=1)
        protocol = asyncio.StreamReaderProtocol(reader, loop=self.loop)
        transport, _ = self.loop.run_until_complete(
            self.loop.connect_read_pipe(lambda: protocol, pipe))

        watcher = asyncio.SafeChildWatcher()
        watcher.attach_loop(self.loop)
        try:
            asyncio.set_child_watcher(watcher)
            create = asyncio.create_subprocess_exec(*args,
                                                    pass_fds={wfd},
                                                    loop=self.loop)
            proc = self.loop.run_until_complete(create)
            self.loop.run_until_complete(proc.wait())
        finally:
            asyncio.set_child_watcher(None)

        os.close(wfd)
        data = self.loop.run_until_complete(reader.read(-1))
        self.assertEqual(data, b'data') 

def _wait_for_data(self, func_name):
    """Wait until feed_data() or feed_eof() is called.

    If stream was paused, automatically resume it.
    """
    # StreamReader uses a future to link the protocol feed_data() method
    # to a read coroutine. Running two read coroutines at the same time
    # would have an unexpected behaviour. It would not possible to know
    # which coroutine would get the next data.
    if self._waiter is not None:
        raise RuntimeError('%s() called while another coroutine is '
                           'already waiting for incoming data' % func_name)

    assert not self._eof, '_wait_for_data after EOF'

    # Waiting for data while paused will make deadlock, so prevent it.
    if self._paused:
        self._paused = False
        self._transport.resume_reading()

    self._waiter = asyncio.futures.Future(loop=self._loop)
    try:
        await self._waiter
    finally:
        self._waiter = None 

def create_pipe_streams_pair():
    path = tempfile.mktemp()
    loop = asyncio.get_event_loop()
    server_side = asyncio.Future()

    def factory():

        def client_connected_cb(reader, writer):
            server_side.set_result((reader, writer))
        reader = asyncio.StreamReader(loop=loop)
        return asyncio.StreamReaderProtocol(reader, client_connected_cb, loop=loop)

    server = yield from loop.create_unix_server(factory, path)

    r1 = asyncio.StreamReader(loop=loop)
    protocol = asyncio.StreamReaderProtocol(r1, loop=loop)
    transport, _ = yield from loop.create_unix_connection(
        lambda: protocol, path)
    w1 = asyncio.StreamWriter(transport, protocol, r1, loop)

    r2, w2 = yield from server_side
    server.close()
    return (r1, w1), (r2, w2) 

def test_it_logs_messages(self):
        asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
        loop = asyncio.get_event_loop()

        async def test():
            reader = asyncio.StreamReader(loop=loop)
            protocol = asyncio.StreamReaderProtocol(reader)

            transport, _ = await loop.connect_read_pipe(
                lambda: protocol, self.read_pipe
            )

            logger = Logger.with_default_handlers()
            await logger.info("Xablau")

            logged_content = await reader.readline()
            self.assertEqual(logged_content, b"Xablau\n")

            transport.close()
            await logger.shutdown()

        loop.run_until_complete(test()) 

def __call__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        """Main entry point for an interpreter session with a single client."""
        self.reader = reader
        self.writer = writer
        self.running = True

        if self.banner:
            writer.write(self.banner)
            await writer.drain()

        while self.running:
            try:
                await self.handle_one_command()
            except ConnectionResetError:
                writer.close()
                self.running = False
                break
            except Exception:
                log.exception("Exception in manhole REPL")
                self.writer.write(traceback.format_exc())
                await self.writer.drain() 

def test_exception_cancel(self):
        stream = asyncio.StreamReader(loop=self.loop)

        t = asyncio.Task(stream.readline(), loop=self.loop)
        test_utils.run_briefly(self.loop)
        t.cancel()
        test_utils.run_briefly(self.loop)
        # The following line fails if set_exception() isn't careful.
        stream.set_exception(RuntimeError('message'))
        test_utils.run_briefly(self.loop)
        self.assertIs(stream._waiter, None) 

def test_readexactly_zero_or_less(self):
        # Read exact number of bytes (zero or less).
        stream = asyncio.StreamReader(loop=self.loop)
        stream.feed_data(self.DATA)

        data = self.loop.run_until_complete(stream.readexactly(0))
        self.assertEqual(b'', data)
        self.assertEqual(self.DATA, stream._buffer)

        data = self.loop.run_until_complete(stream.readexactly(-1))
        self.assertEqual(b'', data)
        self.assertEqual(self.DATA, stream._buffer) 

def __init__(
        self,
        local_type: NodeType,
        connection_type: Optional[NodeType],
        sr: StreamReader,
        sw: StreamWriter,
        server_port: int,
        on_connect: OnConnectFunc,
        log: logging.Logger,
    ):
        self.local_type = local_type
        self.connection_type = connection_type
        self.reader = sr
        self.writer = sw
        socket = self.writer.get_extra_info("socket")
        self.local_host = socket.getsockname()[0]
        self.local_port = server_port
        self.peer_host = self.writer.get_extra_info("peername")[0]
        self.peer_port = self.writer.get_extra_info("peername")[1]
        self.peer_server_port: Optional[int] = None
        self.node_id = None
        self.on_connect = on_connect
        self.log = log

        # Connection metrics
        self.creation_time = time.time()
        self.bytes_read = 0
        self.bytes_written = 0
        self.last_message_time: float = 0
        self._cached_peer_name = self.writer.get_extra_info("peername") 

def __call__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
                       ) -> None:
        sock = writer.transport.get_extra_info("socket")
        # TODO support non-linux OSes
        # I think FreeBSD uses SCM_CREDS
        creds = sock.getsockopt(SOL_SOCKET, SO_PEERCRED, struct.calcsize('3i'))
        pid, uid, gid = struct.unpack('3i', creds)
        user_info = pwd.getpwuid(uid)
        username = f"{user_info.pw_name} ({uid})" if user_info and user_info.pw_name else uid
        if len(self.whitelist) > 0 and uid not in self.whitelist:
            writer.write(b"You are not whitelisted to use the manhole.")
            log.warning(f"Non-whitelisted user {username} tried to connect from PID {pid}")
            await writer.drain()
            writer.close()
            return

        namespace = {**self.namespace}
        if AWAIT_TRANSFORM:
            namespace[AWAIT_FUNC_NAME] = AWAIT_FALLBACK
        interpreter = self.interpreter_class(namespace=namespace, banner=self.banner,
                                             loop=self.loop)
        namespace["exit"] = interpreter.close
        self.clients.append(interpreter)
        conn_id = self.conn_id

        log.info(f"Manhole connection OPENED: {conn_id} from PID {pid} by {username}")
        await asyncio.ensure_future(interpreter(reader, writer))
        log.info(f"Manhole connection CLOSED: {conn_id} from PID {pid} by {username}")
        self.clients.remove(interpreter) 

def __init__(self, config: Dict, constants: Dict):
        self.constants = constants
        self.config: Dict = config
        self.ips_estimate = {
            k: v
            for k, v in list(
                zip(
                    self.config["vdf_clients"]["ip"],
                    self.config["vdf_clients"]["ips_estimate"],
                )
            )
        }
        self.lock: Lock = Lock()
        self.active_discriminants: Dict[bytes32, Tuple[StreamWriter, uint64, str]] = {}
        self.best_weight_three_proofs: int = -1
        self.active_discriminants_start_time: Dict = {}
        self.pending_iters: Dict = {}
        self.submitted_iters: Dict = {}
        self.done_discriminants: List[bytes32] = []
        self.proofs_to_write: List[OutboundMessage] = []
        self.seen_discriminants: List[bytes32] = []
        self.proof_count: Dict = {}
        self.avg_ips: Dict = {}
        self.discriminant_queue: List[Tuple[bytes32, uint128]] = []
        self.max_connection_time = self.config["max_connection_time"]
        self.potential_free_clients: List = []
        self.free_clients: List[Tuple[str, StreamReader, StreamWriter]] = []
        self.server: Optional[ChiaServer] = None
        self._is_shutdown = False 

def _handle_client(self, reader: StreamReader, writer: StreamWriter):
        async with self.lock:
            client_ip = writer.get_extra_info("peername")[0]
            log.info(f"New timelord connection from client: {client_ip}.")
            if client_ip in self.ips_estimate.keys():
                self.free_clients.append((client_ip, reader, writer))
                log.info(f"Added new VDF client {client_ip}.")
                for ip, end_time in list(self.potential_free_clients):
                    if ip == client_ip:
                        self.potential_free_clients.remove((ip, end_time))
                        break 

def start_server(
    self: "ChiaServer", on_connect: OnConnectFunc = None
) -> asyncio.AbstractServer:
    """
    Launches a listening server on host and port specified, to connect to NodeType nodes. On each
    connection, the on_connect asynchronous generator will be called, and responses will be sent.
    Whenever a new TCP connection is made, a new srwt tuple is sent through the pipeline.
    """
    require_cert = self._local_type not in (NodeType.FULL_NODE, NodeType.INTRODUCER)
    ssl_context = ssl_context_for_server(
        self.root_path, self.config, require_cert=require_cert
    )

    server, aiter = await start_server_aiter(
        self._port, host=None, reuse_address=True, ssl=ssl_context
    )

    def add_connection_type(
        srw: Tuple[asyncio.StreamReader, asyncio.StreamWriter]
    ) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter, OnConnectFunc]:
        ssl_object = srw[1].get_extra_info(name="ssl_object")
        peer_cert = ssl_object.getpeercert()
        self.log.info(f"Client authed as {peer_cert}")
        return (srw[0], srw[1], on_connect)

    srwt_aiter = map_aiter(add_connection_type, aiter)

    # Push aiters that come from the server into the pipeline
    if not self._srwt_aiter.is_stopped():
        self._srwt_aiter.push(srwt_aiter)

    self.log.info(f"Server started on port {self._port}")
    return server 

def test_feed_empty_data(self):
        stream = asyncio.StreamReader(loop=self.loop)

        stream.feed_data(b'')
        self.assertEqual(b'', stream._buffer) 

def test_read_zero(self):
        # Read zero bytes.
        stream = asyncio.StreamReader(loop=self.loop)
        stream.feed_data(self.DATA)

        data = self.loop.run_until_complete(stream.read(0))
        self.assertEqual(b'', data)
        self.assertEqual(self.DATA, stream._buffer) 

def test_feed_nonempty_data(self):
        stream = asyncio.StreamReader(loop=self.loop)

        stream.feed_data(self.DATA)
        self.assertEqual(self.DATA, stream._buffer) 

def test_exception(self):
        stream = asyncio.StreamReader(loop=self.loop)
        self.assertIsNone(stream.exception())

        exc = ValueError()
        stream.set_exception(exc)
        self.assertIs(stream.exception(), exc) 

def test_readexactly_exception(self):
        stream = asyncio.StreamReader(loop=self.loop)
        stream.feed_data(b'line\n')

        data = self.loop.run_until_complete(stream.readexactly(2))
        self.assertEqual(b'li', data)

        stream.set_exception(ValueError())
        self.assertRaises(
            ValueError, self.loop.run_until_complete, stream.readexactly(2)) 

def test___repr__nondefault_limit(self):
        stream = asyncio.StreamReader(loop=self.loop, limit=123)
        self.assertEqual("<StreamReader l=123>", repr(stream)) 

def test_readline_read_byte_count(self):
        stream = asyncio.StreamReader(loop=self.loop)
        stream.feed_data(self.DATA)

        self.loop.run_until_complete(stream.readline())

        data = self.loop.run_until_complete(stream.read(7))

        self.assertEqual(b'line2\nl', data)
        self.assertEqual(b'ine3\n', stream._buffer) 

def test_readline_empty_eof(self):
        stream = asyncio.StreamReader(loop=self.loop)
        stream.feed_eof()

        line = self.loop.run_until_complete(stream.readline())
        self.assertEqual(b'', line) 

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

        line = self.loop.run_until_complete(stream.readline())
        self.assertEqual(b'some data', line) 

def test_readline_nolimit_nowait(self):
        # All needed data for the first 'readline' call will be
        # in the buffer.
        stream = asyncio.StreamReader(loop=self.loop)
        stream.feed_data(self.DATA[:6])
        stream.feed_data(self.DATA[6:])

        line = self.loop.run_until_complete(stream.readline())

        self.assertEqual(b'line1\n', line)
        self.assertEqual(b'line2\nline3\n', stream._buffer) 

def test_readline_limit(self):
        # Read one line. StreamReaders are fed with data after
        # their 'readline' methods are called.

        stream = asyncio.StreamReader(limit=7, loop=self.loop)
        def cb():
            stream.feed_data(b'chunk1')
            stream.feed_data(b'chunk2')
            stream.feed_data(b'chunk3\n')
            stream.feed_eof()
        self.loop.call_soon(cb)

        self.assertRaises(
            ValueError, self.loop.run_until_complete, stream.readline())
        # The buffer had just one line of data, and after raising
        # a ValueError it should be empty.
        self.assertEqual(b'', stream._buffer)

        stream = asyncio.StreamReader(limit=7, loop=self.loop)
        def cb():
            stream.feed_data(b'chunk1')
            stream.feed_data(b'chunk2\n')
            stream.feed_data(b'chunk3\n')
            stream.feed_eof()
        self.loop.call_soon(cb)

        self.assertRaises(
            ValueError, self.loop.run_until_complete, stream.readline())
        self.assertEqual(b'chunk3\n', stream._buffer) 

def test_readline_limit_with_existing_data(self):
        # Read one line. The data is in StreamReader's buffer
        # before the event loop is run.

        stream = asyncio.StreamReader(limit=3, loop=self.loop)
        stream.feed_data(b'li')
        stream.feed_data(b'ne1\nline2\n')

        self.assertRaises(
            ValueError, self.loop.run_until_complete, stream.readline())
        # The buffer should contain the remaining data after exception
        self.assertEqual(b'line2\n', stream._buffer)

        stream = asyncio.StreamReader(limit=3, loop=self.loop)
        stream.feed_data(b'li')
        stream.feed_data(b'ne1')
        stream.feed_data(b'li')

        self.assertRaises(
            ValueError, self.loop.run_until_complete, stream.readline())
        # No b'\n' at the end. The 'limit' is set to 3. So before
        # waiting for the new data in buffer, 'readline' will consume
        # the entire buffer, and since the length of the consumed data
        # is more than 3, it will raise a ValueError. The buffer is
        # expected to be empty now.
        self.assertEqual(b'', stream._buffer) 

def test_read_exception(self):
        stream = asyncio.StreamReader(loop=self.loop)
        stream.feed_data(b'line\n')

        data = self.loop.run_until_complete(stream.read(2))
        self.assertEqual(b'li', data)

        stream.set_exception(ValueError())
        self.assertRaises(
            ValueError, self.loop.run_until_complete, stream.read(2)) 

def test_read_until_eof(self):
        # Read all bytes until eof.
        stream = asyncio.StreamReader(loop=self.loop)
        read_task = asyncio.Task(stream.read(-1), loop=self.loop)

        def cb():
            stream.feed_data(b'chunk1\n')
            stream.feed_data(b'chunk2')
            stream.feed_eof()
        self.loop.call_soon(cb)

        data = self.loop.run_until_complete(read_task)

        self.assertEqual(b'chunk1\nchunk2', data)
        self.assertEqual(b'', stream._buffer)