feat: identify identify/push raw-format fix and tests

Author: acul71

libp2p/identity/identify_push/identify_push.py

@@ -91,12 +91,24 @@ async def handle_identify_push(stream: INetStream) -> None:
                     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)

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

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