mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2025-12-31 20:36:24 +00:00
104 lines
3.6 KiB
Python
104 lines
3.6 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):
|
|
# conn is now a list, get the first connection
|
|
if isinstance(conn, list):
|
|
conn = conn[0]
|
|
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)
|