Show useful 502 message for not running/existent sandbox

packages/api/internal/orchestrator/orchestrator.go

@@ -12,7 +12,7 @@ import (
 
 	analyticscollector "github.com/e2b-dev/infra/packages/api/internal/analytics_collector"
 	"github.com/e2b-dev/infra/packages/api/internal/cache/instance"
-	"github.com/e2b-dev/infra/packages/shared/pkg/dns"
+	"github.com/e2b-dev/infra/packages/api/internal/dns"
 	"github.com/e2b-dev/infra/packages/shared/pkg/env"
 	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
 )

packages/nomad/proxies/client.conf

@@ -73,7 +73,7 @@ server {
   location / {
     if ($node_ip = "") {
       # If you set any text, the header will be set to `application/octet-stream` and then browser won't be able to render the content
-      return 404;
+      return 404; # Invalid sandbox url
     }
 
 
@@ -85,6 +85,14 @@ server {
   }
 }
 
+# Mock for sandbox server when the sandbox is not running, 127.0.0.1 is returned by the DNS resolver
+server {
+  listen 3003;
+
+  default_type text/plain;
+  return 502 'Sandbox does not exist.';
+}
+
 server {
   listen 3001;
   location /health {

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

@@ -12,10 +12,10 @@ import (
 
 	"go.opentelemetry.io/otel"
 
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/dns"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/network"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/template"
-	"github.com/e2b-dev/infra/packages/shared/pkg/dns"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logs"
 )

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

@@ -13,10 +13,10 @@ import (
 	"go.opentelemetry.io/otel"
 	"golang.org/x/sync/errgroup"
 
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/dns"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/network"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/template"
-	"github.com/e2b-dev/infra/packages/shared/pkg/dns"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logs"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"

packages/orchestrator/internal/dns/server.go

@@ -0,0 +1,96 @@
+package dns
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"strings"
+	"sync"
+
+	resolver "github.com/miekg/dns"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
+)
+
+const ttl = 0
+
+const defaultRoutingIP = "127.0.0.1"
+
+type DNS struct {
+	mu      sync.Mutex
+	records *smap.Map[string]
+}
+
+func New() *DNS {
+	return &DNS{
+		records: smap.New[string](),
+	}
+}
+
+func (d *DNS) Add(sandboxID, ip string) {
+	d.records.Insert(d.hostname(sandboxID), ip)
+}
+
+func (d *DNS) Remove(sandboxID, ip string) {
+	d.records.RemoveCb(d.hostname(sandboxID), func(key string, v string, exists bool) bool {
+		return v == ip
+	})
+}
+
+func (d *DNS) get(hostname string) (string, bool) {
+	return d.records.Get(hostname)
+}
+
+func (*DNS) hostname(sandboxID string) string {
+	return fmt.Sprintf("%s.", sandboxID)
+}
+
+func (d *DNS) handleDNSRequest(w resolver.ResponseWriter, r *resolver.Msg) {
+	m := new(resolver.Msg)
+	m.SetReply(r)
+	m.Compress = false
+	m.Authoritative = true
+
+	for _, q := range m.Question {
+		if q.Qtype == resolver.TypeA {
+			a := &resolver.A{
+				Hdr: resolver.RR_Header{
+					Name:   q.Name,
+					Rrtype: resolver.TypeA,
+					Class:  resolver.ClassINET,
+					Ttl:    ttl,
+				},
+			}
+
+			sandboxID := strings.Split(q.Name, "-")[0]
+			ip, found := d.get(sandboxID)
+			if found {
+				a.A = net.ParseIP(ip).To4()
+			} else {
+				a.A = net.ParseIP(defaultRoutingIP).To4()
+			}
+
+			m.Answer = append(m.Answer, a)
+		}
+	}
+
+	err := w.WriteMsg(m)
+	if err != nil {
+		log.Printf("Failed to write message: %s\n", err.Error())
+	}
+}
+
+func (d *DNS) Start(address string, port int) error {
+	mux := resolver.NewServeMux()
+
+	mux.HandleFunc(".", d.handleDNSRequest)
+
+	server := resolver.Server{Addr: fmt.Sprintf("%s:%d", address, port), Net: "udp", Handler: mux}
+
+	err := server.ListenAndServe()
+	if err != nil {
+		return fmt.Errorf("failed to start DNS server: %w", err)
+	}
+
+	return nil
+}

packages/orchestrator/internal/sandbox/sandbox.go

@@ -15,14 +15,14 @@ import (
 	"golang.org/x/mod/semver"
 	"golang.org/x/sys/unix"
 
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/dns"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/build"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/fc"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/network"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/rootfs"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/stats"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/template"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd"
-	"github.com/e2b-dev/infra/packages/shared/pkg/dns"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logs"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"

packages/orchestrator/internal/server/main.go

@@ -14,10 +14,10 @@ import (
 	"google.golang.org/grpc/health"
 	"google.golang.org/grpc/health/grpc_health_v1"
 
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/dns"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/network"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/template"
-	"github.com/e2b-dev/infra/packages/shared/pkg/dns"
 	e2bgrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 	"github.com/e2b-dev/infra/packages/shared/pkg/smap"