From db6ee959dfd81d14bd7d89c3ca6cb4ea7b83216f Mon Sep 17 00:00:00 2001
From: Ayush Ranjan <ayushranjan@google.com>
Date: Mon, 6 May 2024 17:33:42 -0700
Subject: [PATCH] Add GoferClientProvider to devutil.

Also introduce CtxDevGoferClientProvider, which is provided in restore context.

PiperOrigin-RevId: 631237657
---
 pkg/devutil/context.go            | 13 ++++++++++++-
 pkg/devutil/devutil.go            |  5 +++++
 pkg/sentry/kernel/kernel.go       |  6 ++++--
 pkg/sentry/kernel/task_context.go |  2 +-
 runsc/boot/restore.go             |  2 ++
 5 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/pkg/devutil/context.go b/pkg/devutil/context.go
index 6095dad7e..50b227415 100644
--- a/pkg/devutil/context.go
+++ b/pkg/devutil/context.go
@@ -14,7 +14,7 @@
 
 package devutil
 
-import "gvisor.dev/gvisor/pkg/context"
+import "context"
 
 // contextID is this package's type for context.Context.Value keys.
 type contextID int
@@ -22,6 +22,9 @@ type contextID int
 const (
 	// CtxDevGoferClient is a Context.Value key for a /dev gofer client.
 	CtxDevGoferClient contextID = iota
+
+	// CtxDevGoferClientProvider is a Context.Value key for GoferClientProvider.
+	CtxDevGoferClientProvider
 )
 
 // GoferClientFromContext returns the device gofer client used by ctx.
@@ -31,3 +34,11 @@ func GoferClientFromContext(ctx context.Context) *GoferClient {
 	}
 	return nil
 }
+
+// GoferClientProviderFromContext returns the GoferClientProvider used by ctx.
+func GoferClientProviderFromContext(ctx context.Context) GoferClientProvider {
+	if v := ctx.Value(CtxDevGoferClientProvider); v != nil {
+		return v.(GoferClientProvider)
+	}
+	return nil
+}
diff --git a/pkg/devutil/devutil.go b/pkg/devutil/devutil.go
index edc1478d4..4e21a4fa4 100644
--- a/pkg/devutil/devutil.go
+++ b/pkg/devutil/devutil.go
@@ -122,3 +122,8 @@ func (g *GoferClient) OpenAt(ctx context.Context, name string, flags uint32) (in
 	client.CloseFD(ctx, childOpenFD, true /* flush */)
 	return childHostFD, nil
 }
+
+// GoferClientProvider provides a GoferClient for a given container.
+type GoferClientProvider interface {
+	GetDevGoferClient(contName string) *GoferClient
+}
diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go
index ccfc9e18b..59bf69a4b 100644
--- a/pkg/sentry/kernel/kernel.go
+++ b/pkg/sentry/kernel/kernel.go
@@ -943,7 +943,7 @@ func (ctx *createProcessContext) Value(key any) any {
 		mntns.IncRef()
 		return mntns
 	case devutil.CtxDevGoferClient:
-		return ctx.kernel.getDevGoferClient(ctx.kernel.ContainerName(ctx.args.ContainerID))
+		return ctx.kernel.GetDevGoferClient(ctx.kernel.ContainerName(ctx.args.ContainerID))
 	case inet.CtxStack:
 		return ctx.kernel.RootNetworkNamespace().Stack()
 	case ktime.CtxRealtimeClock:
@@ -2004,7 +2004,9 @@ func (k *Kernel) RemoveDevGofer(contName string) {
 	delete(k.devGofers, contName)
 }
 
-func (k *Kernel) getDevGoferClient(contName string) *devutil.GoferClient {
+// GetDevGoferClient implements
+// devutil.GoferClientProviderFromContext.GetDevGoferClient.
+func (k *Kernel) GetDevGoferClient(contName string) *devutil.GoferClient {
 	k.devGofersMu.Lock()
 	defer k.devGofersMu.Unlock()
 	return k.devGofers[contName]
diff --git a/pkg/sentry/kernel/task_context.go b/pkg/sentry/kernel/task_context.go
index afca312ca..6b7b9a98f 100644
--- a/pkg/sentry/kernel/task_context.go
+++ b/pkg/sentry/kernel/task_context.go
@@ -105,7 +105,7 @@ func (t *Task) contextValue(key any, isTaskGoroutine bool) any {
 		t.mountNamespace.IncRef()
 		return t.mountNamespace
 	case devutil.CtxDevGoferClient:
-		return t.k.getDevGoferClient(t.k.ContainerName(t.containerID))
+		return t.k.GetDevGoferClient(t.k.ContainerName(t.containerID))
 	case inet.CtxStack:
 		return t.NetworkContext()
 	case ktime.CtxRealtimeClock:
diff --git a/runsc/boot/restore.go b/runsc/boot/restore.go
index ece168f83..3ea078891 100644
--- a/runsc/boot/restore.go
+++ b/runsc/boot/restore.go
@@ -25,6 +25,7 @@ import (
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	"gvisor.dev/gvisor/pkg/cleanup"
 	"gvisor.dev/gvisor/pkg/context"
+	"gvisor.dev/gvisor/pkg/devutil"
 	"gvisor.dev/gvisor/pkg/fd"
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/sentry/control"
@@ -226,6 +227,7 @@ func (r *restorer) restore(l *Loader) error {
 	ctx = context.WithValue(ctx, vfs.CtxRestoreFilesystemFDMap, fdmap)
 	log.Debugf("Restore using mfmap: %v", fdmap)
 	ctx = context.WithValue(ctx, pgalloc.CtxMemoryFileMap, mfmap)
+	ctx = context.WithValue(ctx, devutil.CtxDevGoferClientProvider, l.k)
 
 	// Load the state.
 	loadOpts := state.LoadOpts{Source: r.stateFile, PagesMetadata: r.pagesMetadata, PagesFile: r.pagesFile}