fips/testing/chaos/sim/topology.py

"""Topology generation: random graphs with connectivity guarantees."""

from __future__ import annotations

import math
import random
from collections import deque
from dataclasses import dataclass, field

from .keys import derive_full
from .naming import name_suffix, veth_token
from .scenario import TopologyConfig

# An edge carried by UDP over a dedicated veth pair, each end an
# interface-bound UDP instance (``transports.udp.<iface>.interface``). The
# harness's stand-in for "wifi and cable, both IP": two UDP instances on
# two interfaces, so a peer reachable over both holds two paths.
UDP_VETH = "udp-veth"

# Port of the interface-bound UDP instances. Not 2121: the bridge instance
# binds the wildcard on that port, and a second wildcard bind on the same
# port would conflict.
UDP_VETH_PORT = 2122

# Second octet of the /24s the veth pairs carry. Clear of docker's default
# pool (172.17-31), the sim's claimed 10.30.x ranges and sidecar's 10.40.x.
_UDP_VETH_NET = "10.222"


@dataclass(frozen=True)
class UdpVethLink:
    """One end of a ``udp-veth`` edge, as a node sees it."""

    peer_id: str
    # The veth interface in this node's container, and the UDP instance
    # name bound to it.
    iface: str
    local_ip: str
    peer_ip: str

    @property
    def instance(self) -> str:
        return self.iface

    @property
    def peer_addr(self) -> str:
        return f"{self.peer_ip}:{UDP_VETH_PORT}"


@dataclass
class SimNode:
    node_id: str  # "n01", "n02", ...
    docker_ip: str  # "172.20.0.10", ...
    nsec: str  # 64-char hex
    npub: str  # bech32 npub1...
    peers: list[str] = field(default_factory=list)
    # MAC addresses for Ethernet veth interfaces, keyed by peer_id
    ethernet_macs: dict[str, str] = field(default_factory=dict)


@dataclass
class SimTopology:
    nodes: dict[str, SimNode] = field(default_factory=dict)
    edges: set[tuple[str, str]] = field(default_factory=set)
    # Per-edge transport type; edges not in this dict default to "udp"
    edge_transport: dict[tuple[str, str], str] = field(default_factory=dict)
    # Edges declared ``ethernet+udp``: an Ethernet veth (found by beacon)
    # *and* a UDP static-peer entry over the bridge, so the pair holds two
    # paths under one session. ``edge_transport`` says ``ethernet`` for
    # these, which is what netem and link flaps act on.
    dual_udp_edges: set[tuple[str, str]] = field(default_factory=set)
    # Suffix scoping globally-visible names to this run and scenario; empty
    # outside the CI harness, which keeps a bare run's names unchanged.
    name_suffix: str = ""

    @property
    def veth_token(self) -> str:
        """Short stand-in for the suffix, for names bound by IFNAMSIZ.

        Derived rather than stored so no caller can build a topology whose
        host names are scoped differently from its container names.
        """
        return veth_token(self.name_suffix)

    def is_dual_udp_edge(self, a: str, b: str) -> bool:
        """Whether the edge also carries a UDP static-peer link."""
        return _make_edge(a, b) in self.dual_udp_edges

    def is_veth_transport(self, transport: str) -> bool:
        """Whether edges of this transport run over a dedicated veth pair."""
        return transport in ("ethernet", UDP_VETH)

    def veth_edges(self) -> list[tuple[str, str]]:
        """Every edge that needs a veth pair: Ethernet and ``udp-veth``."""
        return sorted(
            e for e, t in self.edge_transport.items() if self.is_veth_transport(t)
        )

    def has_veth(self) -> bool:
        return bool(self.veth_edges())

    def udp_veth_edges(self) -> list[tuple[str, str]]:
        """Edges carried by UDP over a veth, in canonical order. The index
        of an edge here is what its /24 is numbered by."""
        return sorted(e for e, t in self.edge_transport.items() if t == UDP_VETH)

    def udp_veth_links(self, node_id: str) -> list[UdpVethLink]:
        """This node's ends of its ``udp-veth`` edges, with addressing.

        Edge ``k`` (in ``udp_veth_edges`` order) is ``10.222.k.0/24``: the
        lower node id is ``.1``, the higher ``.2``.
        """
        links = []
        for k, (a, b) in enumerate(self.udp_veth_edges()):
            if node_id not in (a, b):
                continue
            if k > 255:
                raise ValueError("more than 256 udp-veth edges are not addressable")
            local, peer = (a, b) if node_id == a else (b, a)
            local_ip = f"{_UDP_VETH_NET}.{k}.{1 if node_id == a else 2}"
            peer_ip = f"{_UDP_VETH_NET}.{k}.{2 if node_id == a else 1}"
            links.append(
                UdpVethLink(
                    peer_id=peer,
                    iface=veth_interface_name(local, peer),
                    local_ip=local_ip,
                    peer_ip=peer_ip,
                )
            )
        return links

    def udp_veth_ip(self, node_id: str, peer_id: str) -> str | None:
        """The veth IP ``node_id`` has on its ``udp-veth`` edge to ``peer_id``."""
        for link in self.udp_veth_links(node_id):
            if link.peer_id == peer_id:
                return link.local_ip
        return None

    def transport_for_edge(self, a: str, b: str) -> str:
        """Get the transport type for an edge (defaults to 'udp')."""
        edge = _make_edge(a, b)
        return self.edge_transport.get(edge, "udp")

    def ethernet_edges(self) -> list[tuple[str, str]]:
        """Return all edges using Ethernet transport."""
        return [e for e, t in self.edge_transport.items() if t == "ethernet"]

    def has_ethernet(self) -> bool:
        """Check if any edges use Ethernet transport."""
        return any(t == "ethernet" for t in self.edge_transport.values())

    def tcp_edges(self) -> list[tuple[str, str]]:
        """Return all edges using TCP transport."""
        return [e for e, t in self.edge_transport.items() if t == "tcp"]

    def has_tcp(self) -> bool:
        """Check if any edges use TCP transport."""
        return any(t == "tcp" for t in self.edge_transport.values())

    def tcp_peers(self, node_id: str) -> list[str]:
        """Return peer IDs connected to this node via TCP."""
        peers = []
        for (a, b), transport in self.edge_transport.items():
            if transport != "tcp":
                continue
            if a == node_id:
                peers.append(b)
            elif b == node_id:
                peers.append(a)
        return sorted(peers)

    def ethernet_interfaces(self, node_id: str) -> list[str]:
        """Return the veth interface names for a node's Ethernet edges."""
        ifaces = []
        for (a, b), transport in self.edge_transport.items():
            if transport != "ethernet":
                continue
            if a == node_id:
                ifaces.append(veth_interface_name(a, b))
            elif b == node_id:
                ifaces.append(veth_interface_name(b, a))
        return sorted(ifaces)

    def is_connected(self) -> bool:
        """BFS connectivity check."""
        if len(self.nodes) <= 1:
            return True
        start = next(iter(self.nodes))
        visited = set()
        queue = deque([start])
        while queue:
            node = queue.popleft()
            if node in visited:
                continue
            visited.add(node)
            for peer in self.nodes[node].peers:
                if peer not in visited:
                    queue.append(peer)
        return len(visited) == len(self.nodes)

    def neighbors(self, node_id: str) -> list[str]:
        return self.nodes[node_id].peers

    def would_disconnect(self, edge: tuple[str, str]) -> bool:
        """Check if removing this edge would disconnect the graph."""
        a, b = edge
        # Temporarily remove edge
        self.nodes[a].peers.remove(b)
        self.nodes[b].peers.remove(a)
        connected = self.is_connected()
        # Restore
        self.nodes[a].peers.append(b)
        self.nodes[b].peers.append(a)
        return not connected

    def container_name(self, node_id: str) -> str:
        return f"fips-node-{node_id}{self.name_suffix}"

    def veth_host_name(self, node_a: str, node_b: str, end: str) -> str:
        """Generate the host-namespace veth name for one end of an edge.

        Format: ``vh{token}{NN}{MM}{end}`` (max 15 chars for IFNAMSIZ).
        Host interfaces are global, so the token keeps a scenario from
        deleting a concurrent scenario's pair; it is empty outside the CI
        harness, yielding the same "vh0104a" this has always produced.

        ``node_a`` and ``node_b`` must be in canonical edge order. Unlike
        ``veth_interface_name()`` this is not symmetric: the far end is
        ``end="b"`` on the same ordering, so swapping the arguments names
        an interface that does not exist.
        """
        nn_local = node_a.replace("n", "")
        nn_peer = node_b.replace("n", "")
        name = f"vh{self.veth_token}{nn_local}{nn_peer}{end}"
        if len(name) > 15:
            raise ValueError(f"veth host name too long: {name!r} ({len(name)} > 15)")
        return name

    def directed_outbound(self) -> dict[str, list[str]]:
        """Assign each static-config edge to exactly one node for outbound connection.

        Returns a mapping from node_id to the list of peers that node
        should connect to (outbound only). Every edge appears in exactly
        one direction, ensuring auto-reconnect is testable — if B goes
        down, only A (the outbound owner) will attempt to reconnect.

        Ethernet edges are excluded — they use beacon discovery instead
        of static peer configuration. UDP and TCP edges use static config.

        Strategy: BFS spanning tree edges go parent→child. Non-tree
        edges go from the lower node ID to the higher. This guarantees
        every node is reachable via at least one inbound connection.
        """
        # Consider all edges that use static peer config (not Ethernet/discovery)
        static_edges = {
            e for e in self.edges
            if self.edge_transport.get(e, "udp") != "ethernet"
            or e in self.dual_udp_edges
        }

        outbound: dict[str, list[str]] = {nid: [] for nid in self.nodes}

        # Build static-config adjacency for BFS
        static_adj: dict[str, list[str]] = {nid: [] for nid in self.nodes}
        for a, b in static_edges:
            static_adj[a].append(b)
            static_adj[b].append(a)

        # BFS spanning tree from first node (over static-config edges only)
        root = min(self.nodes)
        visited: set[str] = set()
        tree_edges: set[tuple[str, str]] = set()
        queue = deque([root])
        visited.add(root)
        while queue:
            node = queue.popleft()
            for peer in static_adj[node]:
                if peer not in visited:
                    visited.add(peer)
                    queue.append(peer)
                    tree_edges.add((node, peer))  # parent → child
                    outbound[node].append(peer)

        # Non-tree static-config edges: lower ID → higher ID
        for a, b in static_edges:
            if (a, b) not in tree_edges and (b, a) not in tree_edges:
                outbound[a].append(b)  # a < b by _make_edge convention

        return outbound


def generate_topology(
    config: TopologyConfig,
    rng: random.Random,
    mesh_name: str,
) -> SimTopology:
    """Generate a topology according to the config."""
    n = config.num_nodes
    subnet_base = config.subnet.rsplit(".", 1)[0]  # "172.20.0"

    # Create nodes with IPs and keys.
    #
    # The mesh roots itself at the numerically smallest NodeAddr
    # (`src/tree/state.rs:363-390`), which is a hash of the node's public key
    # and so bears no relation to the node numbering. Every scenario diagram in
    # this tree draws n01 at the top, and before this ordering was applied the
    # root landed on an arbitrary node in most scenarios — which is why the
    # cost-selection scenarios that reasoned about a specific root could never
    # be turned into reliable assertions and were moved to sans-IO unit tests.
    #
    # So derive the identities from the mesh name as before, then *assign* them
    # in NodeAddr order: n01 receives the smallest and is the root, n02 the next,
    # and so on. The keys are unchanged and still deterministic; only which node
    # id holds which one changes. Scenarios that want an arbitrary root set
    # `pin_root: false` and keep exercising election.
    node_ids_ordered = [f"n{i + 1:02d}" for i in range(n)]
    identities = [derive_full(mesh_name, nid) for nid in node_ids_ordered]
    if config.pin_root:
        identities.sort(key=lambda t: t[2])

    nodes: dict[str, SimNode] = {}
    for i, node_id in enumerate(node_ids_ordered):
        docker_ip = f"{subnet_base}.{config.ip_start + i}"
        nsec, npub, _ = identities[i]
        nodes[node_id] = SimNode(
            node_id=node_id,
            docker_ip=docker_ip,
            nsec=nsec,
            npub=npub,
        )

    node_ids = sorted(nodes.keys())

    # Generate edges
    if config.algorithm == "chain":
        edges = _generate_chain(node_ids)
    elif config.algorithm == "random_geometric":
        radius = config.params.get("radius", 0.5)
        edges = _generate_random_geometric(node_ids, radius, rng)
    elif config.algorithm == "erdos_renyi":
        p = config.params.get("p", 0.3)
        edges = _generate_erdos_renyi(node_ids, p, rng)
    elif config.algorithm == "explicit":
        adjacency = config.params.get("adjacency")
        if not adjacency:
            raise ValueError("explicit topology requires params.adjacency")
        edges, edge_transport, dual_udp_edges = _generate_explicit(
            adjacency, config.default_transport
        )
        # Validate all referenced nodes exist
        for a, b in edges:
            if a not in nodes:
                raise ValueError(f"explicit adjacency references unknown node {a}")
            if b not in nodes:
                raise ValueError(f"explicit adjacency references unknown node {b}")
    else:
        raise ValueError(f"Unknown algorithm: {config.algorithm}")

    # Assign transport types to edges
    if config.algorithm != "explicit":
        edge_transport = _assign_edge_transports(edges, config, rng)
        dual_udp_edges = set()

    # Build peer lists from edges
    for a, b in edges:
        nodes[a].peers.append(b)
        nodes[b].peers.append(a)

    # Read the environment once, here, so every name a run produces comes
    # from the same value.
    topo = SimTopology(
        nodes=nodes,
        edges=edges,
        edge_transport=edge_transport,
        dual_udp_edges=dual_udp_edges,
        name_suffix=name_suffix(),
    )

    # Connectivity check with retry
    if config.ensure_connected:
        max_retries = 50
        attempt = 0
        while not topo.is_connected() and attempt < max_retries:
            attempt += 1
            # Clear and regenerate
            for node in nodes.values():
                node.peers.clear()

            if config.algorithm == "random_geometric":
                edges = _generate_random_geometric(node_ids, radius, rng)
            elif config.algorithm == "erdos_renyi":
                edges = _generate_erdos_renyi(node_ids, p, rng)
            else:
                break  # chain is always connected

            for a, b in edges:
                nodes[a].peers.append(b)
                nodes[b].peers.append(a)

            topo.edges = edges
            topo.edge_transport = _assign_edge_transports(edges, config, rng)

        if not topo.is_connected():
            raise RuntimeError(
                f"Failed to generate connected topology after {max_retries} attempts"
            )

    return topo


def _generate_chain(node_ids: list[str]) -> set[tuple[str, str]]:
    """Linear topology: n01-n02-n03-..."""
    edges = set()
    for i in range(len(node_ids) - 1):
        edge = _make_edge(node_ids[i], node_ids[i + 1])
        edges.add(edge)
    return edges


def _generate_random_geometric(
    node_ids: list[str],
    radius: float,
    rng: random.Random,
) -> set[tuple[str, str]]:
    """Place nodes randomly in [0,1]^2, connect if distance < radius."""
    positions = {nid: (rng.random(), rng.random()) for nid in node_ids}
    edges = set()
    for i, a in enumerate(node_ids):
        for b in node_ids[i + 1 :]:
            ax, ay = positions[a]
            bx, by = positions[b]
            dist = math.sqrt((ax - bx) ** 2 + (ay - by) ** 2)
            if dist < radius:
                edges.add(_make_edge(a, b))
    return edges


def _generate_erdos_renyi(
    node_ids: list[str],
    p: float,
    rng: random.Random,
) -> set[tuple[str, str]]:
    """Include each edge with probability p."""
    edges = set()
    for i, a in enumerate(node_ids):
        for b in node_ids[i + 1 :]:
            if rng.random() < p:
                edges.add(_make_edge(a, b))
    return edges


def _generate_explicit(
    adjacency: list, default_transport: str = "udp"
) -> tuple[set[tuple[str, str]], dict[tuple[str, str], str], set[tuple[str, str]]]:
    """Build edges from an explicit adjacency list.

    Each entry is a 2-element list ``[nodeA, nodeB]`` (uses default
    transport) or a 3-element list ``[nodeA, nodeB, transport]``. The
    transport ``ethernet+udp`` declares a dual edge: an Ethernet veth and
    a UDP static-peer link between the same two nodes, so the pair holds
    two paths under one session. ``udp-veth+udp`` is the all-IP dual edge:
    UDP over a dedicated veth (an interface-bound UDP instance at each
    end) and UDP over the bridge.

    Returns ``(edges, edge_transport, dual_udp_edges)`` where
    ``edge_transport`` maps each edge to its transport type (``ethernet``
    for a dual edge) and ``dual_udp_edges`` is the set of dual edges.
    """
    edges = set()
    edge_transport: dict[tuple[str, str], str] = {}
    dual_udp_edges: set[tuple[str, str]] = set()
    for i, entry in enumerate(adjacency):
        if not isinstance(entry, (list, tuple)) or len(entry) not in (2, 3):
            raise ValueError(
                f"explicit adjacency[{i}]: expected [nodeA, nodeB] or "
                f"[nodeA, nodeB, transport], got {entry}"
            )
        edge = _make_edge(str(entry[0]), str(entry[1]))
        edges.add(edge)
        transport = str(entry[2]) if len(entry) == 3 else default_transport
        if transport in ("ethernet+udp", f"{UDP_VETH}+udp"):
            transport = transport[: -len("+udp")]
            dual_udp_edges.add(edge)
        edge_transport[edge] = transport
    return edges, edge_transport, dual_udp_edges


def _assign_edge_transports(
    edges: set[tuple[str, str]],
    config: TopologyConfig,
    rng: random.Random,
) -> dict[tuple[str, str], str]:
    """Assign transport types to edges.

    If ``config.transport_mix`` is set, each edge is randomly assigned
    a transport based on the mix weights. Otherwise all edges use
    ``config.default_transport``.
    """
    if config.transport_mix is None:
        return {e: config.default_transport for e in edges}

    transports = list(config.transport_mix.keys())
    weights = [config.transport_mix[t] for t in transports]
    assignments = rng.choices(transports, weights=weights, k=len(edges))
    return dict(zip(sorted(edges), assignments))


def veth_interface_name(local: str, peer: str) -> str:
    """Generate the veth interface name inside a container.

    Format: ``ve-{local}-{peer}`` (max 15 chars for IFNAMSIZ).
    For typical node IDs like "n01", this yields "ve-n01-n02" (10 chars).
    """
    name = f"ve-{local}-{peer}"
    if len(name) > 15:
        raise ValueError(f"veth interface name too long: {name!r} ({len(name)} > 15)")
    return name


def _make_edge(a: str, b: str) -> tuple[str, str]:
    """Canonical edge representation (sorted)."""
    return (min(a, b), max(a, b))