thoxmesh_node/tunnel.py

"""Public tunnel for nodes that have no public address of their own.

A Colab VM is not reachable from the internet, so the node's OpenAI surface has
to be published through a tunnel before anything in the mesh can call it. A Space
already has a public origin and needs none of this.

``cloudflared`` is the default because its quick tunnels need no account, no
token and no interactive approval — which matters when the goal is a notebook
that runs top to bottom without the operator stopping to paste an auth token.
ngrok is supported for operators who already have a token and want a stable
hostname.
"""

from __future__ import annotations

import logging
import os
import platform
import re
import shutil
import subprocess
import threading
import time
import urllib.request
from dataclasses import dataclass

from .errors import TunnelError

logger = logging.getLogger("thox.tunnel")

_CLOUDFLARED_URL_RE = re.compile(r"https://[a-z0-9-]+\.trycloudflare\.com")

_CLOUDFLARED_RELEASES = {
    ("Linux", "x86_64"): "https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64",
    ("Linux", "aarch64"): "https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-arm64",
    ("Darwin", "arm64"): "https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-darwin-arm64.tgz",
}


@dataclass
class Tunnel:
    """A running tunnel and the public URL it exposes."""

    public_url: str
    process: subprocess.Popen | None

    def close(self) -> None:
        """Terminate the tunnel. Best effort; never raises."""
        if self.process is None:
            return
        try:
            self.process.terminate()
            self.process.wait(timeout=10)
        except Exception:  # noqa: BLE001
            try:
                self.process.kill()
            except Exception:  # noqa: BLE001
                logger.debug("tunnel kill failed", exc_info=True)


def ensure_cloudflared(install_dir: str = "/usr/local/bin") -> str:
    """Return a path to ``cloudflared``, downloading the binary if needed.

    Only the Linux targets are auto-installed, since that covers Colab and any
    container; elsewhere the operator is expected to have it on PATH, and we say
    so rather than guessing at a package manager.
    """
    existing = shutil.which("cloudflared")
    if existing:
        return existing

    key = (platform.system(), platform.machine())
    url = _CLOUDFLARED_RELEASES.get(key)
    if url is None or not url.endswith(("amd64", "arm64")):
        raise TunnelError(
            f"cloudflared is not on PATH and no automatic install is available for {key}. "
            "Install cloudflared manually, or set THOX_PUBLIC_URL if the node is already reachable."
        )

    target = os.path.join(install_dir if os.access(install_dir, os.W_OK) else "/tmp", "cloudflared")
    logger.info("downloading cloudflared to %s", target)
    try:
        with urllib.request.urlopen(url, timeout=120) as response, open(target, "wb") as handle:
            shutil.copyfileobj(response, handle)
        os.chmod(target, 0o755)
    except Exception as exc:  # noqa: BLE001
        raise TunnelError(f"could not install cloudflared: {exc}") from exc
    return target


def start_cloudflared(port: int, *, timeout: int = 90) -> Tunnel:
    """Start a quick tunnel to ``localhost:port`` and return its public URL.

    cloudflared prints the assigned hostname to stderr once and never repeats it,
    so stderr is drained on a reader thread. Polling the process output after the
    fact would race with the buffer and intermittently miss the URL.
    """
    binary = ensure_cloudflared()
    try:
        process = subprocess.Popen(
            [binary, "tunnel", "--no-autoupdate", "--url", f"http://127.0.0.1:{port}"],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
        )
    except (OSError, subprocess.SubprocessError) as exc:
        # Popen raises rather than returning a code when the binary is missing,
        # not executable, or built for the wrong architecture. Without this the
        # node dies with a bare OSError traceback instead of a usable message.
        raise TunnelError(
            f"could not start cloudflared at {binary}: {exc}. "
            "Check the binary is executable and matches this machine's architecture, "
            "or set THOX_PUBLIC_URL if the node is already reachable."
        ) from exc

    found: dict[str, str] = {}
    done = threading.Event()

    def _read() -> None:
        assert process.stdout is not None
        for line in process.stdout:
            logger.debug("cloudflared: %s", line.rstrip())
            match = _CLOUDFLARED_URL_RE.search(line)
            if match and "url" not in found:
                found["url"] = match.group(0)
                done.set()
        done.set()

    threading.Thread(target=_read, daemon=True, name="cloudflared-reader").start()

    if not done.wait(timeout=timeout) or "url" not in found:
        try:
            process.terminate()
        except Exception:  # noqa: BLE001
            pass
        raise TunnelError(f"cloudflared did not report a public URL within {timeout}s")

    url = found["url"]
    logger.info("tunnel up: %s", url)
    _wait_for_tunnel(url)
    return Tunnel(public_url=url, process=process)


def _wait_for_tunnel(url: str, *, attempts: int = 10, delay: float = 3.0) -> None:
    """Block until the tunnel actually serves our health endpoint.

    Cloudflare prints the hostname before edge propagation finishes. Registering
    in that gap publishes a URL that 530s for the first few seconds, and the
    controller's health check would immediately mark the fresh node unhealthy.
    """
    for attempt in range(1, attempts + 1):
        try:
            with urllib.request.urlopen(f"{url}/health", timeout=10) as response:
                if response.status == 200:
                    return
        except Exception as exc:  # noqa: BLE001
            logger.debug("tunnel probe %d/%d failed: %s", attempt, attempts, exc)
        time.sleep(delay)
    logger.warning("tunnel %s did not pass a health probe; registering anyway", url)


def start_ngrok(port: int, token: str | None = None) -> Tunnel:
    """Start an ngrok tunnel via ``pyngrok``. Requires an authtoken for stability."""
    try:
        from pyngrok import conf, ngrok
    except ImportError as exc:  # pragma: no cover - env dependent
        raise TunnelError("pyngrok is not installed; `pip install pyngrok` or use cloudflared") from exc
    auth = token or os.environ.get("NGROK_AUTHTOKEN", "")
    if auth:
        conf.get_default().auth_token = auth
    try:
        listener = ngrok.connect(port, "http")
    except Exception as exc:  # noqa: BLE001
        raise TunnelError(f"ngrok failed to start: {exc}") from exc
    url = listener.public_url.replace("http://", "https://")
    logger.info("ngrok tunnel up: %s", url)
    _wait_for_tunnel(url)
    return Tunnel(public_url=url, process=None)


def start_tunnel(provider: str, port: int) -> Tunnel:
    """Start the requested tunnel provider."""
    normalised = (provider or "cloudflared").strip().lower()
    if normalised in ("cloudflared", "cloudflare"):
        return start_cloudflared(port)
    if normalised == "ngrok":
        return start_ngrok(port)
    raise TunnelError(f"unknown tunnel provider {provider!r}; use 'cloudflared' or 'ngrok'")