fix(client-proxy): bound per-sandbox ingress requests without fragmenting the pool

packages/client-proxy/internal/proxy/proxy.go

@@ -16,6 +16,7 @@ import (
 	"go.uber.org/zap"
 	"google.golang.org/grpc/codes"
 
+	"github.com/e2b-dev/infra/packages/shared/pkg/connlimit"
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	proxygrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc/proxy"
@@ -241,6 +242,18 @@ func NewClientProxy(meterProvider metric.MeterProvider, serviceName string, port
 		}))
 	}
 
+	// Per-sandbox ingress limiter. The client-proxy shares one keepalive pool
+	// (ConnectionKey constant) across sandboxes, so the limiter is keyed by the
+	// Destination's LimitKey (the sandbox ID) instead, bounding per-sandbox
+	// in-flight requests without fragmenting the pool. Default flag -1 keeps the
+	// limiter inert until an operator sets a finite cap.
+	connLimitConfig := &reverseproxy.ConnectionLimitConfig{
+		Limiter: connlimit.NewConnectionLimiter(),
+		GetMaxLimit: func(ctx context.Context) int {
+			return featureFlagsClient.IntFlag(ctx, featureflags.ClientProxyMaxConnectionsPerSandbox)
+		},
+	}
+
 	getTargetFromRequest := reverseproxy.GetTargetFromRequest()
 	forwardedHostDomains = normalizeForwardedHostDomains(forwardedHostDomains)
 	proxy := reverseproxy.New(
@@ -342,7 +355,11 @@ func NewClientProxy(meterProvider metric.MeterProvider, serviceName string, port
 				RequestLogger:       l,
 				SandboxPort:         port,
 				ConnectionKey:       connectionKey,
-				Url:                 url,
+				// Share the pool across sandboxes via ConnectionKey, but account
+				// the ingress limiter per sandbox via LimitKey, so one sandbox
+				// cannot occupy the shared edge without bound.
+				LimitKey: sandboxId,
+				Url:      url,
 				MaskRequestHost: clientProxyMaskRequestHost(
 					ctx,
 					featureFlagsClient,
@@ -352,7 +369,7 @@ func NewClientProxy(meterProvider metric.MeterProvider, serviceName string, port
 				),
 			}, nil
 		},
-		nil,
+		connLimitConfig,
 		false,
 		proxyOptions...,
 	)

packages/shared/pkg/featureflags/flags.go

@@ -887,6 +887,13 @@ var (
 	// connections allowed per sandbox. Negative means no limit.
 	SandboxMaxIncomingConnections = NewIntFlag("sandbox-max-incoming-connections", -1)
 
+	// ClientProxyMaxConnectionsPerSandbox is the maximum number of concurrent
+	// client-proxy ingress requests accounted per sandbox. The client-proxy
+	// shares one keepalive connection pool across sandboxes, so without this the
+	// only bound on per-sandbox in-flight requests is the process-wide pool
+	// ceiling. Negative means no limit; the default preserves existing behavior.
+	ClientProxyMaxConnectionsPerSandbox = NewIntFlag("client-proxy-max-connections-per-sandbox", -1)
+
 	// BuildBaseRootfsSizeLimitMB is the maximum size of the base rootfs filesystem created from the OCI image, in MB.
 	BuildBaseRootfsSizeLimitMB = NewIntFlag("build-base-rootfs-size-limit-mb", 25000)
 

packages/shared/pkg/proxy/handler.go

@@ -207,7 +207,7 @@ func handler(p *pool.ProxyPool, getDestination func(r *http.Request) (*pool.Dest
 		// Connection limiting
 		if connLimitConfig != nil {
 			maxLimit := connLimitConfig.GetMaxLimit(ctx)
-			count, acquired := connLimitConfig.Limiter.TryAcquire(d.ConnectionKey, maxLimit)
+			count, acquired := connLimitConfig.Limiter.TryAcquire(d.LimiterKey(), maxLimit)
 			if !acquired {
 				l.Warn(ctx, "sandbox too many incoming connections",
 					zap.String("host", r.Host),
@@ -239,7 +239,7 @@ func handler(p *pool.ProxyPool, getDestination func(r *http.Request) (*pool.Dest
 
 			start := time.Now()
 			defer func() {
-				connLimitConfig.Limiter.Release(d.ConnectionKey)
+				connLimitConfig.Limiter.Release(d.LimiterKey())
 				if connLimitConfig.OnConnectionReleased != nil {
 					connLimitConfig.OnConnectionReleased(ctx, time.Since(start).Milliseconds())
 				}

packages/shared/pkg/proxy/pool/destination.go

@@ -23,13 +23,17 @@ type Destination struct {
 	// Should we return the error about closed port if there is a problem with a connection to upstream?
 	DefaultToPortError bool
 	RequestLogger      logger.Logger
-	// ConnectionKey uniquely identifies a single sandbox lifecycle. It is
-	// used for two purposes:
-	//   1. keepalive connection pool isolation, so connections to a reused
-	//      IP:port pair are not accidentally shared across sandboxes;
-	//   2. per-sandbox ingress connection limiter accounting.
-	// Embedders should pick a value that is unique per lifecycle.
-	ConnectionKey                      string
+	// ConnectionKey identifies the keepalive connection pool entry: connections
+	// to a reused IP:port pair are not accidentally shared across sandboxes.
+	// Embedders that isolate per lifecycle pick a value unique per lifecycle;
+	// embedders that intentionally share one pool (client-proxy) use a constant.
+	ConnectionKey string
+	// LimitKey scopes the ingress connection limiter accounting. When empty the
+	// limiter falls back to ConnectionKey (see LimiterKey), so a lifecycle-keyed
+	// pool is limited per lifecycle with no extra field. An embedder that shares
+	// one ConnectionKey across sandboxes (for pool reuse) sets LimitKey to the
+	// per-sandbox value so limiting stays per sandbox without fragmenting the pool.
+	LimitKey                           string
 	IncludeSandboxIdInProxyErrorLogger bool
 	// MaskRequestHost is used to mask the request host.
 	MaskRequestHost *string
@@ -41,3 +45,14 @@ type Destination struct {
 	// ConnectionKey entry. Closing it ends them.
 	Owner *Owner
 }
+
+// LimiterKey returns the key the ingress connection limiter accounts against:
+// LimitKey when set, otherwise ConnectionKey. This lets an embedder share one
+// ConnectionKey for pool reuse while still limiting per sandbox via LimitKey.
+func (d Destination) LimiterKey() string {
+	if d.LimitKey != "" {
+		return d.LimitKey
+	}
+
+	return d.ConnectionKey
+}

packages/shared/pkg/proxy/pool/destination_limiterkey_test.go

@@ -0,0 +1,28 @@
+package pool
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// TestDestination_LimiterKey verifies the limiter key is decoupled from the
+// pool key: LimitKey when set, else ConnectionKey. This is what lets the
+// client-proxy share one ConnectionKey (pool reuse) while limiting per sandbox.
+func TestDestination_LimiterKey(t *testing.T) {
+	t.Parallel()
+
+	// No LimitKey: falls back to ConnectionKey (orchestrator behavior — limiter
+	// scoped per lifecycle, identical to before this field existed).
+	lifecycle := Destination{ConnectionKey: "lifecycle-abc"}
+	assert.Equal(t, "lifecycle-abc", lifecycle.LimiterKey())
+
+	// LimitKey set: limiter is keyed per sandbox while the pool key stays shared
+	// (client-proxy behavior — one pool, per-sandbox limit).
+	shared := Destination{ConnectionKey: "client-proxy", LimitKey: "sbx-123"}
+	assert.Equal(t, "sbx-123", shared.LimiterKey())
+	assert.Equal(t, "client-proxy", shared.ConnectionKey, "pool key must stay shared, not fragmented per sandbox")
+
+	// Empty both: empty (admission with an empty key is the caller's concern).
+	assert.Equal(t, "", Destination{}.LimiterKey())
+}