nest/uvm/network.py

"""Host bridge, TAP, NAT, and address allocation for local microVMs."""

from __future__ import annotations

import json
import re
from collections.abc import Iterable
from ipaddress import IPv4Address

from .config import Settings
from .domain import VmRecord
from .errors import TapCreationError, UvmError, ValidationError
from .system import CommandRunner


class NetworkManager:
    """Manage the shared bridge and one TAP device per persisted VM."""

    def __init__(self, settings: Settings, runner: CommandRunner) -> None:
        self._settings = settings
        self._runner = runner

    def allocate_ip(
        self,
        vms: Iterable[VmRecord],
        requested: IPv4Address | None,
    ) -> IPv4Address:
        used = {vm.guest_ip for vm in vms}
        if requested is not None:
            if requested not in self._settings.network or requested in {
                self._settings.network.network_address,
                self._settings.network.broadcast_address,
                self._settings.gateway,
            }:
                raise ValidationError(
                    "guest IP must be an unused address inside "
                    f"{self._settings.network}, excluding {self._settings.gateway}"
                )
            if str(requested) in used:
                raise ValidationError(f"IP already allocated: {requested}")
            return requested

        for host in self._settings.network.hosts():
            if host == self._settings.gateway:
                continue
            if str(host) not in used:
                return host
        raise UvmError(f"no free IPs in {self._settings.network}")

    @staticmethod
    def mac_for(index: int) -> str:
        if index < 1 or index > 0xFFFFFF:
            raise UvmError("no free locally administered MAC addresses remain")
        return (
            f"02:fc:00:{(index >> 16) & 255:02x}:"
            f"{(index >> 8) & 255:02x}:{index & 255:02x}"
        )

    @staticmethod
    def tap_name(vm_id: str) -> str:
        suffix = re.sub(r"[^a-zA-Z0-9]", "", vm_id)[-11:]
        if not suffix:
            raise UvmError("cannot derive a TAP name from an empty VM ID")
        return f"uvm-{suffix}"[:15]

    def ensure_bridge(self) -> None:
        bridge = self._settings.bridge
        existing = self._runner.run(("ip", "link", "show", bridge), check=False, capture=True)
        if existing.returncode != 0:
            self._runner.run(("ip", "link", "add", bridge, "type", "bridge"))
        else:
            self._ensure_existing_link_is_bridge(bridge)
        self._runner.run(
            (
                "ip",
                "addr",
                "replace",
                f"{self._settings.gateway}/{self._settings.network.prefixlen}",
                "dev",
                bridge,
            )
        )
        self._runner.run(("ip", "link", "set", bridge, "up"))
        self._runner.run(("sysctl", "-w", "net.ipv4.ip_forward=1"))
        self._ensure_masquerade()

    def _ensure_existing_link_is_bridge(self, bridge: str) -> None:
        details = self._runner.run(
            ("ip", "-j", "-d", "link", "show", "dev", bridge), capture=True
        )
        try:
            links = json.loads(details.stdout)
            kind = links[0]["linkinfo"]["info_kind"]
        except (IndexError, KeyError, TypeError, json.JSONDecodeError) as error:
            raise UvmError(f"could not determine whether existing interface {bridge} is a bridge") from error
        if kind != "bridge":
            raise UvmError(f"configured bridge {bridge} exists but is not a Linux bridge")

    def create_tap(self, name: str) -> None:
        existing = self._runner.run(("ip", "link", "show", name), check=False, capture=True)
        if existing.returncode == 0:
            raise UvmError(f"TAP device already exists: {name}")
        created = False
        try:
            self._runner.run(("ip", "tuntap", "add", "dev", name, "mode", "tap"))
            created = True
            self._runner.run(("ip", "link", "set", name, "master", self._settings.bridge))
            self._runner.run(("ip", "link", "set", name, "up"))
        except BaseException as error:
            if created:
                try:
                    self.delete_tap(name)
                except UvmError as cleanup_error:
                    raise TapCreationError(
                        f"could not configure TAP device {name}; cleanup also failed: {cleanup_error}",
                        tap_created=True,
                    ) from error
            raise

    def delete_tap(self, name: str) -> None:
        deleted = self._runner.run(("ip", "link", "del", name), check=False, capture=True)
        if deleted.returncode == 0:
            return
        remaining = self._runner.run(("ip", "link", "show", name), check=False, capture=True)
        if remaining.returncode != 0:
            return
        detail = deleted.stderr.strip() or deleted.stdout.strip()
        suffix = f": {detail}" if detail else ""
        raise UvmError(f"could not delete TAP device {name}{suffix}")

    def _ensure_masquerade(self) -> None:
        route = self._runner.run(("ip", "route", "show", "default"), capture=True)
        match = re.search(r"\bdev\s+(\S+)", route.stdout)
        if match is None:
            return
        uplink = match.group(1)
        rule = (
            "iptables",
            "-t",
            "nat",
            "-C",
            "POSTROUTING",
            "-s",
            str(self._settings.network),
            "-o",
            uplink,
            "-j",
            "MASQUERADE",
        )
        present = self._runner.run(rule, check=False)
        if present.returncode != 0:
            self._runner.run(
                (
                    "iptables",
                    "-t",
                    "nat",
                    "-A",
                    "POSTROUTING",
                    "-s",
                    str(self._settings.network),
                    "-o",
                    uplink,
                    "-j",
                    "MASQUERADE",
                )
            )
        self._ensure_iptables_rule(
            (
                "iptables",
                "-C",
                "FORWARD",
                "-i",
                self._settings.bridge,
                "-o",
                uplink,
                "-s",
                str(self._settings.network),
                "-j",
                "ACCEPT",
            ),
            (
                "iptables",
                "-A",
                "FORWARD",
                "-i",
                self._settings.bridge,
                "-o",
                uplink,
                "-s",
                str(self._settings.network),
                "-j",
                "ACCEPT",
            ),
        )
        self._ensure_iptables_rule(
            (
                "iptables",
                "-C",
                "FORWARD",
                "-i",
                uplink,
                "-o",
                self._settings.bridge,
                "-d",
                str(self._settings.network),
                "-m",
                "conntrack",
                "--ctstate",
                "ESTABLISHED,RELATED",
                "-j",
                "ACCEPT",
            ),
            (
                "iptables",
                "-A",
                "FORWARD",
                "-i",
                uplink,
                "-o",
                self._settings.bridge,
                "-d",
                str(self._settings.network),
                "-m",
                "conntrack",
                "--ctstate",
                "ESTABLISHED,RELATED",
                "-j",
                "ACCEPT",
            ),
        )

    def _ensure_iptables_rule(
        self,
        check_rule: tuple[str, ...],
        add_rule: tuple[str, ...],
    ) -> None:
        present = self._runner.run(check_rule, check=False)
        if present.returncode != 0:
            self._runner.run(add_rule)