diff --git a/README.md b/README.md index d80333a6..cf18493e 100644 --- a/README.md +++ b/README.md @@ -289,6 +289,56 @@ Provides chunked streaming file reading and writing within the environment conta | `WriteFile` | Streams raw binary or text chunks directly to a target file | +## Reaching a service inside an environment + +A daemon that listens on a port inside an environment, such as one a runtime layer starts next to +the guest ([docs/rootfs-images/RUNTIMES.md](docs/rootfs-images/RUNTIMES.md)), is reachable through +the API's listener. A plain HTTP request that carries `X-Env-Id` and `X-Env-Port` (and +`X-Env-Atespace` outside `ate-env`) is forwarded through the Substrate router to that port, whatever +its path; the API strips the three headers. The router resumes a suspended environment first, +and the proxy waits for the wake up to `--wake-budget` (the same wait as a guest call) before it +answers 503. + +```bash +curl -H 'X-Env-Id: dev1' -H 'X-Env-Port: 6000' localhost:7777/health +``` + +Clients that cannot set headers, such as WebSocket clients and browsers, put the target in the path +instead: `/v1alpha/envs/{id}/ports/{port}/...`, or `/v1alpha/atespaces/{atespace}/envs/{id}/ports/{port}/...` +outside `ate-env`. The API strips that prefix and forwards the rest, so a client given +`http://localhost:7777/v1alpha/envs/dev1/ports/6000` as its base URL works unchanged. WebSocket upgrades pass +through either form; gRPC to the daemon is not supported yet. + +```bash +curl localhost:7777/v1alpha/envs/dev1/ports/6000/health +``` + +The path form reaches the daemon exactly as sent: a path without the trailing slash is the daemon's +root, and `//` or `..` after the prefix are not cleaned or redirected. + +**Who can use it.** The forwarding does not authenticate callers. Whoever can reach the API's +listener can reach any TCP port of any actor the Substrate router serves, in any atespace, +including actors that ate-env did not create. The guest API already reaches each actor that runs +the guest; the forwarding extends that reach to every port. Keep the listener inside the cluster +(the manifest's Service is `ClusterIP`), restrict it with a NetworkPolicy where other workloads +share the cluster, and make each daemon check its own credential. + +**Browsers.** The path form refuses a state-changing request or a WebSocket handshake that a browser +sends from a page of another origin, so a web page cannot drive a daemon through a `kubectl +port-forward`. A page served by the forwarded daemon itself is same-origin and passes. All daemons +behind one listener share that origin, thus a page from one daemon can call another. +`--port-proxy-trusted-origins` adds origins, such as a UI on its own host. The header form needs no +check: a browser sends custom headers to another origin only after a CORS preflight, which the API +never approves. DNS rebinding, where a page on an attacker's name resolves to the forwarded +listener, is not blocked yet; it needs a check of the `Host` header across the whole listener, the +MCP endpoint included. + +| Flag | Default | Effect | +|---|---|---| +| `--atenet-connect` | `atenet-router.ate-system.svc.cluster.local:8081` | The router's CONNECT listener. An empty value turns the forwarding off. | +| `--port-proxy-trusted-origins` | none | Comma-separated browser origins allowed to use the path form. | +| `--port-proxy-response-header-timeout` | `0`, no limit | Bound on the wait for a daemon's response headers. Without it, the caller's deadline ends the wait. | + ## Built-in MCP Server The API exposes a streamable MCP endpoint at `POST /v1alpha/envs/{id}/mcp` serving built-in tools. diff --git a/cmd/ate-env-api/main.go b/cmd/ate-env-api/main.go index ad1a5f0b..666543ff 100644 --- a/cmd/ate-env-api/main.go +++ b/cmd/ate-env-api/main.go @@ -23,12 +23,14 @@ import ( "net/http" "os" "os/signal" + "strings" "syscall" "time" "github.com/agent-substrate/env/internal/apiservice" "github.com/agent-substrate/env/internal/ate" "github.com/agent-substrate/env/internal/mcp" + "github.com/agent-substrate/env/internal/portproxy" ateenvv1alpha "github.com/agent-substrate/env/proto/ateenv/v1alpha" "google.golang.org/grpc" ) @@ -38,10 +40,13 @@ func main() { listen = flag.String("listen", "0.0.0.0:7777", "address to serve the API on") ateapi = flag.String("ateapi", ate.DefaultControlAddr, "address of the ateapi gRPC control plane") atenet = flag.String("atenet", ate.DefaultRouterAddr, "address of the atenet HTTP router") + atenetConn = flag.String("atenet-connect", ate.DefaultRouterConnectAddr, "address of the atenet router's plaintext CONNECT listener; HTTP requests carrying X-Env-Id and X-Env-Port, or to /v1alpha/envs/{id}/ports/{port}/..., are forwarded through it to that port in the environment. The forwarding does not authenticate callers and reaches any port of any actor the router serves. Empty disables it") + ppOrigins = flag.String("port-proxy-trusted-origins", "", "comma-separated browser origins (scheme://host[:port]) whose pages may use the port proxy's path form, next to the listener's own origin") + ppHeaderTimeout = flag.Duration("port-proxy-response-header-timeout", 0, "how long the port proxy waits for a daemon's response headers; 0 means no limit, and the caller's deadline applies") hostSuffix = flag.String("host-suffix", ate.DefaultHostSuffix, "atenet router host suffix for actor routing") skipVerify = flag.Bool("skip-verify", true, "skip TLS certificate verification on the control plane connection") tokenFile = flag.String("ateapi-token-file", "", "file with a bearer token for the control plane, e.g. a projected ServiceAccount token with the ateapi audience") - wakeBudget = flag.Duration("wake-budget", apiservice.DefaultWakeBudget, "how long a guest call waits for a suspended environment to wake before failing; cold image pulls and cross-region snapshot reads exceed the router's own parking budget") + wakeBudget = flag.Duration("wake-budget", apiservice.DefaultWakeBudget, "how long a guest call or a port-proxy request waits for a suspended environment to wake before failing; cold image pulls and cross-region snapshot reads exceed the router's own parking budget") wakeRunningBudget = flag.Duration("wake-running-budget", apiservice.DefaultWakeRunningBudget, "how long a guest call waits when the environment is already RUNNING but the router still answers Unavailable (guest not yet listening vs. path down)") ) flag.Parse() @@ -86,6 +91,31 @@ func main() { mux.Handle("/v1alpha/envs/{id}/mcp", mcp.NewHandler(client)) mux.Handle("/", grpcServer) + // Requests that name an environment port go to the port proxy before + // anything else, whatever their path: the path belongs to the daemon. + var handler http.Handler = mux + if *atenetConn != "" { + var origins []string + for o := range strings.SplitSeq(*ppOrigins, ",") { + if o = strings.TrimSpace(o); o != "" { + origins = append(origins, o) + } + } + pp, err := portproxy.New(portproxy.Options{ + ConnectAddr: *atenetConn, + HostSuffix: *hostSuffix, + DefaultAtespace: ate.DefaultAtespace, + ResponseHeaderTimeout: *ppHeaderTimeout, + TrustedOrigins: origins, + WakeBudget: *wakeBudget, + }) + if err != nil { + log.Fatalf("--port-proxy-trusted-origins: %v", err) + } + // Header form and path form, /v1alpha/envs/{id}/ports/{port}/... + handler = portproxy.Dispatch(pp, mux) + } + // Enable both HTTP/1.1 (for /healthz and /v1alpha/ HTTP endpoints) and // unencrypted HTTP/2 (h2c, for gRPC services) on the same plaintext TCP listener. var protocols http.Protocols @@ -94,7 +124,7 @@ func main() { srv := &http.Server{ Addr: *listen, - Handler: mux, + Handler: handler, Protocols: &protocols, } @@ -107,7 +137,7 @@ func main() { _ = srv.Shutdown(context.Background()) }() - log.Printf("ate-env-api listening on %s (ateapi %s, atenet %s)", *listen, *ateapi, *atenet) + log.Printf("ate-env-api listening on %s (ateapi %s, atenet %s, atenet CONNECT %q)", *listen, *ateapi, *atenet, *atenetConn) if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { log.Fatalf("server error: %v", err) } diff --git a/docs/rootfs-images/DESIGN.md b/docs/rootfs-images/DESIGN.md index b5402167..fd9dfcf6 100644 --- a/docs/rootfs-images/DESIGN.md +++ b/docs/rootfs-images/DESIGN.md @@ -193,10 +193,10 @@ one golden snapshot per image. Neither is reclaimed automatically. with distinct images, not with environments. The plan is a last-used annotation updated on each create, collection by age or by last use, and the golden snapshot removed alongside. -- **Reaching a second runtime from outside.** Layers and sidecars put a - runtime such as OpenSandbox's `execd` inside the actor and gate readiness on - it; an endpoint lookup on `ate-env-api` that names the runtime's port - through the router is what an external SDK still needs. +- **gRPC and raw TCP to a second runtime.** `ate-env-api` forwards HTTP to a + runtime's port through the router's CONNECT listener (see + [RUNTIMES.md](RUNTIMES.md#who-can-reach-the-second-runtime)); a runtime that + speaks gRPC or a non-HTTP protocol needs an h2c upstream or a raw tunnel. - **Per-create overrides** for resources and workspace, so one base can serve images with different needs. - **Image validation at create time**, so a bad reference fails the call diff --git a/docs/rootfs-images/RUNTIMES.md b/docs/rootfs-images/RUNTIMES.md index fa662d1f..cd50f31d 100644 --- a/docs/rootfs-images/RUNTIMES.md +++ b/docs/rootfs-images/RUNTIMES.md @@ -84,14 +84,38 @@ sidecar, run under two different rootfs images. away. An agent framework whose tool calls run inside the sandbox is served. - **The guest's own APIs** are unaffected: exec, files and MCP keep going through `ate-env-api` on port 80. -- **Clients outside the environment** currently reach only port 80. The atenet - router forwards to the actor's port 80 unless the request is a CONNECT whose - authority names another port, and `ate-env-api` proxies only the guest. An - SDK that speaks HTTP to the second runtime therefore needs one of two things, - neither of which is in this branch: an endpoint lookup on `ate-env-api` that - returns a router address naming the runtime's port, or a reverse proxy. This - is the next step for the OpenSandbox integration, whose lifecycle server - resolves sandbox endpoints exactly this way. +- **Clients outside the environment** reach it through `ate-env-api`. A plain + HTTP request to the API's listener that carries `X-Env-Id`, `X-Env-Port` and + optionally `X-Env-Atespace` goes to that port in that environment, whatever + its path: `ate-env-api` opens a CONNECT tunnel through the atenet router's + CONNECT listener (`--atenet-connect`, port 8081 on a stock install), the + router resumes the actor if needed and forwards to the port, and the request + and its streamed response pass through unchanged apart from the routing + headers, which are stripped. Idle tunnels are reused, one pool for each + environment and port. + + ```bash + curl -H 'X-Env-Id: py1' -H 'X-Env-Port: 44772' localhost:7777/ready + curl localhost:7777/v1alpha/envs/py1/ports/44772/ready # the same, for clients without headers + ``` + + The path form, `/v1alpha/envs/{id}/ports/{port}/...` (or + `/v1alpha/atespaces/{atespace}/envs/...`), is for clients that cannot set + headers: the prefix is stripped, so it works as a base URL. The tunnel + carries HTTP/1.1, not raw TCP. Streamed responses are flushed as they arrive + and WebSocket upgrades pass through; gRPC to the runtime is not supported + yet. The NeMo Gym + provider hands this out as its `endpoint(handle, port)` + ([integrations/nemo-gym](../../integrations/nemo-gym/README.md)). +- **The runtime must listen on `0.0.0.0`.** The router reaches the port on + the actor's address. A runtime bound to `127.0.0.1` serves only processes + inside the environment, and requests from outside fail with 502. +- **Nobody is authenticated on the way in.** `ate-env-api` forwards for any + caller that reaches its listener, to any port of any actor the router + serves. A runtime that executes commands or edits files must check its own + credential, as OpenSandbox's `execd` token and the OpenHands agent-server's + session key do. See the README section "Reaching a service inside an + environment" for the browser checks and the flags. ## Suspend, resume and restarts diff --git a/integrations/nemo-gym/README.md b/integrations/nemo-gym/README.md index 68b247dd..77cb0eaa 100644 --- a/integrations/nemo-gym/README.md +++ b/integrations/nemo-gym/README.md @@ -206,6 +206,85 @@ To share one `ate_env.Client` between the provider and other code, pass it in: Unknown keys are rejected at `create()` time, as are unknown keys anywhere in the config block. +### Services inside the sandbox: an injected daemon + +Some NeMo components talk to a service inside the sandbox instead of running commands through +`exec`: an exec server, a code-execution sandbox server, a tool daemon. The task images do not +contain it, and rebuilding hundreds of thousands of SWE-bench images to add it is not an option. +ate-env attaches such a daemon at runtime the same way it attaches its own guest: as a read-only +OCI image volume, a *runtime layer*, that the guest starts and supervises next to itself +([docs/rootfs-images/RUNTIMES.md](../../docs/rootfs-images/RUNTIMES.md)). The provider implements +nemo-gym 0.7's `SupportsSandboxEndpoint`, so Gym reaches the daemon from outside the cluster's +actor network. + +1. **Put the daemon on the base template, once.** Every environment created from this base with a + task image (`rootfs_image`) gets the daemon; the task images stay unmodified. + + ```bash + ate-env manifest template --template gym-base \ + --layer nemo=/nemo-daemon@sha256:=/opt/nemo \ + --sidecar '/opt/nemo/daemon --port 6000' \ + --sidecar-readyz http://127.0.0.1:6000/health \ + --snapshots-bucket gs:///ate-env/ | kubectl-ate create actor-template -f - + ``` + + The actor reports ready only once the daemon answers its readiness URL, so `create()` returns + a sandbox whose daemon is up. + +2. **Select the base and declare the port.** + + ```yaml + sandbox: + substrate: + connection: + api_url: ate-env-api.ate-env:7777 + # endpoint_url: http://ate-env-api.ate-env:7777 # default: api_url + create: + template: gym-base + ``` + + ```python + handle = await provider.create(SandboxSpec(image=task_image, ports=(6000,))) + ``` + +3. **Resolve the endpoint and call the daemon.** + + ```python + ep = await provider.endpoint(handle, 6000) + async with httpx.AsyncClient(base_url=ep.endpoint, headers=ep.headers) as http: + r = await http.post("/execute", json={"generated_code": "print(1)", "language": "python"}) + ``` + + `ep.endpoint` is ate-env-api, and `ep.headers` (`X-Env-Id`, `X-Env-Atespace`, `X-Env-Port`) + route each request through the Substrate router to port 6000 in this sandbox. `endpoint()` does + not touch the guest: the first request through the proxy resumes a parked sandbox, and the proxy + waits for the wake (ate-env-api's `--wake-budget`, two minutes by default, the same wait as a + guest call). It checks only that the sandbox exists, and raises `RuntimeError` at once for one + that is deleted or crashed. A port that `SandboxSpec.ports` did not declare is a `ValueError`, + as in Gym's Docker provider; `create()` rejects a spec whose ports are outside 1-65535 or name + the guest's port 80. Set `connection.endpoint_url` when the code that calls the daemon reaches + ate-env-api at another address than the provider does. + +Limits: + +- **HTTP/1.1 only.** The router tunnels HTTP to the port, not raw TCP. Streamed responses are + flushed as they arrive and WebSocket upgrades pass through; gRPC to the daemon is not supported + yet. +- **The daemon must be self-contained.** A layer cannot rely on the task image's libraries: ship a + static binary, or put the interpreter and its packages in the layer. NeMo Skills' sandbox server + (Flask on CPython, port 6000, `POST /execute`, `GET /health`) fits only with a relocatable Python + inside its layer; Gym's ECS exec server needs `python3` in the task image. +- **The daemon must listen on `0.0.0.0`.** The router reaches the port on the actor's address; a + daemon bound to `127.0.0.1` only answers processes inside the sandbox, and requests through the + endpoint fail with 502. +- **No authentication in front of the daemon.** Anyone who can reach ate-env-api can call any port + in any sandbox through it. Keep ate-env-api inside the cluster, and give the daemon its own + credential when the sandbox runs untrusted code. +- Port 80 belongs to the guest. The daemon shares the task image's user, environment and + filesystem. +- Each distinct task image yields one derived ActorTemplate on its base, created on first use and + reused after; the daemon adds no templates. + ## Testing ```bash @@ -275,7 +354,6 @@ Known limits (tracked in the RL-on-Substrate proposal): - `spec.ttl_s` is **not enforced** — substrate has no server-side TTL yet; clean up via `close()`. A crashed harness leaks actors until deleted. - `exec(user=…)` raises: commands run as the actor's configured user. -- `spec.ports` / `SandboxEndpoint` exposure is not implemented in this demo. ## Troubleshooting diff --git a/integrations/nemo-gym/src/nemo_gym_substrate/_compat.py b/integrations/nemo-gym/src/nemo_gym_substrate/_compat.py index 2d21be8d..c11eb635 100644 --- a/integrations/nemo-gym/src/nemo_gym_substrate/_compat.py +++ b/integrations/nemo-gym/src/nemo_gym_substrate/_compat.py @@ -82,10 +82,41 @@ class SandboxCreateVerificationError(SandboxCreateError): pass +# SandboxEndpoint is newer than the rest (nemo-gym 0.7): an installed 0.5 +# has the base module without it, so it falls back on its own. +try: # pragma: no cover - exercised only with nemo-gym >= 0.7 + from nemo_gym.sandbox.providers.base import SandboxEndpoint # type: ignore +except ImportError: # pragma: no cover - the in-repo test path and nemo-gym 0.5 + from collections.abc import Mapping as _Mapping + from dataclasses import dataclass as _dataclass + from dataclasses import field as _field + from urllib.parse import urlsplit as _urlsplit + + @_dataclass(frozen=True) + class SandboxEndpoint: # type: ignore[no-redef] + endpoint: str + headers: dict[str, str] = _field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.endpoint, str) or not self.endpoint.strip(): + raise ValueError("Sandbox endpoint must be a non-empty absolute URL") + endpoint = self.endpoint.strip() + parsed = _urlsplit(endpoint) + if not parsed.scheme or not parsed.netloc: + raise ValueError("Sandbox endpoint must be a non-empty absolute URL") + if not isinstance(self.headers, _Mapping): + raise TypeError("Sandbox endpoint headers must be a mapping") + object.__setattr__(self, "endpoint", endpoint) + object.__setattr__( + self, "headers", {str(k): str(v) for k, v in self.headers.items()} + ) + + __all__ = [ "NEMO_GYM_AVAILABLE", "SandboxCreateError", "SandboxCreateVerificationError", + "SandboxEndpoint", "SandboxExecResult", "SandboxHandle", "SandboxSpec", diff --git a/integrations/nemo-gym/src/nemo_gym_substrate/provider.py b/integrations/nemo-gym/src/nemo_gym_substrate/provider.py index 0b99a2d2..89940311 100644 --- a/integrations/nemo-gym/src/nemo_gym_substrate/provider.py +++ b/integrations/nemo-gym/src/nemo_gym_substrate/provider.py @@ -46,6 +46,7 @@ from ._compat import ( SandboxCreateError, SandboxCreateVerificationError, + SandboxEndpoint, SandboxExecResult, SandboxHandle, SandboxSpec, @@ -69,6 +70,35 @@ # Exit code the guest reports for a process group it killed with SIGKILL. _SIGKILL_EXIT = 137 +# Routing headers of ate-env-api's port proxy (internal/portproxy): a plain +# HTTP request carrying them is forwarded to that port inside the environment. +_ENV_ID_HEADER = "X-Env-Id" +_ENV_ATESPACE_HEADER = "X-Env-Atespace" +_ENV_PORT_HEADER = "X-Env-Port" +# The ate-env guest's own port in every environment: not a daemon port, and +# the port proxy refuses it too. +_GUEST_PORT = 80 + + +def _validate_ports(ports: object) -> list[int]: + """``SandboxSpec.ports`` as a list of TCP ports a daemon can listen on. + + Checked before the actor starts, so a bad spec fails at once rather than + at ``endpoint()`` time. ``bool`` is excluded explicitly because it is an + ``int`` subclass in Python (``True == 1``).""" + out: list[int] = [] + for p in ports or (): + if isinstance(p, bool) or not isinstance(p, int): + raise TypeError(f"SandboxSpec.ports entries must be integers, got {p!r}") + if not 1 <= p <= 65535: + raise ValueError(f"SandboxSpec.ports entry {p} is not a TCP port (1-65535)") + if p == _GUEST_PORT: + raise ValueError( + f"SandboxSpec.ports entry {p} is the ate-env guest's port; a daemon listens on another port" + ) + out.append(p) + return out + # Env lifecycle status -> Gym status. Parked (suspended / paused) actors report # RUNNING: the router resumes them on the next guest call, which is what a # rollout loop wants to see (see README, "Resource mapping"). @@ -116,17 +146,31 @@ class ConnectionConfig: api_url: str = "http://127.0.0.1:7777" # Deadline for each lifecycle RPC (create / get / delete). request_timeout_s: float = 30.0 + # Base URL that endpoint() hands out for services inside a sandbox. Unset + # means api_url as an http(s) URL: ate-env-api serves its port proxy on the + # same listener. Set it when the code that calls the service reaches + # ate-env-api at another address than the provider does. + endpoint_url: str | None = None - _ALLOWED = frozenset({"api_url", "request_timeout_s"}) + _ALLOWED = frozenset({"api_url", "request_timeout_s", "endpoint_url"}) @classmethod def from_mapping(cls, options: Mapping[str, Any]) -> "ConnectionConfig": _require_keys(options, cls._ALLOWED, "sandbox.substrate.connection") + endpoint_url = options.get("endpoint_url") return cls( api_url=str(options.get("api_url", cls.api_url)).rstrip("/"), request_timeout_s=float(options.get("request_timeout_s", cls.request_timeout_s)), + endpoint_url=None if endpoint_url is None else str(endpoint_url).rstrip("/"), ) + def endpoint_base(self) -> str: + """The absolute URL of ate-env-api's port proxy: ``endpoint_url`` when + set, else ``api_url``; either gets ``http://`` when it names no scheme, + because ``SandboxEndpoint`` rejects a bare ``host:port``.""" + base = self.endpoint_url or self.api_url + return base if "://" in base else f"http://{base}" + @dataclass(frozen=True) class CreateConfig: @@ -214,6 +258,9 @@ def __init__( async def create(self, spec: SandboxSpec) -> SandboxHandle: opts = SubstrateProviderOptions.from_mapping(spec.provider_options or {}) + # Ports a service inside the sandbox listens on (nemo-gym >= 0.7); + # endpoint() resolves only these. + ports = _validate_ports(getattr(spec, "ports", None)) template, rootfs_image = self._resolve_template(spec, opts) atespace = opts.atespace or self._create.atespace env_id = f"gym-{uuid.uuid4().hex[:10]}" @@ -256,6 +303,7 @@ async def create(self, spec: SandboxSpec) -> SandboxHandle: "atespace": atespace, "workdir": spec.workdir, "env": dict(spec.env or {}), + "ports": ports, }, ) @@ -347,6 +395,56 @@ async def status(self, handle: SandboxHandle) -> SandboxStatus: return SandboxStatus.UNKNOWN return _STATUS_MAP.get(info.status, SandboxStatus.UNKNOWN) + async def endpoint(self, handle: SandboxHandle, port: int) -> SandboxEndpoint: + """Resolve a declared service port to a caller-reachable endpoint + (nemo-gym's ``SupportsSandboxEndpoint``). + + The service is a daemon inside the sandbox, typically one a runtime + layer on the base template mounts into the task image and the guest + starts next to itself (see README). The endpoint is ate-env-api's port + proxy; its headers route each plain HTTP request through the Substrate + router to ``port`` in this sandbox. Nothing here touches the guest: a + parked sandbox is resumed by the first request through the proxy, + which waits for the wake (ate-env-api's ``--wake-budget``). Only the + sandbox's existence is checked, from the control plane, so a deleted + or crashed sandbox raises ``RuntimeError`` here rather than on the + first request. + """ + raw = handle.raw or {} + declared = list(raw.get("ports") or ()) + if isinstance(port, bool) or port not in declared: + raise ValueError( + f"Sandbox port {port!r} was not declared in SandboxSpec.ports; " + f"declared ports: {declared!r}" + ) + atespace = self._atespace_of(raw) + try: + info = await asyncio.wait_for( + self._ate().get(handle.sandbox_id, atespace=atespace), + timeout=self._connection.request_timeout_s, + ) + except NotFoundError as exc: + raise RuntimeError(f"substrate env {handle.sandbox_id!r} does not exist") from exc + except EnvError as exc: + raise RuntimeError( + f"substrate env {handle.sandbox_id!r}: {_describe(exc)}" + ) from exc + except TimeoutError as exc: + raise RuntimeError( + f"substrate env {handle.sandbox_id!r}: no answer from " + f"{self._connection.api_url} within {self._connection.request_timeout_s}s" + ) from exc + if info.status == EnvironmentStatus.CRASHED: + raise RuntimeError(f"substrate env {handle.sandbox_id!r} crashed") + return SandboxEndpoint( + endpoint=self._connection.endpoint_base(), + headers={ + _ENV_ID_HEADER: handle.sandbox_id, + _ENV_ATESPACE_HEADER: atespace, + _ENV_PORT_HEADER: str(port), + }, + ) + async def close(self, handle: SandboxHandle) -> None: await self._best_effort_delete(handle.sandbox_id, self._atespace_of(handle.raw or {})) @@ -401,10 +499,20 @@ def _resolve_template( return self._create.template, spec.image return self._create.template, None - async def _wait_ready(self, env: Env, timeout_s: float) -> None: + async def _wait_ready( + self, + env: Env, + timeout_s: float, + *, + error: type[Exception] = SandboxCreateVerificationError, + stop_on_not_found: bool = False, + ) -> None: """Poll a trivial command until the guest answers it. A fresh actor may still be booting, and a guest call made before it serves fails with a - transport-level error; both are retried until the deadline.""" + transport-level error; both are retried until the deadline, then + ``error`` is raised. The first call also resumes a parked actor. With + ``stop_on_not_found``, a sandbox that no longer exists raises ``error`` + at once instead of after the deadline.""" deadline = time.monotonic() + timeout_s last_error = "no probe attempted" while True: @@ -413,6 +521,10 @@ async def _wait_ready(self, env: Env, timeout_s: float) -> None: if result.exit_code == 0: return last_error = f"probe exited {result.exit_code}: {result.stderr.strip()[:200]}" + except NotFoundError as exc: + if stop_on_not_found: + raise error(f"substrate env {env.id!r} does not exist") from exc + last_error = _describe(exc) except EnvError as exc: last_error = _describe(exc) except asyncio.TimeoutError: @@ -420,9 +532,7 @@ async def _wait_ready(self, env: Env, timeout_s: float) -> None: if time.monotonic() >= deadline: break await asyncio.sleep(self._create.ready_poll_interval_s) - raise SandboxCreateVerificationError( - f"substrate env {env.id!r} not ready after {timeout_s}s: {last_error}" - ) + raise error(f"substrate env {env.id!r} not ready after {timeout_s}s: {last_error}") @staticmethod async def _collect(process: Process) -> SandboxExecResult: diff --git a/integrations/nemo-gym/tests/test_provider.py b/integrations/nemo-gym/tests/test_provider.py index c41e2bc8..134d0968 100644 --- a/integrations/nemo-gym/tests/test_provider.py +++ b/integrations/nemo-gym/tests/test_provider.py @@ -26,6 +26,7 @@ import asyncio import contextlib import logging +import time from pathlib import Path import grpc @@ -36,6 +37,7 @@ from nemo_gym_substrate._compat import ( SandboxCreateError, SandboxCreateVerificationError, + SandboxEndpoint, SandboxSpec, SandboxStatus, ) @@ -45,11 +47,16 @@ @contextlib.asynccontextmanager -async def stack(create: dict | None = None, *, state: FakeAteEnv | None = None): +async def stack( + create: dict | None = None, + *, + state: FakeAteEnv | None = None, + connection: dict | None = None, +): """A provider dialing a fake ate-env-api; yields ``(provider, fake)``.""" async with serve(state) as (target, fake): config = { - "connection": {"api_url": target, "request_timeout_s": 5.0}, + "connection": {"api_url": target, "request_timeout_s": 5.0, **(connection or {})}, "create": { "ready_poll_interval_s": 0.01, "image_templates": {"img:1": "tpl-img1"}, @@ -444,3 +451,139 @@ async def t(): assert provider._channel is None # noqa: SLF001 run(t()) + + +# -- endpoint(): services inside the sandbox --------------------------------- + + +def test_endpoint_routes_a_declared_port_through_the_api(): + async def t(): + async with stack() as (provider, fake): + handle = await provider.create(SandboxSpec(ports=(6000,))) + assert handle.raw["ports"] == [6000] + probes = _probe_count(fake) + ep = await provider.endpoint(handle, 6000) + assert isinstance(ep, SandboxEndpoint) + # The port proxy shares ate-env-api's listener. + assert ep.endpoint == provider._connection.endpoint_base() + assert ep.endpoint.startswith("http://") and ep.endpoint.endswith( + provider._connection.api_url.split("://")[-1] + ) + assert ep.headers == { + "X-Env-Id": handle.sandbox_id, + "X-Env-Atespace": "ate-env", + "X-Env-Port": "6000", + } + # Nothing touches the guest: the first request through the proxy + # wakes a parked sandbox, and the proxy waits for it. + assert _probe_count(fake) == probes + + run(t()) + + +def test_endpoint_rejects_an_undeclared_port(): + async def t(): + async with stack() as (provider, _fake): + handle = await provider.create(SandboxSpec(ports=(6000,))) + with pytest.raises(ValueError, match="not declared in SandboxSpec.ports"): + await provider.endpoint(handle, 8080) + bare = await provider.create(SandboxSpec()) + with pytest.raises(ValueError, match=r"declared ports: \[\]"): + await provider.endpoint(bare, 6000) + + run(t()) + + +def test_endpoint_uses_the_configured_url_and_the_sandbox_atespace(): + async def t(): + async with stack( + {"atespace": "team-a"}, + connection={"endpoint_url": "https://gym-proxy.example.com/"}, + ) as (provider, _fake): + handle = await provider.create(SandboxSpec(ports=[6000])) + ep = await provider.endpoint(handle, 6000) + assert ep.endpoint == "https://gym-proxy.example.com" + assert ep.headers["X-Env-Atespace"] == "team-a" + + run(t()) + + +def test_endpoint_raises_for_a_crashed_sandbox(): + async def t(): + async with stack() as (provider, fake): + handle = await provider.create(SandboxSpec(ports=(6000,))) + fake.set_status(handle.sandbox_id, EnvironmentStatus.CRASHED) + with pytest.raises(RuntimeError, match="crashed"): + await provider.endpoint(handle, 6000) + + run(t()) + + +def test_endpoint_does_not_wake_a_parked_sandbox(): + async def t(): + async with stack() as (provider, fake): + handle = await provider.create(SandboxSpec(ports=(6000,))) + fake.set_status(handle.sandbox_id, EnvironmentStatus.SUSPENDED) + before = len(fake.starts) + ep = await provider.endpoint(handle, 6000) + assert ep.headers["X-Env-Port"] == "6000" + assert len(fake.starts) == before # no probe exec, so no wake-up + + run(t()) + + +def test_create_rejects_invalid_ports(): + async def t(): + async with stack() as (provider, fake): + cases = [ + ((0,), ValueError, "not a TCP port"), + ((70000,), ValueError, "not a TCP port"), + ((80,), ValueError, "guest"), + ((True,), TypeError, "integers"), # bool is an int subclass + (("6000",), TypeError, "integers"), + ] + for ports, error, why in cases: + before = len(fake.envs) + with pytest.raises(error, match=why): + await provider.create(SandboxSpec(ports=ports)) + assert len(fake.envs) == before, "a bad spec must fail before the actor starts" + + run(t()) + + +def test_endpoint_rejects_a_bool_port(): + async def t(): + async with stack() as (provider, _fake): + handle = await provider.create(SandboxSpec(ports=(1,))) + with pytest.raises(ValueError, match="not declared"): + await provider.endpoint(handle, True) # True == 1 in Python + + run(t()) + + +def test_endpoint_stops_at_once_for_a_deleted_sandbox(): + async def t(): + async with stack({"ready_timeout_s": 30}) as (provider, _fake): + handle = await provider.create(SandboxSpec(ports=(6000,))) + await provider.close(handle) + started = time.monotonic() + with pytest.raises(RuntimeError, match="does not exist"): + await provider.endpoint(handle, 6000) + assert time.monotonic() - started < 5, "endpoint() waited out the ready timeout" + + run(t()) + + +def test_endpoint_base_defaults_to_the_api_url(): + cfg = provider_mod.ConnectionConfig.from_mapping({"api_url": "ate-env-api.ate-env:7777"}) + assert cfg.endpoint_base() == "http://ate-env-api.ate-env:7777" + cfg = provider_mod.ConnectionConfig.from_mapping({"api_url": "https://env.example.com/"}) + assert cfg.endpoint_base() == "https://env.example.com" + # A bare host:port endpoint_url gets a scheme too; SandboxEndpoint rejects it otherwise. + cfg = provider_mod.ConnectionConfig.from_mapping( + {"api_url": "ate-env-api.ate-env:7777", "endpoint_url": "gym-proxy.example.com:8080"} + ) + assert cfg.endpoint_base() == "http://gym-proxy.example.com:8080" + SandboxEndpoint(endpoint=cfg.endpoint_base(), headers={}) + with pytest.raises(ValueError, match="unknown option"): + provider_mod.ConnectionConfig.from_mapping({"endpoint": "x"}) diff --git a/internal/ate/client.go b/internal/ate/client.go index c434e753..344e953e 100644 --- a/internal/ate/client.go +++ b/internal/ate/client.go @@ -60,6 +60,9 @@ const DefaultHostSuffix = "actors.resources.substrate.ate.dev" const ( DefaultControlAddr = "api.ate-system.svc.cluster.local:443" DefaultRouterAddr = "atenet-router.ate-system.svc.cluster.local:80" + // DefaultRouterConnectAddr is the router's plaintext CONNECT listener, which + // tunnels to a port on an actor other than 80. + DefaultRouterConnectAddr = "atenet-router.ate-system.svc.cluster.local:8081" ) // DefaultAtespace is the Substrate atespace every environment actor lives in. diff --git a/internal/internaltest/fakerouter/fakerouter.go b/internal/internaltest/fakerouter/fakerouter.go index bf0b45ad..8a3ba9f9 100644 --- a/internal/internaltest/fakerouter/fakerouter.go +++ b/internal/internaltest/fakerouter/fakerouter.go @@ -16,12 +16,17 @@ // resolves the target environment from the request's ate-target-actor header // (or, failing that, its Host header) and forwards to that environment's guest // handler, returning 404 for an invalid or duplicated actor reference and 503 -// when the environment is not running. +// when the environment is not running. Like the real router's CONNECT +// listener, it also tunnels an HTTP CONNECT whose authority names a port to +// the handler registered for that environment and port. package fakerouter import ( + "bufio" + "net" "net/http" "net/http/httptest" + "strconv" "strings" "sync" @@ -36,12 +41,26 @@ type Router struct { mu sync.Mutex guests map[string]http.Handler + ports map[portKey]http.Handler +} + +type portKey struct { + id string + port int } // New returns a Router that considers every actor running unless a Running // callback is set. func New() *Router { - return &Router{guests: make(map[string]http.Handler)} + return &Router{guests: make(map[string]http.Handler), ports: make(map[portKey]http.Handler)} +} + +// RegisterPort installs the handler that serves CONNECT tunnels to port on +// environment id, standing in for a daemon listening there. +func (r *Router) RegisterPort(id string, port int, h http.Handler) { + r.mu.Lock() + defer r.mu.Unlock() + r.ports[portKey{id, port}] = h } // Register installs the guest handler serving a environment ID. @@ -57,6 +76,10 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) { http.Error(w, "invalid actor reference", http.StatusNotFound) return } + if req.Method == http.MethodConnect { + r.serveConnect(w, req, id) + return + } r.mu.Lock() guest := r.guests[id] r.mu.Unlock() @@ -102,3 +125,80 @@ func (r *Router) Serve() (addr string, stop func()) { srv.Start() return strings.TrimPrefix(srv.URL, "http://"), srv.Close } + +// serveConnect answers a CONNECT like the real router: the actor comes from +// ate-target-actor, the port from the authority. On success the hijacked +// connection carries HTTP to the registered port handler. +func (r *Router) serveConnect(w http.ResponseWriter, req *http.Request, id string) { + _, portStr, err := net.SplitHostPort(req.Host) + port, perr := strconv.Atoi(portStr) + if err != nil || perr != nil { + http.Error(w, "CONNECT authority names no port", http.StatusBadRequest) + return + } + r.mu.Lock() + h := r.ports[portKey{id, port}] + r.mu.Unlock() + if h == nil { + http.Error(w, "no actor "+id+" listening on "+portStr, http.StatusNotFound) + return + } + if r.Running != nil && !r.Running(id) { + http.Error(w, "actor "+id+" is not running", http.StatusServiceUnavailable) + return + } + conn, rw, err := http.NewResponseController(w).Hijack() + if err != nil { + http.Error(w, "hijack: "+err.Error(), http.StatusInternalServerError) + return + } + if _, err := conn.Write([]byte("HTTP/1.1 200 OK\r\n\r\n")); err != nil { + conn.Close() + return + } + l := &oneConnListener{conn: &tunnelConn{Conn: conn, r: rw.Reader}, done: make(chan struct{})} + _ = (&http.Server{Handler: h}).Serve(l) +} + +// tunnelConn reads what the hijacked reader buffered before the connection. +type tunnelConn struct { + net.Conn + r *bufio.Reader + closeFn func() +} + +func (c *tunnelConn) Read(b []byte) (int, error) { return c.r.Read(b) } + +func (c *tunnelConn) Close() error { + if c.closeFn != nil { + c.closeFn() + } + return c.Conn.Close() +} + +// oneConnListener hands out one connection, then blocks until it is closed. +type oneConnListener struct { + conn *tunnelConn + once sync.Once + closed sync.Once + done chan struct{} +} + +func (l *oneConnListener) Accept() (net.Conn, error) { + var c net.Conn + l.once.Do(func() { + l.conn.closeFn = func() { l.closed.Do(func() { close(l.done) }) } + c = l.conn + }) + if c != nil { + return c, nil + } + <-l.done + return nil, net.ErrClosed +} + +func (l *oneConnListener) Close() error { + l.closed.Do(func() { close(l.done) }) + return nil +} +func (l *oneConnListener) Addr() net.Addr { return l.conn.LocalAddr() } diff --git a/internal/portproxy/portproxy.go b/internal/portproxy/portproxy.go new file mode 100644 index 00000000..759ec158 --- /dev/null +++ b/internal/portproxy/portproxy.go @@ -0,0 +1,503 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package portproxy forwards HTTP requests from callers outside the actor +// network to a TCP port inside an environment, such as a daemon that a runtime +// layer starts next to the guest (docs/rootfs-images/RUNTIMES.md). +// +// A request names its target with three headers, EnvIDHeader, +// EnvAtespaceHeader (optional) and EnvPortHeader, or, for clients that cannot +// set headers (WebSocket clients, browsers), with a path prefix that the proxy +// strips: /v1alpha/envs/{id}/ports/{port}/... in the default atespace, or +// /v1alpha/atespaces/{atespace}/envs/{id}/ports/{port}/... . The proxy opens +// an HTTP CONNECT tunnel through the atenet router's CONNECT listener, naming +// the actor in ate-target-actor and the port in the CONNECT authority. The +// router resumes a suspended actor before it answers, then forwards the tunnel +// to that port. The proxy speaks HTTP/1.1 to the daemon through the tunnel and +// keeps idle tunnels for reuse, one pool for each environment and port. +// +// Trust model: the proxy does not authenticate callers. Whoever can reach the +// listener it is served on can reach any TCP port of any actor the router +// serves, in any atespace, including actors that ate-env did not create. A +// daemon behind it must authenticate its own clients. The path form refuses +// requests that a browser sends from another origin (see Options.TrustedOrigins), +// so a web page cannot drive a daemon through a port-forwarded listener. +package portproxy + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "log" + "net" + "net/http" + "net/http/httputil" + "net/url" + "regexp" + "strconv" + "strings" + "time" + + "github.com/agent-substrate/env/internal/ate" +) + +// The routing headers. They match the x-env-id and x-env-atespace metadata +// keys of the gRPC API, so a caller uses one vocabulary for both. +const ( + EnvIDHeader = "X-Env-Id" + EnvAtespaceHeader = "X-Env-Atespace" + EnvPortHeader = "X-Env-Port" +) + +// DefaultHandshakeTimeout bounds dialing the router and waiting for its +// answer to CONNECT, which includes resuming a suspended actor. +const DefaultHandshakeTimeout = 2 * time.Minute + +// DefaultWakeBudget bounds how long a request waits for a suspended +// environment to wake. The router's CONNECT is what resumes a suspended +// actor, and the router holds it only for its own parking budget, a few +// seconds; a wake that pulls an image or reads a snapshot across a region +// takes longer. The same two minutes as ate-env-api's guest calls +// (apiservice.DefaultWakeBudget). +const DefaultWakeBudget = 2 * time.Minute + +// Idle tunnels: each one is a TCP stream from the API to an actor, held open +// by the router and the worker. Keep few, and not for long, so an idle +// environment is not held awake and a large fleet does not pin thousands +// of sockets in the API pod. +const ( + idleTunnelTimeout = 10 * time.Second + idleTunnelsPerPort = 4 + idleTunnelsTotal = 1024 +) + +const ( + wakeRetryInitialBackoff = 200 * time.Millisecond + wakeRetryMaxBackoff = 2 * time.Second +) + +// guestPort is where the ate-env guest listens in every environment. It is +// not a daemon port: the API already serves it, and a daemon cannot bind it. +const guestPort = 80 + +// Options configures a Proxy. +type Options struct { + // ConnectAddr is the router's plaintext CONNECT listener, host:port. + ConnectAddr string + // HostSuffix is the router's actor host suffix; the CONNECT authority is + // .:. The router reads only the port from it. + HostSuffix string + // DefaultAtespace applies when a request carries no EnvAtespaceHeader. + DefaultAtespace string + // HandshakeTimeout bounds the CONNECT handshake. Zero means + // DefaultHandshakeTimeout. + HandshakeTimeout time.Duration + // ResponseHeaderTimeout bounds the wait for the daemon's response headers + // after the request is sent. Zero means no limit: a daemon that answers a + // long command only when it finishes is served, and the caller's own + // deadline ends a wait on a daemon that never answers. + ResponseHeaderTimeout time.Duration + // TrustedOrigins are browser origins (scheme://host[:port]) allowed to use + // the path form from a page they serve, next to the listener's own origin. + TrustedOrigins []string + // WakeBudget bounds how long a request waits while the router answers + // 503 for an environment that is suspended or still waking. Zero means + // DefaultWakeBudget. + WakeBudget time.Duration +} + +// Proxy is an http.Handler that forwards each request to the environment port +// its headers name. +type Proxy struct { + opts Options + rp *httputil.ReverseProxy + tr *http.Transport + cop *http.CrossOriginProtection +} + +// New returns a Proxy. It fails only on a malformed trusted origin. +func New(opts Options) (*Proxy, error) { + if opts.HostSuffix == "" { + opts.HostSuffix = ate.DefaultHostSuffix + } + if opts.HandshakeTimeout <= 0 { + opts.HandshakeTimeout = DefaultHandshakeTimeout + } + if opts.WakeBudget <= 0 { + opts.WakeBudget = DefaultWakeBudget + } + p := &Proxy{opts: opts, cop: http.NewCrossOriginProtection()} + for _, o := range opts.TrustedOrigins { + if err := p.cop.AddTrustedOrigin(o); err != nil { + return nil, fmt.Errorf("trusted origin: %w", err) + } + } + p.tr = &http.Transport{ + DialContext: p.dial, + MaxIdleConnsPerHost: idleTunnelsPerPort, + MaxIdleConns: idleTunnelsTotal, + IdleConnTimeout: idleTunnelTimeout, + ResponseHeaderTimeout: opts.ResponseHeaderTimeout, + // Pass encodings through untouched; the daemon and the caller + // negotiate them. + DisableCompression: true, + } + p.rp = &httputil.ReverseProxy{ + Rewrite: p.rewrite, + Transport: p.tr, + // Flush each write at once, so streamed responses (server-sent events, + // chunked output) reach the caller as the daemon produces them. + FlushInterval: -1, + ErrorHandler: p.handleError, + } + return p, nil +} + +// Wants reports whether r uses the header form: it carries EnvPortHeader and +// is not a gRPC call. Callers of the gRPC API never send the header, and the +// tunnel carries HTTP/1.1 only. +// Close drops the idle tunnels. Requests in flight are not affected. +func (p *Proxy) Close() { p.tr.CloseIdleConnections() } + +func Wants(r *http.Request) bool { + return r.Header.Get(EnvPortHeader) != "" && + !strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") +} + +// Dispatch returns a handler that sends header-form requests to p.ServeHTTP, +// path-form requests to p.ServePath, and all others to next. The path form is +// matched here rather than on a ServeMux, so the daemon gets the path exactly +// as sent: a mux would redirect a missing trailing slash, "//" and "..", and +// WebSocket clients do not follow redirects. +func Dispatch(p *Proxy, next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case Wants(r): + p.ServeHTTP(w, r) + case isPathForm(r.URL.EscapedPath()): + p.ServePath(w, r) + default: + next.ServeHTTP(w, r) + } + }) +} + +type target struct { + atespace string + id string + port int + // path and rawPath, when set, replace the request's path upstream (the + // path form strips its routing prefix). + path, rawPath string +} + +type targetKey struct{} + +// label is a DNS-1123 label, the shape of environment ids and atespaces. The +// values go into a CONNECT request line and header, so nothing else may pass. +var label = regexp.MustCompile(`^[a-z0-9]([-a-z0-9]{0,61}[a-z0-9])?$`) + +// canonicalPort is a TCP port in decimal: no sign, no leading zero. +var canonicalPort = regexp.MustCompile(`^[1-9][0-9]{0,4}$`) + +func parseTarget(r *http.Request, defaultAtespace string) (target, error) { + for _, h := range []string{EnvIDHeader, EnvAtespaceHeader, EnvPortHeader} { + if len(r.Header.Values(h)) > 1 { + return target{}, fmt.Errorf("%s appears more than once", h) + } + } + if r.Header.Get(EnvIDHeader) == "" { + return target{}, fmt.Errorf("%s is required next to %s", EnvIDHeader, EnvPortHeader) + } + return makeTarget(r.Header.Get(EnvIDHeader), r.Header.Get(EnvAtespaceHeader), r.Header.Get(EnvPortHeader), defaultAtespace) +} + +// makeTarget validates a target from either form. +func makeTarget(id, atespace, portStr, defaultAtespace string) (target, error) { + t := target{id: id, atespace: atespace} + if !label.MatchString(t.id) { + return t, fmt.Errorf("environment id %q is not valid (lowercase alphanumerics and '-', at most 63)", t.id) + } + if t.atespace == "" { + t.atespace = defaultAtespace + } + if !label.MatchString(t.atespace) { + return t, fmt.Errorf("atespace %q is not valid", t.atespace) + } + n, err := strconv.Atoi(portStr) + if !canonicalPort.MatchString(portStr) || err != nil || n > 65535 { + return t, fmt.Errorf("port %q is not a TCP port (1-65535)", portStr) + } + if n == guestPort { + return t, fmt.Errorf("port %d is the ate-env guest, which ate-env-api serves itself; a daemon listens on another port", guestPort) + } + t.port = n + return t, nil +} + +// pathPrefix splits an escaped path of the path form into its routing +// segments and the rest. n is 0 when the path does not have that shape. The +// split is by segment on the escaped path, the way a ServeMux matches, so an +// escaped routing segment (%65nv1 for env1) still ends the prefix where it +// should. +func pathPrefix(escaped string) (id, atespace, portStr, rest string, n int) { + segs := strings.Split(escaped, "/") + switch { + case len(segs) >= 6 && segs[0] == "" && segs[1] == "v1alpha" && segs[2] == "envs" && segs[4] == "ports": + id, portStr, n = segs[3], segs[5], 6 + case len(segs) >= 8 && segs[0] == "" && segs[1] == "v1alpha" && segs[2] == "atespaces" && segs[4] == "envs" && segs[6] == "ports": + atespace, id, portStr, n = segs[3], segs[5], segs[7], 8 + default: + return "", "", "", "", 0 + } + return id, atespace, portStr, "/" + strings.Join(segs[n:], "/"), n +} + +func isPathForm(escaped string) bool { + _, _, _, _, n := pathPrefix(escaped) + return n > 0 +} + +func parsePath(escaped, defaultAtespace string) (target, error) { + id, atespace, portStr, rest, n := pathPrefix(escaped) + if n == 0 { + return target{}, errors.New("path is not /v1alpha/[atespaces/{atespace}/]envs/{id}/ports/{port}[/rest]") + } + // Validation admits only labels and digits, thus an escaped routing + // segment is unescaped first and the value that passes reads the same both + // ways. + for _, s := range []*string{&id, &atespace, &portStr} { + v, err := url.PathUnescape(*s) + if err != nil { + return target{}, fmt.Errorf("bad path escaping: %w", err) + } + *s = v + } + t, err := makeTarget(id, atespace, portStr, defaultAtespace) + if err != nil { + return t, err + } + // The rest keeps its own escaping. The bare prefix, with or without a + // trailing slash, is the daemon's root. + t.rawPath = rest + if t.path, err = url.PathUnescape(rest); err != nil { + return t, fmt.Errorf("bad path escaping: %w", err) + } + if t.rawPath == t.path { + t.rawPath = "" + } + return t, nil +} + +// ServePath forwards a path-form request to the environment port its path +// names, after it strips the routing prefix. It refuses a request that a +// browser sends from another origin: a state-changing method, or a WebSocket +// handshake, which is a GET that opens a channel both ways. The header form +// needs no such check, because a browser sends custom headers to another +// origin only after a CORS preflight, which the proxy never approves. +func (p *Proxy) ServePath(w http.ResponseWriter, r *http.Request) { + t, err := parsePath(r.URL.EscapedPath(), p.opts.DefaultAtespace) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + check := r + if r.Header.Get("Upgrade") != "" { + // CrossOriginProtection lets safe methods through; judge a handshake + // as the state-changing request it is. + check = r.WithContext(r.Context()) + check.Method = http.MethodPost + } + if err := p.cop.Check(check); err != nil { + http.Error(w, "cross-origin request to an environment port refused; "+ + "add the page's origin to the port proxy's trusted origins", http.StatusForbidden) + return + } + p.rp.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), targetKey{}, t))) +} + +// ServeHTTP forwards r to the environment port its headers name. +func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { + t, err := parseTarget(r, p.opts.DefaultAtespace) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + p.rp.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), targetKey{}, t))) +} + +func (p *Proxy) rewrite(pr *httputil.ProxyRequest) { + t := pr.In.Context().Value(targetKey{}).(target) + pr.SetXForwarded() + // The URL host is the transport's pool key, not an address: one pool for + // each environment and port. dial reads the target back from it. + pr.Out.URL.Scheme = "http" + pr.Out.URL.Host = net.JoinHostPort(t.id+"."+t.atespace, strconv.Itoa(t.port)) + // What a client inside the environment would send. + pr.Out.Host = net.JoinHostPort("127.0.0.1", strconv.Itoa(t.port)) + if t.path != "" { + pr.Out.URL.Path, pr.Out.URL.RawPath = t.path, t.rawPath + } + for _, h := range []string{EnvIDHeader, EnvAtespaceHeader, EnvPortHeader} { + pr.Out.Header.Del(h) + } +} + +// routerError is a CONNECT the router refused. +type routerError struct { + status int + msg string +} + +func (e *routerError) Error() string { + return fmt.Sprintf("router answered CONNECT with %d: %s", e.status, e.msg) +} + +// dial opens a CONNECT tunnel to the environment port encoded in addr (see +// rewrite). While the router answers 503, which is what a suspended actor +// and one still waking both look like, it tries again with backoff until +// WakeBudget is spent: the first CONNECT is what resumes the actor, and the +// router holds it only for its own parking budget. +func (p *Proxy) dial(ctx context.Context, _, addr string) (net.Conn, error) { + host, portStr, err := net.SplitHostPort(addr) + if err != nil { + return nil, err + } + id, atespace, _ := strings.Cut(host, ".") + deadline := time.Now().Add(p.opts.WakeBudget) + backoff := wakeRetryInitialBackoff + for attempt := 1; ; attempt++ { + conn, err := p.connect(ctx, id, atespace, portStr) + var re *routerError + if !errors.As(err, &re) || re.status != http.StatusServiceUnavailable { + return conn, err + } + remaining := time.Until(deadline) + if remaining <= 0 { + // Logged here: handleError keeps router 503s out of the log, and + // a spent wake budget is the one 503 an operator wants to see, + // as the close of the "waiting" line from the first attempt. + err = fmt.Errorf("environment %s/%s did not wake within %s (%d attempts): %w", atespace, id, p.opts.WakeBudget, attempt, err) + log.Printf("port proxy: %v", err) + return nil, err + } + if attempt == 1 { + log.Printf("port proxy: waiting for environment %s/%s to wake (wake budget %s)", atespace, id, p.opts.WakeBudget) + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(min(backoff, remaining)): + } + backoff = min(backoff*2, wakeRetryMaxBackoff) + } +} + +// connect sends one CONNECT for the environment port and returns the tunnel. +func (p *Proxy) connect(ctx context.Context, id, atespace, portStr string) (net.Conn, error) { + ctx, cancel := context.WithTimeout(ctx, p.opts.HandshakeTimeout) + defer cancel() + var d net.Dialer + conn, err := d.DialContext(ctx, "tcp", p.opts.ConnectAddr) + if err != nil { + return nil, fmt.Errorf("dialing the router's CONNECT listener %s: %w", p.opts.ConnectAddr, err) + } + // A context that ends mid-handshake unblocks the write or read below. + stop := context.AfterFunc(ctx, func() { _ = conn.SetDeadline(time.Unix(1, 0)) }) + + authority := net.JoinHostPort(id+"."+p.opts.HostSuffix, portStr) + req := &http.Request{ + Method: http.MethodConnect, + URL: &url.URL{Opaque: authority}, + Host: authority, + Header: http.Header{ate.TargetActorHeader: {atespace + "/" + id}}, + } + br := bufio.NewReader(conn) + resp, err := func() (*http.Response, error) { + if err := req.Write(conn); err != nil { + return nil, err + } + return http.ReadResponse(br, req) + }() + if !stop() { + err = errors.Join(ctx.Err(), err) + } + if err != nil { + conn.Close() + return nil, fmt.Errorf("CONNECT to %s/%s port %s: %w", atespace, id, portStr, err) + } + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) + conn.Close() + return nil, &routerError{status: resp.StatusCode, msg: strings.TrimSpace(string(body))} + } + if br.Buffered() > 0 { + return &bufferedConn{Conn: conn, r: br}, nil + } + return conn, nil +} + +// bufferedConn reads what the handshake reader buffered past the CONNECT +// response before it reads the connection again. +type bufferedConn struct { + net.Conn + r *bufio.Reader +} + +func (c *bufferedConn) Read(b []byte) (int, error) { return c.r.Read(b) } + +// CloseWrite half-closes the tunnel when the connection can: after an +// upgrade, httputil.ReverseProxy half-closes the backend once the client +// side is done, and embedding net.Conn would otherwise hide the method. +func (c *bufferedConn) CloseWrite() error { + if cw, ok := c.Conn.(interface{ CloseWrite() error }); ok { + return cw.CloseWrite() + } + return nil +} + +// handleError answers a failed forward. The body names the target and a +// short cause; the full error, which can name the router's address, goes to +// the log only. +func (p *Proxy) handleError(w http.ResponseWriter, r *http.Request, err error) { + t, _ := r.Context().Value(targetKey{}).(target) + var re *routerError + var ne net.Error + status, cause := http.StatusBadGateway, "is unreachable; the ate-env-api log has the cause" + switch { + case errors.Is(err, context.Canceled): + return // the caller went away; nobody reads an answer + case errors.As(err, &re): + // 404 (no such environment) and 503 (not running) mean the same to + // the caller as to the router; anything else is a gateway failure. + switch re.status { + case http.StatusNotFound: + status, cause = re.status, "was not found" + case http.StatusServiceUnavailable: + status, cause = re.status, "is not available (not running, or still waking)" + case http.StatusBadGateway: + cause = "refused the connection; a daemon must listen on 0.0.0.0, not only on 127.0.0.1" + } + case errors.Is(err, context.DeadlineExceeded), errors.As(err, &ne) && ne.Timeout(): + status, cause = http.StatusGatewayTimeout, "did not answer in time" + } + // A 404 or 503 from the router is the caller's answer, not an operator's + // problem; a log line for each would flood the log when a fleet wakes. + if status != http.StatusNotFound && status != http.StatusServiceUnavailable { + log.Printf("port proxy: %s/%s port %d: %v", t.atespace, t.id, t.port, err) + } + http.Error(w, fmt.Sprintf("environment %s/%s port %d %s", t.atespace, t.id, t.port, cause), status) +} diff --git a/internal/portproxy/portproxy_test.go b/internal/portproxy/portproxy_test.go new file mode 100644 index 00000000..3d2b7c47 --- /dev/null +++ b/internal/portproxy/portproxy_test.go @@ -0,0 +1,651 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package portproxy_test + +import ( + "bufio" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/agent-substrate/env/internal/ate" + "github.com/agent-substrate/env/internal/internaltest/fakerouter" + "github.com/agent-substrate/env/internal/portproxy" +) + +// recordingRouter wraps the fake router and records each CONNECT's actor +// reference and authority. +type recordingRouter struct { + *fakerouter.Router + mu sync.Mutex + connects []string +} + +func (r *recordingRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) { + if req.Method == http.MethodConnect { + r.mu.Lock() + r.connects = append(r.connects, req.Header.Get(ate.TargetActorHeader)+" "+req.Host) + r.mu.Unlock() + } + r.Router.ServeHTTP(w, req) +} + +func (r *recordingRouter) connectCount() int { + r.mu.Lock() + defer r.mu.Unlock() + return len(r.connects) +} + +type setup struct { + router *recordingRouter + routerAddr string + proxy *httptest.Server +} + +func newSetup(t *testing.T) *setup { + t.Helper() + rr := &recordingRouter{Router: fakerouter.New()} + srv := httptest.NewServer(rr) + t.Cleanup(srv.Close) + addr := strings.TrimPrefix(srv.URL, "http://") + p := mustNew(t, portproxy.Options{ + ConnectAddr: addr, + DefaultAtespace: "ate-env", + HandshakeTimeout: 5 * time.Second, + WakeBudget: 2 * time.Second, + }) + proxy := httptest.NewServer(p) + t.Cleanup(proxy.Close) + return &setup{router: rr, routerAddr: addr, proxy: proxy} +} + +// mustNew builds a proxy and drops its idle tunnels when the test ends, so +// the fake router's tunnel goroutines end with the test. +func mustNew(t *testing.T, opts portproxy.Options) *portproxy.Proxy { + t.Helper() + p, err := portproxy.New(opts) + if err != nil { + t.Fatal(err) + } + t.Cleanup(p.Close) + return p +} + +func (s *setup) do(t *testing.T, method, path, body string, headers map[string]string) (*http.Response, string) { + t.Helper() + req, err := http.NewRequest(method, s.proxy.URL+path, strings.NewReader(body)) + if err != nil { + t.Fatal(err) + } + for k, v := range headers { + req.Header.Set(k, v) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("%s %s: %v", method, path, err) + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + return resp, string(b) +} + +func target(id string, port int) map[string]string { + return map[string]string{portproxy.EnvIDHeader: id, portproxy.EnvPortHeader: fmt.Sprint(port)} +} + +func TestForwardsToEnvironmentPort(t *testing.T) { + s := newSetup(t) + var seen *http.Request + var seenBody string + s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seen = r + b, _ := io.ReadAll(r.Body) + seenBody = string(b) + w.Header().Set("X-Daemon", "yes") + w.WriteHeader(http.StatusCreated) + io.WriteString(w, "daemon says hi") + })) + + h := target("env1", 8000) + h["X-Custom"] = "kept" + resp, body := s.do(t, http.MethodPost, "/execute?lang=python", `{"code":"1+1"}`, h) + + if resp.StatusCode != http.StatusCreated || body != "daemon says hi" || resp.Header.Get("X-Daemon") != "yes" { + t.Fatalf("got %d %q (X-Daemon %q), want 201 from the daemon", resp.StatusCode, body, resp.Header.Get("X-Daemon")) + } + if seen.Method != http.MethodPost || seen.URL.RequestURI() != "/execute?lang=python" || seenBody != `{"code":"1+1"}` { + t.Errorf("daemon saw %s %s %q", seen.Method, seen.URL.RequestURI(), seenBody) + } + if seen.Host != "127.0.0.1:8000" { + t.Errorf("daemon saw Host %q, want 127.0.0.1:8000 (what an in-sandbox client sends)", seen.Host) + } + for _, hdr := range []string{portproxy.EnvIDHeader, portproxy.EnvAtespaceHeader, portproxy.EnvPortHeader} { + if v := seen.Header.Get(hdr); v != "" { + t.Errorf("daemon saw routing header %s=%q; the proxy must strip it", hdr, v) + } + } + if seen.Header.Get("X-Custom") != "kept" { + t.Error("other request headers must reach the daemon") + } + if got := s.router.connects; len(got) != 1 || got[0] != "ate-env/env1 env1."+ate.DefaultHostSuffix+":8000" { + t.Errorf("router saw CONNECTs %q, want one for ate-env/env1 on port 8000", got) + } +} + +func TestAtespaceHeaderNamesTheActor(t *testing.T) { + s := newSetup(t) + s.router.RegisterPort("env1", 9000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + h := target("env1", 9000) + h[portproxy.EnvAtespaceHeader] = "team-a" + if resp, body := s.do(t, http.MethodGet, "/", "", h); resp.StatusCode != http.StatusOK { + t.Fatalf("got %d %q", resp.StatusCode, body) + } + if got := s.router.connects; len(got) != 1 || !strings.HasPrefix(got[0], "team-a/env1 ") { + t.Errorf("router saw %q, want the actor team-a/env1", got) + } +} + +func TestStreamsResponses(t *testing.T) { + s := newSetup(t) + release := make(chan struct{}) + s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + io.WriteString(w, "first\n") + w.(http.Flusher).Flush() + <-release + io.WriteString(w, "second\n") + })) + req, _ := http.NewRequest(http.MethodGet, s.proxy.URL+"/events", nil) + for k, v := range target("env1", 8000) { + req.Header.Set(k, v) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + buf := make([]byte, len("first\n")) + if _, err := io.ReadFull(resp.Body, buf); err != nil || string(buf) != "first\n" { + t.Fatalf("first chunk %q, %v: the proxy must flush before the daemon finishes", buf, err) + } + close(release) + rest, _ := io.ReadAll(resp.Body) + if string(rest) != "second\n" { + t.Errorf("rest %q", rest) + } +} + +func TestRejectsBadTargets(t *testing.T) { + s := newSetup(t) + for name, h := range map[string]map[string]string{ + "no id": {portproxy.EnvPortHeader: "8000"}, + "port not number": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "http"}, + "port zero": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "0"}, + "port too large": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "70000"}, + "id uppercase": {portproxy.EnvIDHeader: "Env1", portproxy.EnvPortHeader: "8000"}, + "id with slash": {portproxy.EnvIDHeader: "a/b", portproxy.EnvPortHeader: "8000"}, + "port with sign": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "+8000"}, + "port zero-led": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "08000"}, + "bad atespace": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "8000", portproxy.EnvAtespaceHeader: "x y"}, + "guest port": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "80"}, + } { + if resp, body := s.do(t, http.MethodGet, "/", "", h); resp.StatusCode != http.StatusBadRequest { + t.Errorf("%s: got %d %q, want 400", name, resp.StatusCode, body) + } + } + // A routing header sent twice is ambiguous. + req, _ := http.NewRequest(http.MethodGet, s.proxy.URL+"/", nil) + req.Header.Add(portproxy.EnvIDHeader, "env1") + req.Header.Add(portproxy.EnvIDHeader, "env2") + req.Header.Set(portproxy.EnvPortHeader, "8000") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("duplicate %s: got %d, want 400", portproxy.EnvIDHeader, resp.StatusCode) + } + if n := s.router.connectCount(); n != 0 { + t.Errorf("router saw %d CONNECTs for rejected requests, want 0", n) + } +} + +func TestPassesRouterRefusals(t *testing.T) { + s := newSetup(t) + s.router.RegisterPort("asleep", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + s.router.Running = func(id string) bool { return id != "asleep" } + + if resp, body := s.do(t, http.MethodGet, "/", "", target("missing", 8000)); resp.StatusCode != http.StatusNotFound { + t.Errorf("unknown environment: got %d %q, want 404", resp.StatusCode, body) + } + if resp, body := s.do(t, http.MethodGet, "/", "", target("asleep", 8000)); resp.StatusCode != http.StatusServiceUnavailable || !strings.Contains(body, "ate-env/asleep port 8000") { + t.Errorf("not running: got %d %q, want 503 naming the target", resp.StatusCode, body) + } +} + +// A suspended environment answers 503 at the router until it is back; the +// proxy tries again within its wake budget instead of passing the first 503 +// on, the way ate-env-api's guest calls do. +func TestRetriesWhileTheEnvironmentWakes(t *testing.T) { + s := newSetup(t) + s.router.RegisterPort("waking", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "awake") })) + var mu sync.Mutex + refusals := 0 + s.router.Running = func(id string) bool { + mu.Lock() + defer mu.Unlock() + if refusals < 2 { + refusals++ + return false + } + return true + } + resp, body := s.do(t, http.MethodGet, "/", "", target("waking", 8000)) + if resp.StatusCode != http.StatusOK || body != "awake" { + t.Fatalf("got %d %q, want the daemon's answer after the wake", resp.StatusCode, body) + } + if n := s.router.connectCount(); n != 3 { + t.Errorf("router saw %d CONNECTs, want 3 (two refused while waking, one accepted)", n) + } +} + +func TestWakeBudgetBoundsTheWait(t *testing.T) { + s := newSetup(t) + p := mustNew(t, portproxy.Options{ConnectAddr: s.routerAddr, DefaultAtespace: "ate-env", HandshakeTimeout: 5 * time.Second, WakeBudget: 300 * time.Millisecond}) + srv := httptest.NewServer(p) + t.Cleanup(srv.Close) + s.router.RegisterPort("asleep", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + s.router.Running = func(string) bool { return false } + req, _ := http.NewRequest(http.MethodGet, srv.URL+"/", nil) + for k, v := range target("asleep", 8000) { + req.Header.Set(k, v) + } + start := time.Now() + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + b, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusServiceUnavailable || !strings.Contains(string(b), "not available") { + t.Errorf("got %d %q, want 503 once the wake budget is spent", resp.StatusCode, b) + } + if took := time.Since(start); took < 300*time.Millisecond || took > 3*time.Second { + t.Errorf("waited %s, want about the 300ms budget", took) + } + if n := s.router.connectCount(); n < 2 { + t.Errorf("router saw %d CONNECTs, want at least 2 (the budget allows retries)", n) + } +} + +func TestPoolsConnectionsPerTarget(t *testing.T) { + s := newSetup(t) + for _, id := range []string{"env1", "env2"} { + id := id + s.router.RegisterPort(id, 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, id) })) + } + for i := 0; i < 3; i++ { + for _, id := range []string{"env1", "env2"} { + if _, body := s.do(t, http.MethodGet, "/", "", target(id, 8000)); body != id { + t.Fatalf("request for %s answered by %q: tunnels must not be shared across environments", id, body) + } + } + } + if n := s.router.connectCount(); n != 2 { + t.Errorf("router saw %d CONNECTs for 6 sequential requests to 2 targets, want 2 (idle tunnels reused)", n) + } +} + +func TestDispatch(t *testing.T) { + s := newSetup(t) + s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "proxied") })) + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "api") }) + h := portproxy.Dispatch(mustNew(t, portproxy.Options{ConnectAddr: s.routerAddr, DefaultAtespace: "ate-env"}), next) + + rec := httptest.NewRecorder() + h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/healthz", nil)) + if rec.Body.String() != "api" { + t.Errorf("request without %s went to %q, want the API", portproxy.EnvPortHeader, rec.Body.String()) + } + srv := httptest.NewServer(h) + defer srv.Close() + req, _ := http.NewRequest(http.MethodGet, srv.URL+"/healthz", nil) + for k, v := range target("env1", 8000) { + req.Header.Set(k, v) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if b, _ := io.ReadAll(resp.Body); string(b) != "proxied" { + t.Errorf("request with %s got %q, want the daemon's answer", portproxy.EnvPortHeader, b) + } + + // A gRPC call is never proxied, even with the header: the tunnel carries + // HTTP/1.1 only. + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPost, "/ateenv.v1alpha.EnvironmentService/Get", nil) + req.Header.Set("Content-Type", "application/grpc") + req.Header.Set(portproxy.EnvPortHeader, "8000") + h.ServeHTTP(rec, req) + if rec.Body.String() != "api" { + t.Errorf("gRPC request with %s went to %q, want the API", portproxy.EnvPortHeader, rec.Body.String()) + } +} + +// pathSetup serves the proxy the way ate-env-api does: both forms dispatched +// in front of the API's mux. +func pathSetup(t *testing.T, origins ...string) (*setup, string) { + t.Helper() + s := newSetup(t) + p := mustNew(t, portproxy.Options{ConnectAddr: s.routerAddr, DefaultAtespace: "ate-env", HandshakeTimeout: 5 * time.Second, TrustedOrigins: origins}) + mux := http.NewServeMux() + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "api") }) + srv := httptest.NewServer(portproxy.Dispatch(p, mux)) + t.Cleanup(srv.Close) + return s, srv.URL +} + +func TestPathFormStripsPrefix(t *testing.T) { + s, base := pathSetup(t) + var seen *http.Request + s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seen = r + io.WriteString(w, "daemon") + })) + resp, err := http.Get(base + "/v1alpha/envs/env1/ports/8000/api/file/download?path=/etc/os-release") + if err != nil { + t.Fatal(err) + } + b, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if string(b) != "daemon" { + t.Fatalf("got %d %q", resp.StatusCode, b) + } + if got := seen.URL.RequestURI(); got != "/api/file/download?path=/etc/os-release" { + t.Errorf("daemon saw %q, want the path after the routing prefix with its query", got) + } + if got := s.router.connects; len(got) != 1 || !strings.HasPrefix(got[0], "ate-env/env1 ") { + t.Errorf("router saw %q", got) + } + + // Escaping after the prefix is kept as sent. + resp, err = http.Get(base + "/v1alpha/envs/env1/ports/8000/files/a%2Fb") + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if got := seen.URL.EscapedPath(); got != "/files/a%2Fb" { + t.Errorf("daemon saw escaped path %q, want /files/a%%2Fb", got) + } + + // The bare prefix with a trailing slash is the daemon's root. + resp, err = http.Get(base + "/v1alpha/envs/env1/ports/8000/") + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if seen.URL.Path != "/" { + t.Errorf("daemon saw %q for the bare prefix, want /", seen.URL.Path) + } +} + +func TestPathFormAtespaceAndValidation(t *testing.T) { + s, base := pathSetup(t) + s.router.RegisterPort("env1", 9000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, r.URL.Path) })) + resp, err := http.Get(base + "/v1alpha/atespaces/team-a/envs/env1/ports/9000/x") + if err != nil { + t.Fatal(err) + } + b, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if string(b) != "/x" { + t.Errorf("got %d %q, want the daemon's /x", resp.StatusCode, b) + } + if got := s.router.connects; len(got) != 1 || !strings.HasPrefix(got[0], "team-a/env1 ") { + t.Errorf("router saw %q, want team-a/env1", got) + } + for _, path := range []string{ + "/v1alpha/envs/Env1/ports/9000/x", + "/v1alpha/envs/env1/ports/0/x", + "/v1alpha/envs/env1/ports/http/x", + "/v1alpha/envs/env1/ports/09000/x", + "/v1alpha/envs//ports/9000/x", + "/v1alpha/envs/a%2Fb/ports/9000/x", + "/v1alpha/atespaces/A_B/envs/env1/ports/9000/x", + } { + resp, err := http.Get(base + path) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("%s: got %d, want 400", path, resp.StatusCode) + } + } + // Other paths still reach the API. + resp, err = http.Get(base + "/v1alpha/envs/env1/mcp") + if err != nil { + t.Fatal(err) + } + b, _ = io.ReadAll(resp.Body) + resp.Body.Close() + if string(b) != "api" { + t.Errorf("/v1alpha/envs/env1/mcp went to %q, want the API", b) + } +} + +// An HTTP upgrade (what a WebSocket handshake is) passes through the proxy and +// the tunnel, after which bytes flow both ways. +func TestUpgradePassesThrough(t *testing.T) { + s, base := pathSetup(t) + s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Upgrade") != "echo" { + http.Error(w, "want Upgrade: echo", http.StatusBadRequest) + return + } + conn, rw, err := http.NewResponseController(w).Hijack() + if err != nil { + return + } + defer conn.Close() + rw.WriteString("HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nUpgrade: echo\r\n\r\n") + rw.Flush() + buf := make([]byte, 4) + if _, err := io.ReadFull(rw, buf); err == nil { + rw.Write(append([]byte("re:"), buf...)) + rw.Flush() + } + })) + addr := strings.TrimPrefix(base, "http://") + for name, req := range map[string]string{ + "path form": "GET /v1alpha/envs/env1/ports/8000/sockets/events/c1 HTTP/1.1\r\nHost: " + addr + "\r\n", + "header form": "GET /sockets/events/c1 HTTP/1.1\r\nHost: " + addr + "\r\nX-Env-Id: env1\r\nX-Env-Port: 8000\r\n", + } { + conn, err := net.Dial("tcp", addr) + if err != nil { + t.Fatal(err) + } + conn.SetDeadline(time.Now().Add(10 * time.Second)) + fmt.Fprint(conn, req+"Connection: Upgrade\r\nUpgrade: echo\r\n\r\n") + br := bufio.NewReader(conn) + resp, err := http.ReadResponse(br, nil) + if err != nil || resp.StatusCode != http.StatusSwitchingProtocols { + t.Fatalf("%s: upgrade answered %v, %v", name, resp, err) + } + conn.Write([]byte("ping")) + got := make([]byte, 7) + if _, err := io.ReadFull(br, got); err != nil || string(got) != "re:ping" { + t.Errorf("%s: after upgrade got %q, %v; want re:ping", name, got, err) + } + conn.Close() + } +} + +// The path form reaches the daemon exactly as sent: an escaped routing segment +// still ends the prefix, and nothing is redirected or cleaned. +func TestPathFormKeepsPathAsSent(t *testing.T) { + s, base := pathSetup(t) + var seen *http.Request + s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seen = r + io.WriteString(w, "daemon") + })) + noRedirect := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + for _, c := range []struct{ method, path, want string }{ + {http.MethodGet, "/v1alpha/envs/%65nv1/ports/8000/a", "/a"}, + {http.MethodGet, "/v1alpha/envs/env1/ports/%38000/a", "/a"}, + {http.MethodGet, "/v1alpha/envs/env1/ports/8000", "/"}, + {http.MethodPost, "/v1alpha/envs/env1/ports/8000", "/"}, + {http.MethodPost, "/v1alpha/atespaces/ate-env/envs/env1/ports/8000", "/"}, + {http.MethodGet, "/v1alpha/envs/env1/ports/8000/a//b/../c?x=1", "/a//b/../c?x=1"}, + } { + seen = nil + req, _ := http.NewRequest(c.method, base+c.path, nil) + resp, err := noRedirect.Do(req) + if err != nil { + t.Fatal(err) + } + b, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK || string(b) != "daemon" || seen == nil { + t.Errorf("%s %s: got %d %q, want the daemon's 200", c.method, c.path, resp.StatusCode, b) + continue + } + if got := seen.URL.RequestURI(); got != c.want { + t.Errorf("%s %s: daemon saw %q, want %q", c.method, c.path, got, c.want) + } + } +} + +// A browser page from another origin cannot drive a daemon through the path +// form: no state-changing request, no WebSocket handshake. Reads, non-browser +// clients, the listener's own origin and trusted origins pass. +func TestPathFormRefusesCrossOrigin(t *testing.T) { + s, base := pathSetup(t, "https://ui.example") + s.router.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Upgrade") != "" { + w.Header().Set("Connection", "Upgrade") + w.Header().Set("Upgrade", "echo") + w.WriteHeader(http.StatusSwitchingProtocols) + return + } + io.WriteString(w, "daemon") + })) + self := base // the listener's own origin, as a browser on it sends + path := base + "/v1alpha/envs/env1/ports/8000/run" + for _, c := range []struct { + name string + method string + headers map[string]string + want int + }{ + {"cross-site POST", http.MethodPost, map[string]string{"Sec-Fetch-Site": "cross-site", "Origin": "https://evil.example"}, http.StatusForbidden}, + {"foreign Origin POST", http.MethodPost, map[string]string{"Origin": "https://evil.example"}, http.StatusForbidden}, + {"foreign Origin upgrade", http.MethodGet, map[string]string{"Origin": "https://evil.example", "Connection": "Upgrade", "Upgrade": "echo"}, http.StatusForbidden}, + {"cross-site upgrade", http.MethodGet, map[string]string{"Sec-Fetch-Site": "cross-site", "Connection": "Upgrade", "Upgrade": "echo"}, http.StatusForbidden}, + {"cross-site GET", http.MethodGet, map[string]string{"Sec-Fetch-Site": "cross-site", "Origin": "https://evil.example"}, http.StatusOK}, + {"no browser headers", http.MethodPost, nil, http.StatusOK}, + {"own origin POST", http.MethodPost, map[string]string{"Origin": self}, http.StatusOK}, + {"same-origin upgrade", http.MethodGet, map[string]string{"Sec-Fetch-Site": "same-origin", "Connection": "Upgrade", "Upgrade": "echo"}, http.StatusSwitchingProtocols}, + {"trusted origin POST", http.MethodPost, map[string]string{"Sec-Fetch-Site": "cross-site", "Origin": "https://ui.example"}, http.StatusOK}, + } { + req, _ := http.NewRequest(c.method, path, nil) + for k, v := range c.headers { + req.Header.Set(k, v) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("%s: %v", c.name, err) + } + resp.Body.Close() + if resp.StatusCode != c.want { + t.Errorf("%s: got %d, want %d", c.name, resp.StatusCode, c.want) + } + } + if _, err := portproxy.New(portproxy.Options{ConnectAddr: s.routerAddr, TrustedOrigins: []string{"not an origin"}}); err == nil { + t.Error("New accepted a malformed trusted origin") + } +} + +// An error body names the target and a short cause. The router's message and +// address stay in the log. +func TestErrorBodiesNameCauseOnly(t *testing.T) { + refuser := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "dial tcp 10.1.2.3:8000: connect: connection refused", http.StatusBadGateway) + })) + defer refuser.Close() + p := mustNew(t, portproxy.Options{ConnectAddr: strings.TrimPrefix(refuser.URL, "http://"), DefaultAtespace: "ate-env"}) + srv := httptest.NewServer(p) + defer srv.Close() + req, _ := http.NewRequest(http.MethodGet, srv.URL+"/", nil) + for k, v := range target("env1", 8000) { + req.Header.Set(k, v) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + b, _ := io.ReadAll(resp.Body) + resp.Body.Close() + body := string(b) + if resp.StatusCode != http.StatusBadGateway || !strings.Contains(body, "ate-env/env1 port 8000") || !strings.Contains(body, "0.0.0.0") { + t.Errorf("got %d %q, want 502 naming the target and the listen rule", resp.StatusCode, body) + } + if strings.Contains(body, "10.1.2.3") || strings.Contains(body, refuser.URL[len("http://"):]) { + t.Errorf("body %q leaks an internal address", body) + } +} + +func TestResponseHeaderTimeout(t *testing.T) { + rr := &recordingRouter{Router: fakerouter.New()} + router := httptest.NewServer(rr) + defer router.Close() + release := make(chan struct{}) + defer close(release) + rr.RegisterPort("env1", 8000, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-release: + case <-r.Context().Done(): + } + })) + p := mustNew(t, portproxy.Options{ + ConnectAddr: strings.TrimPrefix(router.URL, "http://"), + DefaultAtespace: "ate-env", + ResponseHeaderTimeout: 200 * time.Millisecond, + }) + srv := httptest.NewServer(p) + defer srv.Close() + req, _ := http.NewRequest(http.MethodGet, srv.URL+"/", nil) + for k, v := range target("env1", 8000) { + req.Header.Set(k, v) + } + client := &http.Client{Timeout: 10 * time.Second} // fail, not hang, without the limit + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusGatewayTimeout { + t.Errorf("silent daemon: got %d, want 504", resp.StatusCode) + } +}