From 08ed5f1bc5d0d4f6b7b4fb96b4f52c816939b3ae Mon Sep 17 00:00:00 2001 From: Tomer Glottman <83163454+tomergee@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:44:19 -0700 Subject: [PATCH 1/9] ate-env-api: forward HTTP to a port inside an environment through the router A daemon that a runtime layer starts next to the guest (an exec server, a code-execution server, OpenSandbox's execd) was reachable only from inside its environment: ate-env-api proxied the guest alone. The atenet router already tunnels an HTTP CONNECT to any port on an actor, naming the actor in ate-target-actor and the port in the authority, and resumes the actor first. internal/portproxy uses that. A plain HTTP request carrying X-Env-Id, X-Env-Port and optionally X-Env-Atespace goes to that port in that environment, whatever its path. The proxy validates the target (DNS labels, 1-65535), strips the routing headers, opens a CONNECT tunnel through the router's CONNECT listener, speaks HTTP/1.1 to the daemon, flushes streamed responses at once, and keeps idle tunnels in one pool per environment and port. A router 404 or 503 reaches the caller as such; other failures are 502, a handshake timeout 504. ate-env-api serves it on its existing listener: requests with X-Env-Port go to the proxy before the gRPC and MCP handlers, so no new port or manifest change. --atenet-connect names the router's CONNECT listener (default atenet-router.ate-system.svc.cluster.local:8081, the stock install); empty turns forwarding off. The fake router answers CONNECT the way the real one does, for tests. --- README.md | 15 + cmd/ate-env-api/main.go | 17 +- docs/rootfs-images/DESIGN.md | 8 +- docs/rootfs-images/RUNTIMES.md | 27 +- internal/ate/client.go | 3 + .../internaltest/fakerouter/fakerouter.go | 101 ++++++- internal/portproxy/portproxy.go | 274 ++++++++++++++++++ internal/portproxy/portproxy_test.go | 256 ++++++++++++++++ 8 files changed, 685 insertions(+), 16 deletions(-) create mode 100644 internal/portproxy/portproxy.go create mode 100644 internal/portproxy/portproxy_test.go diff --git a/README.md b/README.md index d80333a6..bea6ddcc 100644 --- a/README.md +++ b/README.md @@ -289,6 +289,21 @@ 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. + +```bash +curl -H 'X-Env-Id: dev1' -H 'X-Env-Port: 6000' localhost:7777/health +``` + +`--atenet-connect` names the router's CONNECT listener (default +`atenet-router.ate-system.svc.cluster.local:8081`); an empty value turns the forwarding off. + ## 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..b1d2d6ee 100644 --- a/cmd/ate-env-api/main.go +++ b/cmd/ate-env-api/main.go @@ -29,6 +29,7 @@ import ( "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,6 +39,7 @@ 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 are forwarded through it to that port in the environment. Empty disables the port proxy") 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") @@ -86,6 +88,17 @@ 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 != "" { + handler = portproxy.Dispatch(portproxy.New(portproxy.Options{ + ConnectAddr: *atenetConn, + HostSuffix: *hostSuffix, + DefaultAtespace: ate.DefaultAtespace, + }), 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 +107,7 @@ func main() { srv := &http.Server{ Addr: *listen, - Handler: mux, + Handler: handler, Protocols: &protocols, } @@ -107,7 +120,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..a2accc9b 100644 --- a/docs/rootfs-images/RUNTIMES.md +++ b/docs/rootfs-images/RUNTIMES.md @@ -84,14 +84,25 @@ 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 + ``` + + The tunnel carries HTTP/1.1, not raw TCP. Streamed responses are flushed as + they arrive; gRPC and WebSocket to the runtime are not supported or tested + yet. The NeMo Gym + provider hands this out as its `endpoint(handle, port)` + ([integrations/nemo-gym](../../integrations/nemo-gym/README.md)). ## Suspend, resume and restarts 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..71bcb698 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,77 @@ 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 { 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..d5d6dc10 --- /dev/null +++ b/internal/portproxy/portproxy.go @@ -0,0 +1,274 @@ +// 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. 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. +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 + +// 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 +} + +// Proxy is an http.Handler that forwards each request to the environment port +// its headers name. +type Proxy struct { + opts Options + rp *httputil.ReverseProxy +} + +// New returns a Proxy. +func New(opts Options) *Proxy { + if opts.HostSuffix == "" { + opts.HostSuffix = ate.DefaultHostSuffix + } + if opts.HandshakeTimeout <= 0 { + opts.HandshakeTimeout = DefaultHandshakeTimeout + } + p := &Proxy{opts: opts} + p.rp = &httputil.ReverseProxy{ + Rewrite: p.rewrite, + Transport: &http.Transport{ + DialContext: p.dial, + MaxIdleConnsPerHost: 16, + IdleConnTimeout: 90 * time.Second, + // Pass encodings through untouched; the daemon and the caller + // negotiate them. + DisableCompression: true, + }, + // 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 +} + +// Wants reports whether r is addressed to the proxy: it carries +// EnvPortHeader. Callers of the gRPC API never send it. +func Wants(r *http.Request) bool { + return r.Header.Get(EnvPortHeader) != "" +} + +// Dispatch returns a handler that sends requests the proxy Wants to p and all +// others to next. +func Dispatch(p *Proxy, next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if Wants(r) { + p.ServeHTTP(w, r) + return + } + next.ServeHTTP(w, r) + }) +} + +type target struct { + atespace string + id string + port int +} + +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])?$`) + +func parseTarget(r *http.Request, defaultAtespace string) (target, error) { + t := target{ + id: r.Header.Get(EnvIDHeader), + atespace: r.Header.Get(EnvAtespaceHeader), + } + if t.id == "" { + return t, fmt.Errorf("%s is required next to %s", EnvIDHeader, EnvPortHeader) + } + if !label.MatchString(t.id) { + return t, fmt.Errorf("%s %q is not a valid environment id (lowercase alphanumerics and '-', at most 63)", EnvIDHeader, t.id) + } + if t.atespace == "" { + t.atespace = defaultAtespace + } + if !label.MatchString(t.atespace) { + return t, fmt.Errorf("%s %q is not a valid atespace", EnvAtespaceHeader, t.atespace) + } + port, err := strconv.Atoi(r.Header.Get(EnvPortHeader)) + if err != nil || port < 1 || port > 65535 { + return t, fmt.Errorf("%s %q is not a TCP port (1-65535)", EnvPortHeader, r.Header.Get(EnvPortHeader)) + } + t.port = port + return t, nil +} + +// 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)) + 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). +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, ".") + + 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) } + +func (p *Proxy) handleError(w http.ResponseWriter, r *http.Request, err error) { + t, _ := r.Context().Value(targetKey{}).(target) + var re *routerError + status := http.StatusBadGateway + 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. + if re.status == http.StatusNotFound || re.status == http.StatusServiceUnavailable { + status = re.status + } + case errors.Is(err, context.DeadlineExceeded): + status = http.StatusGatewayTimeout + } + 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: %v", t.atespace, t.id, t.port, err), status) +} diff --git a/internal/portproxy/portproxy_test.go b/internal/portproxy/portproxy_test.go new file mode 100644 index 00000000..813d1ab1 --- /dev/null +++ b/internal/portproxy/portproxy_test.go @@ -0,0 +1,256 @@ +// 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 ( + "fmt" + "io" + "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 := portproxy.New(portproxy.Options{ + ConnectAddr: addr, + DefaultAtespace: "ate-env", + HandshakeTimeout: 5 * time.Second, + }) + proxy := httptest.NewServer(p) + t.Cleanup(proxy.Close) + return &setup{router: rr, routerAddr: addr, proxy: proxy} +} + +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"}, + "bad atespace": {portproxy.EnvIDHeader: "env1", portproxy.EnvPortHeader: "8000", portproxy.EnvAtespaceHeader: "x y"}, + } { + 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) + } + } + 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) + } +} + +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(portproxy.New(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) + } +} From 2ac45d511e4ece92368cf58bdeb42e456f3896d4 Mon Sep 17 00:00:00 2001 From: Tomer Glottman <83163454+tomergee@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:44:20 -0700 Subject: [PATCH 2/9] integrations/nemo-gym: endpoint() for services inside the sandbox Implements nemo-gym 0.7's SupportsSandboxEndpoint. create() keeps SandboxSpec.ports on the handle; endpoint(handle, port) refuses an undeclared port with ValueError, as Gym's Docker provider does, wakes the sandbox with one guest probe, and returns ate-env-api's port proxy as the URL with the X-Env-Id, X-Env-Atespace and X-Env-Port routing headers. connection.endpoint_url overrides the URL when callers reach ate-env-api at another address than the provider. _compat gains a SandboxEndpoint mirror for nemo-gym 0.5 and the hermetic tests. The README documents the injected-daemon setup: a runtime layer and sidecar on the base template, ports on the spec, endpoint() at call time, and the limits (HTTP/1.1, self-contained daemon, port 80 is the guest's). --- integrations/nemo-gym/README.md | 68 +++++++++++++++ .../src/nemo_gym_substrate/_compat.py | 31 +++++++ .../src/nemo_gym_substrate/provider.py | 68 +++++++++++++-- integrations/nemo-gym/tests/test_provider.py | 85 ++++++++++++++++++- 4 files changed, 244 insertions(+), 8 deletions(-) diff --git a/integrations/nemo-gym/README.md b/integrations/nemo-gym/README.md index 68b247dd..30023df1 100644 --- a/integrations/nemo-gym/README.md +++ b/integrations/nemo-gym/README.md @@ -206,6 +206,74 @@ 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()` + wakes a parked sandbox before it returns. A port that `SandboxSpec.ports` did not declare is + a `ValueError`, as in Gym's Docker provider. 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; gRPC and WebSocket to the daemon are not supported or tested 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. +- 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 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..eb1dbb4d 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,12 @@ # 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" + # 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 +123,29 @@ 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: + if self.endpoint_url: + return self.endpoint_url + return self.api_url if "://" in self.api_url else f"http://{self.api_url}" + @dataclass(frozen=True) class CreateConfig: @@ -256,6 +275,9 @@ async def create(self, spec: SandboxSpec) -> SandboxHandle: "atespace": atespace, "workdir": spec.workdir, "env": dict(spec.env or {}), + # Ports a service inside the sandbox listens on (nemo-gym + # >= 0.7); endpoint() resolves only these. + "ports": [int(p) for p in (getattr(spec, "ports", None) or ())], }, ) @@ -347,6 +369,35 @@ 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. The sandbox is woken first, so the + endpoint answers at once even after Substrate parked it. + """ + raw = handle.raw or {} + declared = list(raw.get("ports") or ()) + if port not in declared: + raise ValueError( + f"Sandbox port {port} was not declared in SandboxSpec.ports; " + f"declared ports: {declared!r}" + ) + env = self._env(handle) + await self._wait_ready(env, self._create.ready_timeout_s, error=RuntimeError) + return SandboxEndpoint( + endpoint=self._connection.endpoint_base(), + headers={ + _ENV_ID_HEADER: handle.sandbox_id, + _ENV_ATESPACE_HEADER: self._atespace_of(raw), + _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 +452,17 @@ 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, + ) -> 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.""" deadline = time.monotonic() + timeout_s last_error = "no probe attempted" while True: @@ -420,9 +478,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..63899da6 100644 --- a/integrations/nemo-gym/tests/test_provider.py +++ b/integrations/nemo-gym/tests/test_provider.py @@ -36,6 +36,7 @@ from nemo_gym_substrate._compat import ( SandboxCreateError, SandboxCreateVerificationError, + SandboxEndpoint, SandboxSpec, SandboxStatus, ) @@ -45,11 +46,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 +450,78 @@ 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", + } + # The sandbox is woken (one guest probe) before the endpoint is handed out. + assert _probe_count(fake) == probes + 1 + + 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_when_the_sandbox_does_not_answer(): + async def t(): + async with stack({"ready_timeout_s": 0.2}) as (provider, fake): + handle = await provider.create(SandboxSpec(ports=(6000,))) + fake.ready_after_probes = 10_000 + _probe_count(fake) + fake._probes = 0 + with pytest.raises(RuntimeError, match="not ready"): + await provider.endpoint(handle, 6000) + + 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" + with pytest.raises(ValueError, match="unknown option"): + provider_mod.ConnectionConfig.from_mapping({"endpoint": "x"}) From 90b3d81159b1a860293563112ad483e6e0481dfa Mon Sep 17 00:00:00 2001 From: Tomer Glottman <83163454+tomergee@users.noreply.github.com> Date: Thu, 15 Oct 2026 12:04:40 -0700 Subject: [PATCH 3/9] portproxy: path form /v1alpha/envs/{id}/ports/{port}/... for clients without headers WebSocket clients and browsers cannot set the routing headers. The OpenHands SDK, for one, opens its event WebSocket at /sockets/events/ with no custom headers, but keeps the host's path. The path form carries the target in a prefix the proxy strips before forwarding: /v1alpha/envs/{id}/ports/{port}/... in the default atespace and /v1alpha/atespaces/{atespace}/envs/{id}/ports/{port}/... otherwise. Both forms share one validation; escaping after the prefix is kept as sent. Tests: prefix stripping with query and escaped segments, the bare prefix, the atespace form, bad ids, ports and atespaces, other /v1alpha paths still reaching the API, and an HTTP upgrade (a WebSocket handshake) through the proxy and the tunnel in both forms. --- cmd/ate-env-api/main.go | 6 +- internal/portproxy/portproxy.go | 68 +++++++++--- internal/portproxy/portproxy_test.go | 149 +++++++++++++++++++++++++++ 3 files changed, 207 insertions(+), 16 deletions(-) diff --git a/cmd/ate-env-api/main.go b/cmd/ate-env-api/main.go index b1d2d6ee..fe656a48 100644 --- a/cmd/ate-env-api/main.go +++ b/cmd/ate-env-api/main.go @@ -92,11 +92,13 @@ func main() { // anything else, whatever their path: the path belongs to the daemon. var handler http.Handler = mux if *atenetConn != "" { - handler = portproxy.Dispatch(portproxy.New(portproxy.Options{ + pp := portproxy.New(portproxy.Options{ ConnectAddr: *atenetConn, HostSuffix: *hostSuffix, DefaultAtespace: ate.DefaultAtespace, - }), mux) + }) + pp.Register(mux) // the path form, /v1alpha/envs/{id}/ports/{port}/... + handler = portproxy.Dispatch(pp, mux) } // Enable both HTTP/1.1 (for /healthz and /v1alpha/ HTTP endpoints) and diff --git a/internal/portproxy/portproxy.go b/internal/portproxy/portproxy.go index d5d6dc10..fae865e8 100644 --- a/internal/portproxy/portproxy.go +++ b/internal/portproxy/portproxy.go @@ -16,8 +16,11 @@ // 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. The proxy opens an HTTP CONNECT tunnel through +// 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 @@ -126,6 +129,9 @@ 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{} @@ -135,30 +141,61 @@ type targetKey struct{} var label = regexp.MustCompile(`^[a-z0-9]([-a-z0-9]{0,61}[a-z0-9])?$`) func parseTarget(r *http.Request, defaultAtespace string) (target, error) { - t := target{ - id: r.Header.Get(EnvIDHeader), - atespace: r.Header.Get(EnvAtespaceHeader), - } - if t.id == "" { - return t, fmt.Errorf("%s is required next to %s", EnvIDHeader, EnvPortHeader) + 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, port, defaultAtespace string) (target, error) { + t := target{id: id, atespace: atespace} if !label.MatchString(t.id) { - return t, fmt.Errorf("%s %q is not a valid environment id (lowercase alphanumerics and '-', at most 63)", EnvIDHeader, 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("%s %q is not a valid atespace", EnvAtespaceHeader, t.atespace) + return t, fmt.Errorf("atespace %q is not valid", t.atespace) } - port, err := strconv.Atoi(r.Header.Get(EnvPortHeader)) - if err != nil || port < 1 || port > 65535 { - return t, fmt.Errorf("%s %q is not a TCP port (1-65535)", EnvPortHeader, r.Header.Get(EnvPortHeader)) + n, err := strconv.Atoi(port) + if err != nil || n < 1 || n > 65535 { + return t, fmt.Errorf("port %q is not a TCP port (1-65535)", port) } - t.port = port + t.port = n return t, nil } +// Register adds the path form of the proxy to mux. +func (p *Proxy) Register(mux *http.ServeMux) { + mux.HandleFunc("/v1alpha/envs/{id}/ports/{port}/{rest...}", p.servePath) + mux.HandleFunc("/v1alpha/atespaces/{atespace}/envs/{id}/ports/{port}/{rest...}", p.servePath) +} + +func (p *Proxy) servePath(w http.ResponseWriter, r *http.Request) { + t, err := makeTarget(r.PathValue("id"), r.PathValue("atespace"), r.PathValue("port"), p.opts.DefaultAtespace) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + prefix := "/v1alpha/envs/" + r.PathValue("id") + "/ports/" + r.PathValue("port") + if a := r.PathValue("atespace"); a != "" { + prefix = "/v1alpha/atespaces/" + a + prefix[len("/v1alpha"):] + } + // The prefix is made of validated labels and digits, so it reads the same + // escaped and unescaped; what follows keeps its own escaping. + t.rawPath = strings.TrimPrefix(r.URL.EscapedPath(), prefix) + if t.path, err = url.PathUnescape(t.rawPath); err != nil { + http.Error(w, "bad path escaping: "+err.Error(), http.StatusBadRequest) + return + } + if t.rawPath == t.path { + t.rawPath = "" + } + 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) @@ -178,6 +215,9 @@ func (p *Proxy) rewrite(pr *httputil.ProxyRequest) { 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) } diff --git a/internal/portproxy/portproxy_test.go b/internal/portproxy/portproxy_test.go index 813d1ab1..8926b1f8 100644 --- a/internal/portproxy/portproxy_test.go +++ b/internal/portproxy/portproxy_test.go @@ -15,8 +15,10 @@ package portproxy_test import ( + "bufio" "fmt" "io" + "net" "net/http" "net/http/httptest" "strings" @@ -254,3 +256,150 @@ func TestDispatch(t *testing.T) { t.Errorf("request with %s got %q, want the daemon's answer", portproxy.EnvPortHeader, b) } } + +// pathSetup serves the proxy the way ate-env-api does: path form on a mux, +// header form dispatched in front of it. +func pathSetup(t *testing.T) (*setup, string) { + t.Helper() + s := newSetup(t) + p := portproxy.New(portproxy.Options{ConnectAddr: s.routerAddr, DefaultAtespace: "ate-env", HandshakeTimeout: 5 * time.Second}) + mux := http.NewServeMux() + p.Register(mux) + 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/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() + } +} From d6d8e78f26332fd81af7495a803d59184d212c77 Mon Sep 17 00:00:00 2001 From: Tomer Glottman <83163454+tomergee@users.noreply.github.com> Date: Thu, 15 Oct 2026 12:13:22 -0700 Subject: [PATCH 4/9] docs: the port proxy's path form, and WebSocket upgrades through it The path form (/v1alpha/envs/{id}/ports/{port}/...) is documented next to the header form in the README and RUNTIMES.md. WebSocket upgrades pass through the proxy and the router: verified live with the OpenHands agent-server's event stream (see the follow-up example). gRPC stays out. --- README.md | 10 ++++++++++ docs/rootfs-images/RUNTIMES.md | 8 ++++++-- integrations/nemo-gym/README.md | 3 ++- 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index bea6ddcc..46f85385 100644 --- a/README.md +++ b/README.md @@ -301,6 +301,16 @@ its path; the API strips the three headers. The router resumes a suspended envir 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 +``` + `--atenet-connect` names the router's CONNECT listener (default `atenet-router.ate-system.svc.cluster.local:8081`); an empty value turns the forwarding off. diff --git a/docs/rootfs-images/RUNTIMES.md b/docs/rootfs-images/RUNTIMES.md index a2accc9b..77482216 100644 --- a/docs/rootfs-images/RUNTIMES.md +++ b/docs/rootfs-images/RUNTIMES.md @@ -96,10 +96,14 @@ sidecar, run under two different rootfs images. ```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 tunnel carries HTTP/1.1, not raw TCP. Streamed responses are flushed as - they arrive; gRPC and WebSocket to the runtime are not supported or tested + 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)). diff --git a/integrations/nemo-gym/README.md b/integrations/nemo-gym/README.md index 30023df1..0254b6f0 100644 --- a/integrations/nemo-gym/README.md +++ b/integrations/nemo-gym/README.md @@ -264,7 +264,8 @@ actor network. Limits: - **HTTP/1.1 only.** The router tunnels HTTP to the port, not raw TCP. Streamed responses are - flushed as they arrive; gRPC and WebSocket to the daemon are not supported or tested yet. + 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 From 7621f273b5f65efc78716f73513d3dbf5c348c4b Mon Sep 17 00:00:00 2001 From: Tomer Glottman <83163454+tomergee@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:31:06 -0700 Subject: [PATCH 5/9] portproxy: segment-matched path form, cross-origin checks, stricter targets - Match the path form in Dispatch by segment on the escaped path instead of on the ServeMux. An escaped routing segment (%65nv1) now ends the prefix where it should; before, the daemon got the whole path. A path without the trailing slash is the daemon's root, and "//" or ".." are passed as sent: the mux redirected all three, and WebSocket clients do not follow redirects. Register is gone. - Refuse a state-changing request or a WebSocket handshake that a browser sends to the path form from another origin (http.CrossOriginProtection, with the handshake judged as a POST). --port-proxy-trusted-origins adds origins. The header form needs no check: custom headers to another origin need a CORS preflight, which nothing on the listener approves. - Reject a routing header sent twice, and a port with a sign or a leading zero. A gRPC call that carries X-Env-Port is not proxied. - Error bodies name the target and a short cause; the router's message and address go to the log only. A router 502 says the daemon must listen on 0.0.0.0. A transport timeout is a 504. - --port-proxy-response-header-timeout bounds the wait for a daemon's response headers; the default stays no limit, because a daemon can answer a long command only when it ends. - Document the trust model in the package doc and the flag help: no authentication, any port of any actor the router serves. --- cmd/ate-env-api/main.go | 26 +++- internal/portproxy/portproxy.go | 195 +++++++++++++++++++------- internal/portproxy/portproxy_test.go | 200 ++++++++++++++++++++++++++- 3 files changed, 360 insertions(+), 61 deletions(-) diff --git a/cmd/ate-env-api/main.go b/cmd/ate-env-api/main.go index fe656a48..2966f835 100644 --- a/cmd/ate-env-api/main.go +++ b/cmd/ate-env-api/main.go @@ -23,6 +23,7 @@ import ( "net/http" "os" "os/signal" + "strings" "syscall" "time" @@ -39,7 +40,9 @@ 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 are forwarded through it to that port in the environment. Empty disables the port proxy") + 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") @@ -92,12 +95,23 @@ func main() { // anything else, whatever their path: the path belongs to the daemon. var handler http.Handler = mux if *atenetConn != "" { - pp := portproxy.New(portproxy.Options{ - ConnectAddr: *atenetConn, - HostSuffix: *hostSuffix, - DefaultAtespace: ate.DefaultAtespace, + 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, }) - pp.Register(mux) // the path form, /v1alpha/envs/{id}/ports/{port}/... + 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) } diff --git a/internal/portproxy/portproxy.go b/internal/portproxy/portproxy.go index fae865e8..d6951999 100644 --- a/internal/portproxy/portproxy.go +++ b/internal/portproxy/portproxy.go @@ -20,12 +20,19 @@ // 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. +// /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 ( @@ -71,6 +78,14 @@ type Options struct { // 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 } // Proxy is an http.Handler that forwards each request to the environment port @@ -78,23 +93,30 @@ type Options struct { type Proxy struct { opts Options rp *httputil.ReverseProxy + cop *http.CrossOriginProtection } -// New returns a Proxy. -func New(opts Options) *Proxy { +// 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 } - p := &Proxy{opts: opts} + 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.rp = &httputil.ReverseProxy{ Rewrite: p.rewrite, Transport: &http.Transport{ - DialContext: p.dial, - MaxIdleConnsPerHost: 16, - IdleConnTimeout: 90 * time.Second, + DialContext: p.dial, + MaxIdleConnsPerHost: 16, + IdleConnTimeout: 90 * time.Second, + ResponseHeaderTimeout: opts.ResponseHeaderTimeout, // Pass encodings through untouched; the daemon and the caller // negotiate them. DisableCompression: true, @@ -104,24 +126,32 @@ func New(opts Options) *Proxy { FlushInterval: -1, ErrorHandler: p.handleError, } - return p + return p, nil } -// Wants reports whether r is addressed to the proxy: it carries -// EnvPortHeader. Callers of the gRPC API never send it. +// 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. func Wants(r *http.Request) bool { - return r.Header.Get(EnvPortHeader) != "" + return r.Header.Get(EnvPortHeader) != "" && + !strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") } -// Dispatch returns a handler that sends requests the proxy Wants to p and all -// others to next. +// 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) { - if Wants(r) { + switch { + case Wants(r): p.ServeHTTP(w, r) - return + case isPathForm(r.URL.EscapedPath()): + p.ServePath(w, r) + default: + next.ServeHTTP(w, r) } - next.ServeHTTP(w, r) }) } @@ -140,7 +170,15 @@ type targetKey struct{} // 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) } @@ -148,7 +186,7 @@ func parseTarget(r *http.Request, defaultAtespace string) (target, error) { } // makeTarget validates a target from either form. -func makeTarget(id, atespace, port, defaultAtespace string) (target, error) { +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) @@ -159,40 +197,92 @@ func makeTarget(id, atespace, port, defaultAtespace string) (target, error) { if !label.MatchString(t.atespace) { return t, fmt.Errorf("atespace %q is not valid", t.atespace) } - n, err := strconv.Atoi(port) - if err != nil || n < 1 || n > 65535 { - return t, fmt.Errorf("port %q is not a TCP port (1-65535)", port) + 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) } t.port = n return t, nil } -// Register adds the path form of the proxy to mux. -func (p *Proxy) Register(mux *http.ServeMux) { - mux.HandleFunc("/v1alpha/envs/{id}/ports/{port}/{rest...}", p.servePath) - mux.HandleFunc("/v1alpha/atespaces/{atespace}/envs/{id}/ports/{port}/{rest...}", p.servePath) +// 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 } -func (p *Proxy) servePath(w http.ResponseWriter, r *http.Request) { - t, err := makeTarget(r.PathValue("id"), r.PathValue("atespace"), r.PathValue("port"), p.opts.DefaultAtespace) +// 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 } - prefix := "/v1alpha/envs/" + r.PathValue("id") + "/ports/" + r.PathValue("port") - if a := r.PathValue("atespace"); a != "" { - prefix = "/v1alpha/atespaces/" + a + prefix[len("/v1alpha"):] + 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 } - // The prefix is made of validated labels and digits, so it reads the same - // escaped and unescaped; what follows keeps its own escaping. - t.rawPath = strings.TrimPrefix(r.URL.EscapedPath(), prefix) - if t.path, err = url.PathUnescape(t.rawPath); err != nil { - http.Error(w, "bad path escaping: "+err.Error(), http.StatusBadRequest) + 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 } - if t.rawPath == t.path { - t.rawPath = "" - } p.rp.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), targetKey{}, t))) } @@ -293,22 +383,31 @@ type bufferedConn struct { func (c *bufferedConn) Read(b []byte) (int, error) { return c.r.Read(b) } +// 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 - status := http.StatusBadGateway + 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. - if re.status == http.StatusNotFound || re.status == http.StatusServiceUnavailable { - status = re.status + 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): - status = http.StatusGatewayTimeout + case errors.Is(err, context.DeadlineExceeded), errors.As(err, &ne) && ne.Timeout(): + status, cause = http.StatusGatewayTimeout, "did not answer in time" } 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: %v", t.atespace, t.id, t.port, err), status) + 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 index 8926b1f8..1b76f3a8 100644 --- a/internal/portproxy/portproxy_test.go +++ b/internal/portproxy/portproxy_test.go @@ -66,7 +66,7 @@ func newSetup(t *testing.T) *setup { srv := httptest.NewServer(rr) t.Cleanup(srv.Close) addr := strings.TrimPrefix(srv.URL, "http://") - p := portproxy.New(portproxy.Options{ + p := mustNew(t, portproxy.Options{ ConnectAddr: addr, DefaultAtespace: "ate-env", HandshakeTimeout: 5 * time.Second, @@ -76,6 +76,15 @@ func newSetup(t *testing.T) *setup { return &setup{router: rr, routerAddr: addr, proxy: proxy} } +func mustNew(t *testing.T, opts portproxy.Options) *portproxy.Proxy { + t.Helper() + p, err := portproxy.New(opts) + if err != nil { + t.Fatal(err) + } + 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)) @@ -188,12 +197,27 @@ func TestRejectsBadTargets(t *testing.T) { "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"}, } { 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) } @@ -234,7 +258,7 @@ 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(portproxy.New(portproxy.Options{ConnectAddr: s.routerAddr, DefaultAtespace: "ate-env"}), next) + 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)) @@ -255,16 +279,26 @@ func TestDispatch(t *testing.T) { 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: path form on a mux, -// header form dispatched in front of it. -func pathSetup(t *testing.T) (*setup, 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 := portproxy.New(portproxy.Options{ConnectAddr: s.routerAddr, DefaultAtespace: "ate-env", HandshakeTimeout: 5 * time.Second}) + p := mustNew(t, portproxy.Options{ConnectAddr: s.routerAddr, DefaultAtespace: "ate-env", HandshakeTimeout: 5 * time.Second, TrustedOrigins: origins}) mux := http.NewServeMux() - p.Register(mux) mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "api") }) srv := httptest.NewServer(portproxy.Dispatch(p, mux)) t.Cleanup(srv.Close) @@ -334,6 +368,9 @@ func TestPathFormAtespaceAndValidation(t *testing.T) { "/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) @@ -403,3 +440,152 @@ func TestUpgradePassesThrough(t *testing.T) { 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) + } +} From 27ed10551b352f465a5340d1a7d5a00a61ff8b04 Mon Sep 17 00:00:00 2001 From: Tomer Glottman <83163454+tomergee@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:31:06 -0700 Subject: [PATCH 6/9] integrations/nemo-gym: endpoint() stops at once for a deleted sandbox endpoint() polled a sandbox that no longer exists for the full ready timeout (120 s by default). It now raises RuntimeError on the first NotFound. The create path keeps retrying as before. The README states the listen rule (0.0.0.0, or 502 through the endpoint) and that nothing authenticates callers in front of the daemon. --- integrations/nemo-gym/README.md | 11 +++++++++-- .../nemo-gym/src/nemo_gym_substrate/provider.py | 13 +++++++++++-- integrations/nemo-gym/tests/test_provider.py | 14 ++++++++++++++ 3 files changed, 34 insertions(+), 4 deletions(-) diff --git a/integrations/nemo-gym/README.md b/integrations/nemo-gym/README.md index 0254b6f0..71c2261d 100644 --- a/integrations/nemo-gym/README.md +++ b/integrations/nemo-gym/README.md @@ -257,8 +257,9 @@ actor network. `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()` - wakes a parked sandbox before it returns. A port that `SandboxSpec.ports` did not declare is - a `ValueError`, as in Gym's Docker provider. Set `connection.endpoint_url` when the code that + wakes a parked sandbox before it returns, and raises `RuntimeError` at once for a sandbox that + no longer exists. A port that `SandboxSpec.ports` did not declare is a `ValueError`, as in Gym's + Docker provider. Set `connection.endpoint_url` when the code that calls the daemon reaches ate-env-api at another address than the provider does. Limits: @@ -270,6 +271,12 @@ Limits: 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 diff --git a/integrations/nemo-gym/src/nemo_gym_substrate/provider.py b/integrations/nemo-gym/src/nemo_gym_substrate/provider.py index eb1dbb4d..53324e21 100644 --- a/integrations/nemo-gym/src/nemo_gym_substrate/provider.py +++ b/integrations/nemo-gym/src/nemo_gym_substrate/provider.py @@ -388,7 +388,9 @@ async def endpoint(self, handle: SandboxHandle, port: int) -> SandboxEndpoint: f"declared ports: {declared!r}" ) env = self._env(handle) - await self._wait_ready(env, self._create.ready_timeout_s, error=RuntimeError) + await self._wait_ready( + env, self._create.ready_timeout_s, error=RuntimeError, stop_on_not_found=True + ) return SandboxEndpoint( endpoint=self._connection.endpoint_base(), headers={ @@ -458,11 +460,14 @@ async def _wait_ready( 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, then - ``error`` is raised. The first call also resumes a parked actor.""" + ``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: @@ -471,6 +476,10 @@ async def _wait_ready( 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: diff --git a/integrations/nemo-gym/tests/test_provider.py b/integrations/nemo-gym/tests/test_provider.py index 63899da6..7d63ecd4 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 @@ -518,6 +519,19 @@ async def t(): 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" From 37fe29a222c55f063337e1d8347e870e3739c025 Mon Sep 17 00:00:00 2001 From: Tomer Glottman <83163454+tomergee@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:31:06 -0700 Subject: [PATCH 7/9] docs: who can reach a port through ate-env-api, and the new flags The README section on reaching a service now states the trust model (no authentication; any port of any actor the router serves, in any atespace), what to do about it, the browser checks on the path form, the DNS-rebinding gap that needs a Host check across the whole listener, and a table of the three port-proxy flags. RUNTIMES.md adds the listen rule and the authentication note to "Who can reach the second runtime". --- README.md | 27 +++++++++++++++++++++++++-- docs/rootfs-images/RUNTIMES.md | 9 +++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 46f85385..369c1343 100644 --- a/README.md +++ b/README.md @@ -311,8 +311,31 @@ through either form; gRPC to the daemon is not supported yet. curl localhost:7777/v1alpha/envs/dev1/ports/6000/health ``` -`--atenet-connect` names the router's CONNECT listener (default -`atenet-router.ate-system.svc.cluster.local:8081`); an empty value turns the forwarding off. +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 diff --git a/docs/rootfs-images/RUNTIMES.md b/docs/rootfs-images/RUNTIMES.md index 77482216..cd50f31d 100644 --- a/docs/rootfs-images/RUNTIMES.md +++ b/docs/rootfs-images/RUNTIMES.md @@ -107,6 +107,15 @@ sidecar, run under two different rootfs images. 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 From 98e641f622a11e82b3cb1377536028217639e422 Mon Sep 17 00:00:00 2001 From: Tomer Glottman <83163454+tomergee@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:17:38 -0700 Subject: [PATCH 8/9] port proxy: wait for a suspended environment to wake, refuse the guest port, bound idle tunnels; provider validates ports and no longer wakes on endpoint() Review follow-ups on #92. Proxy (internal/portproxy): - dial repeats the CONNECT while the router answers 503, with backoff, until --wake-budget is spent: the first CONNECT resumes a suspended actor and the router holds it only for its own parking budget, so a cold image pull or a cross-region snapshot read used to surface as 503. Same budget and flag as the guest calls from #86; the body still names the target when the budget runs out. - Port 80 is refused with 400: it is the ate-env guest, which the API serves itself, not a daemon port. - Idle tunnels: 10 s idle timeout, 4 per environment port, 1024 in total (the zero value meant no limit across hosts), so idle environments are not held open by the API and a large fleet does not pin sockets. - Close drops the idle tunnels; bufferedConn forwards CloseWrite, which httputil.ReverseProxy calls on the backend after an upgrade and the embedded net.Conn hid. - No log line for a router 404 or 503; those are the caller's answer. - Fake router: the one-connection listener's Close ends its Accept, and the tests drop idle tunnels in cleanup, so no tunnel goroutine outlives its test. NeMo Gym provider: - endpoint_base prepends http:// to a bare host:port endpoint_url too; SandboxEndpoint rejects a URL without a scheme. - create() validates SandboxSpec.ports before the actor starts: integers (bool excluded), 1-65535, not 80. - endpoint() no longer probes the guest. It checks that the sandbox exists from the control plane (deleted or crashed raise RuntimeError at once) and returns the static endpoint; the first request through the proxy wakes a parked sandbox and the proxy waits for it. - README: the stale ports/SandboxEndpoint limit is gone; the endpoint paragraph describes the new behaviour. Top-level README notes the wait. --- README.md | 4 +- cmd/ate-env-api/main.go | 3 +- integrations/nemo-gym/README.md | 14 ++- .../src/nemo_gym_substrate/provider.py | 75 +++++++++--- integrations/nemo-gym/tests/test_provider.py | 62 ++++++++-- .../internaltest/fakerouter/fakerouter.go | 5 +- internal/portproxy/portproxy.go | 109 ++++++++++++++++-- internal/portproxy/portproxy_test.go | 60 ++++++++++ 8 files changed, 289 insertions(+), 43 deletions(-) diff --git a/README.md b/README.md index 369c1343..cf18493e 100644 --- a/README.md +++ b/README.md @@ -295,7 +295,9 @@ A daemon that listens on a port inside an environment, such as one a runtime lay 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. +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 diff --git a/cmd/ate-env-api/main.go b/cmd/ate-env-api/main.go index 2966f835..666543ff 100644 --- a/cmd/ate-env-api/main.go +++ b/cmd/ate-env-api/main.go @@ -46,7 +46,7 @@ func main() { 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() @@ -107,6 +107,7 @@ func main() { DefaultAtespace: ate.DefaultAtespace, ResponseHeaderTimeout: *ppHeaderTimeout, TrustedOrigins: origins, + WakeBudget: *wakeBudget, }) if err != nil { log.Fatalf("--port-proxy-trusted-origins: %v", err) diff --git a/integrations/nemo-gym/README.md b/integrations/nemo-gym/README.md index 71c2261d..77cb0eaa 100644 --- a/integrations/nemo-gym/README.md +++ b/integrations/nemo-gym/README.md @@ -256,11 +256,14 @@ actor network. ``` `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()` - wakes a parked sandbox before it returns, and raises `RuntimeError` at once for a sandbox that - no longer exists. A port that `SandboxSpec.ports` did not declare is a `ValueError`, as in Gym's - Docker provider. Set `connection.endpoint_url` when the code that - calls the daemon reaches ate-env-api at another address than the provider does. + 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: @@ -351,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/provider.py b/integrations/nemo-gym/src/nemo_gym_substrate/provider.py index 53324e21..89940311 100644 --- a/integrations/nemo-gym/src/nemo_gym_substrate/provider.py +++ b/integrations/nemo-gym/src/nemo_gym_substrate/provider.py @@ -75,6 +75,29 @@ _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 @@ -142,9 +165,11 @@ def from_mapping(cls, options: Mapping[str, Any]) -> "ConnectionConfig": ) def endpoint_base(self) -> str: - if self.endpoint_url: - return self.endpoint_url - return self.api_url if "://" in self.api_url else f"http://{self.api_url}" + """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) @@ -233,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]}" @@ -275,9 +303,7 @@ async def create(self, spec: SandboxSpec) -> SandboxHandle: "atespace": atespace, "workdir": spec.workdir, "env": dict(spec.env or {}), - # Ports a service inside the sandbox listens on (nemo-gym - # >= 0.7); endpoint() resolves only these. - "ports": [int(p) for p in (getattr(spec, "ports", None) or ())], + "ports": ports, }, ) @@ -377,25 +403,44 @@ async def endpoint(self, handle: SandboxHandle, port: int) -> SandboxEndpoint: 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. The sandbox is woken first, so the - endpoint answers at once even after Substrate parked it. + 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 port not in declared: + if isinstance(port, bool) or port not in declared: raise ValueError( - f"Sandbox port {port} was not declared in SandboxSpec.ports; " + f"Sandbox port {port!r} was not declared in SandboxSpec.ports; " f"declared ports: {declared!r}" ) - env = self._env(handle) - await self._wait_ready( - env, self._create.ready_timeout_s, error=RuntimeError, stop_on_not_found=True - ) + 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: self._atespace_of(raw), + _ENV_ATESPACE_HEADER: atespace, _ENV_PORT_HEADER: str(port), }, ) diff --git a/integrations/nemo-gym/tests/test_provider.py b/integrations/nemo-gym/tests/test_provider.py index 7d63ecd4..134d0968 100644 --- a/integrations/nemo-gym/tests/test_provider.py +++ b/integrations/nemo-gym/tests/test_provider.py @@ -474,8 +474,9 @@ async def t(): "X-Env-Atespace": "ate-env", "X-Env-Port": "6000", } - # The sandbox is woken (one guest probe) before the endpoint is handed out. - assert _probe_count(fake) == probes + 1 + # 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()) @@ -507,18 +508,59 @@ async def t(): run(t()) -def test_endpoint_raises_when_the_sandbox_does_not_answer(): +def test_endpoint_raises_for_a_crashed_sandbox(): async def t(): - async with stack({"ready_timeout_s": 0.2}) as (provider, fake): + async with stack() as (provider, fake): handle = await provider.create(SandboxSpec(ports=(6000,))) - fake.ready_after_probes = 10_000 + _probe_count(fake) - fake._probes = 0 - with pytest.raises(RuntimeError, match="not ready"): + 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): @@ -537,5 +579,11 @@ def test_endpoint_base_defaults_to_the_api_url(): 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/internaltest/fakerouter/fakerouter.go b/internal/internaltest/fakerouter/fakerouter.go index 71bcb698..8a3ba9f9 100644 --- a/internal/internaltest/fakerouter/fakerouter.go +++ b/internal/internaltest/fakerouter/fakerouter.go @@ -197,5 +197,8 @@ func (l *oneConnListener) Accept() (net.Conn, error) { return nil, net.ErrClosed } -func (l *oneConnListener) Close() error { return nil } +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 index d6951999..fe119219 100644 --- a/internal/portproxy/portproxy.go +++ b/internal/portproxy/portproxy.go @@ -66,6 +66,33 @@ const ( // 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. @@ -86,6 +113,10 @@ type Options struct { // 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 @@ -93,6 +124,7 @@ type Options struct { type Proxy struct { opts Options rp *httputil.ReverseProxy + tr *http.Transport cop *http.CrossOriginProtection } @@ -104,23 +136,28 @@ func New(opts Options) (*Proxy, error) { 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: &http.Transport{ - DialContext: p.dial, - MaxIdleConnsPerHost: 16, - IdleConnTimeout: 90 * time.Second, - ResponseHeaderTimeout: opts.ResponseHeaderTimeout, - // Pass encodings through untouched; the daemon and the caller - // negotiate them. - DisableCompression: true, - }, + 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, @@ -132,6 +169,9 @@ func New(opts Options) (*Proxy, error) { // 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") @@ -201,6 +241,9 @@ func makeTarget(id, atespace, portStr, defaultAtespace string) (target, error) { 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 } @@ -324,14 +367,42 @@ func (e *routerError) Error() string { } // dial opens a CONNECT tunnel to the environment port encoded in addr (see -// rewrite). +// 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 { + return nil, fmt.Errorf("environment %s/%s did not wake within %s (%d attempts): %w", atespace, id, p.opts.WakeBudget, attempt, 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 @@ -383,6 +454,16 @@ type bufferedConn struct { 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. @@ -408,6 +489,10 @@ func (p *Proxy) handleError(w http.ResponseWriter, r *http.Request, err error) { case errors.Is(err, context.DeadlineExceeded), errors.As(err, &ne) && ne.Timeout(): status, cause = http.StatusGatewayTimeout, "did not answer in time" } - log.Printf("port proxy: %s/%s port %d: %v", t.atespace, t.id, t.port, err) + // 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 index 1b76f3a8..3d2b7c47 100644 --- a/internal/portproxy/portproxy_test.go +++ b/internal/portproxy/portproxy_test.go @@ -70,18 +70,22 @@ func newSetup(t *testing.T) *setup { 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 } @@ -200,6 +204,7 @@ func TestRejectsBadTargets(t *testing.T) { "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) @@ -236,6 +241,61 @@ func TestPassesRouterRefusals(t *testing.T) { } } +// 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"} { From b1e61b21baec72182614f245955e95b85d8bcb1f Mon Sep 17 00:00:00 2001 From: Tomer Glottman <83163454+tomergee@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:40:46 -0700 Subject: [PATCH 9/9] port proxy: log once when the wake budget is spent handleError keeps router 503s out of the log, and the budget-exhaustion error wraps the last 503, so a request that waited the whole budget for an environment that never woke left no trace after the first attempt's "waiting" line. dial now logs the exhaustion itself before it returns. --- internal/portproxy/portproxy.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/internal/portproxy/portproxy.go b/internal/portproxy/portproxy.go index fe119219..759ec158 100644 --- a/internal/portproxy/portproxy.go +++ b/internal/portproxy/portproxy.go @@ -387,7 +387,12 @@ func (p *Proxy) dial(ctx context.Context, _, addr string) (net.Conn, error) { } remaining := time.Until(deadline) if remaining <= 0 { - return nil, fmt.Errorf("environment %s/%s did not wake within %s (%d attempts): %w", atespace, id, p.opts.WakeBudget, attempt, err) + // 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)