fix(proxy): track connections through shutdown

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
GitOrigin-RevId: bdf063393e4dbadae8d8af4d99170df450e0bbb7

Author: 2 people

docs/ARCHITECTURE.md

@@ -246,6 +246,9 @@ Key mechanisms (all under `pkg/sandbox/`):
 - **Sandbox proxy** (:5007, `pkg/proxy/`): reverse-proxies incoming traffic from client-proxy to
   the sandbox's slot IP and requested port over HTTP or configured HTTPS, enforcing per-sandbox
   traffic access tokens. HTTPS backends may use self-signed certificates.
+  Shared proxy shutdown waits for H2C and upgraded connections to close and includes them in
+  forced connection cleanup. Failed H2C handshakes release their hijacked connections when the
+  handshake handler returns.
 - Writes sandbox lifecycle **events** and cgroup **host stats** to ClickHouse; exports metrics via
   OTel. Sandbox and template-build log writes go through a flag-resolved HTTP route: the legacy
   collector remains the fallback primary destination, and configured shadow destinations can mirror

packages/orchestrator/go.mod

@@ -73,6 +73,7 @@ require (
 	go.opentelemetry.io/otel/sdk/metric v1.46.0
 	go.opentelemetry.io/otel/trace v1.46.0
 	go.uber.org/zap v1.28.0
+	golang.org/x/net v0.58.0
 	golang.org/x/sync v0.23.0
 	golang.org/x/sys v0.48.0
 	golang.org/x/term v0.45.0
@@ -318,7 +319,6 @@ require (
 	golang.org/x/crypto v0.55.0 // indirect
 	golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f // indirect
 	golang.org/x/mod v0.41.0 // indirect
-	golang.org/x/net v0.58.0 // indirect
 	golang.org/x/oauth2 v0.37.0 // indirect
 	golang.org/x/text v0.41.0 // indirect
 	golang.org/x/time v0.15.0 // indirect

packages/orchestrator/pkg/proxy/shutdown_h2c_test.go

@@ -0,0 +1,120 @@
+//go:build linux
+
+package proxy
+
+import (
+	"context"
+	"crypto/tls"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/net/http2"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/connlimit"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	reverseproxy "github.com/e2b-dev/infra/packages/shared/pkg/proxy"
+	"github.com/e2b-dev/infra/packages/shared/pkg/proxy/pool"
+)
+
+func TestSandboxProxyCloseInterruptsStalledH2CUpload(t *testing.T) {
+	t.Parallel()
+
+	for _, forced := range []bool{false, true} {
+		t.Run(fmt.Sprintf("forced=%t", forced), func(t *testing.T) {
+			t.Parallel()
+
+			backendStarted := make(chan error, 1)
+			backend := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
+				_, err := io.ReadFull(r.Body, make([]byte, 1))
+				backendStarted <- err
+				_, _ = io.Copy(io.Discard, r.Body)
+			}))
+			t.Cleanup(backend.Close)
+			backendURL, err := url.Parse(backend.URL)
+			require.NoError(t, err)
+
+			handlerDone := make(chan struct{})
+			proxy := &SandboxProxy{proxy: reverseproxy.New(0, reverseproxy.SandboxProxyRetries, time.Minute,
+				func(*http.Request) (*pool.Destination, error) {
+					return &pool.Destination{
+						Url:           backendURL,
+						ConnectionKey: "upload-lifecycle",
+						RequestLogger: logger.NewNopLogger(),
+					}, nil
+				}, &reverseproxy.ConnectionLimitConfig{
+					Limiter:              connlimit.NewConnectionLimiter(),
+					GetMaxLimit:          func(context.Context) int { return 10 },
+					OnConnectionReleased: func(context.Context, int64) { close(handlerDone) },
+				}, true)}
+			var listenConfig net.ListenConfig
+			listener, err := listenConfig.Listen(t.Context(), "tcp", "127.0.0.1:0")
+			require.NoError(t, err)
+			serveDone := make(chan error, 1)
+			go func() { serveDone <- proxy.proxy.Serve(listener) }()
+			t.Cleanup(func() {
+				assert.NoError(t, proxy.proxy.Close())
+				assert.ErrorIs(t, <-serveDone, http.ErrServerClosed)
+			})
+
+			client := &http.Client{Transport: &http2.Transport{
+				AllowHTTP: true,
+				DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
+					return (&net.Dialer{}).DialContext(ctx, network, addr)
+				},
+			}}
+			t.Cleanup(client.CloseIdleConnections)
+			reader, writer := io.Pipe()
+			req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, "http://"+listener.Addr().String()+"/upload", reader)
+			require.NoError(t, err)
+			req.ContentLength = 8193
+			clientDone := make(chan struct{})
+			go func() {
+				defer close(clientDone)
+				resp, err := client.Do(req)
+				if err == nil {
+					_ = resp.Body.Close()
+				}
+			}()
+			t.Cleanup(func() {
+				_ = writer.Close()
+				_ = reader.Close()
+				select {
+				case <-clientDone:
+				case <-time.After(5 * time.Second):
+					t.Error("upload client did not exit")
+				}
+			})
+			_, err = io.WriteString(writer, strings.Repeat("x", 8192))
+			require.NoError(t, err)
+			select {
+			case err := <-backendStarted:
+				require.NoError(t, err)
+			case <-time.After(5 * time.Second):
+				t.Fatal("backend did not receive the upload")
+			}
+			backend.CloseClientConnections()
+
+			ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
+			defer cancel()
+			if forced {
+				cancel()
+			}
+			require.NoError(t, proxy.Close(ctx))
+			select {
+			case <-handlerDone:
+			case <-time.After(5 * time.Second):
+				t.Fatal("shutdown left the upload handler blocked")
+			}
+			assert.Zero(t, proxy.proxy.CurrentServerConnections())
+		})
+	}
+}

packages/shared/pkg/httpserver/h2c.go

@@ -1,6 +1,8 @@
 package httpserver
 
 import (
+	"bufio"
+	"net"
 	"net/http"
 	"time"
 
@@ -26,19 +28,54 @@ func ConfigureH2C(server *http.Server) {
 	}
 
 	h2cHandler := h2c.NewHandler(handler, h2Server)
-	limitedH2CHandler := http.MaxBytesHandler(h2cHandler, h2cUpgradeBodyLimit)
+	trackedH2CHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		writer := &h2cResponseWriter{ResponseWriter: w}
+		defer func() {
+			// h2c can return a preface read error without closing the hijacked socket.
+			if writer.conn != nil {
+				_ = writer.conn.Close()
+			}
+		}()
+		h2cHandler.ServeHTTP(writer, r)
+	})
+	// MaxBytesHandler needs the original writer's private requestTooLarge hook.
+	limitedH2CHandler := http.MaxBytesHandler(trackedH2CHandler, h2cUpgradeBodyLimit)
 
 	server.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if isH2CUpgrade(r.Header) {
 			limitedH2CHandler.ServeHTTP(w, r)
 
 			return
 		}
+		if r.Method == "PRI" && len(r.Header) == 0 && r.URL.Path == "*" && r.Proto == "HTTP/2.0" {
+			trackedH2CHandler.ServeHTTP(w, r)
+
+			return
+		}
 
 		h2cHandler.ServeHTTP(w, r)
 	})
 }
 
+type h2cResponseWriter struct {
+	http.ResponseWriter
+
+	conn net.Conn
+}
+
+func (w *h2cResponseWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}
+
+func (w *h2cResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	conn, buffered, err := http.NewResponseController(w.ResponseWriter).Hijack()
+	if err == nil {
+		w.conn = conn
+	}
+
+	return conn, buffered, err
+}
+
 func newHTTP2Server() *http2.Server {
 	return &http2.Server{
 		MaxConcurrentStreams:         100,

packages/shared/pkg/httpserver/h2c_test.go

@@ -1,8 +1,10 @@
 package httpserver
 
 import (
+	"bufio"
 	"context"
 	"crypto/tls"
+	"fmt"
 	"io"
 	"net"
 	"net/http"
@@ -15,6 +17,33 @@ import (
 	"golang.org/x/net/http2"
 )
 
+func TestConfigureH2CRejectsIncompleteOversizedUpgrade(t *testing.T) {
+	t.Parallel()
+
+	server := httptest.NewUnstartedServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
+		t.Error("oversized upgrade reached the application handler")
+	}))
+	ConfigureH2C(server.Config)
+	server.Start()
+	t.Cleanup(server.Close)
+
+	var dialer net.Dialer
+	conn, err := dialer.DialContext(t.Context(), "tcp", server.Listener.Addr().String())
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	require.NoError(t, conn.SetDeadline(time.Now().Add(5*time.Second)))
+	_, err = fmt.Fprintf(conn, "POST /upload HTTP/1.1\r\nHost: sandbox.test\r\nConnection: Upgrade, HTTP2-Settings\r\nUpgrade: h2c\r\nHTTP2-Settings: AAMAAABkAAQAAP__\r\nContent-Length: %d\r\n\r\n", h2cUpgradeBodyLimit+2)
+	require.NoError(t, err)
+	_, err = io.WriteString(conn, strings.Repeat("x", h2cUpgradeBodyLimit+1))
+	require.NoError(t, err)
+
+	resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	require.Equal(t, http.StatusInternalServerError, resp.StatusCode)
+	require.True(t, resp.Close)
+}
+
 func TestConfigureH2CAcceptsHTTP2AndHTTP1(t *testing.T) {
 	t.Parallel()
 

packages/shared/pkg/proxy/proxy.go

@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net"
 	"net/http"
@@ -12,6 +13,7 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/httpserver"
 	"github.com/e2b-dev/infra/packages/shared/pkg/proxy/pool"
 	"github.com/e2b-dev/infra/packages/shared/pkg/proxy/tracking"
+	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
 )
 
 // ConnectionLimitConfig bundles connection limiting and associated metric callbacks.
@@ -33,6 +35,7 @@ type Proxy struct {
 	http.Server
 
 	pool                      *pool.ProxyPool
+	serverConnections         *smap.Map[*tracking.Connection]
 	currentServerConnsCounter atomic.Int64
 }
 
@@ -69,7 +72,8 @@ func New(
 			ReadHeaderTimeout: 0,
 			Handler:           handler(p, getDestination, connLimitConfig),
 		},
-		pool: p,
+		pool:              p,
+		serverConnections: smap.New[*tracking.Connection](),
 	}
 	httpserver.ConfigureH2C(&proxy.Server)
 
@@ -109,5 +113,35 @@ func (p *Proxy) ListenAndServe(ctx context.Context) error {
 }
 
 func (p *Proxy) Serve(l net.Listener) error {
-	return p.Server.Serve(tracking.NewListener(l, &p.currentServerConnsCounter))
+	return p.Server.Serve(tracking.NewListener(l, &p.currentServerConnsCounter, p.serverConnections))
+}
+
+func (p *Proxy) Shutdown(ctx context.Context) error {
+	if err := p.Server.Shutdown(ctx); err != nil {
+		return err
+	}
+
+	// net/http does not wait for hijacked H2C and upgraded connections.
+	ticker := time.NewTicker(10 * time.Millisecond)
+	defer ticker.Stop()
+	for p.CurrentServerConnections() != 0 {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-ticker.C:
+		}
+	}
+
+	return nil
+}
+
+func (p *Proxy) Close() error {
+	err := p.Server.Close()
+	for _, conn := range p.serverConnections.Items() {
+		if closeErr := conn.Close(); closeErr != nil && !errors.Is(closeErr, net.ErrClosed) {
+			err = errors.Join(err, closeErr)
+		}
+	}
+
+	return err
 }