From 8739b58b61d11564e5bcfa8f0b6781cf1cbbec05 Mon Sep 17 00:00:00 2001
From: Jing Chen <chjing@google.com>
Date: Thu, 21 Mar 2024 01:24:15 -0700
Subject: [PATCH] Implement pass through ioctl command
 VFIO_DEVICE_GET_REGION_INFO.

PiperOrigin-RevId: 617760440
---
 pkg/abi/linux/vfio.go                         | 39 ++++++++++++++++---
 pkg/sentry/devices/tpuproxy/seccomp_filter.go |  4 ++
 pkg/sentry/devices/tpuproxy/tpu.go            | 26 ++++++++++++-
 3 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/pkg/abi/linux/vfio.go b/pkg/abi/linux/vfio.go
index 4d52ff406..15fe9c1a2 100644
--- a/pkg/abi/linux/vfio.go
+++ b/pkg/abi/linux/vfio.go
@@ -49,13 +49,26 @@ const (
 	VFIO_DEVICE_FLAGS_CDX
 )
 
+// VFIO region info flags.
+const (
+	// Region supports read.
+	VFIO_REGION_INFO_FLAG_READ = 1 << iota
+	// Region supports write.
+	VFIO_REGION_INFO_FLAG_WRITE
+	// Region supports mmap.
+	VFIO_REGION_INFO_FLAG_MMAP
+	// Info supports caps.
+	VFIO_REGION_INFO_FLAG_CAPS
+)
+
 // IOCTLs for VFIO file descriptor from include/uapi/linux/vfio.h.
 var (
-	VFIO_CHECK_EXTENSION     = IO(VFIO_TYPE, VFIO_BASE+1)
-	VFIO_SET_IOMMU           = IO(VFIO_TYPE, VFIO_BASE+2)
-	VFIO_GROUP_SET_CONTAINER = IO(VFIO_TYPE, VFIO_BASE+4)
-	VFIO_GROUP_GET_DEVICE_FD = IO(VFIO_TYPE, VFIO_BASE+6)
-	VFIO_DEVICE_GET_INFO     = IO(VFIO_TYPE, VFIO_BASE+7)
+	VFIO_CHECK_EXTENSION        = IO(VFIO_TYPE, VFIO_BASE+1)
+	VFIO_SET_IOMMU              = IO(VFIO_TYPE, VFIO_BASE+2)
+	VFIO_GROUP_SET_CONTAINER    = IO(VFIO_TYPE, VFIO_BASE+4)
+	VFIO_GROUP_GET_DEVICE_FD    = IO(VFIO_TYPE, VFIO_BASE+6)
+	VFIO_DEVICE_GET_INFO        = IO(VFIO_TYPE, VFIO_BASE+7)
+	VFIO_DEVICE_GET_REGION_INFO = IO(VFIO_TYPE, VFIO_BASE+8)
 )
 
 // VFIODeviceInfo is analogous to vfio_device_info
@@ -73,3 +86,19 @@ type VFIODeviceInfo struct {
 	CapOffset uint32
 	pad       uint32
 }
+
+// VFIORegionInfo is analogous to vfio_region_info
+// from include/uapi/linux/vfio.h.
+//
+// +marshal
+type VFIORegionInfo struct {
+	Argsz uint32
+	Flags uint32
+	Index uint32
+	// Offset within info struct of first cap.
+	capOffset uint32
+	// Region size in bytes.
+	Size uint64
+	// Region offset from start of device fd.
+	Offset uint64
+}
diff --git a/pkg/sentry/devices/tpuproxy/seccomp_filter.go b/pkg/sentry/devices/tpuproxy/seccomp_filter.go
index 3d94f9dee..6920450bd 100644
--- a/pkg/sentry/devices/tpuproxy/seccomp_filter.go
+++ b/pkg/sentry/devices/tpuproxy/seccomp_filter.go
@@ -73,6 +73,10 @@ func Filters() seccomp.SyscallRules {
 				seccomp.NonNegativeFD{},
 				seccomp.EqualTo(linux.VFIO_DEVICE_GET_INFO),
 			},
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(linux.VFIO_DEVICE_GET_REGION_INFO),
+			},
 		},
 	})
 }
diff --git a/pkg/sentry/devices/tpuproxy/tpu.go b/pkg/sentry/devices/tpuproxy/tpu.go
index c87bce874..96beb8bce 100644
--- a/pkg/sentry/devices/tpuproxy/tpu.go
+++ b/pkg/sentry/devices/tpuproxy/tpu.go
@@ -224,18 +224,42 @@ func (fd *pciDeviceFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr,
 		panic("Ioctl should be called from a task context")
 	}
 	switch cmd {
+	// TODO(b/299303493): consider making VFIO's GET_INFO commands more generic.
 	case linux.VFIO_DEVICE_GET_INFO:
 		return fd.vfioDeviceInfo(ctx, t, args[2].Pointer())
+	case linux.VFIO_DEVICE_GET_REGION_INFO:
+		return fd.vfioRegionInfo(ctx, t, args[2].Pointer())
 	}
 	return 0, linuxerr.ENOSYS
 }
 
+// Retrieve the host TPU device's region information, which could be used by
+// vfio driver to setup mappings.
+func (fd *pciDeviceFD) vfioRegionInfo(ctx context.Context, t *kernel.Task, arg hostarch.Addr) (uintptr, error) {
+	var regionInfo linux.VFIORegionInfo
+	if _, err := regionInfo.CopyIn(t, arg); err != nil {
+		return 0, err
+	}
+	if regionInfo.Argsz == 0 {
+		return 0, linuxerr.EINVAL
+	}
+	ret, err := IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_GET_REGION_INFO, &regionInfo)
+	if err != nil {
+		return 0, err
+	}
+	if _, err := regionInfo.CopyOut(t, arg); err != nil {
+		return 0, err
+	}
+	return ret, nil
+}
+
+// Retrieve the host TPU device's information.
 func (fd *pciDeviceFD) vfioDeviceInfo(ctx context.Context, t *kernel.Task, arg hostarch.Addr) (uintptr, error) {
 	var deviceInfo linux.VFIODeviceInfo
 	if _, err := deviceInfo.CopyIn(t, arg); err != nil {
 		return 0, err
 	}
-	// Callers must set VFIODevice.Argsz.
+	// Callers must set VFIODeviceInfo.Argsz.
 	if deviceInfo.Argsz == 0 {
 		return 0, linuxerr.EINVAL
 	}