utun2/tests/stream_netns/nat_scenario.py

"""Проверка детектора uTun на настоящем NAT; контроль не создаёт контакт с помощником."""

import json

from topology import NATS, NODES


def profile(statistics, protocol):
    return next(item for item in statistics["address_profiles"] if item["socket"] == protocol)


def check_profiles(scenario):
    agents = scenario.agents
    public = [node for node in NODES if not node.nat]
    transport = "tcp" if scenario.args.tcp_only else "udp"
    connected = set()

    def connect(source, target):
        pair = (source.name, target.name)
        if pair in connected:
            return
        snapshot = agents[target.name].rpc("snapshot")["snapshot"]
        agents[source.name].rpc("nat_link", snapshot=snapshot, socket=transport)
        target_id = agents[target.name].identity
        scenario.wait(f"nat-link-{source.name}-{target.name}",
                      lambda data: target_id in data[source.name]["direct_up"], 20)
        connected.add(pair)

    def store_owner(owner, observer):
        agents[observer.name].rpc("nat_store", snapshot=agents[owner.name].rpc("snapshot")["snapshot"])

    def measure(owner, observer, protocol):
        before = profile(agents[owner.name].rpc("stats"), protocol)
        agents[owner.name].rpc("nat_check", socket=protocol, observer=agents[observer.name].identity)
        data = scenario.wait(f"nat-profile-{owner.name}-{protocol}", lambda statistics:
                            profile(statistics[owner.name], protocol)["rounds"] > before["rounds"] and
                            profile(statistics[owner.name], protocol)["phase"] == 0, 40)
        return profile(data[owner.name], protocol)

    # Отсутствие второго наблюдателя не превращается в STRICT/DIRECT.
    first, second = public[:2]
    connect(first, second)
    if not scenario.args.tcp_only:
        incomplete = measure(first, second, "udp")
        if (incomplete["mapping"], incomplete["inbound"], incomplete["reason"]) != ("UNKNOWN", "UNKNOWN", "no_helper"):
            raise RuntimeError(f"NAT no-helper test: {incomplete}")

    verified = {}
    for owner in public + [node for node in NODES if node.nat]:
        observer = second if owner.name == first.name else first
        if not scenario.args.tcp_only:
            contacts = set(profile(agents[owner.name].rpc("stats"), "udp")["contacts"])
            helper = next((node for node in public if node.name not in {owner.name, observer.name} and node.ip not in contacts), None)
            if not helper:
                raise RuntimeError(f"{owner.name}: no independent helper remains before cold measurement; contacts={contacts}")
            connect(observer, helper)
        connect(owner, observer)
        store_owner(owner, observer)
        protocols = ["tcp"] if scenario.args.tcp_only else ["udp", "tcp"] if scenario.args.tcp else ["udp"]
        observed = {}
        for protocol in protocols:
            if protocol == "tcp" and owner.nat:
                # Приватный TCP listener не имеет обнаруженного публичного endpoint.
                agents[owner.name].rpc("nat_check", socket="tcp", observer=agents[observer.name].identity)
                data = scenario.wait(f"nat-private-tcp-{owner.name}", lambda statistics:
                                    profile(statistics[owner.name], "tcp")["reason"] == "no_listener", 5)
                actual = profile(data[owner.name], "tcp")
                expected = ("UNKNOWN", "UNKNOWN")
            else:
                actual = measure(owner, observer, protocol)
                if protocol == "tcp" or not owner.nat:
                    expected = ("UNCHANGED", "OPEN")
                else:
                    gateway = next(item for item in NATS if item.name == owner.nat)
                    expected = {"cone": ("EIM", "OPEN"), "restricted": ("EIM", "CONTACT_REQUIRED"),
                                "strict": ("DEPENDENT", "UNKNOWN")}[gateway.mode]
            if (actual["mapping"], actual["inbound"]) != expected:
                raise RuntimeError(f"{owner.name}/{protocol}: expected {expected}, detector returned {actual}")
            if expected[0] == "EIM":
                gateway = next(item for item in NATS if item.name == owner.nat)
                external = f"{gateway.ip}:{owner.mapped_port}"
            elif expected[0] == "UNCHANGED":
                external = f"{owner.ip}:32000"
            else:
                external = ""
            if actual["external"] != external:
                raise RuntimeError(f"{owner.name}/{protocol}: foreign/invalid external endpoint: {actual}, expected {external}")
            observed[protocol] = actual
            print(f"NAT {owner.name}/{protocol}: {actual['mapping']} / {actual['inbound']}"
                  f" external={actual['external'] or '-'} reason={actual['reason']}", flush=True)
        verified[owner.name] = observed
        (scenario.network.directory / "nat-profiles.json").write_text(json.dumps(verified, indent=2) + "\n")

    # Повтор на использовавшемся порту сохраняет mapping, но не обновляет OPEN
    # по контактировавшим IP. Проверка запускается после всех холодных измерений.
    first_profile = verified[first.name].get("udp")
    if first_profile:
        again = measure(first, second, "udp")
        if again["mapping"] != "UNCHANGED" or again["inbound"] not in {"OPEN", "UNKNOWN"}:
            raise RuntimeError(f"{first.name}: warm measurement produced invalid classification: {again}")
        if again["inbound"] == "UNKNOWN" and again["reason"] != "contact_history":
            raise RuntimeError(f"{first.name}: UNKNOWN without contact-history explanation: {again}")

    result = {"nodes_checked": len(verified), "protocols": protocols, "no_helper_checked": not scenario.args.tcp_only,
              "warm_recheck_checked": bool(first_profile), "verified": verified}
    (scenario.network.directory / "nat-result.json").write_text(json.dumps(result, indent=2) + "\n")
    return result