Add finalizer on AtomicRefCount to check for leaks.

PiperOrigin-RevId: 255711454

23 files changed, 145 insertions, 35 deletions

diff --git a/pkg/refs/BUILD b/pkg/refs/BUILD
index 0652fbbe6..9c08452fc 100644
--- a/pkg/refs/BUILD
+++ b/pkg/refs/BUILD
@@ -24,6 +24,9 @@ go_library(
     ],
     importpath = "gvisor.dev/gvisor/pkg/refs",
     visibility = ["//:sandbox"],
+    deps = [
+        "//pkg/log",
+    ],
 )
 
 go_test(
diff --git a/pkg/refs/refcounter.go b/pkg/refs/refcounter.go
index 20f515391..6d6c3a782 100644
--- a/pkg/refs/refcounter.go
+++ b/pkg/refs/refcounter.go
@@ -18,8 +18,11 @@ package refs
 
 import (
 	"reflect"
+	"runtime"
 	"sync"
 	"sync/atomic"
+
+	"gvisor.dev/gvisor/pkg/log"
 )
 
 // RefCounter is the interface to be implemented by objects that are reference
@@ -189,6 +192,11 @@ type AtomicRefCount struct {
 	// how these fields are used.
 	refCount int64
 
+	// name is the name of the type which owns this ref count.
+	//
+	// name is immutable after EnableLeakCheck is called.
+	name string
+
 	// mu protects the list below.
 	mu sync.Mutex `state:"nosave"`
 
@@ -196,6 +204,67 @@ type AtomicRefCount struct {
 	weakRefs weakRefList `state:"nosave"`
 }
 
+// LeakMode configures the leak checker.
+type LeakMode uint32
+
+const (
+	// uninitializedLeakChecking indicates that the leak checker has not yet been initialized.
+	uninitializedLeakChecking LeakMode = iota
+
+	// NoLeakChecking indicates that no effort should be made to check for
+	// leaks.
+	NoLeakChecking
+
+	// LeaksLogWarning indicates that a warning should be logged when leaks
+	// are found.
+	LeaksLogWarning
+)
+
+// leakMode stores the current mode for the reference leak checker.
+//
+// Values must be one of the LeakMode values.
+//
+// leakMode must be accessed atomically.
+var leakMode uint32
+
+// SetLeakMode configures the reference leak checker.
+func SetLeakMode(mode LeakMode) {
+	atomic.StoreUint32(&leakMode, uint32(mode))
+}
+
+func (r *AtomicRefCount) finalize() {
+	var note string
+	switch LeakMode(atomic.LoadUint32(&leakMode)) {
+	case NoLeakChecking:
+		return
+	case uninitializedLeakChecking:
+		note = "(Leak checker uninitialized): "
+	}
+	if n := r.ReadRefs(); n != 0 {
+		log.Warningf("%sAtomicRefCount %p owned by %q garbage collected with ref count of %d (want 0)", note, r, r.name, n)
+	}
+}
+
+// EnableLeakCheck checks for reference leaks when the AtomicRefCount gets
+// garbage collected.
+//
+// This function adds a finalizer to the AtomicRefCount, so the AtomicRefCount
+// must be at the beginning of its parent.
+//
+// name is a friendly name that will be listed as the owner of the
+// AtomicRefCount in logs. It should be the name of the parent type, including
+// package.
+func (r *AtomicRefCount) EnableLeakCheck(name string) {
+	if name == "" {
+		panic("invalid name")
+	}
+	if LeakMode(atomic.LoadUint32(&leakMode)) == NoLeakChecking {
+		return
+	}
+	r.name = name
+	runtime.SetFinalizer(r, (*AtomicRefCount).finalize)
+}
+
 // ReadRefs returns the current number of references. The returned count is
 // inherently racy and is unsafe to use without external synchronization.
 func (r *AtomicRefCount) ReadRefs() int64 {
diff --git a/pkg/sentry/fs/dirent.go b/pkg/sentry/fs/dirent.go
index 28651e58b..fbca06761 100644
--- a/pkg/sentry/fs/dirent.go
+++ b/pkg/sentry/fs/dirent.go
@@ -229,11 +229,13 @@ func newDirent(inode *Inode, name string) *Dirent {
 	if inode != nil {
 		inode.MountSource.IncDirentRefs()
 	}
-	return &Dirent{
+	d := Dirent{
 		Inode:    inode,
 		name:     name,
 		children: make(map[string]*refs.WeakRef),
 	}
+	d.EnableLeakCheck("fs.Dirent")
+	return &d
 }
 
 // NewNegativeDirent returns a new root negative Dirent. Otherwise same as NewDirent.
diff --git a/pkg/sentry/fs/file.go b/pkg/sentry/fs/file.go
index 8e1f5674d..bb8117f89 100644
--- a/pkg/sentry/fs/file.go
+++ b/pkg/sentry/fs/file.go
@@ -130,14 +130,15 @@ type File struct {
 // to false respectively.
 func NewFile(ctx context.Context, dirent *Dirent, flags FileFlags, fops FileOperations) *File {
 	dirent.IncRef()
-	f := &File{
+	f := File{
 		UniqueID:       uniqueid.GlobalFromContext(ctx),
 		Dirent:         dirent,
 		FileOperations: fops,
 		flags:          flags,
 	}
 	f.mu.Init()
-	return f
+	f.EnableLeakCheck("fs.File")
+	return &f
 }
 
 // DecRef destroys the File when it is no longer referenced.
diff --git a/pkg/sentry/fs/file_overlay.go b/pkg/sentry/fs/file_overlay.go
index c6cbd5631..9820f0b13 100644
--- a/pkg/sentry/fs/file_overlay.go
+++ b/pkg/sentry/fs/file_overlay.go
@@ -347,13 +347,14 @@ func (*overlayFileOperations) ConfigureMMap(ctx context.Context, file *File, opt
 		// preventing us from saving a proper inode mapping for the
 		// file.
 		file.IncRef()
-		id := &overlayMappingIdentity{
+		id := overlayMappingIdentity{
 			id:          opts.MappingIdentity,
 			overlayFile: file,
 		}
+		id.EnableLeakCheck("fs.overlayMappingIdentity")
 
 		// Swap out the old MappingIdentity for the wrapped one.
-		opts.MappingIdentity = id
+		opts.MappingIdentity = &id
 		return nil
 	}
 
diff --git a/pkg/sentry/fs/gofer/handles.go b/pkg/sentry/fs/gofer/handles.go
index b87c4f150..27eeae3d9 100644
--- a/pkg/sentry/fs/gofer/handles.go
+++ b/pkg/sentry/fs/gofer/handles.go
@@ -79,11 +79,12 @@ func newHandles(ctx context.Context, file contextFile, flags fs.FileFlags) (*han
 		newFile.close(ctx)
 		return nil, err
 	}
-	h := &handles{
+	h := handles{
 		File: newFile,
 		Host: hostFile,
 	}
-	return h, nil
+	h.EnableLeakCheck("gofer.handles")
+	return &h, nil
 }
 
 type handleReadWriter struct {
diff --git a/pkg/sentry/fs/gofer/path.go b/pkg/sentry/fs/gofer/path.go
index b91386909..8c17603f8 100644
--- a/pkg/sentry/fs/gofer/path.go
+++ b/pkg/sentry/fs/gofer/path.go
@@ -145,16 +145,17 @@ func (i *inodeOperations) Create(ctx context.Context, dir *fs.Inode, name string
 	defer d.DecRef()
 
 	// Construct the new file, caching the handles if allowed.
-	h := &handles{
+	h := handles{
 		File: newFile,
 		Host: hostFile,
 	}
+	h.EnableLeakCheck("gofer.handles")
 	if iops.fileState.canShareHandles() {
 		iops.fileState.handlesMu.Lock()
-		iops.fileState.setSharedHandlesLocked(flags, h)
+		iops.fileState.setSharedHandlesLocked(flags, &h)
 		iops.fileState.handlesMu.Unlock()
 	}
-	return NewFile(ctx, d, name, flags, iops, h), nil
+	return NewFile(ctx, d, name, flags, iops, &h), nil
 }
 
 // CreateLink uses Create to create a symlink between oldname and newname.
diff --git a/pkg/sentry/fs/gofer/session.go b/pkg/sentry/fs/gofer/session.go
index 9f7660ed1..69d08a627 100644
--- a/pkg/sentry/fs/gofer/session.go
+++ b/pkg/sentry/fs/gofer/session.go
@@ -241,7 +241,7 @@ func Root(ctx context.Context, dev string, filesystem fs.Filesystem, superBlockF
 	}
 
 	// Construct the session.
-	s := &session{
+	s := session{
 		connID:          dev,
 		msize:           o.msize,
 		version:         o.version,
@@ -250,13 +250,14 @@ func Root(ctx context.Context, dev string, filesystem fs.Filesystem, superBlockF
 		superBlockFlags: superBlockFlags,
 		mounter:         mounter,
 	}
+	s.EnableLeakCheck("gofer.session")
 
 	if o.privateunixsocket {
 		s.endpoints = newEndpointMaps()
 	}
 
 	// Construct the MountSource with the session and superBlockFlags.
-	m := fs.NewMountSource(ctx, s, filesystem, superBlockFlags)
+	m := fs.NewMountSource(ctx, &s, filesystem, superBlockFlags)
 
 	// Given that gofer files can consume host FDs, restrict the number
 	// of files that can be held by the cache.
@@ -290,7 +291,7 @@ func Root(ctx context.Context, dev string, filesystem fs.Filesystem, superBlockF
 		return nil, err
 	}
 
-	sattr, iops := newInodeOperations(ctx, s, s.attach, qid, valid, attr, false)
+	sattr, iops := newInodeOperations(ctx, &s, s.attach, qid, valid, attr, false)
 	return fs.NewInode(ctx, iops, m, sattr), nil
 }
 
diff --git a/pkg/sentry/fs/gofer/session_state.go b/pkg/sentry/fs/gofer/session_state.go
index 29a79441e..d045e04ff 100644
--- a/pkg/sentry/fs/gofer/session_state.go
+++ b/pkg/sentry/fs/gofer/session_state.go
@@ -111,5 +111,4 @@ func (s *session) afterLoad() {
 			panic("failed to restore endpoint maps: " + err.Error())
 		}
 	}
-
 }
diff --git a/pkg/sentry/fs/host/socket.go b/pkg/sentry/fs/host/socket.go
index 7fedc88bc..44c4ee5f2 100644
--- a/pkg/sentry/fs/host/socket.go
+++ b/pkg/sentry/fs/host/socket.go
@@ -47,12 +47,12 @@ const maxSendBufferSize = 8 << 20
 //
 // +stateify savable
 type ConnectedEndpoint struct {
-	queue *waiter.Queue
-	path  string
-
 	// ref keeps track of references to a connectedEndpoint.
 	ref refs.AtomicRefCount
 
+	queue *waiter.Queue
+	path  string
+
 	// If srfd >= 0, it is the host FD that file was imported from.
 	srfd int `state:"wait"`
 
@@ -133,6 +133,8 @@ func NewConnectedEndpoint(ctx context.Context, file *fd.FD, queue *waiter.Queue,
 	// AtomicRefCounters start off with a single reference. We need two.
 	e.ref.IncRef()
 
+	e.ref.EnableLeakCheck("host.ConnectedEndpoint")
+
 	return &e, nil
 }
 
diff --git a/pkg/sentry/fs/inode.go b/pkg/sentry/fs/inode.go
index e4aae1135..f4ddfa406 100644
--- a/pkg/sentry/fs/inode.go
+++ b/pkg/sentry/fs/inode.go
@@ -86,12 +86,14 @@ type LockCtx struct {
 // NewInode takes a reference on msrc.
 func NewInode(ctx context.Context, iops InodeOperations, msrc *MountSource, sattr StableAttr) *Inode {
 	msrc.IncRef()
-	return &Inode{
+	i := Inode{
 		InodeOperations: iops,
 		StableAttr:      sattr,
 		Watches:         newWatches(),
 		MountSource:     msrc,
 	}
+	i.EnableLeakCheck("fs.Inode")
+	return &i
 }
 
 // DecRef drops a reference on the Inode.
diff --git a/pkg/sentry/fs/mount.go b/pkg/sentry/fs/mount.go
index 912495528..7a9692800 100644
--- a/pkg/sentry/fs/mount.go
+++ b/pkg/sentry/fs/mount.go
@@ -138,12 +138,14 @@ func NewMountSource(ctx context.Context, mops MountSourceOperations, filesystem
 	if filesystem != nil {
 		fsType = filesystem.Name()
 	}
-	return &MountSource{
+	msrc := MountSource{
 		MountSourceOperations: mops,
 		Flags:                 flags,
 		FilesystemType:        fsType,
 		fscache:               NewDirentCache(DefaultDirentCacheSize),
 	}
+	msrc.EnableLeakCheck("fs.MountSource")
+	return &msrc
 }
 
 // DirentRefs returns the current mount direntRefs.
diff --git a/pkg/sentry/fs/mounts.go b/pkg/sentry/fs/mounts.go
index 281364dfc..ce7ffeed2 100644
--- a/pkg/sentry/fs/mounts.go
+++ b/pkg/sentry/fs/mounts.go
@@ -181,12 +181,14 @@ func NewMountNamespace(ctx context.Context, root *Inode) (*MountNamespace, error
 		d: newRootMount(1, d),
 	}
 
-	return &MountNamespace{
+	mns := MountNamespace{
 		userns:  creds.UserNamespace,
 		root:    d,
 		mounts:  mnts,
 		mountID: 2,
-	}, nil
+	}
+	mns.EnableLeakCheck("fs.MountNamespace")
+	return &mns, nil
 }
 
 // UserNamespace returns the user namespace associated with this mount manager.
diff --git a/pkg/sentry/fs/tty/terminal.go b/pkg/sentry/fs/tty/terminal.go
index 8290f2530..b7cecb2ed 100644
--- a/pkg/sentry/fs/tty/terminal.go
+++ b/pkg/sentry/fs/tty/terminal.go
@@ -38,9 +38,11 @@ type Terminal struct {
 
 func newTerminal(ctx context.Context, d *dirInodeOperations, n uint32) *Terminal {
 	termios := linux.DefaultSlaveTermios
-	return &Terminal{
+	t := Terminal{
 		d:  d,
 		n:  n,
 		ld: newLineDiscipline(termios),
 	}
+	t.EnableLeakCheck("tty.Terminal")
+	return &t
 }
diff --git a/pkg/sentry/kernel/fd_map.go b/pkg/sentry/kernel/fd_map.go
index 786936a7d..1b84bfe14 100644
--- a/pkg/sentry/kernel/fd_map.go
+++ b/pkg/sentry/kernel/fd_map.go
@@ -98,11 +98,13 @@ func (f *FDMap) ID() uint64 {
 
 // NewFDMap allocates a new FDMap that may be used by tasks in k.
 func (k *Kernel) NewFDMap() *FDMap {
-	return &FDMap{
+	f := FDMap{
 		k:     k,
 		files: make(map[kdefs.FD]descriptor),
 		uid:   atomic.AddUint64(&k.fdMapUids, 1),
 	}
+	f.EnableLeakCheck("kernel.FDMap")
+	return &f
 }
 
 // destroy removes all of the file descriptors from the map.
diff --git a/pkg/sentry/kernel/fs_context.go b/pkg/sentry/kernel/fs_context.go
index 938239aeb..ded27d668 100644
--- a/pkg/sentry/kernel/fs_context.go
+++ b/pkg/sentry/kernel/fs_context.go
@@ -51,11 +51,13 @@ type FSContext struct {
 func newFSContext(root, cwd *fs.Dirent, umask uint) *FSContext {
 	root.IncRef()
 	cwd.IncRef()
-	return &FSContext{
+	f := FSContext{
 		root:  root,
 		cwd:   cwd,
 		umask: umask,
 	}
+	f.EnableLeakCheck("kernel.FSContext")
+	return &f
 }
 
 // destroy is the destructor for an FSContext.
diff --git a/pkg/sentry/kernel/sessions.go b/pkg/sentry/kernel/sessions.go
index 355984140..81fcd8258 100644
--- a/pkg/sentry/kernel/sessions.go
+++ b/pkg/sentry/kernel/sessions.go
@@ -294,6 +294,7 @@ func (tg *ThreadGroup) createSession() error {
 		id:     SessionID(id),
 		leader: tg,
 	}
+	s.refs.EnableLeakCheck("kernel.Session")
 
 	// Create a new ProcessGroup, belonging to that Session.
 	// This also has a single reference (assigned below).
@@ -307,6 +308,7 @@ func (tg *ThreadGroup) createSession() error {
 		session:    s,
 		ancestors:  0,
 	}
+	pg.refs.EnableLeakCheck("kernel.ProcessGroup")
 
 	// Tie them and return the result.
 	s.processGroups.PushBack(pg)
@@ -378,11 +380,13 @@ func (tg *ThreadGroup) CreateProcessGroup() error {
 	// We manually adjust the ancestors if the parent is in the same
 	// session.
 	tg.processGroup.session.incRef()
-	pg := &ProcessGroup{
+	pg := ProcessGroup{
 		id:         ProcessGroupID(id),
 		originator: tg,
 		session:    tg.processGroup.session,
 	}
+	pg.refs.EnableLeakCheck("kernel.ProcessGroup")
+
 	if tg.leader.parent != nil && tg.leader.parent.tg.processGroup.session == pg.session {
 		pg.ancestors++
 	}
@@ -390,20 +394,20 @@ func (tg *ThreadGroup) CreateProcessGroup() error {
 	// Assign the new process group; adjust children.
 	oldParentPG := tg.parentPG()
 	tg.forEachChildThreadGroupLocked(func(childTG *ThreadGroup) {
-		childTG.processGroup.incRefWithParent(pg)
+		childTG.processGroup.incRefWithParent(&pg)
 		childTG.processGroup.decRefWithParent(oldParentPG)
 	})
 	tg.processGroup.decRefWithParent(oldParentPG)
-	tg.processGroup = pg
+	tg.processGroup = &pg
 
 	// Add the new process group to the session.
-	pg.session.processGroups.PushBack(pg)
+	pg.session.processGroups.PushBack(&pg)
 
 	// Ensure this translation is added to all namespaces.
 	for ns := tg.pidns; ns != nil; ns = ns.parent {
 		local := ns.tgids[tg]
-		ns.pgids[pg] = ProcessGroupID(local)
-		ns.processGroups[ProcessGroupID(local)] = pg
+		ns.pgids[&pg] = ProcessGroupID(local)
+		ns.processGroups[ProcessGroupID(local)] = &pg
 	}
 
 	return nil
diff --git a/pkg/sentry/kernel/shm/shm.go b/pkg/sentry/kernel/shm/shm.go
index 3e9fe70e2..5bd610f68 100644
--- a/pkg/sentry/kernel/shm/shm.go
+++ b/pkg/sentry/kernel/shm/shm.go
@@ -224,6 +224,7 @@ func (r *Registry) newShm(ctx context.Context, pid int32, key Key, creator fs.Fi
 		creatorPID:    pid,
 		changeTime:    ktime.NowFromContext(ctx),
 	}
+	shm.EnableLeakCheck("kernel.Shm")
 
 	// Find the next available ID.
 	for id := r.lastIDUsed + 1; id != r.lastIDUsed; id++ {
diff --git a/pkg/sentry/mm/aio_context.go b/pkg/sentry/mm/aio_context.go
index c9a942023..1b746d030 100644
--- a/pkg/sentry/mm/aio_context.go
+++ b/pkg/sentry/mm/aio_context.go
@@ -213,7 +213,9 @@ func newAIOMappable(mfp pgalloc.MemoryFileProvider) (*aioMappable, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &aioMappable{mfp: mfp, fr: fr}, nil
+	m := aioMappable{mfp: mfp, fr: fr}
+	m.EnableLeakCheck("mm.aioMappable")
+	return &m, nil
 }
 
 // DecRef implements refs.RefCounter.DecRef.
diff --git a/pkg/sentry/mm/special_mappable.go b/pkg/sentry/mm/special_mappable.go
index 2f8651a0a..ea2d7af74 100644
--- a/pkg/sentry/mm/special_mappable.go
+++ b/pkg/sentry/mm/special_mappable.go
@@ -45,7 +45,9 @@ type SpecialMappable struct {
 //
 // Preconditions: fr.Length() != 0.
 func NewSpecialMappable(name string, mfp pgalloc.MemoryFileProvider, fr platform.FileRange) *SpecialMappable {
-	return &SpecialMappable{mfp: mfp, fr: fr, name: name}
+	m := SpecialMappable{mfp: mfp, fr: fr, name: name}
+	m.EnableLeakCheck("mm.SpecialMappable")
+	return &m
 }
 
 // DecRef implements refs.RefCounter.DecRef.
diff --git a/pkg/sentry/socket/unix/transport/connectioned.go b/pkg/sentry/socket/unix/transport/connectioned.go
index e4c416233..73d2df15d 100644
--- a/pkg/sentry/socket/unix/transport/connectioned.go
+++ b/pkg/sentry/socket/unix/transport/connectioned.go
@@ -143,7 +143,9 @@ func NewPair(ctx context.Context, stype linux.SockType, uid UniqueIDProvider) (E
 	}
 
 	q1 := &queue{ReaderQueue: a.Queue, WriterQueue: b.Queue, limit: initialLimit}
+	q1.EnableLeakCheck("transport.queue")
 	q2 := &queue{ReaderQueue: b.Queue, WriterQueue: a.Queue, limit: initialLimit}
+	q2.EnableLeakCheck("transport.queue")
 
 	if stype == linux.SOCK_STREAM {
 		a.receiver = &streamQueueReceiver{queueReceiver: queueReceiver{q1}}
@@ -294,12 +296,14 @@ func (e *connectionedEndpoint) BidirectionalConnect(ctx context.Context, ce Conn
 	}
 
 	readQueue := &queue{ReaderQueue: ce.WaiterQueue(), WriterQueue: ne.Queue, limit: initialLimit}
+	readQueue.EnableLeakCheck("transport.queue")
 	ne.connected = &connectedEndpoint{
 		endpoint:   ce,
 		writeQueue: readQueue,
 	}
 
 	writeQueue := &queue{ReaderQueue: ne.Queue, WriterQueue: ce.WaiterQueue(), limit: initialLimit}
+	writeQueue.EnableLeakCheck("transport.queue")
 	if e.stype == linux.SOCK_STREAM {
 		ne.receiver = &streamQueueReceiver{queueReceiver: queueReceiver{readQueue: writeQueue}}
 	} else {
diff --git a/pkg/sentry/socket/unix/transport/connectionless.go b/pkg/sentry/socket/unix/transport/connectionless.go
index e987519f0..c591f261d 100644
--- a/pkg/sentry/socket/unix/transport/connectionless.go
+++ b/pkg/sentry/socket/unix/transport/connectionless.go
@@ -41,7 +41,9 @@ var (
 // NewConnectionless creates a new unbound dgram endpoint.
 func NewConnectionless(ctx context.Context) Endpoint {
 	ep := &connectionlessEndpoint{baseEndpoint{Queue: &waiter.Queue{}}}
-	ep.receiver = &queueReceiver{readQueue: &queue{ReaderQueue: ep.Queue, WriterQueue: &waiter.Queue{}, limit: initialLimit}}
+	q := queue{ReaderQueue: ep.Queue, WriterQueue: &waiter.Queue{}, limit: initialLimit}
+	q.EnableLeakCheck("transport.queue")
+	ep.receiver = &queueReceiver{readQueue: &q}
 	return ep
 }
 
diff --git a/pkg/sentry/socket/unix/unix.go b/pkg/sentry/socket/unix/unix.go
index 6190de0c5..bf7d2cfa2 100644
--- a/pkg/sentry/socket/unix/unix.go
+++ b/pkg/sentry/socket/unix/unix.go
@@ -69,10 +69,13 @@ func New(ctx context.Context, endpoint transport.Endpoint, stype linux.SockType)
 
 // NewWithDirent creates a new unix socket using an existing dirent.
 func NewWithDirent(ctx context.Context, d *fs.Dirent, ep transport.Endpoint, stype linux.SockType, flags fs.FileFlags) *fs.File {
-	return fs.NewFile(ctx, d, flags, &SocketOperations{
+	s := SocketOperations{
 		ep:    ep,
 		stype: stype,
-	})
+	}
+	s.EnableLeakCheck("unix.SocketOperations")
+
+	return fs.NewFile(ctx, d, flags, &s)
 }
 
 // DecRef implements RefCounter.DecRef.