Python源码示例:asyncio.StreamReader()

示例1
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 
示例2
def __init__(
        self,
        app: ASGIFramework,
        loop: asyncio.AbstractEventLoop,
        config: Config,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        self.app = app
        self.config = config
        self.loop = loop
        self.protocol: ProtocolWrapper
        self.reader = reader
        self.writer = writer
        self.send_lock = asyncio.Lock()
        self.timeout_lock = asyncio.Lock()

        self._keep_alive_timeout_handle: Optional[asyncio.Task] = None 
示例3
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()) 
示例4
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) 
示例5
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) 
示例6
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) 
示例7
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) 
示例8
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 
示例9
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() 
示例10
def open_serial_connection(*,
                           loop=None,
                           limit=asyncio.streams._DEFAULT_LIMIT,
                           **kwargs):
    """A wrapper for create_serial_connection() returning a (reader,
    writer) pair.

    The reader returned is a StreamReader instance; the writer is a
    StreamWriter instance.

    The arguments are all the usual arguments to Serial(). Additional
    optional keyword arguments are loop (to set the event loop instance
    to use) and limit (to set the buffer limit passed to the
    StreamReader.

    This function is a coroutine.
    """
    if loop is None:
        loop = asyncio.get_event_loop()
    reader = asyncio.StreamReader(limit=limit, loop=loop)
    protocol = asyncio.StreamReaderProtocol(reader, loop=loop)
    transport, _ = yield from create_serial_connection(
        loop=loop,
        protocol_factory=lambda: protocol,
        **kwargs)
    writer = asyncio.StreamWriter(transport, protocol, reader, loop)
    return reader, writer


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# test 
示例11
def open_connection(host=None, port=None, *,
                          limit, loop=None,
                          parser=None, **kwds):
    # XXX: parser is not used (yet)
    if loop is not None and sys.version_info >= (3, 8):
        warnings.warn("The loop argument is deprecated",
                      DeprecationWarning)
    reader = StreamReader(limit=limit)
    protocol = asyncio.StreamReaderProtocol(reader)
    transport, _ = await get_event_loop().create_connection(
        lambda: protocol, host, port, **kwds)
    writer = asyncio.StreamWriter(transport, protocol, reader,
                                  loop=get_event_loop())
    return reader, writer 
示例12
def open_unix_connection(address, *,
                               limit, loop=None,
                               parser=None, **kwds):
    # XXX: parser is not used (yet)
    if loop is not None and sys.version_info >= (3, 8):
        warnings.warn("The loop argument is deprecated",
                      DeprecationWarning)
    reader = StreamReader(limit=limit)
    protocol = asyncio.StreamReaderProtocol(reader)
    transport, _ = await get_event_loop().create_unix_connection(
        lambda: protocol, address, **kwds)
    writer = asyncio.StreamWriter(transport, protocol, reader,
                                  loop=get_event_loop())
    return reader, writer 
示例13
def _handshake(
    initiator: "HandshakeInitiator",
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
    token: CancelToken,
) -> Tuple[bytes, bytes, BasePreImage, BasePreImage]:
    """See the handshake() function above.

    This code was factored out into this helper so that we can create Peers with directly
    connected readers/writers for our tests.
    """
    initiator_nonce = keccak(os.urandom(HASH_LEN))
    auth_msg = initiator.create_auth_message(initiator_nonce)
    auth_init = initiator.encrypt_auth_message(auth_msg)
    writer.write(auth_init)

    auth_ack = await token.cancellable_wait(
        reader.read(ENCRYPTED_AUTH_ACK_LEN), timeout=REPLY_TIMEOUT
    )

    if reader.at_eof():
        # This is what happens when Parity nodes have blacklisted us
        # (https://github.com/ethereum/py-evm/issues/901).
        raise HandshakeDisconnectedFailure(
            "%s disconnected before sending auth ack", repr(initiator.remote)
        )

    ephemeral_pubkey, responder_nonce = initiator.decode_auth_ack_message(auth_ack)
    aes_secret, mac_secret, egress_mac, ingress_mac = initiator.derive_secrets(
        initiator_nonce, responder_nonce, ephemeral_pubkey, auth_init, auth_ack
    )

    return aes_secret, mac_secret, egress_mac, ingress_mac 
示例14
def connect(self) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter]:
        return await self.cancel_token.cancellable_wait(
            asyncio.open_connection(
                host=self.remote.address.ip, port=self.remote.address.tcp_port
            ),
            timeout=REPLY_TIMEOUT,
        ) 
示例15
def receive_handshake(
        self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
    ) -> None:
        ip, socket, *_ = writer.get_extra_info("peername")
        remote_address = Address(ip, socket)
        if self.peer_pool.chk_dialin_blacklist(remote_address):
            Logger.info_every_n(
                "{} has been blacklisted, refusing connection".format(remote_address),
                100,
            )
            reader.feed_eof()
            writer.close()
        expected_exceptions = (
            TimeoutError,
            PeerConnectionLost,
            HandshakeFailure,
            asyncio.IncompleteReadError,
            HandshakeDisconnectedFailure,
        )
        try:
            await self._receive_handshake(reader, writer)
        except expected_exceptions as e:
            self.logger.debug("Could not complete handshake: %s", e)
            Logger.error_every_n("Could not complete handshake: {}".format(e), 100)
            reader.feed_eof()
            writer.close()
        except OperationCancelled:
            self.logger.error("OperationCancelled")
            reader.feed_eof()
            writer.close()
        except Exception as e:
            self.logger.exception("Unexpected error handling handshake")
            reader.feed_eof()
            writer.close() 
示例16
def get_directly_connected_streams() -> TConnectedStreams:
    bob_reader = asyncio.StreamReader()
    alice_reader = asyncio.StreamReader()
    # Link the alice's writer to the bob's reader, and the bob's writer to the
    # alice's reader.
    bob_writer = MockStreamWriter(alice_reader.feed_data)
    alice_writer = MockStreamWriter(bob_reader.feed_data)
    return (
        (alice_reader, cast(asyncio.StreamWriter, alice_writer)),
        (bob_reader, cast(asyncio.StreamWriter, bob_writer)),
    ) 
示例17
def __init__(self, num_tokens_per_source, input_queue_maxsize, write, read):
        super().__init__(num_tokens_per_source)
        self._input_queue = asyncio.Queue(input_queue_maxsize)
        self._write = write

        loop = asyncio.get_event_loop()
        self._stream_reader = asyncio.StreamReader()
        def protocol_factory():
            return asyncio.StreamReaderProtocol(self._stream_reader)
        pipe = os.fdopen(read, mode='r')
        self._transport, _ = loop.run_until_complete(
            loop.connect_read_pipe(protocol_factory, pipe)) 
示例18
def get_producer_wrapper(self):
        async def receiver():
            stream_reader = asyncio.StreamReader()
            def protocol_factory():
                return asyncio.StreamReaderProtocol(stream_reader)
            transport = await asyncio.get_event_loop().connect_read_pipe(
                protocol_factory, os.fdopen(self._read, mode='r'))

            # TODO do we need to close transport when filter gets garbage
            # collected?

            while True:
                size_bytes = await stream_reader.readexactly(
                    _NUM_BYTES_FOR_SIZE)
                size_of_message = int.from_bytes(size_bytes, _BYTEORDER)

                input_frame = gabriel_pb2.InputFrame()
                input_frame.ParseFromString(
                    await stream_reader.readexactly(size_of_message))
                self._latest_input_frame = input_frame
                self._frame_available.set()

        async def producer():
            if not self._started_receiver:
                self._started_receiver = True
                assert os.getpid() == self._constructor_pid
                asyncio.ensure_future(receiver())

            await self._frame_available.wait()

            # Clear because we are sending self._latest_input_frame
            self._frame_available.clear()

            return self._latest_input_frame

        return ProducerWrapper(producer=producer, source_name=self._source_name) 
示例19
def connect(self):
        loop = self._loop
        self._closed_fut = create_future(loop=loop)
        if self._security_protocol in ["PLAINTEXT", "SASL_PLAINTEXT"]:
            ssl = None
        else:
            assert self._security_protocol in ["SSL", "SASL_SSL"]
            assert self._ssl_context is not None
            ssl = self._ssl_context
        # Create streams same as `open_connection`, but using custom protocol
        reader = asyncio.StreamReader(limit=READER_LIMIT, loop=loop)
        protocol = AIOKafkaProtocol(self._closed_fut, reader, loop=loop)
        transport, _ = await asyncio.wait_for(
            loop.create_connection(
                lambda: protocol, self.host, self.port, ssl=ssl),
            loop=loop, timeout=self._request_timeout)
        writer = asyncio.StreamWriter(transport, protocol, reader, loop)
        self._reader, self._writer, self._protocol = reader, writer, protocol

        # Start reader task.
        self._read_task = self._create_reader_task()

        # Start idle checker
        if self._max_idle_ms is not None:
            self._idle_handle = self._loop.call_soon(
                self._idle_check, weakref.ref(self))

        if self._version_hint and self._version_hint >= (0, 10):
            await self._do_version_lookup()

        if self._security_protocol in ["SASL_SSL", "SASL_PLAINTEXT"]:
            await self._do_sasl_handshake()

        return reader, writer 
示例20
def authenticate(
            self,
            reader: asyncio.StreamReader,
            writer: asyncio.StreamWriter,
    ) -> bool:
        """(async) Authenticate a connecting client.

        Returns:
            True if authentication is successful and False otherwise. The
            caller is responsible for closing the connection in case of
            failure.
        """
        writer.write(enums.ExtOrPortAuthTypes.SAFE_COOKIE
                     + enums.ExtOrPortAuthTypes.END_AUTH_TYPES)
        client_auth_type = await reader.readexactly(1)
        if client_auth_type != enums.ExtOrPortAuthTypes.SAFE_COOKIE:
            return False
        client_nonce = await reader.readexactly(self.nonce_len)
        server_nonce = secrets.token_bytes(self.nonce_len)
        server_hash = self.hash(b''.join((
            self.server_hash_header, client_nonce, server_nonce)))
        writer.write(server_hash + server_nonce)
        client_hash = await reader.readexactly(self.hash_len)
        result = hmac.compare_digest(client_hash, self.hash(b''.join((
            self.client_hash_header, client_nonce, server_nonce))))
        writer.write(int(result).to_bytes(1, 'big'))
        return result 
示例21
def _read_ext_msg(
            reader: asyncio.StreamReader,
    ) -> Tuple[bytes, bytes]:
        command = await reader.readexactly(2)
        body_len = int.from_bytes(await reader.readexactly(2), 'big')
        body = await reader.readexactly(body_len)
        return command, body 
示例22
def handle_client_connection(
        adapter: adapters.ClientAdapter,
        transport: str,
        upstream_host: str,
        upstream_port: int,
        args: Dict[str, str],
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
) -> None:
    handler_logger.debug(
        'Accepted connection for transport %s from %r on %r',
        transport,
        writer.get_extra_info('peername'), writer.get_extra_info('sockname'))
    async with contexts.log_unhandled_exc(handler_logger), \
               contexts.aclosing_multiple_writers(writer) as writers:
        try:
            ureader, uwriter = await adapter.open_transport_connection(
                transport, upstream_host, upstream_port, args)
        except exceptions.PTConnectError as e:
            handler_logger.warning(
                'PT reported error while connecting to upstream '
                '(%r, %r): %r', upstream_host, upstream_port, e)
            writer.transport.abort()
            return
        writers.add(uwriter)
        logname = (f'{writer.get_extra_info("peername")!r} ==> '
                   f'({upstream_host!r}, {upstream_port})')
        handler_logger.info('[%s] %s', transport, logname)
        try:
            await relays.relay(reader, writer, ureader, uwriter)
        except OSError as e:
            handler_logger.warning(
                '[%s] %s caught %r', transport, logname, e) 
示例23
def negotiate_socks4_userid(
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
        host: Union[str, ipaddress.IPv4Address, ipaddress.IPv6Address],
        port: int,
        args: Optional[Dict[str, str]],
) -> None:
    try:
        host = ipaddress.IPv4Address(host)
    except ValueError:
        raise ValueError('SOCKS4 only supports IPv4 address')
    if args:
        args_bytes = encode_args(args)
    else:
        args_bytes = b''
    writer.write(b''.join((
        b'\x04',  # ver
        enums.SOCKS4Command.CONNECT,
        port.to_bytes(2, 'big'),
        host.packed,
        args_bytes,
        b'\0',
    )))
    buf = await reader.readexactly(8)
    assert buf[0] == 0, 'Invalid SOCKS4 reply version'
    reply = enums.SOCKS4Reply(buf[1:2])
    if reply is not enums.SOCKS4Reply.GRANTED:
        raise exceptions.PTSOCKS4ConnectError(reply) 
示例24
def _relay_data_side(
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
) -> None:
    """Pass data and EOF from reader to writer."""
    while True:
        buf = await reader.read(BUF_SIZE)
        if not buf:  # EOF
            break
        writer.write(buf)
        await writer.drain()
    writer.write_eof()
    await writer.drain() 
示例25
def relay(
        dreader: asyncio.StreamReader,
        dwriter: asyncio.StreamWriter,
        ureader: asyncio.StreamReader,
        uwriter: asyncio.StreamWriter,
) -> None:
    """Pass data/EOF from dreader to uwriter, and ureader to dwriter.

    Both writers are ensured to be closed upon exiting this function.
    """
    _logger.debug(
        'Relaying %r <=> %r', dwriter.get_extra_info('peername'),
        uwriter.get_extra_info('peername'))
    utask = asyncio.create_task(_relay_data_side(dreader, uwriter))
    dtask = asyncio.create_task(_relay_data_side(ureader, dwriter))
    async with contexts.aclosing_multiple_writers(dwriter, uwriter):
        try:
            await asyncio.gather(utask, dtask)
            _logger.debug(
                'Relay %r <=> %r ended normally',
                dwriter.get_extra_info('peername'),
                uwriter.get_extra_info('peername'))
        except:
            dtask.cancel()
            utask.cancel()
            raise
        finally:
            await asyncio.wait({dtask, utask})
            for t in (dtask, utask):
                if t.exception():
                    _logger.debug(
                        'Relay task %r caught exception %r', t, t.exception()) 
示例26
def get_reader(self):
        if self._reader is None:
            self._reader = asyncio.StreamReader()
            protocol = asyncio.StreamReaderProtocol(self._reader)
            loop = asyncio.get_event_loop()
            await loop.connect_read_pipe(lambda: protocol, sys.stdin)
        return self._reader 
示例27
def __init__(self, proxy, proxy_auth, dst, app_protocol_factory, waiter, *,
                 remote_resolve=True, loop=None, ssl=False,
                 server_hostname=None, negotiate_done_cb=None,
                 reader_limit=DEFAULT_LIMIT):
        if not isinstance(dst, (tuple, list)) or len(dst) != 2:
            raise ValueError(
                'Invalid dst format, tuple("dst_host", dst_port))'
            )

        self._proxy = proxy
        self._auth = proxy_auth
        self._dst_host, self._dst_port = dst
        self._remote_resolve = remote_resolve
        self._waiter = waiter
        self._ssl = ssl
        self._server_hostname = server_hostname
        self._negotiate_done_cb = negotiate_done_cb
        self._loop = loop or asyncio.get_event_loop()

        self._transport = None
        self._negotiate_done = False
        self._proxy_peername = None
        self._proxy_sockname = None

        if app_protocol_factory:
            self._app_protocol = app_protocol_factory()
        else:
            self._app_protocol = self

        reader = asyncio.StreamReader(loop=self._loop, limit=reader_limit)

        super().__init__(stream_reader=reader,
                         client_connected_cb=self.negotiate, loop=self._loop) 
示例28
def get_pipe_reader(self, fd_reader):
        reader = asyncio.StreamReader()
        protocol = asyncio.StreamReaderProtocol(reader)
        try:
            await self.loop.connect_read_pipe(lambda: protocol, fd_reader)
        except:
            return None
        return reader 
示例29
def get_pipe_reader(self, fd_reader):
        reader = asyncio.StreamReader()
        protocol = asyncio.StreamReaderProtocol(reader)
        try:
            await self.loop.connect_read_pipe(lambda: protocol, fd_reader)
        except:
            return None
        return reader 
示例30
def __init__(self) -> None:
        self.host_ = '127.0.0.1'
        """ The local EPMD is always located on the local host. """

        self.port_ = EPMD_DEFAULT_PORT
        """ We expect local EPMD to be available on the default port. """

        self.reader_ = None  # type: [asyncio.StreamReader, None]
        self.writer_ = None  # type: [asyncio.StreamWriter, None]
        self.n_connection_attempts_ = 5