fix: register sandbox before WaitForEnvd for TCP firewall proxy

The TCP firewall proxy looks up sandboxes by source address to allow
egress traffic. When initEnvd makes outbound connections (e.g., GCS
for volume mount), the proxy couldn't find the sandbox because it
wasn't registered yet - Insert happened after ResumeSandbox returned.

Fix: Move sandbox registration into Factory, right before WaitForEnvd.
Remove the ineffective 1s delay that was added as a workaround.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: omins

packages/orchestrator/benchmark_test.go

@@ -175,7 +175,7 @@ func BenchmarkBaseImageLaunch(b *testing.B) {
 	templateCache.Start(b.Context())
 	b.Cleanup(templateCache.Stop)
 
-	sandboxFactory := sandbox.NewFactory(config.BuilderConfig, networkPool, devicePool, featureFlags)
+	sandboxFactory := sandbox.NewFactory(config.BuilderConfig, networkPool, devicePool, featureFlags, nil)
 
 	dockerhubRepository, err := dockerhub.GetRemoteRepository(b.Context())
 	require.NoError(b, err)

packages/orchestrator/cmd/build-template/main.go

@@ -216,7 +216,7 @@ func buildTemplate(
 		return fmt.Errorf("failed to create build metrics: %w", err)
 	}
 
-	sandboxFactory := sandbox.NewFactory(c.BuilderConfig, networkPool, devicePool, featureFlags)
+	sandboxFactory := sandbox.NewFactory(c.BuilderConfig, networkPool, devicePool, featureFlags, nil)
 
 	builder := build.NewBuilder(
 		builderConfig,

packages/orchestrator/internal/sandbox/envd.go

@@ -126,18 +126,6 @@ func (s *Sandbox) initEnvd(ctx context.Context) (e error) {
 		span.End()
 	}()
 
-	// When volume is configured, add a delay to allow VM network stack to fully initialize.
-	// Without this delay, the envd may try to connect to GCS before network is ready,
-	// resulting in EOF errors during TLS handshake.
-	if s.volumeInitConfig != nil {
-		const networkInitDelay = 300 * time.Millisecond
-		logger.L().Info(ctx, "waiting for VM network initialization before volume mount",
-			logger.WithSandboxID(s.Runtime.SandboxID),
-			zap.Duration("delay", networkInitDelay),
-		)
-		time.Sleep(networkInitDelay)
-	}
-
 	attributes := []attribute.KeyValue{telemetry.WithEnvdVersion(s.Config.Envd.Version), attribute.Int64("timeout_ms", s.internalConfig.EnvdInitRequestTimeout.Milliseconds())}
 	attributesFail := append(attributes, attribute.Bool("success", false))
 	attributesSuccess := append(attributes, attribute.Bool("success", true))

packages/orchestrator/internal/sandbox/sandbox.go

@@ -177,19 +177,22 @@ type Factory struct {
 	featureFlags *featureflags.Client
 	volumes      *VolumesConfig
 	tokenMinter  *gcstoken.Minter
+	sandboxes    *Map
 }
 
 func NewFactory(
 	config cfg.BuilderConfig,
 	networkPool *network.Pool,
 	devicePool *nbd.DevicePool,
 	featureFlags *featureflags.Client,
+	sandboxes *Map,
 ) *Factory {
 	return &Factory{
 		config:       config,
 		networkPool:  networkPool,
 		devicePool:   devicePool,
 		featureFlags: featureFlags,
+		sandboxes:    sandboxes,
 	}
 }
 
@@ -682,11 +685,22 @@ func (f *Factory) ResumeSandbox(
 
 	telemetry.ReportEvent(execCtx, "waiting for envd")
 
+	// Register sandbox before WaitForEnvd so TCP firewall proxy can find it.
+	// This is needed because initEnvd may make outbound connections (e.g., GCS for volume mount)
+	// that go through the TCP firewall proxy, which looks up sandboxes by source address.
+	if f.sandboxes != nil {
+		f.sandboxes.Insert(sbx)
+	}
+
 	err = sbx.WaitForEnvd(
 		ctx,
 		f.config.EnvdTimeout,
 	)
 	if err != nil {
+		// Remove from map on failure
+		if f.sandboxes != nil {
+			f.sandboxes.Remove(sbx.Runtime.SandboxID)
+		}
 		return nil, fmt.Errorf("failed to wait for sandbox start: %w", err)
 	}
 

packages/orchestrator/internal/server/sandboxes.go

@@ -189,7 +189,9 @@ func (s *Server) Create(ctx context.Context, req *orchestrator.SandboxCreateRequ
 		return nil, status.Errorf(codes.Internal, "failed to create sandbox: %s", err)
 	}
 
-	s.sandboxes.Insert(sbx)
+	// Note: sandbox is already inserted into the map by Factory.ResumeSandbox
+	// before WaitForEnvd is called, so TCP firewall proxy can find it.
+
 	go func() {
 		ctx, childSpan := tracer.Start(context.WithoutCancel(ctx), "sandbox-create-stop", trace.WithNewRoot())
 		defer childSpan.End()

packages/orchestrator/main.go

@@ -411,7 +411,7 @@ func run(config cfg.Config) (success bool) {
 	closers = append(closers, closer{"network pool", networkPool.Close})
 
 	// sandbox factory
-	sandboxFactory := sandbox.NewFactory(config.BuilderConfig, networkPool, devicePool, featureFlags)
+	sandboxFactory := sandbox.NewFactory(config.BuilderConfig, networkPool, devicePool, featureFlags, sandboxes)
 
 	// Configure volumes support if enabled
 	if config.VolumesRedisURL != "" {