tractor/tests/discovery/test_multiaddr.py

'''
Multiaddr construction, parsing, and round-trip tests for
`tractor.discovery._multiaddr.mk_maddr()` and
`tractor.discovery._multiaddr.parse_maddr()`.

'''
from pathlib import Path
from types import SimpleNamespace

import pytest
from multiaddr import Multiaddr

from tractor.discovery import (
    TunnelledAddress,
    WGTunnelSpec,
    mb_pubkey,
    mk_wg_maddr,
    parse_wg_maddr,
    tunnels_of,
)
from tractor.ipc._tcp import TCPAddress
from tractor.ipc._uds import UDSAddress
from tractor.discovery._multiaddr import (
    mk_maddr,
    parse_maddr,
    parse_endpoints,
    _tpt_proto_to_maddr,
    _maddr_to_tpt_proto,
)
from tractor.discovery._addr import wrap_address


_WG_PUBKEY: str = (
    'g3x7z0AdV1rM6UQU22CC7IL3/ivn4DzrE7ikDhCZ/Dc='
)
_WG_PUBKEY_2: str = (
    'AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8='
)
_WG_MADDR: str = (
    f'/ip4/192.168.1.50/udp/51820'
    f'/wg/{mb_pubkey(_WG_PUBKEY)}'
    f'/ip4/10.0.11.1/tcp/1616'
)


def test_tpt_proto_to_maddr_mapping():
    '''
    `_tpt_proto_to_maddr` maps all supported `proto_key`
    values to their correct multiaddr protocol names.

    '''
    assert _tpt_proto_to_maddr['tcp'] == 'tcp'
    assert _tpt_proto_to_maddr['uds'] == 'unix'
    assert len(_tpt_proto_to_maddr) == 2


def test_mk_maddr_tcp_ipv4():
    '''
    `mk_maddr()` on a `TCPAddress` with an IPv4 host
    produces the correct `/ip4/<host>/tcp/<port>` multiaddr.

    '''
    addr = TCPAddress('127.0.0.1', 1234)
    result: Multiaddr = mk_maddr(addr)

    assert isinstance(result, Multiaddr)
    assert str(result) == '/ip4/127.0.0.1/tcp/1234'

    protos = result.protocols()
    assert protos[0].name == 'ip4'
    assert protos[1].name == 'tcp'

    assert result.value_for_protocol('ip4') == '127.0.0.1'
    assert result.value_for_protocol('tcp') == '1234'


def test_mk_maddr_tcp_ipv6():
    '''
    `mk_maddr()` on a `TCPAddress` with an IPv6 host
    produces the correct `/ip6/<host>/tcp/<port>` multiaddr.

    '''
    addr = TCPAddress('::1', 5678)
    result: Multiaddr = mk_maddr(addr)

    assert str(result) == '/ip6/::1/tcp/5678'

    protos = result.protocols()
    assert protos[0].name == 'ip6'
    assert protos[1].name == 'tcp'


def test_mk_maddr_uds():
    '''
    `mk_maddr()` on a `UDSAddress` produces a `/unix/<path>`
    multiaddr containing the full socket path.

    '''
    # NOTE, use an absolute `filedir` to match real runtime
    # UDS paths; `mk_maddr()` strips the leading `/` to avoid
    # the double-slash `/unix//run/..` that py-multiaddr
    # rejects as "empty protocol path".
    filedir = '/tmp/tractor_test'
    filename = 'test_sock.sock'
    addr = UDSAddress(
        filedir=filedir,
        filename=filename,
    )
    result: Multiaddr = mk_maddr(addr)

    assert isinstance(result, Multiaddr)

    result_str: str = str(result)
    assert result_str.startswith('/unix/')
    # verify the leading `/` was stripped to avoid double-slash
    assert '/unix/tmp/tractor_test/' in result_str

    sockpath_rel: str = str(
        Path(filedir) / filename
    ).lstrip('/')
    unix_val: str = result.value_for_protocol('unix')
    assert unix_val.endswith(sockpath_rel)


def test_mk_maddr_unsupported_proto_key():
    '''
    `mk_maddr()` raises `ValueError` for an unsupported
    `proto_key`.

    '''
    fake_addr = SimpleNamespace(proto_key='quic')
    with pytest.raises(
        ValueError,
        match='Unsupported proto_key',
    ):
        mk_maddr(fake_addr)


@pytest.mark.parametrize(
    'addr',
    [
        pytest.param(
            TCPAddress('127.0.0.1', 9999),
            id='tcp-ipv4',
        ),
        pytest.param(
            UDSAddress(
                filedir='/tmp/tractor_rt',
                filename='roundtrip.sock',
            ),
            id='uds',
        ),
    ],
)
def test_mk_maddr_roundtrip(addr):
    '''
    `mk_maddr()` output is valid multiaddr syntax that the
    library can re-parse back into an equivalent `Multiaddr`.

    '''
    maddr: Multiaddr = mk_maddr(addr)
    reparsed = Multiaddr(str(maddr))

    assert reparsed == maddr
    assert str(reparsed) == str(maddr)


# ------ parse_maddr() tests ------

def test_maddr_to_tpt_proto_mapping():
    '''
    `_maddr_to_tpt_proto` is the exact inverse of
    `_tpt_proto_to_maddr`.

    '''
    assert _maddr_to_tpt_proto == {
        'tcp': 'tcp',
        'unix': 'uds',
    }


def test_parse_maddr_tcp_ipv4():
    '''
    `parse_maddr()` on an IPv4 TCP multiaddr string
    produce a `TCPAddress` with the correct host and port.

    '''
    result = parse_maddr('/ip4/127.0.0.1/tcp/1234')

    assert isinstance(result, TCPAddress)
    assert result.unwrap() == ('tcp', '127.0.0.1', 1234)


def test_parse_maddr_tcp_ipv6():
    '''
    `parse_maddr()` on an IPv6 TCP multiaddr string
    produce a `TCPAddress` with the correct host and port.

    '''
    result = parse_maddr('/ip6/::1/tcp/5678')

    assert isinstance(result, TCPAddress)
    assert result.unwrap() == ('tcp', '::1', 5678)


def test_parse_maddr_uds():
    '''
    `parse_maddr()` on a `/unix/...` multiaddr string
    produce a `UDSAddress` with the correct dir and filename,
    preserving absolute path semantics.

    '''
    result = parse_maddr('/unix/tmp/tractor_test/test.sock')

    assert isinstance(result, UDSAddress)
    assert result.unwrap() == (
        'unix',
        '/tmp/tractor_test/test.sock',
    )


def test_parse_maddr_unsupported():
    '''
    `parse_maddr()` raise `ValueError` for an unsupported
    protocol combination like UDP.

    '''
    with pytest.raises(
        ValueError,
        match='Unsupported multiaddr protocol combo',
    ):
        parse_maddr('/ip4/127.0.0.1/udp/1234')


def test_parse_wg_maddr():
    '''
    `parse_maddr()` previously rejected the canonical infix `/wg/`
    grammar even though `py-multiaddr` parsed it. Feed a bearer,
    identity, and TCP overlay through both the WG-specific and public
    parsers, then prove they produce the same local-only tunnel
    annotation without changing the bindable overlay.

    '''
    parsed = parse_wg_maddr(_WG_MADDR)

    assert parse_maddr(_WG_MADDR) == parsed
    assert isinstance(parsed, TunnelledAddress)
    assert parsed.tunnel == WGTunnelSpec(
        peer_pubkey=_WG_PUBKEY,
        bearer=('192.168.1.50', 51820),
    )
    assert isinstance(parsed.overlay, TCPAddress)
    assert parsed.overlay.unwrap() == ('tcp', '10.0.11.1', 1616)


def test_mk_wg_maddr_roundtrip():
    '''
    `mk_maddr()` previously saw only the wrapper's delegated TCP
    proto-key and silently dropped all tunnel metadata. Parse the
    canonical maddr, compose it through both public entry points, and
    prove bearer, key, and overlay survive byte-for-byte.

    '''
    parsed = parse_wg_maddr(_WG_MADDR)

    assert str(mk_wg_maddr(parsed)) == _WG_MADDR
    assert str(mk_maddr(parsed)) == _WG_MADDR


def test_nested_wg_maddr_roundtrip():
    '''
    A single first-match lookup confuses nested WG keys and bearers.
    Arrange an IPv4 outer bearer around an IPv6 inner bearer, parse
    from the last `/wg/` outward, and assert tunnel ordering plus an
    exact re-composition of the original stack.

    '''
    nested_maddr: str = (
        f'/ip4/192.168.1.50/udp/51820'
        f'/wg/{mb_pubkey(_WG_PUBKEY)}'
        f'/ip6/2001:db8::2/udp/51821'
        f'/wg/{mb_pubkey(_WG_PUBKEY_2)}'
        f'/ip4/10.0.11.1/tcp/1616'
    )

    parsed = parse_maddr(nested_maddr)
    specs = tunnels_of(parsed)

    assert len(specs) == 2
    assert specs[0].peer_pubkey == _WG_PUBKEY
    assert specs[0].bearer == ('192.168.1.50', 51820)
    assert specs[1].peer_pubkey == _WG_PUBKEY_2
    assert specs[1].bearer == ('2001:db8::2', 51821)
    assert str(mk_maddr(parsed)) == nested_maddr


@pytest.mark.parametrize(
    'maddr, match',
    [
        pytest.param(
            (
                f'/ip4/192.168.1.50/tcp/51820'
                f'/wg/{mb_pubkey(_WG_PUBKEY)}'
                f'/ip4/10.0.11.1/tcp/1616'
            ),
            'Bad `wg` bearer',
            id='non-udp-bearer',
        ),
        pytest.param(
            (
                f'/ip4/192.168.1.50/udp/51820'
                f'/wg/{mb_pubkey(_WG_PUBKEY)}'
            ),
            'no overlay endpoint',
            id='missing-overlay',
        ),
        pytest.param(
            (
                f'/ip4/192.168.1.50/udp/51820'
                f'/wg/{mb_pubkey(_WG_PUBKEY)}'
                f'/ip4/10.0.11.1/udp/1616'
            ),
            'Unsupported `wg` overlay',
            id='non-tcp-overlay',
        ),
    ],
)
def test_parse_wg_maddr_rejects_bad_grammar(
    maddr: str,
    match: str,
):
    '''
    Accepting an invalid bearer or overlay assigns an endpoint to the
    wrong runtime owner. Exercise parseable but unsupported protocol
    combinations and prove each fails before constructing a wrapper,
    with an error identifying the violated WG grammar boundary.

    '''
    with pytest.raises(ValueError, match=match):
        parse_maddr(maddr)


def test_parse_wg_maddr_rejects_malformed_key():
    '''
    A truncated multibase key used to be vulnerable to silent
    identity corruption in hand-written parsers. Give the upstream
    `/wg/` codec a short key and prove `Multiaddr()` rejects it
    before tractor's wrapper parser runs.

    '''
    maddr: str = (
        '/ip4/192.168.1.50/udp/51820'
        '/wg/udG9vIHNob3J0'
        '/ip4/10.0.11.1/tcp/1616'
    )
    with pytest.raises(ValueError):
        parse_maddr(maddr)


def test_parse_wg_maddr_reports_missing_codec(
    monkeypatch,
):
    '''
    Released `multiaddr==0.2.0` does not know `/wg/` and emits an
    opaque unknown-protocol parse error. Simulate that registry and
    prove an actual WG stack reports the dependency action while a
    Unix path containing a `wg` directory remains ordinary UDS data.

    '''
    from multiaddr.exceptions import ProtocolNotFoundError
    from tractor.discovery import _tunnel

    def no_wg_proto(name: str):
        raise ProtocolNotFoundError(name)

    monkeypatch.setattr(
        _tunnel,
        'protocol_with_name',
        no_wg_proto,
    )
    uds = parse_maddr('/unix/tmp/wg/service.sock')
    assert isinstance(uds, UDSAddress)

    with pytest.raises(
        RuntimeError,
        match='py-multiaddr#108',
    ):
        parse_maddr(_WG_MADDR)


def test_mk_wg_maddr_requires_bearer():
    '''
    A key-only tunnel spec relies on local configuration and cannot
    be reconstructed as the canonical bearer-first maddr. Build that
    incomplete annotation and prove composition raises instead of
    emitting a misleading overlay-only address.

    '''
    addr = TunnelledAddress(
        overlay=TCPAddress('10.0.11.1', 1616),
        tunnel=WGTunnelSpec(peer_pubkey=_WG_PUBKEY),
    )
    with pytest.raises(
        ValueError,
        match='without a bearer',
    ):
        mk_maddr(addr)


@pytest.mark.parametrize(
    'addr',
    [
        pytest.param(
            TCPAddress('127.0.0.1', 9999),
            id='tcp-ipv4',
        ),
        pytest.param(
            UDSAddress(
                filedir='/tmp/tractor_rt',
                filename='roundtrip.sock',
            ),
            id='uds',
        ),
    ],
)
def test_parse_maddr_roundtrip(addr):
    '''
    Full round-trip: `addr -> mk_maddr -> str -> parse_maddr`
    produce an `Address` whose `.unwrap()` matches the original.

    '''
    maddr: Multiaddr = mk_maddr(addr)
    maddr_str: str = str(maddr)
    parsed = parse_maddr(maddr_str)

    assert type(parsed) is type(addr)
    assert parsed.unwrap() == addr.unwrap()


def test_wrap_address_maddr_str():
    '''
    `wrap_address()` accept a multiaddr-format string and
    return the correct `Address` type.

    '''
    result = wrap_address('/ip4/127.0.0.1/tcp/9999')

    assert isinstance(result, TCPAddress)
    assert result.unwrap() == ('tcp', '127.0.0.1', 9999)


def test_wrap_address_wg_maddr_str():
    '''
    `wrap_address()` delegates slash-prefixed strings to
    `parse_maddr()`. Pass a canonical WG maddr through that public
    boundary and prove it preserves the tunnel annotation rather than
    rejecting the protocol stack or returning only its TCP overlay.

    '''
    result = wrap_address(_WG_MADDR)

    assert isinstance(result, TunnelledAddress)
    assert result.tunnel.peer_pubkey == _WG_PUBKEY
    assert result.overlay.unwrap() == ('tcp', '10.0.11.1', 1616)


# ------ parse_endpoints() tests ------

def test_parse_endpoints_tcp_only():
    '''
    `parse_endpoints()` with a single TCP maddr per actor
    produce the correct `TCPAddress` instances.

    '''
    table = {
        'registry': ['/ip4/127.0.0.1/tcp/1616'],
        'data_feed': ['/ip4/0.0.0.0/tcp/5555'],
    }
    result = parse_endpoints(table)

    assert set(result.keys()) == {'registry', 'data_feed'}

    reg_addr = result['registry'][0]
    assert isinstance(reg_addr, TCPAddress)
    assert reg_addr.unwrap() == ('tcp', '127.0.0.1', 1616)

    feed_addr = result['data_feed'][0]
    assert isinstance(feed_addr, TCPAddress)
    assert feed_addr.unwrap() == ('tcp', '0.0.0.0', 5555)


def test_parse_endpoints_mixed_tpts():
    '''
    `parse_endpoints()` with both TCP and UDS maddrs for
    the same actor produce the correct mixed `Address` list.

    '''
    table = {
        'broker': [
            '/ip4/127.0.0.1/tcp/4040',
            '/unix/tmp/tractor/broker.sock',
        ],
    }
    result = parse_endpoints(table)
    addrs = result['broker']

    assert len(addrs) == 2
    assert isinstance(addrs[0], TCPAddress)
    assert addrs[0].unwrap() == ('tcp', '127.0.0.1', 4040)

    assert isinstance(addrs[1], UDSAddress)
    assert addrs[1].unwrap() == (
        'unix',
        '/tmp/tractor/broker.sock',
    )


def test_parse_endpoints_wg_maddr():
    '''
    Service endpoint tables previously rejected WG protocol stacks.
    Put a tunnelled maddr beside a plain TCP address and prove
    `parse_endpoints()` retains input order while delivering the
    wrapper needed by the future bindspace lifecycle.

    '''
    table = {
        'registry': [
            _WG_MADDR,
            '/ip4/127.0.0.1/tcp/1616',
        ],
    }
    addrs = parse_endpoints(table)['registry']

    assert isinstance(addrs[0], TunnelledAddress)
    assert addrs[0].tunnel.bearer == (
        '192.168.1.50',
        51820,
    )
    assert isinstance(addrs[1], TCPAddress)


def test_parse_endpoints_unwrapped_tuples():
    '''
    `parse_endpoints()` accept raw `(host, port)` tuples
    and wrap them as `TCPAddress`.

    '''
    table = {
        'ems': [('127.0.0.1', 6666)],
    }
    result = parse_endpoints(table)

    addr = result['ems'][0]
    assert isinstance(addr, TCPAddress)
    assert addr.unwrap() == ('tcp', '127.0.0.1', 6666)


def test_parse_endpoints_mixed_str_and_tuple():
    '''
    `parse_endpoints()` accept a mix of maddr strings and
    raw tuples in the same actor entry list.

    '''
    table = {
        'quoter': [
            '/ip4/127.0.0.1/tcp/7777',
            ('127.0.0.1', 8888),
        ],
    }
    result = parse_endpoints(table)
    addrs = result['quoter']

    assert len(addrs) == 2
    assert isinstance(addrs[0], TCPAddress)
    assert addrs[0].unwrap() == ('tcp', '127.0.0.1', 7777)

    assert isinstance(addrs[1], TCPAddress)
    assert addrs[1].unwrap() == ('tcp', '127.0.0.1', 8888)


def test_parse_endpoints_unsupported_proto():
    '''
    `parse_endpoints()` raise `ValueError` when a maddr
    string uses an unsupported protocol like `/udp/`.

    '''
    table = {
        'bad_actor': ['/ip4/127.0.0.1/udp/9999'],
    }
    with pytest.raises(
        ValueError,
        match='Unsupported multiaddr protocol combo',
    ):
        parse_endpoints(table)


def test_parse_endpoints_empty_table():
    '''
    `parse_endpoints()` on an empty table return an empty
    dict.

    '''
    assert parse_endpoints({}) == {}


def test_parse_endpoints_empty_actor_list():
    '''
    `parse_endpoints()` with an actor mapped to an empty
    list preserve the key with an empty list value.

    '''
    result = parse_endpoints({'x': []})
    assert result == {'x': []}