"""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'")