feat(orch): live-upgrade envd inside a running sandbox at resume

docs/ARCHITECTURE.md

@@ -174,9 +174,18 @@ The agent inside every VM (started by systemd very early in boot), port 49983, c
   stdout/stderr, stdin, signals, PTYs — this is what SDKs use to "run code".
 - **Filesystem service** (`spec/filesystem/filesystem.proto`): stat/list/make/move/remove/watch.
 - **REST**: `/health`, `/metrics`, `/files` upload/download, `/init` (orchestrator pushes env
-  vars, access token, metadata after boot/resume), freeze/thaw hooks used during pause.
+  vars, access token, metadata after boot/resume), `/upgrade` (live self-upgrade, below),
+  freeze/thaw hooks used during pause.
 - **Auth**: `X-Access-Token` header checked against a token delivered via Firecracker MMDS;
   signed URLs for file endpoints.
+- **Live upgrade** (`internal/services/process/upgrade.go`): an authenticated `POST /upgrade` lets
+  the orchestrator swap envd inside a *running* sandbox at resume. It streams the new binary in the
+  request body and envd `syscall.Exec`s into it **with the same PID**, carrying the workload's
+  stdio/PTY fds, process table, recently-retained exit codes and filesystem watchers forward via a
+  tmpfs handover blob. The workload cgroups stay frozen until the post-upgrade `/init` restores the
+  access token (so no re-adopted process runs unauthenticated), and the handover outcome
+  (procs/watchers re-adopted, plus any failures) rides back on that `/init`'s `X-Envd-Handover`
+  header for fleet visibility.
 - Scans guest ports and forwards them so any port a user process opens becomes reachable through
   sandbox URLs. **`pkg/version.go` must be bumped on every behavioral change** — the API and the
   orchestrator gate features on the envd version recorded in each template build.
@@ -290,6 +299,12 @@ sequenceDiagram
 - **Resume**: same path as creation, but placement prefers the **origin node** — if the snapshot
   is still in its local cache, resume avoids any object-storage reads. `Checkpoint` is a
   pause+resume in place used to persist state while keeping the sandbox running.
+- **Envd live-upgrade on resume**: the orchestrator can upgrade the sandbox's envd to a newer
+  node-local build during resume (gated by the `envd-upgrade-target` flag in
+  `packages/shared/pkg/featureflags`), via envd's `POST /upgrade` (see the envd section). It is
+  best-effort — a delivery failure before the `exec` leaves the old envd serving — except an
+  unrecoverable post-`exec` failure (the new envd never re-initializes), which fails the resume
+  rather than return a permanently unusable sandbox.
 - Auto-pause/auto-resume make sandboxes effectively serverless: idle sandboxes pause, traffic
   resumes them (see traffic flow above).
 

packages/envd/internal/api/auth.go

@@ -31,13 +31,27 @@ var authExcludedPaths = []string{
 	"POST/init",
 }
 
+// handoverPreInitAllowedPaths is the MINIMAL set reachable on a live-upgraded
+// envd before its post-upgrade /init has restored the access token: only /init
+// (which restores auth and lifts this gate, self-authenticated via MMDS) and
+// the health check. It deliberately omits /files — unlike authExcludedPaths —
+// so a re-adopted (possibly hostile) guest process can't reach the
+// root-privileged file API unauthenticated in that window. The orchestrator
+// delivers the upgrade over /upgrade's body, not /files, so nothing legitimate
+// needs /files before /init.
+var handoverPreInitAllowedPaths = []string{
+	"GET/health",
+	"POST/init",
+}
+
 func (a *API) WithAuthorization(handler http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
-		if a.accessToken.IsSet() {
-			authHeader := req.Header.Get(accessTokenHeader)
+		// check if this path is allowed without authentication (e.g., health check, endpoints supporting signing)
+		allowedPath := slices.Contains(authExcludedPaths, req.Method+req.URL.Path)
 
-			// check if this path is allowed without authentication (e.g., health check, endpoints supporting signing)
-			allowedPath := slices.Contains(authExcludedPaths, req.Method+req.URL.Path)
+		switch {
+		case a.accessToken.IsSet():
+			authHeader := req.Header.Get(accessTokenHeader)
 
 			if !a.accessToken.Equals(authHeader) && !allowedPath {
 				a.logger.Error().Msg("Trying to access secured envd without correct access token")
@@ -47,6 +61,24 @@ func (a *API) WithAuthorization(handler http.Handler) http.Handler {
 
 				return
 			}
+
+		case a.handover != nil && !a.initialized.Load() && !slices.Contains(handoverPreInitAllowedPaths, req.Method+req.URL.Path):
+			// A live-upgraded envd serves before its post-upgrade /init has
+			// restored the access token — and the fallback thaw may already be
+			// running the re-adopted workload. The sandbox HAD a token (it is a
+			// resume), so treat the unset token as "not yet restored" and fail
+			// CLOSED here rather than falling through to the open path below,
+			// which would let a re-adopted (and possibly hostile) guest process
+			// reach control endpoints unauthenticated in this window. Only the
+			// minimal handoverPreInitAllowedPaths (/init, /health) get through —
+			// notably NOT /files, whose root file API is otherwise unauthenticated
+			// (it is in authExcludedPaths). /init self-authenticates via MMDS and
+			// lifts this gate.
+			a.logger.Warn().Msg("blocking pre-init request on live-upgraded envd (auth not yet restored)")
+
+			jsonError(w, http.StatusUnauthorized, errors.New("envd not initialized"))
+
+			return
 		}
 
 		handler.ServeHTTP(w, req)

packages/envd/internal/api/compose_test.go

@@ -33,7 +33,7 @@ func newComposeTestAPI(t *testing.T) (*API, *user.User) {
 		User:    currentUser.Username,
 	}
 
-	return New(&logger, defaults, nil, false, cgroups.NewNoopManager()), currentUser
+	return New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager())), currentUser
 }
 
 func writeSourceFile(t *testing.T, dir string, name string, data []byte) string {

packages/envd/internal/api/download_test.go

@@ -96,7 +96,7 @@ func TestGetFilesContentDisposition(t *testing.T) {
 				EnvVars: utils.NewEnvVars(),
 				User:    currentUser.Username,
 			}
-			api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+			api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 			// Create request and response recorder
 			req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/files?path="+url.QueryEscape(tempFile), nil)
@@ -145,7 +145,7 @@ func TestGetFilesContentDispositionWithNestedPath(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	// Create request and response recorder
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/files?path="+url.QueryEscape(tempFile), nil)
@@ -188,7 +188,7 @@ func TestGetFiles_GzipEncoding_ExplicitIdentityOffWithRange(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	// Create request and response recorder
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/files?path="+url.QueryEscape(tempFile), nil)
@@ -229,7 +229,7 @@ func TestGetFiles_GzipDownload(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/files?path="+url.QueryEscape(tempFile), nil)
 	req.Header.Set("Accept-Encoding", "gzip")
@@ -294,7 +294,7 @@ func TestPostFiles_GzipUpload(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), &gzBuf)
 	req.Header.Set("Content-Type", mpWriter.FormDataContentType())
@@ -334,7 +334,7 @@ func TestPostFiles_RawBodyUpload(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), bytes.NewReader(originalContent))
 	req.Header.Set("Content-Type", "application/octet-stream")
@@ -372,7 +372,7 @@ func TestPostFiles_RawBodyUploadCreatesDirectories(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), bytes.NewReader(originalContent))
 	req.Header.Set("Content-Type", "application/octet-stream")
@@ -405,7 +405,7 @@ func TestPostFiles_RawBodyUploadRequiresPath(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files", bytes.NewReader([]byte("some content")))
 	req.Header.Set("Content-Type", "application/octet-stream")
@@ -440,7 +440,7 @@ func TestPostFiles_RawBodyUploadOverwritesExisting(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), bytes.NewReader(newContent))
 	req.Header.Set("Content-Type", "application/octet-stream")
@@ -486,7 +486,7 @@ func TestPostFiles_RawBodyGzipUpload(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), &gzBuf)
 	req.Header.Set("Content-Type", "application/octet-stream")
@@ -520,7 +520,7 @@ func TestPostFiles_UnsupportedContentType(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	tempDir := t.TempDir()
 	destPath := filepath.Join(tempDir, "test.txt")
@@ -566,7 +566,7 @@ func TestPostFiles_MultipartStillWorksWithoutContentType(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), &multipartBuf)
 	req.Header.Set("Content-Type", mpWriter.FormDataContentType())
@@ -624,7 +624,7 @@ func TestGzipUploadThenGzipDownload(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	uploadReq := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), &gzBuf)
 	uploadReq.Header.Set("Content-Type", mpWriter.FormDataContentType())

packages/envd/internal/api/init.go

@@ -22,7 +22,7 @@ import (
 	"github.com/e2b-dev/infra/packages/envd/internal/host"
 	"github.com/e2b-dev/infra/packages/envd/internal/logs"
 	"github.com/e2b-dev/infra/packages/envd/internal/logs/ratelimit"
-	"github.com/e2b-dev/infra/packages/envd/internal/services/cgroups"
+	"github.com/e2b-dev/infra/packages/envd/pkg"
 	"github.com/e2b-dev/infra/packages/shared/pkg/keys"
 )
 
@@ -115,6 +115,22 @@ func (a *API) checkMMDSHash(ctx context.Context, requestToken *SecureToken) (boo
 func (a *API) PostInit(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close()
 
+	// Report the running envd version on every /init response (even error/stale
+	// ones) so the orchestrator can read the live version off the resume-path
+	// call it already makes — no extra round-trip. Set before any WriteHeader.
+	w.Header().Set("X-Envd-Version", pkg.Version)
+
+	// If this envd booted from a live-upgrade handover, advertise its outcome on
+	// /init so the orchestrator can record it — the envd-side result (re-adopted
+	// procs, restored retained exits, watcher re-arm success/failures) is
+	// otherwise only logged in-guest and invisible fleet-wide. Set before any
+	// WriteHeader.
+	if a.handover != nil {
+		if b, err := json.Marshal(a.handover); err == nil {
+			w.Header().Set("X-Envd-Handover", string(b))
+		}
+	}
+
 	ctx := r.Context()
 
 	operationID := logs.AssignOperationID()
@@ -167,14 +183,27 @@ func (a *API) PostInit(w http.ResponseWriter, r *http.Request) {
 		// requests still thaw cgroups after pre-pause freeze.
 		defer a.unfreezeUserCgroups(ctx, logger)
 
-		// Update data only if the request is newer or if there's no timestamp at all
+		// Restore the access token (and env) BEFORE marking the envd initialized.
+		// On a live-upgraded envd, WithAuthorization only fails CLOSED while
+		// !initialized; flipping initialized first, with the token not yet
+		// restored, falls through to the fail-OPEN path and lets a re-adopted
+		// (and possibly hostile) guest process reach control endpoints — including
+		// /upgrade — unauthenticated in that window. Update only if the request is
+		// newer (or carries no timestamp); a stale/replayed /init keeps the token
+		// already set by the first one.
 		if initRequest.Timestamp == nil || a.lastSetTime.SetToGreater(initRequest.Timestamp.UnixNano()) {
 			if err := a.SetData(ctx, logger, initRequest); err != nil {
 				writeInitError(w, logger, err)
 
 				return
 			}
 		}
+
+		// Auth passed and token restored: mark the envd initialized so the
+		// live-upgrade /upgrade endpoint and the post-upgrade fallback thaw open
+		// up — a guest process that can't pass auth can't flip this and drive an
+		// unauthenticated upgrade.
+		a.initialized.Store(true)
 	}
 
 	go func() { //nolint:contextcheck // TODO: fix this later
@@ -245,11 +274,6 @@ func (a *API) SetData(ctx context.Context, logger zerolog.Logger, data PostInitJ
 }
 
 // userCgroupsToFreeze is the cgroup set frozen pre-pause and thawed on /init.
-var userCgroupsToFreeze = []cgroups.ProcessType{
-	cgroups.ProcessTypeUser,
-	cgroups.ProcessTypePTY,
-}
-
 // PostFreeze freezes user/pty cgroups directly (no Process.Start / shell).
 // Orchestrator calls this just before pause; the frozen state persists into the
 // snapshot and /init thaws on resume. Best-effort: tries every cgroup even if
@@ -259,22 +283,16 @@ func (a *API) PostFreeze(w http.ResponseWriter, r *http.Request) {
 
 	logger := a.logger.With().Str(string(logs.OperationIDKey), logs.AssignOperationID()).Logger()
 
-	if err := a.freezeLock.Acquire(r.Context(), 1); err != nil {
-		w.WriteHeader(http.StatusServiceUnavailable)
-
-		return
-	}
-	defer a.freezeLock.Release(1)
+	if err := a.workloadFreezer.Freeze(r.Context()); err != nil {
+		// A failed lock acquire means the request ctx was cancelled; a failed
+		// sweep leaves ctx intact.
+		if r.Context().Err() != nil {
+			w.WriteHeader(http.StatusServiceUnavailable)
 
-	var errs []error
-	for _, pt := range userCgroupsToFreeze {
-		if err := a.cgroupManager.Freeze(pt); err != nil {
-			logger.Error().Err(err).Msgf("freeze %s cgroup", pt)
-			errs = append(errs, fmt.Errorf("freeze %s cgroup: %w", pt, err))
+			return
 		}
-	}
-	if len(errs) > 0 {
-		jsonError(w, http.StatusInternalServerError, errors.Join(errs...))
+		logger.Error().Err(err).Msg("freeze workload cgroups")
+		jsonError(w, http.StatusInternalServerError, err)
 
 		return
 	}
@@ -290,25 +308,11 @@ func (a *API) PostFreeze(w http.ResponseWriter, r *http.Request) {
 func (a *API) PostUnfreeze(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close()
 
-	ctx := r.Context()
 	logger := a.logger.With().Str(string(logs.OperationIDKey), logs.AssignOperationID()).Logger()
 
-	if err := a.freezeLock.Acquire(context.WithoutCancel(ctx), 1); err != nil {
-		w.WriteHeader(http.StatusServiceUnavailable)
-
-		return
-	}
-	defer a.freezeLock.Release(1)
-
-	var errs []error
-	for _, pt := range userCgroupsToFreeze {
-		if err := a.cgroupManager.Unfreeze(pt); err != nil {
-			logger.Error().Err(err).Msgf("unfreeze %s cgroup", pt)
-			errs = append(errs, fmt.Errorf("unfreeze %s cgroup: %w", pt, err))
-		}
-	}
-	if len(errs) > 0 {
-		jsonError(w, http.StatusInternalServerError, errors.Join(errs...))
+	if err := a.workloadFreezer.Unfreeze(r.Context()); err != nil {
+		logger.Error().Err(err).Msg("unfreeze workload cgroups")
+		jsonError(w, http.StatusInternalServerError, err)
 
 		return
 	}
@@ -318,15 +322,11 @@ func (a *API) PostUnfreeze(w http.ResponseWriter, r *http.Request) {
 }
 
 // unfreezeUserCgroups unfreezes user/pty cgroups (idempotent if not frozen).
-// Wraps the context with WithoutCancel so the unfreeze always completes.
+// The freezer detaches the wait from ctx cancellation so the unfreeze always
+// completes.
 func (a *API) unfreezeUserCgroups(ctx context.Context, logger zerolog.Logger) {
-	_ = a.freezeLock.Acquire(context.WithoutCancel(ctx), 1)
-	defer a.freezeLock.Release(1)
-
-	for _, pt := range userCgroupsToFreeze {
-		if err := a.cgroupManager.Unfreeze(pt); err != nil {
-			logger.Warn().Err(err).Msgf("unfreeze %s cgroup", pt)
-		}
+	if err := a.workloadFreezer.Unfreeze(ctx); err != nil {
+		logger.Warn().Err(err).Msg("unfreeze workload cgroups")
 	}
 }