async def test_receive_exactly_incomplete(): send_stream, receive_stream = create_memory_object_stream(1) buffered_stream = BufferedByteReceiveStream(receive_stream) await send_stream.send(b'abcd') await send_stream.aclose() with pytest.raises(IncompleteRead): await buffered_stream.receive_exactly(8)
async def handle_tcpros(self, protocol: str, client: SocketStream) -> None: async with client: buffered_receiver = BufferedByteReceiveStream(client) try: header = await read_header(buffered_receiver) try: if "service" in header: handler_getter = self._registry.get_service_handler key = "service" elif "topic" in header: handler_getter = self._registry.get_publication_handler key = "topic" else: await client.send( encode_header( dict(error="no topic or service name detected") ) ) return try: handler = handler_getter(header[key]) except KeyError as exc: raise ProtocolError( f"{key} {header[key]} is not provided by this node" ) from exc else: await handler(protocol, header, client) except ProtocolError as err: await client.send(encode_header(dict(error=err.args[0]))) except anyio.IncompleteRead: # Raised if the connection is closed before the requested amount of # bytes has been read. pass
async def test_receive_exactly(): send_stream, receive_stream = create_memory_object_stream(2) buffered_stream = BufferedByteReceiveStream(receive_stream) await send_stream.send(b'abcd') await send_stream.send(b'efgh') result = await buffered_stream.receive_exactly(8) assert result == b'abcdefgh' assert isinstance(result, bytes)
async def test_receive_exactly() -> None: send_stream, receive_stream = create_memory_object_stream(2) buffered_stream = BufferedByteReceiveStream(receive_stream) await send_stream.send(b"abcd") await send_stream.send(b"efgh") result = await buffered_stream.receive_exactly(8) assert result == b"abcdefgh" assert isinstance(result, bytes)
async def test_receive_until_incomplete(): send_stream, receive_stream = create_memory_object_stream(1) buffered_stream = BufferedByteReceiveStream(receive_stream) await send_stream.send(b'abcd') await send_stream.aclose() with pytest.raises(IncompleteRead): assert await buffered_stream.receive_until(b'de', 10) assert buffered_stream.buffer == b'abcd'
async def _get_client_stream(self) -> SocketStream: tcpros_uri = await self._master.lookup_service(self._service_name) client_stream = await anyio.connect_tcp(*split_tcpros_uri(tcpros_uri)) buffered_receiver = BufferedByteReceiveStream(client_stream) await client_stream.send(encode_header(self._header)) header_dict = await read_header(buffered_receiver) if "error" in header_dict: raise ServiceClientInitError(header_dict["error"]) from None return client_stream
async def test_receive_until(): send_stream, receive_stream = create_memory_object_stream(2) buffered_stream = BufferedByteReceiveStream(receive_stream) await send_stream.send(b'abcd') await send_stream.send(b'efgh') result = await buffered_stream.receive_until(b'de', 10) assert result == b'abc' assert isinstance(result, bytes) result = await buffered_stream.receive_until(b'h', 10) assert result == b'fg' assert isinstance(result, bytes)
async def _handle(self, client): buffered_client = BufferedByteReceiveStream(client) # Support multiple queries while True: self.logger.debug("Reading message from client") # TODO: Issue with infinite loops here with kdig? https://github.com/agronholm/anyio/issues/162 # First 2 bytes of protocol are how much to read len_to_read = int.from_bytes( await buffered_client.receive_exactly(2), "big") query = await buffered_client.receive_exactly(len_to_read) response = await handle_dns_query(query, self.logger) len_to_write = len(response).to_bytes(2, byteorder="big") await client.send(len_to_write + response)
async def _call(self, client_stream: SocketStream, request: abc.ServiceRequestT) -> abc.ServiceResponseT: buffered_receiver = BufferedByteReceiveStream(client_stream) data = await anyio.to_thread.run_sync(self._serializer.serialize, request) await client_stream.send(data) if not await read_byte(buffered_receiver): error = await read_error(buffered_receiver) raise ServiceClientError(error) from None data = await read_data(buffered_receiver) msg = self.response_class() await anyio.to_thread.run_sync(msg.deserialize, data) return msg
async def _subscribe_task(self, publisher: ConnectedPublisher) -> None: with publisher.cancel_scope: protocol, params = await self._get_publisher_connection_params( publisher.xmlrpc_uri, ["UDSROS", "TCPROS"]) client: Optional[SocketStream] = None if protocol == "UDSROS": if params[0] == get_local_address(): logger.debug("Connecting via %s", params[1]) client = await anyio.connect_unix(params[1]) else: # non-local connection => switch to TCPROS protocol, params = await self._get_publisher_connection_params( publisher.xmlrpc_uri, ["TCPROS"]) if client is None: logger.debug("Connecting via %s:%s", params[0], params[1]) client = await anyio.connect_tcp(params[0], int(params[1])) publisher.protocol = protocol publisher.connection_parameter = params try: async with client: publisher.connected = True async with self.condition: self.condition.notify_all() await client.send(encode_header(self.header)) buffered_client = BufferedByteReceiveStream(client) header_dict = await read_header(buffered_client) if "error" in header_dict: # TODO proper exception # raise SubscriberInitError(header_dict["error"]) raise RuntimeError() while not publisher.cancel_scope.cancel_called: msg = self.topic_type() msg.deserialize(await read_data(buffered_client)) for stream in self.subscriptions: stream.send_nowait(msg) except anyio.IncompleteRead: pass finally: self.connected_publishers.pop(publisher.xmlrpc_uri, None)
async def test_terminate(tmp_path): script_path = tmp_path / 'script.py' script_path.write_text(dedent("""\ import signal, sys, time def terminate(signum, frame): sys.exit(2) signal.signal(signal.SIGTERM, terminate) print('ready', flush=True) time.sleep(5) """)) async with await open_process([sys.executable, str(script_path)]) as process: buffered_stdout = BufferedByteReceiveStream(process.stdout) line = await buffered_stdout.receive_until(b'\n', 100) assert line.rstrip() == b'ready' process.terminate() assert await process.wait() == 2
async def handle_tcpros( self, protocol: str, header: abc.Header, client: SocketStream, ) -> None: """Handle incoming service client traffic.""" require_fields(header, "service", "md5sum", "callerid") check_md5sum(header, getattr(self.service_type, "_md5sum")) await client.send(encode_header(self.header)) if header.get("probe") == "1": return persistent = header.get("persistent", "").lower() in ("1", "true") serializer: Serializer[abc.ServiceRequestT] = Serializer() buffered_receiver = BufferedByteReceiveStream(client) while True: request = self.request_class() data = await read_data(buffered_receiver) await anyio.to_thread.run_sync(request.deserialize, data) setattr(request, "_connection_header", header) try: result = await self._handler(request) except anyio.get_cancelled_exc_class(): raise except Exception: # pylint: disable=broad-except logger.error("Exception occured in service callback", exc_info=True) await client.send( encode_byte(0) + encode_str("error processing request")) else: data = await anyio.to_thread.run_sync(serializer.serialize, result) await client.send(encode_byte(1) + data) if not persistent: return
async def __aenter__(self): self.connection_closed = anyio.create_event() self.state = CONNECTING self.version_major = None self.version_minor = None self.server_properties = None self.server_mechanisms = None self.server_locales = None self.server_heartbeat = None self.channels = {} self.server_frame_max = None self.server_channel_max = None self.channels_ids_ceil = 0 self.channels_ids_free = set() self._send_queue_w, self._send_queue_r = anyio.create_memory_object_stream( 1) if self._ssl: if self._ssl is True: ssl_context = ssl.create_default_context() else: ssl_context = self._ssl port = self._port if port is None: if self._ssl: port = 5671 else: port = 5672 if self._ssl: stream = await anyio.connect_tcp(self._host, port, ssl_context=ssl_context, autostart_tls=True) else: stream = await anyio.connect_tcp(self._host, port) stream.extra(SocketAttribute.raw_socket).setsockopt( socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) from anyio.streams.buffered import BufferedByteReceiveStream self._stream = stream self._rstream = BufferedByteReceiveStream(stream) # the writer loop needs to run since the beginning done_here = anyio.create_event() await self._nursery.spawn(self._writer_loop, done_here) await done_here.wait() try: await self._stream.send(amqp_constants.PROTOCOL_HEADER) # Wait 'start' method from the server await self.dispatch_frame() if self.version_major is None: raise RuntimeError("Server didn't start with a START packet") client_properties = { 'capabilities': { 'consumer_cancel_notify': True, 'connection.blocked': False, }, } client_properties.update(self.client_properties) # waiting reply start with credentions and co await self.start_ok(client_properties, 'AMQPLAIN', self._auth, self.server_locales) # wait for a "tune" reponse await self.dispatch_frame() if self.server_channel_max is None: raise RuntimeError("Server didn't send a TUNE packet") tune_ok = { 'channel_max': self.connection_tunning.get('channel_max', self.server_channel_max), 'frame_max': self.connection_tunning.get('frame_max', self.server_frame_max), 'heartbeat': self.connection_tunning.get('heartbeat', self.server_heartbeat), } # "tune" the connexion with max channel, max frame, heartbeat await self.tune_ok(**tune_ok) # update connection tunning values self.server_frame_max = tune_ok['frame_max'] self.server_channel_max = tune_ok['channel_max'] self.server_heartbeat = tune_ok['heartbeat'] # open a virtualhost await self.open(self._virtualhost, capabilities='', insist=self._insist) # wait for open-ok await self.dispatch_frame() if self.state != OPEN: if self._close_reason is not None: raise exceptions.AmqpClosedConnection( self._close_reason['text']) else: raise exceptions.AmqpClosedConnection() # read the other server's responses asynchronously done_here = anyio.create_event() await self._nursery.spawn(self._reader_loop, done_here) await done_here.wait() except BaseException as exc: async with anyio.fail_after(2, shield=True): await self.close(no_wait=True) raise return self