Internal change.

PiperOrigin-RevId: 817500534

Author: nixprime

runsc/boot/controller.go

@@ -17,6 +17,7 @@ package boot
 import (
 	"errors"
 	"fmt"
+	"io"
 	"path"
 	"strconv"
 	"sync"
@@ -489,7 +490,7 @@ func (cm *containerManager) PortForward(opts *PortForwardOpts, _ *struct{}) erro
 
 // RestoreOpts contains options related to restoring a container's file system.
 type RestoreOpts struct {
-	// FilePayload contains the state file to be restored, followed in order by:
+	// FilePayload contains, in order:
 	// 1. checkpoint state file.
 	// 2. optional checkpoint pages metadata file.
 	// 3. optional checkpoint pages file.
@@ -498,6 +499,8 @@ type RestoreOpts struct {
 	HavePagesFile  bool
 	HaveDeviceFile bool
 	Background     bool
+
+	RestoreOptsExtra
 }
 
 // Restore loads a container from a statefile.
@@ -521,27 +524,32 @@ func (cm *containerManager) Restore(o *RestoreOpts, _ *struct{}) error {
 		return fmt.Errorf("at least one file must be passed to Restore")
 	}
 
-	stateFile, err := o.ReleaseFD(0)
+	stateFile, pagesMetadata, pagesFile, err := getRestoreReadersImpl(o)
 	if err != nil {
 		return err
 	}
+	defer func() {
+		if stateFile != nil {
+			stateFile.Close()
+		}
+		if pagesMetadata != nil {
+			pagesMetadata.Close()
+		}
+		if pagesFile != nil {
+			pagesFile.Close()
+		}
+	}()
 
-	var stat unix.Stat_t
-	if err := unix.Fstat(stateFile.FD(), &stat); err != nil {
-		return err
-	}
-	if stat.Size == 0 {
-		return fmt.Errorf("statefile cannot be empty")
-	}
-
-	reader, metadata, err := state.NewStatefileReader(stateFile, nil)
+	reader, metadata, err := state.NewStatefileReader(stateFile /* transfers ownership on success */, nil)
 	if err != nil {
 		return fmt.Errorf("creating statefile reader: %w", err)
 	}
+	stateFile = nil
 
 	// Create the main MemoryFile.
 	mf, err := createMemoryFile(cm.l.root.conf.AppHugePages, cm.l.hostTHP)
 	if err != nil {
+		reader.Close()
 		return fmt.Errorf("creating memory file: %v", err)
 	}
 
@@ -559,42 +567,18 @@ func (cm *containerManager) Restore(o *RestoreOpts, _ *struct{}) error {
 	// Release `cm.l.mu`.
 	cu.Clean()
 
-	fileIdx := 1
 	if o.HavePagesFile {
-		pagesMetadataFD, err := o.ReleaseFD(fileIdx)
-		if err != nil {
-			return err
-		}
-		fileIdx++
-
-		pagesFileFD, err := o.ReleaseFD(fileIdx)
-		if err != nil {
-			return err
-		}
-		fileIdx++
-
-		// //pkg/state/wire reads one byte at a time; buffer these reads to
-		// avoid making one syscall per read. For the state file, this
-		// buffering is handled by statefile.NewReader() => compressio.Reader
-		// or compressio.NewSimpleReader().
-		pagesMetadata := stateio.NewBufioReadCloser(pagesMetadataFD)
-		// TODO: Allow `runsc restore` to override I/O parameters.
-		pagesFile := stateio.NewPagesFileFDReaderDefault(int32(pagesFileFD.Release()))
-
 		// This immediately starts loading the main MemoryFile asynchronously.
-		cm.restorer.asyncMFLoader = kernel.NewAsyncMFLoader(pagesMetadata, pagesFile, cm.restorer.mainMF, timer.Fork("PagesFileLoader"))
+		cm.restorer.asyncMFLoader = kernel.NewAsyncMFLoader(pagesMetadata, pagesFile, cm.restorer.mainMF, timer.Fork("PagesFileLoader")) // transfers ownership
+		pagesMetadata = nil
+		pagesFile = nil
 	}
 
 	if o.HaveDeviceFile {
-		cm.restorer.deviceFile, err = o.ReleaseFD(fileIdx)
+		cm.restorer.deviceFile, err = o.ReleaseFD(len(o.Files) - 1)
 		if err != nil {
 			return err
 		}
-		fileIdx++
-	}
-
-	if fileIdx < len(o.Files) {
-		return fmt.Errorf("more files passed to Restore than expected")
 	}
 	timer.Reached("restorer ok")
 
@@ -634,6 +618,45 @@ func (cm *containerManager) Restore(o *RestoreOpts, _ *struct{}) error {
 	return cm.restorer.restoreContainerInfo(cm.l, &cm.l.root, timer.Fork("cont:root"))
 }
 
+func getRestoreReadersForLocalCheckpointFiles(o *RestoreOpts) (io.ReadCloser, io.ReadCloser, stateio.AsyncReader, error) {
+	stateFile, err := o.ReleaseFD(0)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	cu := cleanup.Make(func() { stateFile.Close() })
+	defer cu.Clean()
+	var stat unix.Stat_t
+	if err := unix.Fstat(stateFile.FD(), &stat); err != nil {
+		return nil, nil, nil, err
+	}
+	if stat.Size == 0 {
+		return nil, nil, nil, fmt.Errorf("statefile cannot be empty")
+	}
+
+	if !o.HavePagesFile {
+		cu.Release()
+		return stateFile, nil, nil, nil
+	}
+	pagesMetadataFile, err := o.ReleaseFD(1)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	cu.Add(func() { pagesMetadataFile.Close() })
+	pagesFile, err := o.ReleaseFD(2)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	cu.Release()
+	// //pkg/state/wire reads one byte at a time; buffer reads from
+	// pagesMetadataFile to avoid making one syscall per read. For the state
+	// file, this buffering is handled by statefile.NewReader() =>
+	// compressio.Reader or compressio.NewSimpleReader().
+	return stateFile,
+		stateio.NewBufioReadCloser(pagesMetadataFile),
+		stateio.NewPagesFileFDReaderDefault(int32(pagesFile.Release())),
+		nil
+}
+
 func (cm *containerManager) onRestoreDone() {
 	cm.l.mu.Lock()
 	cm.l.state = restored

runsc/boot/restore_impl.go

@@ -18,8 +18,11 @@
 package boot
 
 import (
+	"io"
+
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"gvisor.dev/gvisor/pkg/sentry/fsimpl/proc"
+	"gvisor.dev/gvisor/pkg/sentry/state/stateio"
 	"gvisor.dev/gvisor/runsc/config"
 )
 
@@ -30,3 +33,9 @@ func newProcInternalData(conf *config.Config, _ *specs.Spec) *proc.InternalData
 }
 
 func (l *Loader) kernelInitExtra() {}
+
+type RestoreOptsExtra struct{}
+
+func getRestoreReadersImpl(o *RestoreOpts) (io.ReadCloser, io.ReadCloser, stateio.AsyncReader, error) {
+	return getRestoreReadersForLocalCheckpointFiles(o)
+}

runsc/cmd/util/BUILD

@@ -12,8 +12,7 @@ go_library(
         "util.go",
     ],
     visibility = [
-        "//runsc/cli:__subpackages__",
-        "//runsc/cmd:__subpackages__",
+        "//runsc:__subpackages__",
         "//tools:__subpackages__",
     ],
     deps = [

runsc/sandbox/sandbox.go

@@ -502,44 +502,16 @@ func (s *Sandbox) Restore(conf *config.Config, spec *specs.Spec, cid string, ima
 
 	log.Debugf("Restore sandbox %q from path %q", s.ID, imagePath)
 
-	stateFileName := path.Join(imagePath, checkpointfiles.StateFileName)
-	sf, err := os.Open(stateFileName)
-	if err != nil {
-		return fmt.Errorf("opening state file %q failed: %v", stateFileName, err)
-	}
-	defer sf.Close()
-
 	opt := boot.RestoreOpts{
-		FilePayload: urpc.FilePayload{
-			Files: []*os.File{sf},
-		},
 		Background: background,
 	}
-
-	// If the pages file exists, we must pass it in.
-	pagesFileName := path.Join(imagePath, checkpointfiles.PagesFileName)
-	pagesReadFlags := os.O_RDONLY
-	if direct {
-		// The contents are page-aligned, so it can be opened with O_DIRECT.
-		pagesReadFlags |= syscall.O_DIRECT
-	}
-	if pf, err := os.OpenFile(pagesFileName, pagesReadFlags, 0); err == nil {
-		defer pf.Close()
-		pagesMetadataFileName := path.Join(imagePath, checkpointfiles.PagesMetadataFileName)
-		pmf, err := os.Open(pagesMetadataFileName)
-		if err != nil {
-			return fmt.Errorf("opening restore image file %q failed: %v", pagesMetadataFileName, err)
+	defer func() {
+		for _, f := range opt.FilePayload.Files {
+			_ = f.Close()
 		}
-		defer pmf.Close()
-
-		opt.HavePagesFile = true
-		opt.FilePayload.Files = append(opt.FilePayload.Files, pmf, pf)
-		log.Infof("Found page files for sandbox %q. Page metadata: %q, pages: %q", s.ID, pagesMetadataFileName, pagesFileName)
-
-	} else if !os.IsNotExist(err) {
-		return fmt.Errorf("opening restore pages file %q failed: %v", pagesFileName, err)
-	} else {
-		log.Infof("Using single checkpoint file for sandbox %q", s.ID)
+	}()
+	if err := s.setRestoreOptsImpl(conf, imagePath, direct, &opt); err != nil {
+		return err
 	}
 
 	// If the platform needs a device FD we must pass it in.
@@ -575,6 +547,39 @@ func (s *Sandbox) Restore(conf *config.Config, spec *specs.Spec, cid string, ima
 	return nil
 }
 
+func (s *Sandbox) setRestoreOptsForLocalCheckpointFiles(conf *config.Config, imagePath string, direct bool, opt *boot.RestoreOpts) error {
+	stateFileName := path.Join(imagePath, checkpointfiles.StateFileName)
+	sf, err := os.Open(stateFileName)
+	if err != nil {
+		return fmt.Errorf("opening state file %q failed: %w", stateFileName, err)
+	}
+	opt.FilePayload.Files = append(opt.FilePayload.Files, sf)
+
+	// If either the pages metadata file or pages file exist, both must exist,
+	// and we must pass them in.
+	pagesMetadataFileName := path.Join(imagePath, checkpointfiles.PagesMetadataFileName)
+	if pmf, err := os.Open(pagesMetadataFileName); err == nil {
+		opt.FilePayload.Files = append(opt.FilePayload.Files, pmf)
+		pagesFileName := path.Join(imagePath, checkpointfiles.PagesFileName)
+		pagesReadFlags := os.O_RDONLY
+		if direct {
+			// The contents are page-aligned, so it can be opened with O_DIRECT.
+			pagesReadFlags |= syscall.O_DIRECT
+		}
+		pf, err := os.OpenFile(pagesFileName, pagesReadFlags, 0)
+		if err != nil {
+			return fmt.Errorf("opening pages file %q failed: %w", pagesFileName, err)
+		}
+		opt.FilePayload.Files = append(opt.FilePayload.Files, pf)
+		opt.HavePagesFile = true
+	} else if !os.IsNotExist(err) {
+		return fmt.Errorf("opening pages metadata file %q failed: %w", pagesMetadataFileName, err)
+	} else {
+		log.Infof("Using single checkpoint file for sandbox %q", s.ID)
+	}
+	return nil
+}
+
 // RestoreSubcontainer sends the restore call for a sub-container in the sandbox.
 func (s *Sandbox) RestoreSubcontainer(spec *specs.Spec, conf *config.Config, cid string, stdios, goferFiles, goferFilestoreFiles []*os.File, devIOFile *os.File, goferMountConf []boot.GoferMountConf) error {
 	log.Debugf("Restore sub-container %q in sandbox %q, PID: %d", cid, s.ID, s.Pid.Load())

runsc/sandbox/sandbox_impl.go

@@ -18,10 +18,15 @@
 package sandbox
 
 import (
+	"gvisor.dev/gvisor/runsc/boot"
 	"gvisor.dev/gvisor/runsc/config"
 	"gvisor.dev/gvisor/runsc/donation"
 )
 
 func createSandboxProcessExtra(conf *config.Config, args *Args, donations *donation.Agency) error {
 	return nil
 }
+
+func (s *Sandbox) setRestoreOptsImpl(conf *config.Config, imagePath string, direct bool, opt *boot.RestoreOpts) error {
+	return s.setRestoreOptsForLocalCheckpointFiles(conf, imagePath, direct, opt)
+}

tools/embeddedbinary/embeddedbinary_template.go

@@ -38,6 +38,11 @@ var compressedBinary []byte
 
 // Options is the set of options to execute the embedded binary.
 type Options struct {
+	// If TempDir is non-empty, the embedded binary will be extracted to a
+	// subdirectory of TempDir. Otherwise, the embedded binary will be
+	// extracted to a subdirectory of os.TempDir().
+	TempDir string
+
 	// Argv is the set of arguments to exec with.
 	// `Argv[0]` is the name of the binary as invoked.
 	// If Argv is empty, it will default to a single-element slice, with
@@ -50,6 +55,10 @@ type Options struct {
 	// Files is the set of file descriptors to pass to forked processes.
 	// Only used when forking, not pure exec'ing.
 	Files []uintptr
+
+	// SysProcAttr provides OS-specific options to the executed process.
+	// Only used when forking, not pure exec'ing.
+	SysProcAttr *unix.SysProcAttr
 }
 
 // Bogus import to satisfy the compiler that we are using the flate import,
@@ -60,7 +69,7 @@ const _ = flate.NoCompression
 // If fork is true, the binary runs in a separate process, and its PID is
 // returned.
 // Otherwise, the binary is exec'd, so the current process stops executing.
-func run(options Options, fork bool) (int, error) {
+func run(options *Options, fork bool) (int, error) {
 	if len(options.Argv) == 0 {
 		options.Argv = []string{BinaryName}
 	}
@@ -71,7 +80,7 @@ func run(options Options, fork bool) (int, error) {
 	defer runtime.UnlockOSThread()
 	oldMask := unix.Umask(0077)
 	defer unix.Umask(oldMask)
-	tmpDir, err := os.MkdirTemp("", "gvisor.*.tmp")
+	tmpDir, err := os.MkdirTemp(options.TempDir, "gvisor.*.tmp")
 	if err != nil {
 		return 0, fmt.Errorf("cannot create temp directory: %w", err)
 	}
@@ -112,6 +121,7 @@ func run(options Options, fork bool) (int, error) {
 		return syscall.ForkExec(fdPath, options.Argv, &syscall.ProcAttr{
 			Env:   options.Envv,
 			Files: options.Files,
+			Sys:   options.SysProcAttr,
 		})
 	}
 	if err := unix.Exec(fdPath, options.Argv, options.Envv); err != nil {
@@ -123,12 +133,12 @@ func run(options Options, fork bool) (int, error) {
 // Exec execs the embedded binary. The current process is replaced.
 // This function only returns if unsuccessful.
 func Exec(options Options) error {
-	_, err := run(options, false)
+	_, err := run(&options, false)
 	return err
 }
 
 // ForkExec runs the embedded binary in a separate process.
 // Returns the PID of the child process.
 func ForkExec(options Options) (int, error) {
-	return run(options, true)
+	return run(&options, true)
 }