Add finalizer on AtomicRefCount to check for leaks.

PiperOrigin-RevId: 255711454

4 files changed, 17 insertions, 8 deletions

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++ {