From 7137ec8798914009f427fe01a88bf8647447433f Mon Sep 17 00:00:00 2001
From: Fabricio Voznika <fvoznika@google.com>
Date: Tue, 9 Apr 2024 11:43:00 -0700
Subject: [PATCH] Refactor containerMounter.configureRestore

Setup maps that will be part of restore context first and then set
it to the context, to avoid setting the same map multiple times
when there are more than one container being restored.

Updates #1956

PiperOrigin-RevId: 623234955
---
 runsc/boot/restore.go | 11 +++++++----
 runsc/boot/vfs.go     | 13 +++++--------
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/runsc/boot/restore.go b/runsc/boot/restore.go
index 76d931c39..d9640a3ec 100644
--- a/runsc/boot/restore.go
+++ b/runsc/boot/restore.go
@@ -22,6 +22,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/fsimpl/host"
 	"gvisor.dev/gvisor/pkg/sentry/inet"
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
+	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
 	"gvisor.dev/gvisor/pkg/sentry/socket/hostinet"
 	"gvisor.dev/gvisor/pkg/sentry/socket/netstack"
 	"gvisor.dev/gvisor/pkg/sentry/state"
@@ -114,12 +115,12 @@ func (r *restorer) restore(l *Loader) error {
 
 	// TODO(b/298078576): Need to process hints here probably
 	mntr := newContainerMounter(&l.root, l.k, l.mountHints, l.sharedMounts, l.productName, l.sandboxID)
-	ctx, err = mntr.configureRestore(ctx)
-	if err != nil {
+
+	fdmap := make(map[vfs.RestoreID]int)
+	mfmap := make(map[string]*pgalloc.MemoryFile)
+	if err := mntr.configureRestore(fdmap, mfmap); err != nil {
 		return fmt.Errorf("configuring filesystem restore: %v", err)
 	}
-
-	fdmap := vfs.RestoreFilesystemFDMapFromContext(ctx)
 	for appFD, fd := range r.container.stdioFDs {
 		key := host.MakeRestoreID(r.container.containerName, appFD)
 		fdmap[key] = fd.Release()
@@ -128,6 +129,8 @@ func (r *restorer) restore(l *Loader) error {
 		key := host.MakeRestoreID(r.container.containerName, customFD.guest)
 		fdmap[key] = customFD.host.FD()
 	}
+	ctx = context.WithValue(ctx, vfs.CtxRestoreFilesystemFDMap, fdmap)
+	ctx = context.WithValue(ctx, pgalloc.CtxMemoryFileMap, mfmap)
 
 	// Load the state.
 	loadOpts := state.LoadOpts{Source: r.stateFile}
diff --git a/runsc/boot/vfs.go b/runsc/boot/vfs.go
index 9e056024f..94bca9644 100644
--- a/runsc/boot/vfs.go
+++ b/runsc/boot/vfs.go
@@ -1266,26 +1266,23 @@ func (c *containerMounter) makeMountPoint(ctx context.Context, creds *auth.Crede
 
 // configureRestore returns an updated context.Context including filesystem
 // state used by restore defined by conf.
-func (c *containerMounter) configureRestore(ctx context.Context) (context.Context, error) {
+func (c *containerMounter) configureRestore(fdmap map[vfs.RestoreID]int, mfmap map[string]*pgalloc.MemoryFile) error {
 	// Compare createMountNamespace(); rootfs always consumes a gofer FD and a
 	// filestore FD is consumed if the rootfs GoferMountConf indicates so.
-	fdmap := make(map[vfs.RestoreID]int)
-
 	rootKey := vfs.RestoreID{ContainerName: c.containerName, Path: "/"}
 	fdmap[rootKey] = c.goferFDs.remove()
 
-	mfmap := make(map[string]*pgalloc.MemoryFile)
 	if rootfsConf := c.goferMountConfs[0]; rootfsConf.IsFilestorePresent() {
 		mf, err := createPrivateMemoryFile(c.goferFilestoreFDs.removeAsFD().ReleaseToFile("overlay-filestore"), rootKey)
 		if err != nil {
-			return ctx, fmt.Errorf("failed to create private memory file for mount rootfs: %w", err)
+			return fmt.Errorf("failed to create private memory file for mount rootfs: %w", err)
 		}
 		mfmap[rootKey.String()] = mf
 	}
 	// prepareMounts() consumes the remaining FDs for submounts.
 	mounts, err := c.prepareMounts()
 	if err != nil {
-		return ctx, err
+		return err
 	}
 	for i := range mounts {
 		submount := &mounts[i]
@@ -1297,12 +1294,12 @@ func (c *containerMounter) configureRestore(ctx context.Context) (context.Contex
 			key := vfs.RestoreID{ContainerName: c.containerName, Path: submount.mount.Destination}
 			mf, err := createPrivateMemoryFile(submount.filestoreFD.ReleaseToFile("overlay-filestore"), key)
 			if err != nil {
-				return ctx, fmt.Errorf("failed to create private memory file for mount %q: %w", submount.mount.Destination, err)
+				return fmt.Errorf("failed to create private memory file for mount %q: %w", submount.mount.Destination, err)
 			}
 			mfmap[key.String()] = mf
 		}
 	}
-	return context.WithValue(context.WithValue(ctx, vfs.CtxRestoreFilesystemFDMap, fdmap), pgalloc.CtxMemoryFileMap, mfmap), nil
+	return nil
 }
 
 func createDeviceFiles(ctx context.Context, creds *auth.Credentials, info *containerInfo, vfsObj *vfs.VirtualFilesystem, root vfs.VirtualDentry) error {