From 4115d033a8844d52fa9070602aea9637d3a0f783 Mon Sep 17 00:00:00 2001 From: acul71 Date: Wed, 16 Jul 2025 20:20:35 +0200 Subject: [PATCH] feat: identify identify/push raw-format fix and tests --- .../identity/identify_push/identify_push.py | 22 +++- newsfragments/761.internal.rst | 1 + .../identify_push/test_identify_push.py | 101 ++++++++++++++++++ 3 files changed, 119 insertions(+), 5 deletions(-) create mode 100644 newsfragments/761.internal.rst diff --git a/libp2p/identity/identify_push/identify_push.py b/libp2p/identity/identify_push/identify_push.py index 5d6cf163..f13bd970 100644 --- a/libp2p/identity/identify_push/identify_push.py +++ b/libp2p/identity/identify_push/identify_push.py @@ -91,12 +91,24 @@ def identify_push_handler_for( return else: # Read raw protobuf message from the stream + # For raw format, we need to read all data before the stream is closed data = b"" - while True: - chunk = await stream.read(4096) - if not chunk: - break - data += chunk + try: + # Read all available data in a single operation + data = await stream.read() + except StreamClosed: + # Try to read any remaining data + try: + data = await stream.read() + except Exception: + pass + + # If we got no data, log a warning and return + if not data: + logger.warning( + "No data received in raw format from peer %s", peer_id + ) + return identify_msg = Identify() identify_msg.ParseFromString(data) diff --git a/newsfragments/761.internal.rst b/newsfragments/761.internal.rst new file mode 100644 index 00000000..59496ebc --- /dev/null +++ b/newsfragments/761.internal.rst @@ -0,0 +1 @@ +Fix raw format reading in identify/push protocol and add comprehensive test coverage for both varint and raw formats diff --git a/tests/core/identity/identify_push/test_identify_push.py b/tests/core/identity/identify_push/test_identify_push.py index e62bad7a..a1e2e472 100644 --- a/tests/core/identity/identify_push/test_identify_push.py +++ b/tests/core/identity/identify_push/test_identify_push.py @@ -597,3 +597,104 @@ async def test_all_peers_receive_identify_push_with_semaphore_under_high_peer_lo assert peer_id_a in dummy_peerstore.peer_ids() nursery.cancel_scope.cancel() + + +@pytest.mark.trio +async def test_identify_push_default_varint_format(security_protocol): + """ + Test that the identify/push protocol uses varint format by default. + + This test verifies that: + 1. The default behavior uses length-prefixed messages (varint format) + 2. Messages are correctly encoded with varint length prefix + 3. Messages are correctly decoded with varint length prefix + 4. The peerstore is updated correctly with the received information + """ + async with host_pair_factory(security_protocol=security_protocol) as ( + host_a, + host_b, + ): + # Set up the identify/push handlers with default settings + # (use_varint_format=True) + host_b.set_stream_handler(ID_PUSH, identify_push_handler_for(host_b)) + + # Push identify information from host_a to host_b using default settings + success = await push_identify_to_peer(host_a, host_b.get_id()) + assert success, "Identify push should succeed with default varint format" + + # Wait a bit for the push to complete + await trio.sleep(0.1) + + # Get the peerstore from host_b + peerstore = host_b.get_peerstore() + peer_id = host_a.get_id() + + # Verify that the peerstore was updated correctly + assert peer_id in peerstore.peer_ids() + + # Check that addresses have been updated + host_a_addrs = set(host_a.get_addrs()) + peerstore_addrs = set(peerstore.addrs(peer_id)) + assert all(addr in peerstore_addrs for addr in host_a_addrs) + + # Check that protocols have been updated + host_a_protocols = set(host_a.get_mux().get_protocols()) + peerstore_protocols = set(peerstore.get_protocols(peer_id)) + assert all(protocol in peerstore_protocols for protocol in host_a_protocols) + + # Check that the public key has been updated + host_a_public_key = host_a.get_public_key().serialize() + peerstore_public_key = peerstore.pubkey(peer_id).serialize() + assert host_a_public_key == peerstore_public_key + + +@pytest.mark.trio +async def test_identify_push_legacy_raw_format(security_protocol): + """ + Test that the identify/push protocol can use legacy raw format when specified. + + This test verifies that: + 1. When use_varint_format=False, messages are sent without length prefix + 2. Raw protobuf messages are correctly encoded and decoded + 3. The peerstore is updated correctly with the received information + 4. The legacy format is backward compatible + """ + async with host_pair_factory(security_protocol=security_protocol) as ( + host_a, + host_b, + ): + # Set up the identify/push handlers with legacy format (use_varint_format=False) + host_b.set_stream_handler( + ID_PUSH, identify_push_handler_for(host_b, use_varint_format=False) + ) + + # Push identify information from host_a to host_b using legacy format + success = await push_identify_to_peer( + host_a, host_b.get_id(), use_varint_format=False + ) + assert success, "Identify push should succeed with legacy raw format" + + # Wait a bit for the push to complete + await trio.sleep(0.1) + + # Get the peerstore from host_b + peerstore = host_b.get_peerstore() + peer_id = host_a.get_id() + + # Verify that the peerstore was updated correctly + assert peer_id in peerstore.peer_ids() + + # Check that addresses have been updated + host_a_addrs = set(host_a.get_addrs()) + peerstore_addrs = set(peerstore.addrs(peer_id)) + assert all(addr in peerstore_addrs for addr in host_a_addrs) + + # Check that protocols have been updated + host_a_protocols = set(host_a.get_mux().get_protocols()) + peerstore_protocols = set(peerstore.get_protocols(peer_id)) + assert all(protocol in peerstore_protocols for protocol in host_a_protocols) + + # Check that the public key has been updated + host_a_public_key = host_a.get_public_key().serialize() + peerstore_public_key = peerstore.pubkey(peer_id).serialize() + assert host_a_public_key == peerstore_public_key