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