mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2025-12-31 20:36:24 +00:00
cb11f076c8
* feat: base implementation of dcutr for hole-punching * chore: removed circuit-relay imports from __init__ * feat: implemented dcutr protocol * added test suite with mock setup * Fix pre-commit hook issues in DCUtR implementation * usages of CONNECT_TYPE and SYNC_TYPE have been replaced with HolePunch.Type.CONNECT and HolePunch.Type.SYNC * added unit tests for dcutr and nat module and * added multiaddr.get_peer_id() with proper DNS address handling and fixed method signature inconsistencies * added assertions to verify DCUtR hole punch result in integration test --------- Co-authored-by: Manu Sheel Gupta <manusheel.edu@gmail.com>
564 lines
24 KiB
Python
564 lines
24 KiB
Python
"""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)
|