py-libp2p/tests/security/test_secio.py

import pytest
import trio

from libp2p.crypto.secp256k1 import create_new_key_pair
from libp2p.peer.id import ID
from libp2p.security.secio.transport import NONCE_SIZE, create_secure_session
from libp2p.tools.constants import MAX_READ_LEN
from libp2p.tools.factories import raw_conn_factory


@pytest.mark.trio
async def test_create_secure_session(nursery):
    local_nonce = b"\x01" * NONCE_SIZE
    local_key_pair = create_new_key_pair(b"a")
    local_peer = ID.from_pubkey(local_key_pair.public_key)

    remote_nonce = b"\x02" * NONCE_SIZE
    remote_key_pair = create_new_key_pair(b"b")
    remote_peer = ID.from_pubkey(remote_key_pair.public_key)

    async with raw_conn_factory(nursery) as conns:
        local_conn, remote_conn = conns

        local_secure_conn, remote_secure_conn = None, None

        async def local_create_secure_session():
            nonlocal local_secure_conn
            local_secure_conn = await create_secure_session(
                local_nonce,
                local_peer,
                local_key_pair.private_key,
                local_conn,
                remote_peer,
            )

        async def remote_create_secure_session():
            nonlocal remote_secure_conn
            remote_secure_conn = await create_secure_session(
                remote_nonce, remote_peer, remote_key_pair.private_key, remote_conn
            )

        async with trio.open_nursery() as nursery_1:
            nursery_1.start_soon(local_create_secure_session)
            nursery_1.start_soon(remote_create_secure_session)

        msg = b"abc"
        await local_secure_conn.write(msg)
        received_msg = await remote_secure_conn.read(MAX_READ_LEN)
        assert received_msg == msg