Support envd proxying via header as well as host

packages/api/internal/orchestrator/create_instance.go

@@ -184,12 +184,12 @@ func (o *Orchestrator) CreateSandbox(
 	algorithm := o.getPlacementAlgorithm(ctx)
 	node, err = placement.PlaceSandbox(ctx, algorithm, clusterNodes, node, sbxRequest)
 	if err != nil {
-		telemetry.ReportError(ctx, "failed to create sandbox", err)
+		telemetry.ReportError(ctx, "failed to place sandbox", err)
 
 		return sandbox.Sandbox{}, &api.APIError{
 			Code:      http.StatusInternalServerError,
-			ClientMsg: "Failed to create sandbox",
-			Err:       fmt.Errorf("failed to get create sandbox: %w", err),
+			ClientMsg: "Failed to place sandbox",
+			Err:       fmt.Errorf("failed to place sandbox: %w", err),
 		}
 	}
 

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

@@ -11,6 +11,7 @@ import (
 	"go.opentelemetry.io/otel/metric"
 	"go.uber.org/zap"
 
+	"github.com/e2b-dev/infra/packages/shared/pkg/env"
 	l "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"
@@ -50,13 +51,15 @@ func catalogResolution(ctx context.Context, sandboxId string, c catalog.Sandboxe
 }
 
 func NewClientProxy(meterProvider metric.MeterProvider, serviceName string, port uint16, catalog catalog.SandboxesCatalog) (*reverseproxy.Proxy, error) {
+	getUpstreamFromRequest := reverseproxy.GetUpstreamFromRequest(env.IsLocal())
+
 	proxy := reverseproxy.New(
 		port,
 		// Retries that are needed to handle port forwarding delays in sandbox envd are handled by the orchestrator proxy
 		reverseproxy.ClientProxyRetries,
 		idleTimeout,
 		func(r *http.Request) (*pool.Destination, error) {
-			sandboxId, port, err := reverseproxy.ParseHost(r.Host)
+			sandboxId, port, err := getUpstreamFromRequest(r)
 			if err != nil {
 				return nil, err
 			}

packages/orchestrator/internal/proxy/proxy.go

@@ -11,6 +11,7 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox"
+	"github.com/e2b-dev/infra/packages/shared/pkg/env"
 	"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"
@@ -29,13 +30,15 @@ type SandboxProxy struct {
 }
 
 func NewSandboxProxy(meterProvider metric.MeterProvider, port uint16, sandboxes *sandbox.Map) (*SandboxProxy, error) {
+	getUpstreamFromRequest := reverseproxy.GetUpstreamFromRequest(env.IsLocal())
+
 	proxy := reverseproxy.New(
 		port,
 		// Retry 5 times to handle port forwarding delays in sandbox envd.
 		reverseproxy.SandboxProxyRetries,
 		idleTimeout,
 		func(r *http.Request) (*pool.Destination, error) {
-			sandboxId, port, err := reverseproxy.ParseHost(r.Host)
+			sandboxId, port, err := getUpstreamFromRequest(r)
 			if err != nil {
 				return nil, err
 			}

packages/shared/pkg/proxy/errors.go

@@ -0,0 +1,39 @@
+package proxy
+
+import "errors"
+
+var ErrInvalidHost = errors.New("invalid url host")
+
+type InvalidSandboxPortError struct {
+	Port    string
+	wrapped error
+}
+
+func (e InvalidSandboxPortError) Error() string {
+	return "invalid sandbox port"
+}
+
+func (e InvalidSandboxPortError) Is(err error) bool {
+	var other InvalidSandboxPortError
+	ok := errors.As(err, &other)
+
+	return ok
+}
+
+func (e InvalidSandboxPortError) Unwrap() error {
+	return e.wrapped
+}
+
+type SandboxNotFoundError struct {
+	SandboxId string
+}
+
+func NewErrSandboxNotFound(sandboxId string) *SandboxNotFoundError {
+	return &SandboxNotFoundError{
+		SandboxId: sandboxId,
+	}
+}
+
+func (e SandboxNotFoundError) Error() string {
+	return "sandbox not found"
+}

packages/shared/pkg/proxy/errors_test.go

@@ -0,0 +1,28 @@
+package proxy
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInvalidSandboxPortError_Is(t *testing.T) {
+	var e InvalidSandboxPortError
+
+	t.Run("error is nil", func(t *testing.T) {
+		ok := e.Is(nil)
+		assert.False(t, ok)
+	})
+
+	t.Run("error is port error", func(t *testing.T) {
+		var e2 InvalidSandboxPortError
+		ok := e.Is(e2)
+		assert.True(t, ok)
+	})
+
+	t.Run("error is not port error", func(t *testing.T) {
+		var e2 SandboxNotFoundError
+		ok := e.Is(e2)
+		assert.False(t, ok)
+	})
+}

packages/shared/pkg/proxy/handler.go

@@ -13,55 +13,38 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/proxy/template"
 )
 
-type InvalidHostError struct{}
-
-func (e *InvalidHostError) Error() string {
-	return "invalid url host"
-}
-
-type InvalidSandboxPortError struct{}
-
-func (e *InvalidSandboxPortError) Error() string {
-	return "invalid sandbox port"
-}
-
-func NewErrSandboxNotFound(sandboxId string) *SandboxNotFoundError {
-	return &SandboxNotFoundError{
-		SandboxId: sandboxId,
-	}
-}
-
-type SandboxNotFoundError struct {
-	SandboxId string
-}
-
-func (e *SandboxNotFoundError) Error() string {
-	return "sandbox not found"
-}
-
 func handler(p *pool.ProxyPool, getDestination func(r *http.Request) (*pool.Destination, error)) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		d, err := getDestination(r)
 
-		var invalidHostErr *InvalidHostError
-		if errors.As(err, &invalidHostErr) {
+		var mhe MissingHeaderError
+		if errors.As(err, &mhe) {
+			zap.L().Warn("missing header", zap.Error(mhe))
+			http.Error(w, "missing header", http.StatusBadRequest)
+
+			return
+		}
+
+		if errors.Is(err, ErrInvalidHost) {
 			zap.L().Warn("invalid host", zap.String("host", r.Host))
 			http.Error(w, "Invalid host", http.StatusBadRequest)
 
 			return
 		}
 
-		var invalidPortErr *InvalidSandboxPortError
+		var invalidPortErr InvalidSandboxPortError
 		if errors.As(err, &invalidPortErr) {
-			zap.L().Warn("invalid sandbox port", zap.String("host", r.Host))
+			zap.L().Warn("invalid sandbox port", zap.String("host", r.Host), zap.String("port", invalidPortErr.Port))
 			http.Error(w, "Invalid sandbox port", http.StatusBadRequest)
 
 			return
 		}
 
-		var notFoundErr *SandboxNotFoundError
+		var notFoundErr SandboxNotFoundError
 		if errors.As(err, &notFoundErr) {
-			zap.L().Warn("sandbox not found", zap.String("host", r.Host), logger.WithSandboxID(notFoundErr.SandboxId))
+			zap.L().Warn("sandbox not found",
+				zap.String("host", r.Host),
+				logger.WithSandboxID(notFoundErr.SandboxId))
 
 			err := template.
 				NewSandboxNotFoundError(notFoundErr.SandboxId, r.Host).

packages/shared/pkg/proxy/host.go

@@ -1,33 +1,93 @@
 package proxy
 
 import (
+	"fmt"
+	"net/http"
 	"strconv"
 	"strings"
 )
 
-func ParseHost(host string) (sandboxID string, port uint64, err error) {
+func GetUpstreamFromRequest(processHeaders bool) func(r *http.Request) (sandboxId string, port uint64, err error) {
+	return func(r *http.Request) (sandboxId string, port uint64, err error) {
+		if processHeaders {
+			var ok bool
+			sandboxId, port, ok, err = parseHeaders(r.Header)
+			if err != nil {
+				return "", 0, err
+			} else if ok {
+				return sandboxId, port, nil
+			}
+		}
+
+		sandboxId, port, err = parseHost(r.Host)
+		if err != nil {
+			return "", 0, err
+		}
+
+		return sandboxId, port, nil
+	}
+}
+
+func parseHost(host string) (sandboxID string, port uint64, err error) {
 	dot := strings.Index(host, ".")
 
 	// There must be always domain part used
 	if dot == -1 {
-		return "", 0, &InvalidHostError{}
+		return "", 0, ErrInvalidHost
 	}
 
 	// Keep only the left-most subdomain part, i.e. everything before the
 	host = host[:dot]
 
 	hostParts := strings.Split(host, "-")
 	if len(hostParts) < 2 {
-		return "", 0, &InvalidHostError{}
+		return "", 0, ErrInvalidHost
 	}
 
 	sandboxPortString := hostParts[0]
 	sandboxID = hostParts[1]
 
 	sandboxPort, err := strconv.ParseUint(sandboxPortString, 10, 64)
 	if err != nil {
-		return "", 0, &InvalidSandboxPortError{}
+		return "", 0, InvalidSandboxPortError{sandboxPortString, err}
 	}
 
 	return sandboxID, sandboxPort, nil
 }
+
+type MissingHeaderError struct {
+	Header string
+}
+
+func (e MissingHeaderError) Error() string {
+	return fmt.Sprintf("Missing header: %s", e.Header)
+}
+
+const (
+	headerSandboxID   = "X-Sandbox-Id"
+	headerSandboxPort = "X-Sandbox-Port"
+)
+
+func parseHeaders(h http.Header) (sandboxID string, port uint64, ok bool, err error) {
+	sandboxID = h.Get(headerSandboxID)
+	portString := h.Get(headerSandboxPort)
+
+	if sandboxID == "" && portString == "" {
+		return "", 0, false, nil
+	}
+
+	if sandboxID == "" {
+		return "", 0, false, MissingHeaderError{Header: headerSandboxID}
+	}
+
+	if portString == "" {
+		return "", 0, false, MissingHeaderError{Header: headerSandboxPort}
+	}
+
+	port, err = strconv.ParseUint(portString, 10, 64)
+	if err != nil {
+		return "", 0, false, InvalidSandboxPortError{portString, err}
+	}
+
+	return sandboxID, port, true, nil
+}