Replace Nginx based session proxy with orchestrator proxy

main.tf

@@ -219,6 +219,7 @@ module "nomad" {
 
   # Orchestrator
   orchestrator_port           = var.orchestrator_port
+  orchestrator_proxy_port     = var.orchestrator_proxy_port
   fc_env_pipeline_bucket_name = module.buckets.fc_env_pipeline_bucket_name
 
   # Template manager

packages/client-proxy/main.go

@@ -28,12 +28,13 @@ import (
 )
 
 const (
-	ServiceName     = "client-proxy"
-	dnsServer       = "api.service.consul:5353"
-	healthCheckPort = 3001
-	port            = 3002
-	sandboxPort     = 3003
-	maxRetries      = 3
+	ServiceName           = "client-proxy"
+	dnsServer             = "api.service.consul:5353"
+	healthCheckPort       = 3001
+	port                  = 3002
+	sandboxPort           = 3003 // legacy session proxy port
+	orchestratorProxyPort = 5007 // orchestrator proxy port
+	maxRetries            = 3
 )
 
 var commitSHA string

packages/nomad/main.tf

@@ -324,6 +324,7 @@ resource "nomad_job" "orchestrator" {
   jobspec = templatefile("${path.module}/orchestrator.hcl", {
     gcp_zone         = var.gcp_zone
     port             = var.orchestrator_port
+    proxy_port       = var.orchestrator_proxy_port
     environment      = var.environment
     consul_acl_token = var.consul_acl_token_secret
 

packages/nomad/orchestrator.hcl

@@ -25,6 +25,11 @@ job "orchestrator" {
       }
     }
 
+    service {
+      name = "orchestrator-proxy"
+      port = "${proxy_port}"
+    }
+
     task "start" {
       driver = "raw_exec"
 
@@ -45,7 +50,7 @@ job "orchestrator" {
 
       config {
         command = "/bin/bash"
-        args    = ["-c", " chmod +x local/orchestrator && local/orchestrator --port ${port}"]
+        args    = ["-c", " chmod +x local/orchestrator && local/orchestrator --port ${port} --proxy-port ${proxy_port}"]
       }
 
       artifact {

packages/nomad/variables.tf

@@ -185,6 +185,10 @@ variable "orchestrator_port" {
   type = number
 }
 
+variable "orchestrator_proxy_port" {
+  type = number
+}
+
 variable "fc_env_pipeline_bucket_name" {
   type = string
 }

packages/orchestrator/cmd/mock-sandbox/mock.go

@@ -13,6 +13,7 @@ import (
 	"go.opentelemetry.io/otel"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/dns"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/proxy"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/nbd"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/network"
@@ -51,6 +52,8 @@ func main() {
 	}()
 
 	dnsServer := dns.New()
+	proxyServer := proxy.New(3333)
+
 	go func() {
 		log.Printf("Starting DNS server")
 
@@ -87,6 +90,7 @@ func main() {
 			*buildId,
 			*sandboxId+"-"+strconv.Itoa(v),
 			dnsServer,
+			proxyServer,
 			time.Duration(*keepAlive)*time.Second,
 			networkPool,
 			templateCache,
@@ -104,6 +108,7 @@ func mockSandbox(
 	buildId,
 	sandboxId string,
 	dns *dns.DNS,
+	proxy *proxy.SandboxProxy,
 	keepAlive time.Duration,
 	networkPool *network.Pool,
 	templateCache *template.Cache,
@@ -128,6 +133,7 @@ func mockSandbox(
 		childCtx,
 		tracer,
 		dns,
+		proxy,
 		networkPool,
 		templateCache,
 		&orchestrator.SandboxConfig{

packages/orchestrator/cmd/mock-snapshot/mock.go

@@ -14,6 +14,7 @@ import (
 	"golang.org/x/sync/errgroup"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/dns"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/proxy"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/nbd"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/network"
@@ -51,6 +52,7 @@ func main() {
 		cancel()
 	}()
 
+	proxyServer := proxy.New(3333)
 	dnsServer := dns.New()
 	go func() {
 		log.Printf("Starting DNS server")
@@ -90,6 +92,7 @@ func main() {
 			*buildId,
 			*sandboxId+"-"+strconv.Itoa(v),
 			dnsServer,
+			proxyServer,
 			time.Duration(*keepAlive)*time.Second,
 			networkPool,
 			templateCache,
@@ -113,6 +116,7 @@ func mockSnapshot(
 	buildId,
 	sandboxId string,
 	dns *dns.DNS,
+	proxy *proxy.SandboxProxy,
 	keepAlive time.Duration,
 	networkPool *network.Pool,
 	templateCache *template.Cache,
@@ -137,6 +141,7 @@ func mockSnapshot(
 		childCtx,
 		tracer,
 		dns,
+		proxy,
 		networkPool,
 		templateCache,
 		&orchestrator.SandboxConfig{
@@ -242,6 +247,7 @@ func mockSnapshot(
 		childCtx,
 		tracer,
 		dns,
+		proxy,
 		networkPool,
 		templateCache,
 		&orchestrator.SandboxConfig{

packages/orchestrator/internal/proxy/proxy.go

@@ -0,0 +1,196 @@
+package proxy
+
+import (
+	"bytes"
+	"context"
+	_ "embed"
+	"encoding/json"
+	"fmt"
+	"go.uber.org/zap"
+	"html/template"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"regexp"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/meters"
+	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
+)
+
+//go:embed proxy_browser_502.html
+var proxyBrowser502PageHtml string
+
+var browserRegex = regexp.MustCompile(`(?i)mozilla|chrome|safari|firefox|edge|opera|msie`)
+var browserTemplate = template.Must(template.New("template").Parse(proxyBrowser502PageHtml))
+
+type htmlTemplateData struct {
+	SandboxId   string
+	SandboxHost string
+	SandboxPort string
+}
+
+type jsonTemplateData struct {
+	Error     string `json:"error"`
+	SandboxId string `json:"sandboxId"`
+	Port      uint64 `json:"port"`
+}
+
+type SandboxProxy struct {
+	sandboxes *smap.Map[string]
+	server    *http.Server
+}
+
+func New(port uint) *SandboxProxy {
+	server := &http.Server{Addr: fmt.Sprintf(":%d", port)}
+
+	return &SandboxProxy{
+		server:    server,
+		sandboxes: smap.New[string](),
+	}
+}
+
+func (p *SandboxProxy) AddSandbox(sandboxID, ip string) {
+	p.sandboxes.Insert(sandboxID, ip)
+}
+
+func (p *SandboxProxy) RemoveSandbox(sandboxID string, ip string) {
+	p.sandboxes.RemoveCb(sandboxID, func(k string, v string, ok bool) bool { return ok && v == ip })
+}
+
+func (p *SandboxProxy) Start() error {
+	// similar values to our old the nginx configuration
+	serverTransport := &http.Transport{
+		Proxy:                 http.ProxyFromEnvironment,
+		MaxIdleConns:          1024,              // Matches worker_connections
+		MaxIdleConnsPerHost:   8192,              // Matches keepalive_requests
+		IdleConnTimeout:       620 * time.Second, // Matches keepalive_timeout
+		TLSHandshakeTimeout:   10 * time.Second,  // Similar to client_header_timeout
+		ResponseHeaderTimeout: 24 * time.Hour,    // Matches proxy_read_timeout
+		DisableKeepAlives:     true,              // Disable keep-alives, envd doesn't support idle connections
+	}
+
+	p.server.Handler = http.HandlerFunc(p.proxyHandler(serverTransport))
+	return p.server.ListenAndServe()
+}
+
+func (p *SandboxProxy) Shutdown(ctx context.Context) {
+	err := p.server.Shutdown(ctx)
+	if err != nil {
+		zap.L().Error("failed to shutdown proxy server", zap.Error(err))
+	}
+}
+
+func (p *SandboxProxy) proxyHandler(transport *http.Transport) func(w http.ResponseWriter, r *http.Request) {
+	activeConnections, err := meters.GetUpDownCounter(meters.OrchestratorProxyActiveConnectionsCounterMeterName)
+	if err != nil {
+		zap.L().Error("failed to create active connections counter", zap.Error(err))
+	}
+
+	return func(w http.ResponseWriter, r *http.Request) {
+		if activeConnections != nil {
+			activeConnections.Add(r.Context(), 1)
+			defer func() {
+				activeConnections.Add(r.Context(), -1)
+			}()
+		}
+
+		// Extract sandbox id from the host (<port>-<sandbox id>-<old client id>.e2b.dev)
+		hostSplit := strings.Split(r.Host, "-")
+		if len(hostSplit) < 2 {
+			zap.L().Warn("invalid host to proxy", zap.String("host", r.Host))
+			http.Error(w, "Invalid host", http.StatusBadRequest)
+			return
+		}
+
+		sandboxID := hostSplit[1]
+		sandboxPortRaw := hostSplit[0]
+		sandboxPort, sandboxPortErr := strconv.ParseUint(sandboxPortRaw, 10, 64)
+		if sandboxPortErr != nil {
+			zap.L().Warn("invalid sandbox port", zap.String("sandbox_port", sandboxPortRaw))
+			http.Error(w, "Invalid sandbox port", http.StatusBadRequest)
+		}
+
+		sbxIp, sbxFound := p.sandboxes.Get(sandboxID)
+		if !sbxFound {
+			zap.L().Warn("sandbox not found", zap.String("sandbox_id", sandboxID))
+			http.Error(w, "Sandbox not found", http.StatusNotFound)
+			return
+		}
+
+		logger := zap.L().With(zap.String("sandbox_id", sandboxID), zap.String("sandbox_ip", sbxIp), zap.Uint64("sandbox_req_port", sandboxPort), zap.String("sandbox_port_path", r.URL.Path))
+
+		// We've resolved the node to proxy the request to
+		logger.Debug("Proxying request")
+		targetUrl := &url.URL{
+			Scheme: "http",
+			Host:   fmt.Sprintf("%s:%d", sbxIp, sandboxPort),
+		}
+
+		// Proxy the request
+		proxy := httputil.NewSingleHostReverseProxy(targetUrl)
+		proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
+			logger.Error("Reverse proxy error", zap.Error(err))
+
+			if p.isBrowser(r.UserAgent()) {
+				res, resErr := p.buildHtmlClosedPortError(sandboxID, r.Host, sandboxPort)
+				if resErr != nil {
+					logger.Error("Failed to build HTML error response", zap.Error(resErr))
+					w.WriteHeader(http.StatusInternalServerError)
+					return
+				}
+
+				w.WriteHeader(http.StatusBadGateway)
+				w.Header().Add("Content-Type", "text/html")
+				w.Write(res)
+				return
+			}
+
+			w.WriteHeader(http.StatusBadGateway)
+			w.Header().Add("Content-Type", "application/json")
+			w.Write(p.buildJsonClosedPortError(sandboxID, sandboxPort))
+		}
+
+		proxy.ModifyResponse = func(resp *http.Response) error {
+			if resp.StatusCode >= 500 {
+				logger.Error("Backend responded with error", zap.Int("status_code", resp.StatusCode))
+			} else {
+				logger.Info("Backend responded", zap.Int("status_code", resp.StatusCode))
+			}
+
+			return nil
+		}
+
+		proxy.Transport = transport
+		proxy.ServeHTTP(w, r)
+	}
+}
+
+func (p *SandboxProxy) buildHtmlClosedPortError(sandboxId string, host string, port uint64) ([]byte, error) {
+	htmlResponse := new(bytes.Buffer)
+	htmlVars := htmlTemplateData{SandboxId: sandboxId, SandboxHost: host, SandboxPort: strconv.FormatUint(port, 10)}
+
+	err := browserTemplate.Execute(htmlResponse, htmlVars)
+	if err != nil {
+		return nil, err
+	}
+
+	return htmlResponse.Bytes(), nil
+}
+
+func (p *SandboxProxy) buildJsonClosedPortError(sandboxId string, port uint64) []byte {
+	response := jsonTemplateData{
+		Error:     "The sandbox is running but port is not open",
+		SandboxId: sandboxId,
+		Port:      port,
+	}
+
+	responseBytes, _ := json.Marshal(response)
+	return responseBytes
+}
+
+func (p *SandboxProxy) isBrowser(userAgent string) bool {
+	return browserRegex.MatchString(userAgent)
+}

packages/orchestrator/internal/proxy/proxy_browser_502.html

@@ -0,0 +1,23 @@
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width,initial-scale=1">
+    <title>Closed Port Error</title>
+    <style>:root{--brand:#ff8800;--error:#dc2626;--error-light:#fef2f2;--text:#1a1a1a;--background:#ffffff;--border:#e5e7eb;--details-bg:#f9fafb;--code-text:#374151;--muted-text:#6b7280}@media (prefers-color-scheme:dark){:root{--error:#ef4444;--error-light:#2a0f0f;--text:#e5e7eb;--background:#121212;--border:#2f2f2f;--details-bg:#1c1c1c;--code-text:#d1d5db;--muted-text:#9ca3af}}*{margin:0;padding:0;box-sizing:border-box}body{font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,"Helvetica Neue",Arial,sans-serif;background:#f5f5f5;min-height:100vh;display:flex;align-items:center;justify-content:center;padding:1rem;color:var(--text)}@media (prefers-color-scheme:dark){body{background:#0a0a0a}}.error-card{background:var(--background);border-radius:12px;box-shadow:0 4px 6px -1px rgb(0 0 0 / .1),0 2px 4px -2px rgb(0 0 0 / .1);width:100%;max-width:600px;padding:1.5rem 2rem 2rem;position:relative}.logo{position:absolute;top:1rem;right:1.5rem;width:40px;height:40px;border-radius:50%;overflow:hidden}.error-header{margin-bottom:1.5rem;padding-right:3.5rem}.error-title{display:inline-block;color:var(--error);font-size:.9375rem;font-weight:500;margin-bottom:1rem;padding:.25rem .5rem;background:var(--error-light);border-radius:4px}.error-message{font-size:1.125rem;line-height:1.5;color:var(--error);font-weight:400;font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,"Helvetica Neue",Arial,sans-serif}.error-details{background:var(--details-bg);border:1px solid var(--border);border-radius:8px;padding:1rem;margin-top:1.5rem}.error-code{font-family:ui-monospace,SFMono-Regular,Menlo,Monaco,Consolas,monospace;font-size:.875rem;color:var(--code-text)}.sandbox-url{color:var(--muted-text);font-size:.875rem;display:block;margin-bottom:.5rem}.highlight{font-weight:700}.help-text{margin-top:1.5rem;font-size:.875rem;color:var(--muted-text)}.debug-link{display:block;margin-top:2rem;color:var(--brand);text-decoration:none;font-size:.875rem}.debug-link:hover{text-decoration:underline}@media (max-width:640px){.error-card{margin:1rem;padding:1.25rem 1.5rem 1.5rem}.logo{top:.75rem;right:1rem;width:32px;height:32px}.error-header{padding-right:2.5rem}}</style>
+</head>
+<body>
+<main class="error-card">
+    <img src="https://hebbkx1anhila5yf.public.blob.vercel-storage.com/Symbol%20Gradient-Kr5pnWlK3ZhzBcRGf6Am4cNbJvY1Ge.svg" alt="Logo" class="logo">
+    <div class="error-header">
+        <h1 class="error-title">Closed Port Error</h1>
+        <p class="error-message">The sandbox <span class="highlight" id="sandbox-id">{{.SandboxId}}</span> is running but there&#39s no service running on port <span class="highlight" id="port-number">{{.SandboxPort}}</span>.</p>
+    </div>
+    <div class="error-details">
+        <span class="sandbox-url">{{.SandboxHost}}</span>
+        <div class="error-code">Connection refused on port <span class="highlight" id="port-number-code">{{.SandboxPort}}</span></div>
+    </div>
+    <p class="help-text">Please ensure that your service is properly configured and running on the specified port.</p>
+    <a class="debug-link" href="https://e2b.dev/docs/sdk-reference/cli/v1.0.9/sandbox#e2b-sandbox-logs">Check the sandbox logs for more information →</a>
+</main>
+</body>
+</html>