feat: route shared envd host through client proxy

Author: matthewlouisbrockman

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

@@ -62,15 +62,15 @@ func normalizeNodeIP(nodeIP string) (string, error) {
 	return nodeIP, nil
 }
 
-func legacySandboxHostForEnvdRequest(host string, sandboxID string, port uint64) *string {
+func orchestratorSandboxHost(host string, sandboxID string, port uint64) *string {
 	hostname := strings.Split(host, ":")[0]
 	domain, ok := strings.CutPrefix(hostname, "envd.")
 	if !ok || domain == "" {
 		return nil
 	}
 
-	legacyHost := fmt.Sprintf("%d-%s.%s", port, sandboxID, domain)
-	return &legacyHost
+	host := fmt.Sprintf("%d-%s.%s", port, sandboxID, domain)
+	return &host
 }
 
 func catalogResolution(ctx context.Context, sandboxId string, sandboxPort uint64, trafficAccessToken string, envdAccessToken string, c catalog.SandboxesCatalog, pausedChecker PausedSandboxResumer, featureFlags *featureflags.Client) (string, error) {
@@ -143,7 +143,7 @@ func handlePausedSandbox(
 }
 
 func NewClientProxy(meterProvider metric.MeterProvider, serviceName string, port uint16, catalog catalog.SandboxesCatalog, pausedSandboxResumer PausedSandboxResumer, featureFlagsClient *featureflags.Client) (*reverseproxy.Proxy, error) {
-	getTargetFromRequest := reverseproxy.GetTargetFromRequest(true)
+	getTargetFromRequest := reverseproxy.GetTargetFromRequest(reverseproxy.HeaderRoutingEnabled)
 	proxy := reverseproxy.New(
 		port,
 		// Retries that are needed to handle port forwarding delays in sandbox envd are handled by the orchestrator proxy
@@ -208,7 +208,7 @@ func NewClientProxy(meterProvider metric.MeterProvider, serviceName string, port
 				SandboxPort:   port,
 				ConnectionKey: pool.ClientProxyConnectionKey,
 				Url:           url,
-				MaskRequestHost: legacySandboxHostForEnvdRequest(
+				MaskRequestHost: orchestratorSandboxHost(
 					r.Host,
 					sandboxId,
 					port,

packages/orchestrator/pkg/proxy/proxy.go

@@ -15,7 +15,6 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox"
 	"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/env"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	reverseproxy "github.com/e2b-dev/infra/packages/shared/pkg/proxy"
@@ -40,7 +39,7 @@ type SandboxProxy struct {
 }
 
 func NewSandboxProxy(meterProvider metric.MeterProvider, port uint16, sandboxes *sandbox.Map, featureFlags *featureflags.Client) (*SandboxProxy, error) {
-	getTargetFromRequest := reverseproxy.GetTargetFromRequest(env.IsLocal())
+	getTargetFromRequest := reverseproxy.GetTargetFromRequest(reverseproxy.HeaderRoutingDisabled)
 	limiter := connlimit.NewConnectionLimiter()
 	metrics := NewMetrics(meterProvider)
 

packages/shared/pkg/proxy/host.go

@@ -9,9 +9,16 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/id"
 )
 
-func GetTargetFromRequest(processHeaders bool) func(r *http.Request) (sandboxId string, port uint64, err error) {
+type HeaderRoutingMode uint8
+
+const (
+	HeaderRoutingDisabled HeaderRoutingMode = iota
+	HeaderRoutingEnabled
+)
+
+func GetTargetFromRequest(headerRouting HeaderRoutingMode) func(r *http.Request) (sandboxId string, port uint64, err error) {
 	return func(r *http.Request) (sandboxId string, port uint64, err error) {
-		if processHeaders && shouldParseHeaders(r.Host) {
+		if headerRouting == HeaderRoutingEnabled && shouldParseHeaders(r.Host) && hasRoutingHeaders(r.Header) {
 			var ok bool
 			sandboxId, port, ok, err = parseHeaders(r.Header)
 			if err != nil {
@@ -43,6 +50,10 @@ func shouldParseHeaders(host string) bool {
 	return host == "localhost" || strings.HasPrefix(host, "envd.")
 }
 
+func hasRoutingHeaders(h http.Header) bool {
+	return h.Get(headerSandboxID) != "" || h.Get(headerSandboxPort) != ""
+}
+
 func parseHost(host string) (sandboxID string, port uint64, err error) {
 	dot := strings.Index(host, ".")
 

packages/shared/pkg/proxy/host_test.go

@@ -10,7 +10,7 @@ import (
 func TestGetTargetFromRequest(t *testing.T) { //nolint:tparallel // cannot call t.Setenv with t.Parallel
 	t.Setenv("ENVIRONMENT", "local")
 
-	getTargetFromRequest := GetTargetFromRequest(true)
+	getTargetFromRequest := GetTargetFromRequest(HeaderRoutingEnabled)
 
 	tests := []struct {
 		name     string