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