"""Integration tests for DCUtR protocol with real libp2p hosts using circuit relay.""" import logging from unittest.mock import AsyncMock, MagicMock import pytest from multiaddr import Multiaddr import trio from libp2p.relay.circuit_v2.dcutr import ( MAX_HOLE_PUNCH_ATTEMPTS, PROTOCOL_ID, DCUtRProtocol, ) from libp2p.relay.circuit_v2.pb.dcutr_pb2 import ( HolePunch, ) from libp2p.relay.circuit_v2.protocol import ( DEFAULT_RELAY_LIMITS, CircuitV2Protocol, ) from libp2p.tools.async_service import ( background_trio_service, ) from tests.utils.factories import ( HostFactory, ) logger = logging.getLogger(__name__) # Test timeouts SLEEP_TIME = 0.5 # seconds @pytest.mark.trio async def test_dcutr_through_relay_connection(): """ Test DCUtR protocol where peers are connected via relay, then upgrade to direct. """ # Create three hosts: two peers and one relay async with HostFactory.create_batch_and_listen(3) as hosts: peer1, peer2, relay = hosts # Create circuit relay protocol for the relay relay_protocol = CircuitV2Protocol(relay, DEFAULT_RELAY_LIMITS, allow_hop=True) # Create DCUtR protocols for both peers dcutr1 = DCUtRProtocol(peer1) dcutr2 = DCUtRProtocol(peer2) # Track if DCUtR stream handlers were called handler1_called = False handler2_called = False # Override stream handlers to track calls original_handler1 = dcutr1._handle_dcutr_stream original_handler2 = dcutr2._handle_dcutr_stream async def tracked_handler1(stream): nonlocal handler1_called handler1_called = True await original_handler1(stream) async def tracked_handler2(stream): nonlocal handler2_called handler2_called = True await original_handler2(stream) dcutr1._handle_dcutr_stream = tracked_handler1 dcutr2._handle_dcutr_stream = tracked_handler2 # Start all protocols async with background_trio_service(relay_protocol): async with background_trio_service(dcutr1): async with background_trio_service(dcutr2): await relay_protocol.event_started.wait() await dcutr1.event_started.wait() await dcutr2.event_started.wait() # Connect both peers to the relay relay_addrs = relay.get_addrs() # Add relay addresses to both peers' peerstores for addr in relay_addrs: peer1.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) peer2.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) # Connect peers to relay await peer1.connect(relay.get_peerstore().peer_info(relay.get_id())) await peer2.connect(relay.get_peerstore().peer_info(relay.get_id())) await trio.sleep(0.1) # Verify peers are connected to relay assert relay.get_id() in [ peer_id for peer_id in peer1.get_network().connections.keys() ] assert relay.get_id() in [ peer_id for peer_id in peer2.get_network().connections.keys() ] # Verify peers are NOT directly connected to each other assert peer2.get_id() not in [ peer_id for peer_id in peer1.get_network().connections.keys() ] assert peer1.get_id() not in [ peer_id for peer_id in peer2.get_network().connections.keys() ] # Now test DCUtR: peer1 opens a DCUtR stream to peer2 through the # relay # This should trigger the DCUtR protocol for hole punching try: # Create a circuit relay multiaddr for peer2 through the relay relay_addr = relay_addrs[0] circuit_addr = Multiaddr( f"{relay_addr}/p2p-circuit/p2p/{peer2.get_id()}" ) # Add the circuit address to peer1's peerstore peer1.get_peerstore().add_addrs( peer2.get_id(), [circuit_addr], 3600 ) # Open a DCUtR stream from peer1 to peer2 through the relay stream = await peer1.new_stream(peer2.get_id(), [PROTOCOL_ID]) # Send a CONNECT message with observed addresses peer1_addrs = peer1.get_addrs() connect_msg = HolePunch( type=HolePunch.CONNECT, ObsAddrs=[addr.to_bytes() for addr in peer1_addrs[:2]], ) await stream.write(connect_msg.SerializeToString()) # Wait for the message to be processed await trio.sleep(0.2) # Verify that the DCUtR stream handler was called on peer2 assert handler2_called, ( "DCUtR stream handler should have been called on peer2" ) # Close the stream await stream.close() except Exception as e: logger.info( "Expected error when trying to open DCUtR stream through " "relay: %s", e, ) # This might fail because we need more setup, but the important # thing is testing the right scenario # Wait a bit more await trio.sleep(0.1) @pytest.mark.trio async def test_dcutr_relay_to_direct_upgrade(): """Test the complete flow: relay connection -> DCUtR -> direct connection.""" # Create three hosts: two peers and one relay async with HostFactory.create_batch_and_listen(3) as hosts: peer1, peer2, relay = hosts # Create circuit relay protocol for the relay relay_protocol = CircuitV2Protocol(relay, DEFAULT_RELAY_LIMITS, allow_hop=True) # Create DCUtR protocols for both peers dcutr1 = DCUtRProtocol(peer1) dcutr2 = DCUtRProtocol(peer2) # Track messages received messages_received = [] # Override stream handler to capture messages original_handler = dcutr2._handle_dcutr_stream async def message_capturing_handler(stream): try: # Read the message msg_data = await stream.read() hole_punch = HolePunch() hole_punch.ParseFromString(msg_data) messages_received.append(hole_punch) # Send a SYNC response sync_msg = HolePunch(type=HolePunch.SYNC) await stream.write(sync_msg.SerializeToString()) await original_handler(stream) except Exception as e: logger.error(f"Error in message capturing handler: {e}") await stream.close() dcutr2._handle_dcutr_stream = message_capturing_handler # Start all protocols async with background_trio_service(relay_protocol): async with background_trio_service(dcutr1): async with background_trio_service(dcutr2): await relay_protocol.event_started.wait() await dcutr1.event_started.wait() await dcutr2.event_started.wait() # Re-register the handler with the host dcutr2.host.set_stream_handler( PROTOCOL_ID, message_capturing_handler ) # Connect both peers to the relay relay_addrs = relay.get_addrs() # Add relay addresses to both peers' peerstores for addr in relay_addrs: peer1.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) peer2.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) # Connect peers to relay await peer1.connect(relay.get_peerstore().peer_info(relay.get_id())) await peer2.connect(relay.get_peerstore().peer_info(relay.get_id())) await trio.sleep(0.1) # Verify peers are connected to relay but not to each other assert relay.get_id() in [ peer_id for peer_id in peer1.get_network().connections.keys() ] assert relay.get_id() in [ peer_id for peer_id in peer2.get_network().connections.keys() ] assert peer2.get_id() not in [ peer_id for peer_id in peer1.get_network().connections.keys() ] # Try to open a DCUtR stream through the relay try: # Create a circuit relay multiaddr for peer2 through the relay relay_addr = relay_addrs[0] circuit_addr = Multiaddr( f"{relay_addr}/p2p-circuit/p2p/{peer2.get_id()}" ) # Add the circuit address to peer1's peerstore peer1.get_peerstore().add_addrs( peer2.get_id(), [circuit_addr], 3600 ) # Open a DCUtR stream from peer1 to peer2 through the relay stream = await peer1.new_stream(peer2.get_id(), [PROTOCOL_ID]) # Send a CONNECT message with observed addresses peer1_addrs = peer1.get_addrs() connect_msg = HolePunch( type=HolePunch.CONNECT, ObsAddrs=[addr.to_bytes() for addr in peer1_addrs[:2]], ) await stream.write(connect_msg.SerializeToString()) # Wait for the message to be processed await trio.sleep(0.2) # Verify that the CONNECT message was received assert len(messages_received) == 1, ( "Should have received one message" ) assert messages_received[0].type == HolePunch.CONNECT, ( "Should have received CONNECT message" ) assert len(messages_received[0].ObsAddrs) == 2, ( "Should have received 2 observed addresses" ) # Close the stream await stream.close() except Exception as e: logger.info( "Expected error when trying to open DCUtR stream through " "relay: %s", e, ) # Wait a bit more await trio.sleep(0.1) @pytest.mark.trio async def test_dcutr_hole_punch_through_relay(): """Test hole punching when peers are connected through relay.""" # Create three hosts: two peers and one relay async with HostFactory.create_batch_and_listen(3) as hosts: peer1, peer2, relay = hosts # Create circuit relay protocol for the relay relay_protocol = CircuitV2Protocol(relay, DEFAULT_RELAY_LIMITS, allow_hop=True) # Create DCUtR protocols for both peers dcutr1 = DCUtRProtocol(peer1) dcutr2 = DCUtRProtocol(peer2) # Start all protocols async with background_trio_service(relay_protocol): async with background_trio_service(dcutr1): async with background_trio_service(dcutr2): await relay_protocol.event_started.wait() await dcutr1.event_started.wait() await dcutr2.event_started.wait() # Connect both peers to the relay relay_addrs = relay.get_addrs() # Add relay addresses to both peers' peerstores for addr in relay_addrs: peer1.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) peer2.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) # Connect peers to relay await peer1.connect(relay.get_peerstore().peer_info(relay.get_id())) await peer2.connect(relay.get_peerstore().peer_info(relay.get_id())) await trio.sleep(0.1) # Verify peers are connected to relay but not to each other assert relay.get_id() in [ peer_id for peer_id in peer1.get_network().connections.keys() ] assert relay.get_id() in [ peer_id for peer_id in peer2.get_network().connections.keys() ] assert peer2.get_id() not in [ peer_id for peer_id in peer1.get_network().connections.keys() ] # Check if there's already a direct connection (should be False) has_direct = await dcutr1._have_direct_connection(peer2.get_id()) assert not has_direct, "Peers should not have a direct connection" # Try to initiate a hole punch (this should work through the relay # connection) # In a real scenario, this would be called after establishing a # relay connection result = await dcutr1.initiate_hole_punch(peer2.get_id()) # This should attempt hole punching but likely fail due to no public # addresses # The important thing is that the DCUtR protocol logic is executed logger.info( "Hole punch result: %s", result, ) assert result is not None, "Hole punch result should not be None" assert isinstance(result, bool), ( "Hole punch result should be a boolean" ) # Wait a bit more await trio.sleep(0.1) @pytest.mark.trio async def test_dcutr_relay_connection_verification(): """Test that DCUtR works correctly when peers are connected via relay.""" # Create three hosts: two peers and one relay async with HostFactory.create_batch_and_listen(3) as hosts: peer1, peer2, relay = hosts # Create circuit relay protocol for the relay relay_protocol = CircuitV2Protocol(relay, DEFAULT_RELAY_LIMITS, allow_hop=True) # Create DCUtR protocols for both peers dcutr1 = DCUtRProtocol(peer1) dcutr2 = DCUtRProtocol(peer2) # Start all protocols async with background_trio_service(relay_protocol): async with background_trio_service(dcutr1): async with background_trio_service(dcutr2): await relay_protocol.event_started.wait() await dcutr1.event_started.wait() await dcutr2.event_started.wait() # Connect both peers to the relay relay_addrs = relay.get_addrs() # Add relay addresses to both peers' peerstores for addr in relay_addrs: peer1.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) peer2.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) # Connect peers to relay await peer1.connect(relay.get_peerstore().peer_info(relay.get_id())) await peer2.connect(relay.get_peerstore().peer_info(relay.get_id())) await trio.sleep(0.1) # Verify peers are connected to relay assert relay.get_id() in [ peer_id for peer_id in peer1.get_network().connections.keys() ] assert relay.get_id() in [ peer_id for peer_id in peer2.get_network().connections.keys() ] # Verify peers are NOT directly connected to each other assert peer2.get_id() not in [ peer_id for peer_id in peer1.get_network().connections.keys() ] assert peer1.get_id() not in [ peer_id for peer_id in peer2.get_network().connections.keys() ] # Test getting observed addresses (real implementation) observed_addrs1 = await dcutr1._get_observed_addrs() observed_addrs2 = await dcutr2._get_observed_addrs() assert isinstance(observed_addrs1, list) assert isinstance(observed_addrs2, list) # Should contain the hosts' actual addresses assert len(observed_addrs1) > 0, ( "Peer1 should have observed addresses" ) assert len(observed_addrs2) > 0, ( "Peer2 should have observed addresses" ) # Test decoding observed addresses test_addrs = [ Multiaddr("/ip4/127.0.0.1/tcp/1234").to_bytes(), Multiaddr("/ip4/192.168.1.1/tcp/5678").to_bytes(), b"invalid-addr", # This should be filtered out ] decoded = dcutr1._decode_observed_addrs(test_addrs) assert len(decoded) == 2, "Should decode 2 valid addresses" assert all(str(addr).startswith("/ip4/") for addr in decoded) # Wait a bit more await trio.sleep(0.1) @pytest.mark.trio async def test_dcutr_relay_error_handling(): """Test DCUtR error handling when working through relay connections.""" # Create three hosts: two peers and one relay async with HostFactory.create_batch_and_listen(3) as hosts: peer1, peer2, relay = hosts # Create circuit relay protocol for the relay relay_protocol = CircuitV2Protocol(relay, DEFAULT_RELAY_LIMITS, allow_hop=True) # Create DCUtR protocols for both peers dcutr1 = DCUtRProtocol(peer1) dcutr2 = DCUtRProtocol(peer2) # Start all protocols async with background_trio_service(relay_protocol): async with background_trio_service(dcutr1): async with background_trio_service(dcutr2): await relay_protocol.event_started.wait() await dcutr1.event_started.wait() await dcutr2.event_started.wait() # Connect both peers to the relay relay_addrs = relay.get_addrs() # Add relay addresses to both peers' peerstores for addr in relay_addrs: peer1.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) peer2.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) # Connect peers to relay await peer1.connect(relay.get_peerstore().peer_info(relay.get_id())) await peer2.connect(relay.get_peerstore().peer_info(relay.get_id())) await trio.sleep(0.1) # Test with a stream that times out timeout_stream = MagicMock() timeout_stream.muxed_conn.peer_id = peer2.get_id() timeout_stream.read = AsyncMock(side_effect=trio.TooSlowError()) timeout_stream.write = AsyncMock() timeout_stream.close = AsyncMock() # This should not raise an exception, just log and close await dcutr1._handle_dcutr_stream(timeout_stream) # Verify stream was closed assert timeout_stream.close.called # Test with malformed message malformed_stream = MagicMock() malformed_stream.muxed_conn.peer_id = peer2.get_id() malformed_stream.read = AsyncMock(return_value=b"not-a-protobuf") malformed_stream.write = AsyncMock() malformed_stream.close = AsyncMock() # This should not raise an exception, just log and close await dcutr1._handle_dcutr_stream(malformed_stream) # Verify stream was closed assert malformed_stream.close.called # Wait a bit more await trio.sleep(0.1) @pytest.mark.trio async def test_dcutr_relay_attempt_limiting(): """Test DCUtR attempt limiting when working through relay connections.""" # Create three hosts: two peers and one relay async with HostFactory.create_batch_and_listen(3) as hosts: peer1, peer2, relay = hosts # Create circuit relay protocol for the relay relay_protocol = CircuitV2Protocol(relay, DEFAULT_RELAY_LIMITS, allow_hop=True) # Create DCUtR protocols for both peers dcutr1 = DCUtRProtocol(peer1) dcutr2 = DCUtRProtocol(peer2) # Start all protocols async with background_trio_service(relay_protocol): async with background_trio_service(dcutr1): async with background_trio_service(dcutr2): await relay_protocol.event_started.wait() await dcutr1.event_started.wait() await dcutr2.event_started.wait() # Connect both peers to the relay relay_addrs = relay.get_addrs() # Add relay addresses to both peers' peerstores for addr in relay_addrs: peer1.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) peer2.get_peerstore().add_addrs(relay.get_id(), [addr], 3600) # Connect peers to relay await peer1.connect(relay.get_peerstore().peer_info(relay.get_id())) await peer2.connect(relay.get_peerstore().peer_info(relay.get_id())) await trio.sleep(0.1) # Set max attempts reached dcutr1._hole_punch_attempts[peer2.get_id()] = ( MAX_HOLE_PUNCH_ATTEMPTS ) # Try to initiate hole punch - should fail due to max attempts result = await dcutr1.initiate_hole_punch(peer2.get_id()) assert result is False, "Hole punch should fail due to max attempts" # Reset attempts dcutr1._hole_punch_attempts.clear() # Add to direct connections dcutr1._direct_connections.add(peer2.get_id()) # Try to initiate hole punch - should succeed immediately result = await dcutr1.initiate_hole_punch(peer2.get_id()) assert result is True, ( "Hole punch should succeed for already connected peers" ) # Wait a bit more await trio.sleep(0.1)