Enable checkpoint/restore in cases of UDS use.

Previously, processes which used file-system Unix Domain Sockets could not be
checkpoint-ed in runsc because the sockets were saved with their inode
numbers which do not necessarily remain the same upon restore. Now,
the sockets are also saved with their paths so that the new inodes
can be determined for the sockets based on these paths after restoring.
Tests for cases with UDS use are included. Test cleanup to come.

PiperOrigin-RevId: 208268781
Change-Id: Ieaa5d5d9a64914ca105cae199fd8492710b1d7ec

5 files changed, 175 insertions, 41 deletions

diff --git a/pkg/sentry/fs/gofer/gofer_test.go b/pkg/sentry/fs/gofer/gofer_test.go
index 764b530cb..45fdaacfd 100644
--- a/pkg/sentry/fs/gofer/gofer_test.go
+++ b/pkg/sentry/fs/gofer/gofer_test.go
@@ -74,7 +74,7 @@ func root(ctx context.Context, cp cachePolicy, mode p9.FileMode, size uint64) (*
 	}
 
 	rootFile := goodMockFile(mode, size)
-	sattr, rootInodeOperations := newInodeOperations(ctx, s, contextFile{file: rootFile}, p9.QID{}, rootFile.GetAttrMock.Valid, rootFile.GetAttrMock.Attr)
+	sattr, rootInodeOperations := newInodeOperations(ctx, s, contextFile{file: rootFile}, p9.QID{}, rootFile.GetAttrMock.Valid, rootFile.GetAttrMock.Attr, false /* socket */)
 	m := fs.NewMountSource(s, &filesystem{}, fs.MountSourceFlags{})
 	return rootFile, fs.NewInode(rootInodeOperations, m, sattr), nil
 }
diff --git a/pkg/sentry/fs/gofer/path.go b/pkg/sentry/fs/gofer/path.go
index bfeab3833..15e9863fb 100644
--- a/pkg/sentry/fs/gofer/path.go
+++ b/pkg/sentry/fs/gofer/path.go
@@ -57,7 +57,7 @@ func (i *inodeOperations) Lookup(ctx context.Context, dir *fs.Inode, name string
 	}
 
 	// Construct the Inode operations.
-	sattr, node := newInodeOperations(ctx, i.fileState.s, newFile, qids[0], mask, p9attr)
+	sattr, node := newInodeOperations(ctx, i.fileState.s, newFile, qids[0], mask, p9attr, false)
 
 	// Construct a positive Dirent.
 	return fs.NewDirent(fs.NewInode(node, dir.MountSource, sattr), name), nil
@@ -113,7 +113,7 @@ func (i *inodeOperations) Create(ctx context.Context, dir *fs.Inode, name string
 	}
 
 	// Construct the InodeOperations.
-	sattr, iops := newInodeOperations(ctx, i.fileState.s, unopened, qid, mask, p9attr)
+	sattr, iops := newInodeOperations(ctx, i.fileState.s, unopened, qid, mask, p9attr, false)
 
 	// Construct the positive Dirent.
 	d := fs.NewDirent(fs.NewInode(iops, dir.MountSource, sattr), name)
@@ -175,10 +175,10 @@ func (i *inodeOperations) CreateDirectory(ctx context.Context, dir *fs.Inode, s
 	return nil
 }
 
-// Bind implements InodeOperations.
-func (i *inodeOperations) Bind(ctx context.Context, dir *fs.Inode, name string, ep unix.BoundEndpoint, perm fs.FilePermissions) error {
+// Bind implements InodeOperations.Bind.
+func (i *inodeOperations) Bind(ctx context.Context, dir *fs.Inode, name string, ep unix.BoundEndpoint, perm fs.FilePermissions) (*fs.Dirent, error) {
 	if i.session().endpoints == nil {
-		return syscall.EOPNOTSUPP
+		return nil, syscall.EOPNOTSUPP
 	}
 
 	// Create replaces the directory fid with the newly created/opened
@@ -186,7 +186,7 @@ func (i *inodeOperations) Bind(ctx context.Context, dir *fs.Inode, name string,
 	// this node.
 	_, newFile, err := i.fileState.file.walk(ctx, nil)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	// Stabilize the endpoint map while creation is in progress.
@@ -198,7 +198,7 @@ func (i *inodeOperations) Bind(ctx context.Context, dir *fs.Inode, name string,
 	owner := fs.FileOwnerFromContext(ctx)
 	hostFile, err := newFile.create(ctx, name, p9.ReadWrite, p9.FileMode(perm.LinuxMode()), p9.UID(owner.UID), p9.GID(owner.GID))
 	if err != nil {
-		return err
+		return nil, err
 	}
 	// We're not going to use this file.
 	hostFile.Close()
@@ -206,10 +206,10 @@ func (i *inodeOperations) Bind(ctx context.Context, dir *fs.Inode, name string,
 	i.touchModificationTime(ctx, dir)
 
 	// Get the attributes of the file to create inode key.
-	qid, _, attr, err := getattr(ctx, newFile)
+	qid, mask, attr, err := getattr(ctx, newFile)
 	if err != nil {
 		newFile.close(ctx)
-		return err
+		return nil, err
 	}
 
 	key := device.MultiDeviceKey{
@@ -217,9 +217,24 @@ func (i *inodeOperations) Bind(ctx context.Context, dir *fs.Inode, name string,
 		SecondaryDevice: i.session().connID,
 		Inode:           qid.Path,
 	}
-	i.session().endpoints.add(key, ep)
 
-	return nil
+	// Create child dirent.
+
+	// Get an unopened p9.File for the file we created so that it can be
+	// cloned and re-opened multiple times after creation.
+	_, unopened, err := i.fileState.file.walk(ctx, []string{name})
+	if err != nil {
+		newFile.close(ctx)
+		return nil, err
+	}
+
+	// Construct the InodeOperations.
+	sattr, iops := newInodeOperations(ctx, i.fileState.s, unopened, qid, mask, attr, true)
+
+	// Construct the positive Dirent.
+	childDir := fs.NewDirent(fs.NewInode(iops, dir.MountSource, sattr), name)
+	i.session().endpoints.add(key, childDir, ep)
+	return childDir, nil
 }
 
 // CreateFifo implements fs.InodeOperations.CreateFifo. Gofer nodes do not support the
diff --git a/pkg/sentry/fs/gofer/session.go b/pkg/sentry/fs/gofer/session.go
index 648a11435..bfb1154dc 100644
--- a/pkg/sentry/fs/gofer/session.go
+++ b/pkg/sentry/fs/gofer/session.go
@@ -15,6 +15,7 @@
 package gofer
 
 import (
+	"fmt"
 	"sync"
 
 	"gvisor.googlesource.com/gvisor/pkg/p9"
@@ -28,39 +29,60 @@ import (
 )
 
 // +stateify savable
-type endpointMap struct {
+type endpointMaps struct {
+	// mu protexts the direntMap, the keyMap, and the pathMap below.
 	mu sync.RWMutex `state:"nosave"`
-	// TODO: Make map with private unix sockets savable.
-	m map[device.MultiDeviceKey]unix.BoundEndpoint
+
+	// direntMap links sockets to their dirents.
+	// It is filled concurrently with the keyMap and is stored upon save.
+	// Before saving, this map is used to populate the pathMap.
+	direntMap map[unix.BoundEndpoint]*fs.Dirent
+
+	// keyMap links MultiDeviceKeys (containing inode IDs) to their sockets.
+	// It is not stored during save because the inode ID may change upon restore.
+	keyMap map[device.MultiDeviceKey]unix.BoundEndpoint `state:"nosave"`
+
+	// pathMap links the sockets to their paths.
+	// It is filled before saving from the direntMap and is stored upon save.
+	// Upon restore, this map is used to re-populate the keyMap.
+	pathMap map[unix.BoundEndpoint]string
 }
 
-// add adds the endpoint to the map.
+// add adds the endpoint to the maps.
+// A reference is taken on the dirent argument.
 //
-// Precondition: map must have been locked with 'lock'.
-func (e *endpointMap) add(key device.MultiDeviceKey, ep unix.BoundEndpoint) {
-	e.m[key] = ep
+// Precondition: maps must have been locked with 'lock'.
+func (e *endpointMaps) add(key device.MultiDeviceKey, d *fs.Dirent, ep unix.BoundEndpoint) {
+	e.keyMap[key] = ep
+	d.IncRef()
+	e.direntMap[ep] = d
 }
 
-// remove deletes the key from the map.
+// remove deletes the key from the maps.
 //
-// Precondition: map must have been locked with 'lock'.
-func (e *endpointMap) remove(key device.MultiDeviceKey) {
-	delete(e.m, key)
+// Precondition: maps must have been locked with 'lock'.
+func (e *endpointMaps) remove(key device.MultiDeviceKey) {
+	endpoint := e.get(key)
+	delete(e.keyMap, key)
+
+	d := e.direntMap[endpoint]
+	d.DecRef()
+	delete(e.direntMap, endpoint)
 }
 
 // lock blocks other addition and removal operations from happening while
 // the backing file is being created or deleted. Returns a function that unlocks
 // the endpoint map.
-func (e *endpointMap) lock() func() {
+func (e *endpointMaps) lock() func() {
 	e.mu.Lock()
 	return func() { e.mu.Unlock() }
 }
 
-func (e *endpointMap) get(key device.MultiDeviceKey) unix.BoundEndpoint {
-	e.mu.RLock()
-	ep := e.m[key]
-	e.mu.RUnlock()
-	return ep
+// get returns the endpoint mapped to the given key.
+//
+// Precondition: maps must have been locked for reading.
+func (e *endpointMaps) get(key device.MultiDeviceKey) unix.BoundEndpoint {
+	return e.keyMap[key]
 }
 
 // session holds state for each 9p session established during sys_mount.
@@ -115,7 +137,7 @@ type session struct {
 	// TODO: there are few possible races with someone stat'ing the
 	// file and another deleting it concurrently, where the file will not be
 	// reported as socket file.
-	endpoints *endpointMap `state:"wait"`
+	endpoints *endpointMaps `state:"wait"`
 }
 
 // Destroy tears down the session.
@@ -149,7 +171,9 @@ func (s *session) SaveInodeMapping(inode *fs.Inode, path string) {
 
 // newInodeOperations creates a new 9p fs.InodeOperations backed by a p9.File and attributes
 // (p9.QID, p9.AttrMask, p9.Attr).
-func newInodeOperations(ctx context.Context, s *session, file contextFile, qid p9.QID, valid p9.AttrMask, attr p9.Attr) (fs.StableAttr, *inodeOperations) {
+//
+// Endpoints lock must not be held if socket == false.
+func newInodeOperations(ctx context.Context, s *session, file contextFile, qid p9.QID, valid p9.AttrMask, attr p9.Attr, socket bool) (fs.StableAttr, *inodeOperations) {
 	deviceKey := device.MultiDeviceKey{
 		Device:          attr.RDev,
 		SecondaryDevice: s.connID,
@@ -164,10 +188,16 @@ func newInodeOperations(ctx context.Context, s *session, file contextFile, qid p
 	}
 
 	if s.endpoints != nil {
-		// If unix sockets are allowed on this filesystem, check if this file is
-		// supposed to be a socket file.
-		if s.endpoints.get(deviceKey) != nil {
+		if socket {
 			sattr.Type = fs.Socket
+		} else {
+			// If unix sockets are allowed on this filesystem, check if this file is
+			// supposed to be a socket file.
+			unlock := s.endpoints.lock()
+			if s.endpoints.get(deviceKey) != nil {
+				sattr.Type = fs.Socket
+			}
+			unlock()
 		}
 	}
 
@@ -215,7 +245,7 @@ func Root(ctx context.Context, dev string, filesystem fs.Filesystem, superBlockF
 	}
 
 	if o.privateunixsocket {
-		s.endpoints = &endpointMap{m: make(map[device.MultiDeviceKey]unix.BoundEndpoint)}
+		s.endpoints = newEndpointMaps()
 	}
 
 	// Construct the MountSource with the session and superBlockFlags.
@@ -248,6 +278,77 @@ func Root(ctx context.Context, dev string, filesystem fs.Filesystem, superBlockF
 		return nil, err
 	}
 
-	sattr, iops := newInodeOperations(ctx, s, s.attach, qid, valid, attr)
+	sattr, iops := newInodeOperations(ctx, s, s.attach, qid, valid, attr, false)
 	return fs.NewInode(iops, m, sattr), nil
 }
+
+// newEndpointMaps creates a new endpointMaps.
+func newEndpointMaps() *endpointMaps {
+	return &endpointMaps{
+		direntMap: make(map[unix.BoundEndpoint]*fs.Dirent),
+		keyMap:    make(map[device.MultiDeviceKey]unix.BoundEndpoint),
+		pathMap:   make(map[unix.BoundEndpoint]string),
+	}
+}
+
+// fillKeyMap populates key and dirent maps upon restore from saved
+// pathmap.
+func (s *session) fillKeyMap(ctx context.Context) error {
+	unlock := s.endpoints.lock()
+	defer unlock()
+
+	for ep, dirPath := range s.endpoints.pathMap {
+		_, file, err := s.attach.walk(ctx, splitAbsolutePath(dirPath))
+		if err != nil {
+			return fmt.Errorf("error filling endpointmaps, failed to walk to %q: %v", dirPath, err)
+		}
+
+		qid, _, attr, err := file.getAttr(ctx, p9.AttrMaskAll())
+		if err != nil {
+			return fmt.Errorf("failed to get file attributes of %s: %v", dirPath, err)
+		}
+
+		key := device.MultiDeviceKey{
+			Device:          attr.RDev,
+			SecondaryDevice: s.connID,
+			Inode:           qid.Path,
+		}
+
+		s.endpoints.keyMap[key] = ep
+	}
+	return nil
+}
+
+// fillPathMap populates paths for endpoints from dirents in direntMap
+// before save.
+func (s *session) fillPathMap() error {
+	unlock := s.endpoints.lock()
+	defer unlock()
+
+	for ep, dir := range s.endpoints.direntMap {
+		mountRoot := dir.MountRoot()
+		defer mountRoot.DecRef()
+		dirPath, _ := dir.FullName(mountRoot)
+		if dirPath == "" {
+			return fmt.Errorf("error getting path from dirent")
+		}
+		s.endpoints.pathMap[ep] = dirPath
+	}
+	return nil
+}
+
+// restoreEndpointMaps recreates and fills the key and dirent maps.
+func (s *session) restoreEndpointMaps(ctx context.Context) error {
+	// When restoring, only need to create the keyMap because the dirent and path
+	// maps got stored through the save.
+	s.endpoints.keyMap = make(map[device.MultiDeviceKey]unix.BoundEndpoint)
+	if err := s.fillKeyMap(ctx); err != nil {
+		return fmt.Errorf("failed to insert sockets into endpoint map: %v", err)
+	}
+
+	// Re-create pathMap because it can no longer be trusted as socket paths can
+	// change while process continues to run. Empty pathMap will be re-filled upon
+	// next save.
+	s.endpoints.pathMap = make(map[unix.BoundEndpoint]string)
+	return nil
+}
diff --git a/pkg/sentry/fs/gofer/session_state.go b/pkg/sentry/fs/gofer/session_state.go
index 0154810c8..8e6424492 100644
--- a/pkg/sentry/fs/gofer/session_state.go
+++ b/pkg/sentry/fs/gofer/session_state.go
@@ -18,16 +18,17 @@ import (
 	"fmt"
 
 	"gvisor.googlesource.com/gvisor/pkg/p9"
+	"gvisor.googlesource.com/gvisor/pkg/sentry/context"
 	"gvisor.googlesource.com/gvisor/pkg/sentry/fs"
 	"gvisor.googlesource.com/gvisor/pkg/unet"
 )
 
 // beforeSave is invoked by stateify.
-//
-// TODO: Make map with private unix sockets savable.
-func (e *endpointMap) beforeSave() {
-	if len(e.m) != 0 {
-		panic("EndpointMap with existing private unix sockets cannot be saved")
+func (s *session) beforeSave() {
+	if s.endpoints != nil {
+		if err := s.fillPathMap(); err != nil {
+			panic("failed to save paths to endpoint map before saving" + err.Error())
+		}
 	}
 }
 
@@ -72,6 +73,9 @@ func (s *session) afterLoad() {
 	if opts.aname != s.aname {
 		panic(fmt.Sprintf("new attach name %v, want %v", opts.aname, s.aname))
 	}
+
+	// Check if endpointMaps exist when uds sockets are enabled
+	// (only pathmap will actualy have been saved).
 	if opts.privateunixsocket != (s.endpoints != nil) {
 		panic(fmt.Sprintf("new privateunixsocket option %v, want %v", opts.privateunixsocket, s.endpoints != nil))
 	}
@@ -96,4 +100,16 @@ func (s *session) afterLoad() {
 	if err != nil {
 		panic(fmt.Sprintf("failed to attach to aname: %v", err))
 	}
+
+	// If private unix sockets are enabled, create and fill the session's endpoint
+	// maps.
+	if opts.privateunixsocket {
+		// TODO: Context is not plumbed to save/restore.
+		ctx := &dummyClockContext{context.Background()}
+
+		if err = s.restoreEndpointMaps(ctx); err != nil {
+			panic("failed to restore endpoint maps: " + err.Error())
+		}
+	}
+
 }
diff --git a/pkg/sentry/fs/gofer/socket.go b/pkg/sentry/fs/gofer/socket.go
index 406756f5f..8628b9c69 100644
--- a/pkg/sentry/fs/gofer/socket.go
+++ b/pkg/sentry/fs/gofer/socket.go
@@ -30,6 +30,8 @@ func (i *inodeOperations) BoundEndpoint(inode *fs.Inode, path string) unix.Bound
 	}
 
 	if i.session().endpoints != nil {
+		unlock := i.session().endpoints.lock()
+		defer unlock()
 		ep := i.session().endpoints.get(i.fileState.key)
 		if ep != nil {
 			return ep