tractor/tests/discovery/test_wg_iface_lifecycle.py

'''
WireGuard interface policy and owned lifecycle contracts.

'''
from __future__ import annotations

from collections.abc import AsyncIterator
from contextlib import asynccontextmanager as acm
import os
from pathlib import Path
from typing import BinaryIO

import pytest
import trio

from tractor.discovery import (
    Bindspace,
    BindspaceRef,
    BindspaceSpec,
    WGInterfaceConfig,
    WGPeerConfig,
    WGTunnelSpec,
    open_wg_bindspace,
    open_wg_iface,
)
from tractor.discovery import _tunnel


_LOCAL_KEY: str = 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA='
_PEER_KEY: str = 'r1LKM1pqhuY9Z6L4y5jQ2fGX67kJSrq5kRV5Jk2ywEo='


def test_wg_iface_settings_follow_role() -> None:
    '''
    Listen and dial roles interpret the tunnel bearer differently.

    Prove listen derives its local port from the bearer while dial
    applies the bearer as the selected peer's omitted endpoint. Other
    explicit peers retain their own endpoint and routing policy.

    '''
    selected: WGPeerConfig = WGPeerConfig(
        public_key=_PEER_KEY,
        allowed_ips=('10.1.0.0/16',),
    )
    other: WGPeerConfig = WGPeerConfig(
        public_key=_LOCAL_KEY,
        endpoint=('198.51.100.2', 51821),
    )
    config: WGInterfaceConfig = WGInterfaceConfig(
        private_key=_LOCAL_KEY,
        peers=(selected, other),
    )
    spec: WGTunnelSpec = WGTunnelSpec(
        peer_pubkey=_PEER_KEY,
        bearer=('192.0.2.1', 51820),
    )

    listen_port: int|None
    listen_peers: tuple[dict[str, object], ...]
    listen_port, listen_peers = _tunnel._wg_iface_settings(
        spec,
        config,
        'listen',
    )
    assert listen_port == 51820
    # A listen bearer configures the local port, not a peer endpoint.
    assert 'endpoint_addr' not in listen_peers[0]

    dial_port: int|None
    dial_peers: tuple[dict[str, object], ...]
    dial_port, dial_peers = _tunnel._wg_iface_settings(
        spec,
        config,
        'dial',
    )
    # No local dial listen port was declared; the bearer is remote.
    assert dial_port is None
    assert dial_peers[0]['endpoint_addr'] == '192.0.2.1'
    assert dial_peers[0]['endpoint_port'] == 51820
    assert dial_peers[0]['allowed_ips'] == ['10.1.0.0/16']
    assert dial_peers[1]['endpoint_addr'] == '198.51.100.2'
    assert dial_peers[1]['endpoint_port'] == 51821


@pytest.mark.parametrize(
    ('config', 'role', 'match'),
    (
        pytest.param(
            WGInterfaceConfig(
                private_key=_LOCAL_KEY,
            ),
            'dial',
            'not in.*configured peer keys',
            id='missing-dial-peer',
        ),
        pytest.param(
            WGInterfaceConfig(
                private_key=_LOCAL_KEY,
                listen_port=51821,
            ),
            'listen',
            '51821.*51820',
            id='listen-port-51821-vs-bearer-51820',
        ),
        pytest.param(
            WGInterfaceConfig(
                private_key=_LOCAL_KEY,
                peers=(
                    WGPeerConfig(
                        public_key=_PEER_KEY,
                        endpoint=('198.51.100.1', 51820),
                    ),
                ),
            ),
            'dial',
            '198.51.100.1.*192.0.2.1',
            id='dial-endpoint-conflict',
        ),
    ),
)
def test_wg_iface_settings_reject_conflicts(
    config: WGInterfaceConfig,
    role: str,
    match: str,
) -> None:
    '''
    Role-dependent conflicts must fail before pyroute2 side effects.

    '''
    spec: WGTunnelSpec = WGTunnelSpec(
        peer_pubkey=_PEER_KEY,
        bearer=('192.0.2.1', 51820),
    )
    with pytest.raises(ValueError, match=match):
        _tunnel._wg_iface_settings(
            spec,
            config,
            role,  # type: ignore[arg-type]
        )


def test_open_wg_iface_shields_cancelled_cleanup(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    '''
    Cancellation after creation must still remove the owned WG iface.

    Pin a stand-in namespace FD and fake privileged create/remove
    calls. Cancel inside the yielded context and prove shielded
    teardown runs before cancellation leaves the enclosing scope.

    '''
    token_path: Path = tmp_path / 'netns'
    token_path.touch()
    events: list[str] = []

    def create(
        spec: WGTunnelSpec,
        config: WGInterfaceConfig,
        bindspace: Bindspace,
        listen_port: int|None,
        peers: tuple[dict[str, object], ...],
    ) -> None:
        '''
        Record the validated create request.

        '''
        assert bindspace.namespace_fd is not None
        assert listen_port is None
        assert peers[0]['public_key'] == _PEER_KEY
        events.append('create')

    def remove(
        spec: WGTunnelSpec,
        bindspace: Bindspace,
    ) -> None:
        '''
        Record shielded removal after cancellation.

        '''
        events.append('remove')

    monkeypatch.setattr(
        _tunnel,
        '_sync_create_wg_iface',
        create,
    )
    monkeypatch.setattr(
        _tunnel,
        '_sync_remove_wg_iface',
        remove,
    )

    namespace_file: BinaryIO
    with token_path.open('rb') as namespace_file:
        namespace_fd: int = namespace_file.fileno()
        bindspace_spec: BindspaceSpec = BindspaceSpec(
            kind='netns',
            key='tractor-wg0',
        )
        bindspace: Bindspace = Bindspace(
            spec=bindspace_spec,
            ref=BindspaceRef(
                kind='netns',
                key='tractor-wg0',
                inode=os.fstat(namespace_fd).st_ino,
            ),
            namespace_fd=namespace_fd,
            ownership='borrowed',
        )
        peer: WGPeerConfig = WGPeerConfig(
            public_key=_PEER_KEY,
        )
        config: WGInterfaceConfig = WGInterfaceConfig(
            private_key=_LOCAL_KEY,
            peers=(peer,),
        )
        spec: WGTunnelSpec = WGTunnelSpec(
            peer_pubkey=_PEER_KEY,
        )

        async def main() -> None:
            '''
            Cancel while the fake WG iface is owned.

            '''
            with trio.CancelScope() as scope:
                async with open_wg_iface(
                    spec,
                    config,
                    bindspace,
                    'dial',
                ):
                    events.append('yield')
                    scope.cancel()
                    await trio.sleep_forever()

        trio.run(main)

    assert events == ['create', 'yield', 'remove']


def test_open_wg_bindspace_nests_resource_lifetimes(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    '''
    Nested WG interfaces must exit before their bindspace capability.

    Fake two interface layers over one bindspace. Clear the caller's
    mutable layer list at bindspace entry, then cancel from inside
    the yielded application scope and checkpoint. The trace proves
    stack snapshots its declaration before entry, layers enter
    outermost-first, cancellation exits them inside-out, and the live
    bindspace remains available through every interface exit.

    '''
    events: list[str] = []
    calls: list[
        tuple[
            WGTunnelSpec,
            WGInterfaceConfig,
            Bindspace,
            _tunnel.WGRole,
        ]
    ] = []
    bindspace_spec: BindspaceSpec = BindspaceSpec(
        kind='netns',
    )
    bindspace: Bindspace = Bindspace(
        spec=bindspace_spec,
        ref=BindspaceRef(
            kind='netns',
            key=None,
            inode=1,
        ),
        namespace_fd=None,
        ownership='borrowed',
    )
    outer_spec: WGTunnelSpec = WGTunnelSpec(
        peer_pubkey=_PEER_KEY,
        iface='wg-outer',
    )
    inner_spec: WGTunnelSpec = WGTunnelSpec(
        peer_pubkey=_LOCAL_KEY,
        iface='wg-inner',
    )
    outer_config: WGInterfaceConfig = WGInterfaceConfig(
        private_key=_LOCAL_KEY,
    )
    inner_config: WGInterfaceConfig = WGInterfaceConfig(
        private_key=_PEER_KEY,
    )
    layers: list[
        tuple[WGTunnelSpec, WGInterfaceConfig]
    ] = [
        (outer_spec, outer_config),
        (inner_spec, inner_config),
    ]

    @acm
    async def fake_open_bindspace(
        spec: BindspaceSpec,
    ) -> AsyncIterator[Bindspace]:
        '''
        Yield the stand-in bindspace and record its full lifetime.

        '''
        assert spec is bindspace_spec
        events.append('bindspace-enter')
        layers.clear()
        try:
            yield bindspace
        finally:
            events.append('bindspace-exit')

    @acm
    async def fake_open_wg_iface(
        spec: WGTunnelSpec,
        config: WGInterfaceConfig,
        bindspace_arg: Bindspace,
        role: _tunnel.WGRole,
    ) -> AsyncIterator[WGTunnelSpec]:
        '''
        Record one interface's arguments and nested lifetime.

        '''
        calls.append((spec, config, bindspace_arg, role))
        events.append(f'{spec.iface}-enter')
        try:
            yield spec
        finally:
            assert bindspace_arg is bindspace
            events.append(f'{spec.iface}-exit')

    monkeypatch.setattr(
        _tunnel,
        'open_bindspace',
        fake_open_bindspace,
    )
    monkeypatch.setattr(
        _tunnel,
        'open_wg_iface',
        fake_open_wg_iface,
    )

    async def main() -> None:
        '''
        Cancel while both interface layers are live.

        '''
        with trio.CancelScope() as scope:
            async with open_wg_bindspace(
                bindspace_spec,
                layers,
                'dial',
            ) as opened_bindspace:
                assert opened_bindspace is bindspace
                events.append('yield')
                scope.cancel()
                await trio.sleep_forever()

    trio.run(main)

    assert calls == [
        (outer_spec, outer_config, bindspace, 'dial'),
        (inner_spec, inner_config, bindspace, 'dial'),
    ]
    assert events == [
        'bindspace-enter',
        'wg-outer-enter',
        'wg-inner-enter',
        'yield',
        'wg-inner-exit',
        'wg-outer-exit',
        'bindspace-exit',
    ]