Reassociate pma.file to the correct pgalloc.MemoryFile on restore.

Earlier we were always restoring pma.file to mm.mfp.MemoryFile(). However,
d8eb29ed6f ("Add support for saving PMAs referencing tmpfs filestore files.")
added support for saving PMAs that reference "private" pgalloc.MemoryFiles that
are different from mm.mfp.MemoryFile().

We achieve the correct restore by:
- Adding a "RestoreID" field to pgalloc.MemoryFile. Private MemoryFiles set
  this with a vfs.RestoreID.String(). Non-private MemoryFile does not set it.
- MemoryFile struct is not savable by itself, but pma.file field is saved as a
  string. We store the RestoreID string there.
- On restore, if RestoreID is "", then restore using CtxMemoryFile. If it has a
  non-empty RestoreID, then restore using CtxMemoryFileMap.
- Cleanup: vfs.CtxFilesystemMemoryFileMap was moved to pgalloc.CtxMemoryFileMap
  so we can now provide a pgalloc.MemoryFileMapFromContext() method which
  cleans up some code. Also the key to this map (MemoryFileOpts.RestoreID)
  belongs to pgalloc, so it seems like the right place to have this context.

PiperOrigin-RevId: 614903073

pkg/sentry/fsimpl/tmpfs/save_restore.go

@@ -45,12 +45,14 @@ func (fs *filesystem) PrepareSave(ctx context.Context) error {
 	if !fs.privateMF {
 		return nil
 	}
-	mfmapv := ctx.Value(vfs.CtxFilesystemMemoryFileMap)
-	if mfmapv == nil {
-		return fmt.Errorf("CtxFilesystemMemoryFileMap was not provided")
+	mfmap := pgalloc.MemoryFileMapFromContext(ctx)
+	if mfmap == nil {
+		return fmt.Errorf("CtxMemoryFileMap was not provided")
 	}
-	mfmap := mfmapv.(map[vfs.RestoreID]*pgalloc.MemoryFile)
-	mfmap[fs.uniqueID] = fs.mf
+	if _, ok := mfmap[fs.uniqueID.String()]; ok {
+		return fmt.Errorf("memory file for %q already exists in CtxMemoryFileMap", fs.uniqueID)
+	}
+	mfmap[fs.uniqueID.String()] = fs.mf
 	return nil
 }
 
@@ -60,14 +62,13 @@ func (fs *filesystem) CompleteRestore(ctx context.Context, opts vfs.CompleteRest
 	if !fs.privateMF {
 		return nil
 	}
-	mfmapv := ctx.Value(vfs.CtxFilesystemMemoryFileMap)
-	if mfmapv == nil {
-		return fmt.Errorf("CtxFilesystemMemoryFileMap was not provided")
+	mfmap := pgalloc.MemoryFileMapFromContext(ctx)
+	if mfmap == nil {
+		return fmt.Errorf("CtxMemoryFileMap was not provided")
 	}
-	mfmap := mfmapv.(map[vfs.RestoreID]*pgalloc.MemoryFile)
-	mf, ok := mfmap[fs.uniqueID]
+	mf, ok := mfmap[fs.uniqueID.String()]
 	if !ok {
-		return fmt.Errorf("memory file for %q not found in CtxFilesystemMemoryFileMap", fs.uniqueID)
+		return fmt.Errorf("memory file for %q not found in CtxMemoryFileMap", fs.uniqueID)
 	}
 	fs.mf = mf
 	return nil

pkg/sentry/kernel/kernel.go

@@ -522,10 +522,10 @@ func (k *Kernel) Init(args InitKernelArgs) error {
 
 // +stateify savable
 type privateMemoryFileMetadata struct {
-	owners []vfs.RestoreID
+	owners []string
 }
 
-func savePrivateMFs(ctx context.Context, w wire.Writer, mfsToSave map[vfs.RestoreID]*pgalloc.MemoryFile) error {
+func savePrivateMFs(ctx context.Context, w wire.Writer, mfsToSave map[string]*pgalloc.MemoryFile) error {
 	var meta privateMemoryFileMetadata
 	// Generate the order in which private memory files are saved.
 	for fsID := range mfsToSave {
@@ -550,10 +550,7 @@ func loadPrivateMFs(ctx context.Context, r wire.Reader) error {
 	if _, err := state.Load(ctx, r, &meta); err != nil {
 		return err
 	}
-	var mfmap map[vfs.RestoreID]*pgalloc.MemoryFile
-	if mfmapv := ctx.Value(vfs.CtxFilesystemMemoryFileMap); mfmapv != nil {
-		mfmap = mfmapv.(map[vfs.RestoreID]*pgalloc.MemoryFile)
-	}
+	mfmap := pgalloc.MemoryFileMapFromContext(ctx)
 	// Ensure that it is consistent with CtxFilesystemMemoryFileMap.
 	if len(mfmap) != len(meta.owners) {
 		return fmt.Errorf("inconsistent private memory files on restore: savedMFOwners = %v, CtxFilesystemMemoryFileMap = %v", meta.owners, mfmap)
@@ -595,8 +592,8 @@ func (k *Kernel) SaveTo(ctx context.Context, w wire.Writer) error {
 	}
 
 	// Capture all private memory files.
-	mfsToSave := make(map[vfs.RestoreID]*pgalloc.MemoryFile)
-	vfsCtx := context.WithValue(ctx, vfs.CtxFilesystemMemoryFileMap, mfsToSave)
+	mfsToSave := make(map[string]*pgalloc.MemoryFile)
+	vfsCtx := context.WithValue(ctx, pgalloc.CtxMemoryFileMap, mfsToSave)
 	// Prepare filesystems for saving. This must be done after
 	// invalidateUnsavableMappings(), since dropping memory mappings may
 	// affect filesystem state (e.g. page cache reference counts).

pkg/sentry/mm/mm.go

@@ -345,10 +345,10 @@ func (v *vma) copy() vma {
 //
 // +stateify savable
 type pma struct {
-	// file is the file mapped by this pma. Only pmas for which file ==
-	// MemoryManager.mfp.MemoryFile() may be saved. pmas hold a reference to
-	// the corresponding file range while they exist.
-	file memmap.File `state:"nosave"`
+	// file is the file mapped by this pma. Only pmas for which file is of type
+	// pgalloc.MemoryFile may be saved. pmas hold a reference to the
+	// corresponding file range while they exist.
+	file memmap.File `state:".(string)"`
 
 	// off is the offset into file at which this pma begins.
 	off uint64

pkg/sentry/mm/save_restore.go

@@ -37,28 +37,10 @@ func (mm *MemoryManager) InvalidateUnsavable(ctx context.Context) error {
 	return nil
 }
 
-// beforeSave is invoked by stateify.
-func (mm *MemoryManager) beforeSave() {
-	for pseg := mm.pmas.FirstSegment(); pseg.Ok(); pseg = pseg.NextSegment() {
-		if pma := pseg.ValuePtr(); pma.file != nil {
-			if mf, ok := pma.file.(*pgalloc.MemoryFile); ok && mf.IsSavable() {
-				// If the MemoryFile will be saved, then its PMAs are preserved.
-				continue
-			}
-			// InvalidateUnsavable should have caused all such pmas to be
-			// invalidated.
-			panic(fmt.Sprintf("Can't save pma %#v with non-MemoryFile of type %T:\n%s", pseg.Range(), pma.file, mm))
-		}
-	}
-}
-
 // afterLoad is invoked by stateify.
 func (mm *MemoryManager) afterLoad(goContext.Context) {
 	mm.mf = mm.mfp.MemoryFile()
 	mm.haveASIO = mm.p.SupportsAddressSpaceIO()
-	for pseg := mm.pmas.FirstSegment(); pseg.Ok(); pseg = pseg.NextSegment() {
-		pseg.ValuePtr().file = mm.mf
-	}
 }
 
 const (
@@ -148,3 +130,29 @@ func (v *vma) loadRealPerms(_ goContext.Context, b int) {
 		v.growsDown = true
 	}
 }
+
+func (p *pma) saveFile() string {
+	mf, ok := p.file.(*pgalloc.MemoryFile)
+	if !ok {
+		// InvalidateUnsavable should have caused all such pmas to be
+		// invalidated.
+		panic(fmt.Sprintf("Can't save pma with non-MemoryFile of type %T", p.file))
+	}
+	if !mf.IsSavable() {
+		panic(fmt.Sprintf("Can't save pma because its MemoryFile is not savable: %v", mf))
+	}
+	return mf.RestoreID()
+}
+
+func (p *pma) loadFile(ctx goContext.Context, restoreID string) {
+	if restoreID == "" {
+		p.file = pgalloc.MemoryFileFromContext(ctx)
+		return
+	}
+	mfmap := pgalloc.MemoryFileMapFromContext(ctx)
+	mf, ok := mfmap[restoreID]
+	if !ok {
+		panic(fmt.Sprintf("can't restore pma because its MemoryFile's restore ID %q was not found in CtxMemoryFileMap", restoreID))
+	}
+	p.file = mf
+}

pkg/sentry/pgalloc/context.go

@@ -15,7 +15,7 @@
 package pgalloc
 
 import (
-	"gvisor.dev/gvisor/pkg/context"
+	"context"
 )
 
 // contextID is this package's type for context.Context.Value keys.
@@ -30,6 +30,10 @@ const (
 
 	// CtxMemoryCgroupID is the memory cgroup id which the task belongs to.
 	CtxMemoryCgroupID
+
+	// CtxMemoryFileMap is a Context.Value key for mapping
+	// MemoryFileOpts.RestoreID to *MemoryFile. This is used for save/restore.
+	CtxMemoryFileMap
 )
 
 // MemoryFileFromContext returns the MemoryFile used by ctx, or nil if no such
@@ -41,8 +45,8 @@ func MemoryFileFromContext(ctx context.Context) *MemoryFile {
 	return nil
 }
 
-// MemoryFileProviderFromContext returns the MemoryFileProvider used by ctx, or nil if no such
-// MemoryFileProvider exists.
+// MemoryFileProviderFromContext returns the MemoryFileProvider used by ctx, or
+// nil if no such MemoryFileProvider exists.
 func MemoryFileProviderFromContext(ctx context.Context) MemoryFileProvider {
 	if v := ctx.Value(CtxMemoryFileProvider); v != nil {
 		return v.(MemoryFileProvider)
@@ -58,3 +62,12 @@ func MemoryCgroupIDFromContext(ctx context.Context) uint32 {
 	}
 	return 0
 }
+
+// MemoryFileMapFromContext returns the memory file map used by ctx, or nil if
+// no such map exists.
+func MemoryFileMapFromContext(ctx context.Context) map[string]*MemoryFile {
+	if v := ctx.Value(CtxMemoryFileMap); v != nil {
+		return v.(map[string]*MemoryFile)
+	}
+	return nil
+}

pkg/sentry/pgalloc/pgalloc.go

@@ -216,6 +216,10 @@ type MemoryFileOpts struct {
 
 	// DiskBackedFile indicates that the MemoryFile is backed by a file on disk.
 	DiskBackedFile bool
+
+	// RestoreID is an opaque string used to reassociate the MemoryFile with its
+	// replacement during restore.
+	RestoreID string
 }
 
 // DelayedEvictionType is the type of MemoryFileOpts.DelayedEviction.

pkg/sentry/pgalloc/save_restore.go

@@ -129,6 +129,11 @@ func (f *MemoryFile) IsSavable() bool {
 	return f.savable
 }
 
+// RestoreID returns the restore ID for f.
+func (f *MemoryFile) RestoreID() string {
+	return f.opts.RestoreID
+}
+
 // LoadFrom loads MemoryFile state from the given stream.
 func (f *MemoryFile) LoadFrom(ctx context.Context, r wire.Reader) error {
 	// Load metadata.

pkg/sentry/vfs/context.go

@@ -32,10 +32,6 @@ const (
 	// mapping filesystem unique IDs (cf. gofer.InternalFilesystemOptions.UniqueID)
 	// to host FDs.
 	CtxRestoreFilesystemFDMap
-
-	// CtxFilesystemMemoryFileMap is a Context.Value key for mapping tmpfs unique
-	// IDs to private memory files. This is used for save/restore.
-	CtxFilesystemMemoryFileMap
 )
 
 // MountNamespaceFromContext returns the MountNamespace used by ctx. If ctx is

runsc/boot/vfs.go

@@ -637,7 +637,7 @@ func (c *containerMounter) configureOverlay(ctx context.Context, conf *config.Co
 	}
 	if filestoreFD != nil {
 		// Create memory file for disk-backed overlays.
-		mf, err := createPrivateMemoryFile(filestoreFD.ReleaseToFile("overlay-filestore"))
+		mf, err := createPrivateMemoryFile(filestoreFD.ReleaseToFile("overlay-filestore"), vfs.RestoreID{ContainerName: c.containerName, Path: dst})
 		if err != nil {
 			return nil, nil, fmt.Errorf("failed to create memory file for overlay: %v", err)
 		}
@@ -897,7 +897,7 @@ func getMountNameAndOptions(spec *specs.Spec, conf *config.Config, m *mountInfo,
 			return "", nil, err
 		}
 		if m.filestoreFD != nil {
-			mf, err := createPrivateMemoryFile(m.filestoreFD.ReleaseToFile("tmpfs-filestore"))
+			mf, err := createPrivateMemoryFile(m.filestoreFD.ReleaseToFile("tmpfs-filestore"), vfs.RestoreID{ContainerName: containerName, Path: m.mount.Destination})
 			if err != nil {
 				return "", nil, fmt.Errorf("failed to create memory file for tmpfs: %v", err)
 			}
@@ -986,7 +986,7 @@ func parseKeyValue(s string) (string, string, bool) {
 	return strings.TrimSpace(tokens[0]), strings.TrimSpace(tokens[1]), true
 }
 
-func createPrivateMemoryFile(file *os.File) (*pgalloc.MemoryFile, error) {
+func createPrivateMemoryFile(file *os.File, restoreID vfs.RestoreID) (*pgalloc.MemoryFile, error) {
 	mfOpts := pgalloc.MemoryFileOpts{
 		// Private memory files are usually backed by files on disk. Ideally we
 		// would confirm with fstatfs(2) but that is prohibited by seccomp.
@@ -997,6 +997,8 @@ func createPrivateMemoryFile(file *os.File) (*pgalloc.MemoryFile, error) {
 		// pgalloc.IMAWorkAroundForMemFile() uses. Users of private memory files
 		// are expected to have performed the work around outside the sandbox.
 		DisableIMAWorkAround: true,
+		// Private memory files need to be restored correctly using this ID.
+		RestoreID: restoreID.String(),
 	}
 	return pgalloc.NewMemoryFile(file, mfOpts)
 }
@@ -1274,13 +1276,13 @@ func (c *containerMounter) configureRestore(ctx context.Context) (context.Contex
 	rootKey := vfs.RestoreID{ContainerName: c.containerName, Path: "/"}
 	fdmap[rootKey] = c.goferFDs.remove()
 
-	mfmap := make(map[vfs.RestoreID]*pgalloc.MemoryFile)
+	mfmap := make(map[string]*pgalloc.MemoryFile)
 	if rootfsConf := c.goferMountConfs[0]; rootfsConf.IsFilestorePresent() {
-		mf, err := createPrivateMemoryFile(c.goferFilestoreFDs.removeAsFD().ReleaseToFile("overlay-filestore"))
+		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)
 		}
-		mfmap[rootKey] = mf
+		mfmap[rootKey.String()] = mf
 	}
 	// prepareMounts() consumes the remaining FDs for submounts.
 	mounts, err := c.prepareMounts()
@@ -1294,15 +1296,15 @@ func (c *containerMounter) configureRestore(ctx context.Context) (context.Contex
 			fdmap[key] = submount.goferFD.Release()
 		}
 		if submount.filestoreFD != nil {
-			mf, err := createPrivateMemoryFile(submount.filestoreFD.ReleaseToFile("overlay-filestore"))
+			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)
 			}
-			key := vfs.RestoreID{ContainerName: c.containerName, Path: submount.mount.Destination}
-			mfmap[key] = mf
+			mfmap[key.String()] = mf
 		}
 	}
-	return context.WithValue(context.WithValue(ctx, vfs.CtxRestoreFilesystemFDMap, fdmap), vfs.CtxFilesystemMemoryFileMap, mfmap), nil
+	return context.WithValue(context.WithValue(ctx, vfs.CtxRestoreFilesystemFDMap, fdmap), pgalloc.CtxMemoryFileMap, mfmap), nil
 }
 
 func createDeviceFiles(ctx context.Context, creds *auth.Credentials, info *containerInfo, vfsObj *vfs.VirtualFilesystem, root vfs.VirtualDentry) error {