From 7ff0b64d6e72eb879bf8e692181c3cb44411d1f2 Mon Sep 17 00:00:00 2001
From: Jing Chen <chjing@google.com>
Date: Fri, 12 Apr 2024 14:42:12 -0700
Subject: [PATCH] Add pciDeviceFD mmap and initialize it with the corresponding
 host FD.

It removes kernfs.CachedMappable from pciDeviceFD.

PiperOrigin-RevId: 624296998
---
 pkg/sentry/devices/tpuproxy/tpu.go      |  9 ++--
 pkg/sentry/devices/tpuproxy/tpu_mmap.go | 55 ++++++++++++++++++++++++-
 2 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/pkg/sentry/devices/tpuproxy/tpu.go b/pkg/sentry/devices/tpuproxy/tpu.go
index ab135003c..e5d3f60f1 100644
--- a/pkg/sentry/devices/tpuproxy/tpu.go
+++ b/pkg/sentry/devices/tpuproxy/tpu.go
@@ -27,7 +27,6 @@ import (
 	"gvisor.dev/gvisor/pkg/marshal/primitive"
 	"gvisor.dev/gvisor/pkg/sentry/arch"
 	"gvisor.dev/gvisor/pkg/sentry/fsimpl/eventfd"
-	"gvisor.dev/gvisor/pkg/sentry/fsimpl/kernfs"
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
 	"gvisor.dev/gvisor/pkg/sentry/mm"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
@@ -177,7 +176,7 @@ func (fd *tpuFD) getPciDeviceFd(t *kernel.Task, arg hostarch.Addr) (uintptr, fun
 		return 0, cleanup, err
 	}
 	// Initialize a mapping that is backed by a host FD.
-	pciDevFD.CachedMappable.Init(int(hostFD))
+	pciDevFD.memmapFile.fd = pciDevFD
 	return uintptr(newFD), func() {}, nil
 }
 
@@ -187,10 +186,10 @@ type pciDeviceFD struct {
 	vfs.FileDescriptionDefaultImpl
 	vfs.DentryMetadataFileDescriptionImpl
 	vfs.NoLockFD
-	kernfs.CachedMappable
 
-	hostFD int32
-	queue  waiter.Queue
+	hostFD     int32
+	queue      waiter.Queue
+	memmapFile pciDeviceFdMemmapFile
 }
 
 // Release implements vfs.FileDescriptionImpl.Release.
diff --git a/pkg/sentry/devices/tpuproxy/tpu_mmap.go b/pkg/sentry/devices/tpuproxy/tpu_mmap.go
index ad672947c..553b9c09a 100644
--- a/pkg/sentry/devices/tpuproxy/tpu_mmap.go
+++ b/pkg/sentry/devices/tpuproxy/tpu_mmap.go
@@ -85,6 +85,59 @@ func (mf *tpuFDMemmapFile) FD() int {
 
 // ConfigureMMap implements vfs.FileDescriptionImpl.ConfigureMMap.
 func (fd *pciDeviceFD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
-	fd.CachedMappable.InitFileMapperOnce()
 	return vfs.GenericConfigureMMap(&fd.vfsfd, fd, opts)
 }
+
+// AddMapping implements memmap.Mappable.AddMapping.
+func (fd *pciDeviceFD) AddMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) error {
+	return nil
+}
+
+// RemoveMapping implements memmap.Mappable.RemoveMapping.
+func (fd *pciDeviceFD) RemoveMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) {
+}
+
+// CopyMapping implements memmap.Mappable.CopyMapping.
+func (fd *pciDeviceFD) CopyMapping(ctx context.Context, ms memmap.MappingSpace, srcAR, dstAR hostarch.AddrRange, offset uint64, writable bool) error {
+	return nil
+}
+
+// Translate implements memmap.Mappable.Translate.
+func (fd *pciDeviceFD) Translate(ctx context.Context, required, optional memmap.MappableRange, at hostarch.AccessType) ([]memmap.Translation, error) {
+	return []memmap.Translation{
+		{
+			Source: optional,
+			File:   &fd.memmapFile,
+			Offset: optional.Start,
+			Perms:  at,
+		},
+	}, nil
+}
+
+// InvalidateUnsavable implements memmap.Mappable.InvalidateUnsavable.
+func (fd *pciDeviceFD) InvalidateUnsavable(ctx context.Context) error {
+	return nil
+}
+
+type pciDeviceFdMemmapFile struct {
+	fd *pciDeviceFD
+}
+
+// IncRef implements memmap.File.IncRef.
+func (mf *pciDeviceFdMemmapFile) IncRef(memmap.FileRange, uint32) {
+}
+
+// DecRef implements memmap.File.DecRef.
+func (mf *pciDeviceFdMemmapFile) DecRef(fr memmap.FileRange) {
+}
+
+// MapInternal implements memmap.File.MapInternal.
+func (mf *pciDeviceFdMemmapFile) MapInternal(fr memmap.FileRange, at hostarch.AccessType) (safemem.BlockSeq, error) {
+	log.Traceback("tpuproxy: rejecting pciDeviceFdMemmapFile.MapInternal")
+	return safemem.BlockSeq{}, linuxerr.EINVAL
+}
+
+// FD implements memmap.File.FD.
+func (mf *pciDeviceFdMemmapFile) FD() int {
+	return int(mf.fd.hostFD)
+}