From 5fd40fc0be5f79afb73a67e87b721c1b03cff61f Mon Sep 17 00:00:00 2001
From: Jing Chen <chjing@google.com>
Date: Tue, 12 Mar 2024 19:15:19 -0700
Subject: [PATCH] Implement memmap.Mappable and memmap.File for tpuFd.

PiperOrigin-RevId: 615250135
---
 pkg/sentry/devices/tpuproxy/BUILD       |  4 ++
 pkg/sentry/devices/tpuproxy/device.go   |  1 +
 pkg/sentry/devices/tpuproxy/tpu.go      |  7 ++-
 pkg/sentry/devices/tpuproxy/tpu_mmap.go | 84 +++++++++++++++++++++++++
 4 files changed, 93 insertions(+), 3 deletions(-)
 create mode 100644 pkg/sentry/devices/tpuproxy/tpu_mmap.go

diff --git a/pkg/sentry/devices/tpuproxy/BUILD b/pkg/sentry/devices/tpuproxy/BUILD
index 7c1ba7b1c..31ab05871 100644
--- a/pkg/sentry/devices/tpuproxy/BUILD
+++ b/pkg/sentry/devices/tpuproxy/BUILD
@@ -10,6 +10,7 @@ go_library(
         "device.go",
         "seccomp_filter.go",
         "tpu.go",
+        "tpu_mmap.go",
         "vfio.go",
     ],
     visibility = [
@@ -21,9 +22,12 @@ go_library(
         "//pkg/devutil",
         "//pkg/errors/linuxerr",
         "//pkg/fdnotifier",
+        "//pkg/hostarch",
         "//pkg/log",
+        "//pkg/safemem",
         "//pkg/seccomp",
         "//pkg/sentry/arch",
+        "//pkg/sentry/memmap",
         "//pkg/sentry/vfs",
         "//pkg/sync",
         "//pkg/usermem",
diff --git a/pkg/sentry/devices/tpuproxy/device.go b/pkg/sentry/devices/tpuproxy/device.go
index 59c819a42..46e5544a0 100644
--- a/pkg/sentry/devices/tpuproxy/device.go
+++ b/pkg/sentry/devices/tpuproxy/device.go
@@ -78,6 +78,7 @@ func (dev *tpuDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry, o
 		unix.Close(hostFD)
 		return nil, err
 	}
+	fd.memmapFile.fd = fd
 	return &fd.vfsfd, nil
 }
 
diff --git a/pkg/sentry/devices/tpuproxy/tpu.go b/pkg/sentry/devices/tpuproxy/tpu.go
index 51c4b6734..624d2c01f 100644
--- a/pkg/sentry/devices/tpuproxy/tpu.go
+++ b/pkg/sentry/devices/tpuproxy/tpu.go
@@ -37,9 +37,10 @@ type tpuFD struct {
 	vfs.DentryMetadataFileDescriptionImpl
 	vfs.NoLockFD
 
-	hostFD int32
-	device *tpuDevice
-	queue  waiter.Queue
+	hostFD     int32
+	device     *tpuDevice
+	queue      waiter.Queue
+	memmapFile tpuFdMemmapFile
 }
 
 // Release implements vfs.FileDescriptionImpl.Release.
diff --git a/pkg/sentry/devices/tpuproxy/tpu_mmap.go b/pkg/sentry/devices/tpuproxy/tpu_mmap.go
new file mode 100644
index 000000000..f2a47869a
--- /dev/null
+++ b/pkg/sentry/devices/tpuproxy/tpu_mmap.go
@@ -0,0 +1,84 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tpuproxy
+
+import (
+	"gvisor.dev/gvisor/pkg/context"
+	"gvisor.dev/gvisor/pkg/errors/linuxerr"
+	"gvisor.dev/gvisor/pkg/hostarch"
+	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/safemem"
+	"gvisor.dev/gvisor/pkg/sentry/memmap"
+	"gvisor.dev/gvisor/pkg/sentry/vfs"
+)
+
+// ConfigureMMap implements vfs.FileDescriptionImpl.ConfigureMMap.
+func (fd *tpuFD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
+	return vfs.GenericConfigureMMap(&fd.vfsfd, fd, opts)
+}
+
+// AddMapping implements memmap.Mappable.AddMapping.
+func (fd *tpuFD) AddMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) error {
+	return nil
+}
+
+// RemoveMapping implements memmap.Mappable.RemoveMapping.
+func (fd *tpuFD) RemoveMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) {
+}
+
+// CopyMapping implements memmap.Mappable.CopyMapping.
+func (fd *tpuFD) 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 *tpuFD) 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 *tpuFD) InvalidateUnsavable(ctx context.Context) error {
+	return nil
+}
+
+type tpuFdMemmapFile struct {
+	fd *tpuFD
+}
+
+// IncRef implements memmap.File.IncRef.
+func (mf *tpuFdMemmapFile) IncRef(memmap.FileRange, uint32) {
+}
+
+// DecRef implements memmap.File.DecRef.
+func (mf *tpuFdMemmapFile) DecRef(fr memmap.FileRange) {
+}
+
+// MapInternal implements memmap.File.MapInternal.
+func (mf *tpuFdMemmapFile) MapInternal(fr memmap.FileRange, at hostarch.AccessType) (safemem.BlockSeq, error) {
+	log.Traceback("tpuproxy: rejecting tpuFdMemmapFile.MapInternal")
+	return safemem.BlockSeq{}, linuxerr.EINVAL
+}
+
+// FD implements memmap.File.FD.
+func (mf *tpuFdMemmapFile) FD() int {
+	return int(mf.fd.hostFD)
+}