Files
py-libp2p/tests/core/security/test_security_multistream.py
Soham Bhoir 66bd027161 Feat/587-circuit-relay (#611)
* feat: implemented setup of circuit relay and test cases

* chore: remove test files to be rewritten

* added 1 test suite for protocol

* added 1 test suite for discovery

* fixed protocol timeouts and message types to handle reservations and stream operations.

* Resolved merge conflict in libp2p/tools/utils.py by combining timeout approach with retry mechanism

* fix: linting issues

* docs: updated documentation with circuit-relay

* chore: added enums, improved typing, security and examples

* fix: created proper __init__ file to ensure importability

* fix: replace transport_opt with listen_addrs in examples, fixed typing and improved code

* fix type checking issues across relay module and test suite

* regenerated circuit_pb2 file protobuf version 3

* fixed circuit relay example and moved imports to top in test_security_multistream

* chore: moved imports to the top

* chore: fixed linting of test_circuit_v2_transport.py

---------

Co-authored-by: Manu Sheel Gupta <manusheel.edu@gmail.com>
2025-06-18 15:39:39 -06:00

101 lines
3.5 KiB
Python

import pytest
import trio
from libp2p.crypto.rsa import (
create_new_key_pair,
)
from libp2p.security.insecure.transport import (
PLAINTEXT_PROTOCOL_ID,
InsecureSession,
)
from libp2p.security.noise.transport import PROTOCOL_ID as NOISE_PROTOCOL_ID
from libp2p.security.secio.transport import ID as SECIO_PROTOCOL_ID
from libp2p.security.secure_session import (
SecureSession,
)
from libp2p.stream_muxer.mplex.mplex import Mplex
from libp2p.stream_muxer.yamux.yamux import Yamux
from tests.utils.factories import (
host_pair_factory,
)
initiator_key_pair = create_new_key_pair()
noninitiator_key_pair = create_new_key_pair()
async def perform_simple_test(assertion_func, security_protocol):
async with host_pair_factory(security_protocol=security_protocol) as hosts:
# Use a different approach to verify connections
# Wait for both sides to establish connection
for _ in range(5): # Try up to 5 times
try:
# Check if connection established from host0 to host1
conn_0 = hosts[0].get_network().connections.get(hosts[1].get_id())
# Check if connection established from host1 to host0
conn_1 = hosts[1].get_network().connections.get(hosts[0].get_id())
if conn_0 and conn_1:
break
# Wait a bit and retry
await trio.sleep(0.2)
except Exception:
# Wait a bit and retry
await trio.sleep(0.2)
# If we couldn't establish connection after retries,
# the test will fail with clear error
assert conn_0 is not None, "Failed to establish connection from host0 to host1"
assert conn_1 is not None, "Failed to establish connection from host1 to host0"
# Extract the secured connection from either Mplex or Yamux implementation
def get_secured_conn(conn):
muxed_conn = conn.muxed_conn
# Direct attribute access for known implementations
has_secured_conn = hasattr(muxed_conn, "secured_conn")
if isinstance(muxed_conn, (Mplex, Yamux)) and has_secured_conn:
return muxed_conn.secured_conn
# Fallback to _connection attribute if it exists
elif hasattr(muxed_conn, "_connection"):
return muxed_conn._connection
# Last resort - warn but return the muxed_conn itself for type checking
else:
print(f"Warning: Cannot find secured connection in {type(muxed_conn)}")
return muxed_conn
# Get secured connections for both peers
secured_conn_0 = get_secured_conn(conn_0)
secured_conn_1 = get_secured_conn(conn_1)
# Perform assertion on the secured connections
assertion_func(secured_conn_0)
assertion_func(secured_conn_1)
@pytest.mark.trio
@pytest.mark.parametrize(
"security_protocol, transport_type",
(
(PLAINTEXT_PROTOCOL_ID, InsecureSession),
(SECIO_PROTOCOL_ID, SecureSession),
(NOISE_PROTOCOL_ID, SecureSession),
),
)
@pytest.mark.trio
async def test_single_insecure_security_transport_succeeds(
security_protocol, transport_type
):
def assertion_func(conn):
assert isinstance(conn, transport_type)
await perform_simple_test(assertion_func, security_protocol)
@pytest.mark.trio
async def test_default_insecure_security():
def assertion_func(conn):
assert isinstance(conn, InsecureSession)
await perform_simple_test(assertion_func, None)