todo: handled ls command in multiselect.py (#622)

This commit is contained in:
Abhinav Agarwalla
2025-05-28 07:48:37 +05:30
committed by GitHub
parent 5496b2709a
commit 67ca1d7769
5 changed files with 95 additions and 2 deletions

View File

@ -195,6 +195,29 @@ class BasicHost(IHost):
net_stream.set_protocol(selected_protocol)
return net_stream
async def send_command(self, peer_id: ID, command: str) -> list[str]:
"""
Send a multistream-select command to the specified peer and return
the response.
:param peer_id: peer_id that host is connecting
:param command: supported multistream-select command (e.g., "ls)
:raise StreamFailure: If the stream cannot be opened or negotiation fails
:return: list of strings representing the response from peer.
"""
new_stream = await self._network.new_stream(peer_id)
try:
response = await self.multiselect_client.query_multistream_command(
MultiselectCommunicator(new_stream), command
)
except MultiselectClientError as error:
logger.debug("fail to open a stream to peer %s, error=%s", peer_id, error)
await new_stream.reset()
raise StreamFailure(f"failed to open a stream to peer {peer_id}") from error
return response
async def connect(self, peer_info: PeerInfo) -> None:
"""
Ensure there is a connection between this host and the peer

View File

@ -60,8 +60,14 @@ class Multiselect(IMultiselectMuxer):
raise MultiselectError() from error
if command == "ls":
# TODO: handle ls command
pass
supported_protocols = list(self.handlers.keys())
response = "\n".join(supported_protocols) + "\n"
try:
await communicator.write(response)
except MultiselectCommunicatorError as error:
raise MultiselectError() from error
else:
protocol = TProtocol(command)
if protocol in self.handlers:

View File

@ -70,6 +70,36 @@ class MultiselectClient(IMultiselectClient):
raise MultiselectClientError("protocols not supported")
async def query_multistream_command(
self, communicator: IMultiselectCommunicator, command: str
) -> list[str]:
"""
Send a multistream-select command over the given communicator and return
parsed response.
:param communicator: communicator to use to communicate with counterparty
:param command: supported multistream-select command(e.g., ls)
:raise MultiselectClientError: If the communicator fails to process data.
:return: list of strings representing the response from peer.
"""
await self.handshake(communicator)
if command == "ls":
try:
await communicator.write("ls")
except MultiselectCommunicatorError as error:
raise MultiselectClientError() from error
else:
raise ValueError("Command not supported")
try:
response = await communicator.read()
response_list = response.strip().splitlines()
except MultiselectCommunicatorError as error:
raise MultiselectClientError() from error
return response_list
async def try_select(
self, communicator: IMultiselectCommunicator, protocol: TProtocol
) -> TProtocol:

View File

@ -0,0 +1,2 @@
Feature: Support for sending `ls` command over `multistream-select` to list supported protocols from remote peer.
This allows inspecting which protocol handlers a peer supports at runtime.

View File

@ -116,3 +116,35 @@ async def test_multiple_protocol_fails(security_protocol):
await perform_simple_test(
"", protocols_for_client, protocols_for_listener, security_protocol
)
@pytest.mark.trio
async def test_multistream_command(security_protocol):
supported_protocols = [PROTOCOL_ECHO, PROTOCOL_FOO, PROTOCOL_POTATO, PROTOCOL_ROCK]
async with HostFactory.create_batch_and_listen(
2, security_protocol=security_protocol
) as hosts:
listener, dialer = hosts[1], hosts[0]
for protocol in supported_protocols:
listener.set_stream_handler(
protocol, create_echo_stream_handler(ACK_PREFIX)
)
# Ensure dialer knows how to reach the listener
dialer.get_peerstore().add_addrs(listener.get_id(), listener.get_addrs(), 10)
# Dialer asks peer to list the supported protocols using `ls`
response = await dialer.send_command(listener.get_id(), "ls")
# We expect all supported protocols to show up
for protocol in supported_protocols:
assert protocol in response
assert "/does/not/exist" not in response
assert "/foo/bar/1.2.3" not in response
# Dialer asks for unspoorted command
with pytest.raises(ValueError, match="Command not supported"):
await dialer.send_command(listener.get_id(), "random")