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/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
+}