feat: dynamic egress network rule updates for running sandboxes

packages/api/internal/handlers/sandbox_network_update.go

@@ -0,0 +1,57 @@
+package handlers
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+
+	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/utils"
+	"github.com/e2b-dev/infra/packages/auth/pkg/auth"
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+)
+
+func (a *APIStore) PutSandboxesSandboxIDNetwork(
+	c *gin.Context,
+	sandboxID string,
+) {
+	ctx := c.Request.Context()
+
+	var err error
+	sandboxID, err = utils.ShortID(sandboxID)
+	if err != nil {
+		a.sendAPIStoreError(c, http.StatusBadRequest, "Invalid sandbox ID")
+
+		return
+	}
+
+	team := auth.MustGetTeamInfo(c)
+
+	body, err := utils.ParseBody[api.PutSandboxesSandboxIDNetworkJSONBody](ctx, c)
+	if err != nil {
+		a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Error when parsing request: %s", err))
+		telemetry.ReportCriticalError(ctx, "error when parsing request", err)
+
+		return
+	}
+
+	var allowedCIDRs []string
+	if body.AllowOut != nil {
+		allowedCIDRs = *body.AllowOut
+	}
+
+	var deniedCIDRs []string
+	if body.DenyOut != nil {
+		deniedCIDRs = *body.DenyOut
+	}
+
+	if apiErr := a.orchestrator.UpdateSandboxNetworkConfig(ctx, team.ID, sandboxID, allowedCIDRs, deniedCIDRs); apiErr != nil {
+		telemetry.ReportError(ctx, "error updating sandbox network config", apiErr.Err)
+		a.sendAPIStoreError(c, apiErr.Code, apiErr.ClientMsg)
+
+		return
+	}
+
+	c.Status(http.StatusNoContent)
+}

packages/api/internal/orchestrator/update_network.go

@@ -0,0 +1,138 @@
+package orchestrator
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+
+	"github.com/google/uuid"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
+	"github.com/e2b-dev/infra/packages/api/internal/utils"
+	"github.com/e2b-dev/infra/packages/db/pkg/types"
+	orchestratorgrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
+	sandbox_network "github.com/e2b-dev/infra/packages/shared/pkg/sandbox-network"
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+)
+
+func (o *Orchestrator) UpdateSandboxNetworkConfig(
+	ctx context.Context,
+	teamID uuid.UUID,
+	sandboxID string,
+	allowedEntries []string,
+	deniedEntries []string,
+) *api.APIError {
+	// Split allowed entries into CIDRs/IPs and domains, matching the creation
+	// path in buildNetworkConfig.
+	allowedAddresses, allowedDomains := sandbox_network.ParseAddressesAndDomains(allowedEntries)
+
+	// If allowed domains are provided, add the default nameserver so the sandbox
+	// can resolve domain names — same as the creation path.
+	if len(allowedDomains) > 0 {
+		allowedAddresses = append(allowedAddresses, sandbox_network.DefaultNameserver)
+	}
+
+	allowedCIDRs := sandbox_network.AddressStringsToCIDRs(allowedAddresses)
+	deniedCIDRs := sandbox_network.AddressStringsToCIDRs(deniedEntries)
+
+	// Read the sandbox first to validate state and get routing info,
+	// without mutating the store yet.
+	sbx, err := o.sandboxStore.Get(ctx, teamID, sandboxID)
+	if err != nil {
+		var sbxNotFoundErr *sandbox.NotFoundError
+		if errors.As(err, &sbxNotFoundErr) {
+			return &api.APIError{Code: http.StatusNotFound, ClientMsg: utils.SandboxNotFoundMsg(sandboxID), Err: err}
+		}
+
+		return &api.APIError{Code: http.StatusInternalServerError, ClientMsg: "Error reading sandbox", Err: err}
+	}
+
+	if sbx.State != sandbox.StateRunning {
+		return &api.APIError{Code: http.StatusConflict, ClientMsg: utils.SandboxChangingStateMsg(sandboxID, sbx.State), Err: fmt.Errorf("sandbox '%s' is not running (state: %s)", sandboxID, sbx.State)}
+	}
+
+	// Apply the network update on the orchestrator node first.
+	// Only persist to the store after the node update succeeds.
+	if apiErr := o.updateSandboxNetworkOnNode(ctx, sbx, allowedCIDRs, deniedCIDRs, allowedDomains); apiErr != nil {
+		return apiErr
+	}
+
+	// Node update succeeded — now persist the new config in the store.
+	updateFunc := func(sbx sandbox.Sandbox) (sandbox.Sandbox, error) {
+		if sbx.Network == nil {
+			sbx.Network = &types.SandboxNetworkConfig{}
+		}
+
+		sbx.Network.Egress = &types.SandboxNetworkEgressConfig{
+			AllowedAddresses: allowedEntries,
+			DeniedAddresses:  deniedEntries,
+		}
+
+		return sbx, nil
+	}
+
+	if _, err := o.sandboxStore.Update(ctx, teamID, sandboxID, updateFunc); err != nil {
+		telemetry.ReportError(ctx, "network updated on node but failed to persist in store", err)
+
+		return &api.APIError{Code: http.StatusInternalServerError, ClientMsg: "Network rules applied but failed to persist config", Err: err}
+	}
+
+	return nil
+}
+
+func (o *Orchestrator) updateSandboxNetworkOnNode(
+	ctx context.Context,
+	sbx sandbox.Sandbox,
+	allowedCIDRs []string,
+	deniedCIDRs []string,
+	allowedDomains []string,
+) *api.APIError {
+	ctx, span := tracer.Start(ctx, "update-sandbox-network-on-node",
+		trace.WithAttributes(
+			attribute.String("instance.id", sbx.SandboxID),
+		),
+	)
+	defer span.End()
+
+	node := o.GetNode(sbx.ClusterID, sbx.NodeID)
+	if node == nil {
+		return &api.APIError{
+			Code:      http.StatusInternalServerError,
+			ClientMsg: fmt.Sprintf("Node hosting sandbox '%s' not found", sbx.SandboxID),
+			Err:       fmt.Errorf("node '%s' not found for cluster '%s'", sbx.NodeID, sbx.ClusterID),
+		}
+	}
+
+	client, ctx := node.GetClient(ctx)
+	_, err := client.Sandbox.UpdateNetwork(ctx, &orchestratorgrpc.SandboxUpdateNetworkRequest{
+		SandboxId:      sbx.SandboxID,
+		AllowedCidrs:   allowedCIDRs,
+		DeniedCidrs:    deniedCIDRs,
+		AllowedDomains: allowedDomains,
+	})
+	if err != nil {
+		grpcErr, ok := status.FromError(err)
+		if ok && grpcErr.Code() == codes.NotFound {
+			return &api.APIError{Code: http.StatusNotFound, ClientMsg: utils.SandboxNotFoundMsg(sbx.SandboxID), Err: err}
+		}
+
+		err = utils.UnwrapGRPCError(err)
+		telemetry.ReportCriticalError(ctx, "failed to update sandbox network on node", err)
+
+		return &api.APIError{
+			Code:      http.StatusInternalServerError,
+			ClientMsg: "Error applying network config to sandbox",
+			Err:       fmt.Errorf("failed to update sandbox network on node: %w", err),
+		}
+	}
+
+	telemetry.ReportEvent(ctx, "Updated sandbox network on node")
+
+	return nil
+}

packages/orchestrator/internal/proxy/proxy.go

@@ -72,10 +72,8 @@ func NewSandboxProxy(meterProvider metric.MeterProvider, port uint16, sandboxes
 				return nil, reverseproxy.NewErrSandboxNotFound(sandboxId)
 			}
 
-			var accessToken *string = nil
-			if net := sbx.Config.Network; net != nil && net.GetIngress() != nil {
-				accessToken = net.GetIngress().TrafficAccessToken
-			}
+			ingress := sbx.GetNetwork().GetIngress()
+			accessToken := ingress.TrafficAccessToken
 
 			isNonEnvdTraffic := int64(port) != consts.DefaultEnvdServerPort
 
@@ -92,7 +90,7 @@ func NewSandboxProxy(meterProvider metric.MeterProvider, port uint16, sandboxes
 
 			// Handle request host masking only for non-envd traffic.
 			var maskRequestHost *string = nil
-			if h := sbx.Config.Network.GetIngress().GetMaskRequestHost(); isNonEnvdTraffic && h != "" {
+			if h := ingress.GetMaskRequestHost(); isNonEnvdTraffic && h != "" {
 				h = strings.ReplaceAll(h, pool.MaskRequestHostPortPlaceholder, strconv.FormatUint(port, 10))
 				maskRequestHost = &h
 			}

packages/orchestrator/internal/sandbox/network/firewall.go

@@ -3,6 +3,7 @@ package network
 import (
 	"fmt"
 	"net/netip"
+	"slices"
 
 	"github.com/google/nftables"
 	"github.com/google/nftables/binaryutil"
@@ -94,7 +95,7 @@ func NewFirewall(tapIf string, orchestratorInternalIP string) (*Firewall, error)
 	}
 
 	// Populate the sets with initial data
-	err = fw.Reset()
+	err = fw.ReplaceUserRules(nil, nil)
 	if err != nil {
 		return nil, fmt.Errorf("error while configuring initial data: %w", err)
 	}
@@ -240,120 +241,72 @@ func (fw *Firewall) installRules() error {
 	return nil
 }
 
-// AddDeniedCIDR adds a single CIDR to the deny set at runtime.
-func (fw *Firewall) AddDeniedCIDR(cidr string) error {
-	err := addCIDRToSet(fw.conn, fw.userDenySet, cidr)
-	if err != nil {
-		return fmt.Errorf("add denied CIDR to set: %w", err)
+// ReplaceUserRules atomically replaces all firewall sets in a single flush.
+// This avoids any window where rules are partially applied.
+func (fw *Firewall) ReplaceUserRules(allowedCIDRs, deniedCIDRs []string) error {
+	// 1. Reset predefined deny set to default blocked ranges (buffered, no flush).
+	if err := fw.predefinedDenySet.ClearAndAddElements(fw.conn, sandbox_network.DeniedSandboxSetData); err != nil {
+		return fmt.Errorf("reset predefined deny set: %w", err)
 	}
 
-	err = fw.conn.Flush()
+	// 2. Reset predefined allow set to allowedRanges (buffered, no flush).
+	allowedSetData, err := set.AddressStringsToSetData(fw.allowedRanges)
 	if err != nil {
-		return fmt.Errorf("flush add denied changes: %w", err)
-	}
-
-	return nil
-}
-
-// AddAllowedCIDR adds a single CIDR to the allow set at runtime.
-func (fw *Firewall) AddAllowedCIDR(cidr string) error {
-	err := addCIDRToSet(fw.conn, fw.userAllowSet, cidr)
-	if err != nil {
-		return fmt.Errorf("add allowed CIDR to set: %w", err)
+		return fmt.Errorf("parse initial allowed CIDRs: %w", err)
 	}
-
-	err = fw.conn.Flush()
-	if err != nil {
-		return fmt.Errorf("flush add allowed changes: %w", err)
+	if err := fw.predefinedAllowSet.ClearAndAddElements(fw.conn, allowedSetData); err != nil {
+		return fmt.Errorf("reset predefined allow set: %w", err)
 	}
 
-	return nil
-}
-
-func (fw *Firewall) Reset() error {
-	if err := fw.ResetDeniedSets(); err != nil {
-		return fmt.Errorf("clear denied set: %w", err)
-	}
-	if err := fw.ResetAllowedSets(); err != nil {
-		return fmt.Errorf("clear allow set: %w", err)
+	// 3. Replace user deny set with new denied CIDRs (buffered, no flush).
+	if err := clearAndReplaceCIDRs(fw.conn, fw.userDenySet, deniedCIDRs); err != nil {
+		return fmt.Errorf("replace user deny set: %w", err)
 	}
 
-	return nil
-}
-
-// ResetDeniedSets resets the deny set back to original ranges.
-func (fw *Firewall) ResetDeniedSets() error {
-	// Always deny the default ranges
-	if err := fw.predefinedDenySet.ClearAndAddElements(fw.conn, sandbox_network.DeniedSandboxSetData); err != nil {
-		return err
+	// 4. Replace user allow set with new allowed CIDRs (buffered, no flush).
+	if err := clearAndReplaceCIDRs(fw.conn, fw.userAllowSet, allowedCIDRs); err != nil {
+		return fmt.Errorf("replace user allow set: %w", err)
 	}
 
-	// User defined denied ranges
-	if err := fw.userDenySet.ClearAndAddElements(fw.conn, nil); err != nil {
-		return err
-	}
-
-	return fw.conn.Flush()
-}
-
-// ResetAllowedSets resets allow set back to original ranges.
-func (fw *Firewall) ResetAllowedSets() error {
-	// Always allowed ranges
-	initData, err := set.AddressStringsToSetData(fw.allowedRanges)
-	if err != nil {
-		return fmt.Errorf("parse initial allowed CIDRs: %w", err)
-	}
-	if err := fw.predefinedAllowSet.ClearAndAddElements(fw.conn, initData); err != nil {
-		return err
-	}
-
-	// User defined allowed ranges
-	if err := fw.userAllowSet.ClearAndAddElements(fw.conn, nil); err != nil {
-		return err
+	// 5. Single atomic flush.
+	if err := fw.conn.Flush(); err != nil {
+		return fmt.Errorf("flush atomic rule replacement: %w", err)
 	}
 
-	return fw.conn.Flush()
+	return nil
 }
 
-func addCIDRToSet(conn *nftables.Conn, ipset set.Set, cidr string) error {
-	current, err := ipset.Elements(conn)
-	if err != nil {
-		return err
-	}
+// clearAndReplaceCIDRs flushes a set and repopulates it with the given CIDRs.
+// Handles the special 0.0.0.0/0 case which the firewall_toolkit validation
+// rejects (0.0.0.0 is "unspecified") by directly creating nftables elements.
+func clearAndReplaceCIDRs(conn *nftables.Conn, s set.Set, cidrs []string) error {
+	conn.FlushSet(s.Set())
 
-	// The checked range is 0.0.0.0 to 255.255.255.254, because when 255.255.255.255 is added, it's then requested as 255.255.255.254.
-	if len(current) == 1 && current[0].AddressRangeStart == netip.MustParseAddr("0.0.0.0") && current[0].AddressRangeEnd == netip.MustParseAddr("255.255.255.254") {
-		// Because 0.0.0.0/0 is not valid IP per GoLang, we can't add new addresses to the set.
+	if len(cidrs) == 0 {
 		return nil
 	}
 
-	// 0.0.0.0/0 is not valid IP per GoLang, so we handle it as a special case
-	if cidr == sandbox_network.AllInternetTrafficCIDR {
-		conn.FlushSet(ipset.Set())
-
-		toAppend := []nftables.SetElement{
-			{
-				Key: netip.MustParseAddr("0.0.0.0").AsSlice(),
-			},
-			{
-				Key:         netip.MustParseAddr("255.255.255.255").AsSlice(),
-				IntervalEnd: true,
-			},
+	// 0.0.0.0/0 must be handled specially: the firewall_toolkit's
+	// ValidateAddress rejects 0.0.0.0 as "unspecified", so we bypass
+	// the toolkit and create raw nftables interval elements directly.
+	if slices.Contains(cidrs, sandbox_network.AllInternetTrafficCIDR) {
+		elems := []nftables.SetElement{
+			{Key: netip.MustParseAddr("0.0.0.0").AsSlice()},
+			{Key: netip.MustParseAddr("255.255.255.255").AsSlice(), IntervalEnd: true},
 		}
-
-		if err := conn.SetAddElements(ipset.Set(), toAppend); err != nil {
-			return fmt.Errorf("add elements to denied set: %w", err)
+		if err := conn.SetAddElements(s.Set(), elems); err != nil {
+			return fmt.Errorf("add all-traffic elements: %w", err)
 		}
 
 		return nil
 	}
 
-	data, err := set.AddressStringsToSetData([]string{cidr})
+	data, err := set.AddressStringsToSetData(cidrs)
 	if err != nil {
 		return err
 	}
-
-	merged := append(current, data...)
-
-	return ipset.ClearAndAddElements(conn, merged)
+	// We already flushed the set above, so just add the new elements.
+	// Using ClearAndAddElements would flush again which is fine (FlushSet is idempotent
+	// when buffered), but this is more direct.
+	return s.ClearAndAddElements(conn, data)
 }