1 files changed, 57 insertions, 1 deletions

diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go
index 43e9823cb..e7e5ff777 100644
--- a/pkg/sentry/kernel/kernel.go
+++ b/pkg/sentry/kernel/kernel.go
@@ -43,6 +43,7 @@ import (
 	"gvisor.googlesource.com/gvisor/pkg/cpuid"
 	"gvisor.googlesource.com/gvisor/pkg/eventchannel"
 	"gvisor.googlesource.com/gvisor/pkg/log"
+	"gvisor.googlesource.com/gvisor/pkg/refs"
 	"gvisor.googlesource.com/gvisor/pkg/sentry/arch"
 	"gvisor.googlesource.com/gvisor/pkg/sentry/context"
 	"gvisor.googlesource.com/gvisor/pkg/sentry/fs"
@@ -164,7 +165,7 @@ type Kernel struct {
 	// nextInotifyCookie is a monotonically increasing counter used for
 	// generating unique inotify event cookies.
 	//
-	// nextInotifyCookie is mutable, and is accesed using atomic memory
+	// nextInotifyCookie is mutable, and is accessed using atomic memory
 	// operations.
 	nextInotifyCookie uint32
 
@@ -177,6 +178,10 @@ type Kernel struct {
 
 	// danglingEndpoints is used to save / restore tcpip.DanglingEndpoints.
 	danglingEndpoints struct{} `state:".([]tcpip.Endpoint)"`
+
+	// socketTable is used to track all sockets on the system. Protected by
+	// extMu.
+	socketTable map[int]map[*refs.WeakRef]struct{}
 }
 
 // InitKernelArgs holds arguments to Init.
@@ -266,6 +271,7 @@ func (k *Kernel) Init(args InitKernelArgs) error {
 	k.monotonicClock = &timekeeperClock{tk: args.Timekeeper, c: sentrytime.Monotonic}
 	k.futexes = futex.NewManager()
 	k.netlinkPorts = port.New()
+	k.socketTable = make(map[int]map[*refs.WeakRef]struct{})
 
 	return nil
 }
@@ -1051,6 +1057,56 @@ func (k *Kernel) EmitUnimplementedEvent(ctx context.Context) {
 	})
 }
 
+// socketEntry represents a socket recorded in Kernel.socketTable. It implements
+// refs.WeakRefUser for sockets stored in the socket table.
+//
+// +stateify savable
+type socketEntry struct {
+	k      *Kernel
+	sock   *refs.WeakRef
+	family int
+}
+
+// WeakRefGone implements refs.WeakRefUser.WeakRefGone.
+func (s *socketEntry) WeakRefGone() {
+	s.k.extMu.Lock()
+	// k.socketTable is guaranteed to point to a valid socket table for s.family
+	// at this point, since we made sure of the fact when we created this
+	// socketEntry, and we never delete socket tables.
+	delete(s.k.socketTable[s.family], s.sock)
+	s.k.extMu.Unlock()
+}
+
+// RecordSocket adds a socket to the system-wide socket table for tracking.
+//
+// Precondition: Caller must hold a reference to sock.
+func (k *Kernel) RecordSocket(sock *fs.File, family int) {
+	k.extMu.Lock()
+	table, ok := k.socketTable[family]
+	if !ok {
+		table = make(map[*refs.WeakRef]struct{})
+		k.socketTable[family] = table
+	}
+	se := socketEntry{k: k, family: family}
+	se.sock = refs.NewWeakRef(sock, &se)
+	table[se.sock] = struct{}{}
+	k.extMu.Unlock()
+}
+
+// ListSockets returns a snapshot of all sockets of a given family.
+func (k *Kernel) ListSockets(family int) []*refs.WeakRef {
+	k.extMu.Lock()
+	socks := []*refs.WeakRef{}
+	if table, ok := k.socketTable[family]; ok {
+		socks = make([]*refs.WeakRef, 0, len(table))
+		for s, _ := range table {
+			socks = append(socks, s)
+		}
+	}
+	k.extMu.Unlock()
+	return socks
+}
+
 type supervisorContext struct {
 	context.NoopSleeper
 	log.Logger