From 88ee65f3a862e404e42d7481a4ef1217d226879c Mon Sep 17 00:00:00 2001
From: Jing Chen <chjing@google.com>
Date: Thu, 28 Mar 2024 17:43:38 -0700
Subject: [PATCH] Implement pass through ioctl VFIO_IOMMU_UNMAP_DMA.

PiperOrigin-RevId: 620107765
---
 pkg/abi/linux/vfio.go               | 21 ++++++++++++++++++++
 pkg/sentry/devices/tpuproxy/vfio.go | 30 +++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/pkg/abi/linux/vfio.go b/pkg/abi/linux/vfio.go
index 56c1101e5..d77a5a367 100644
--- a/pkg/abi/linux/vfio.go
+++ b/pkg/abi/linux/vfio.go
@@ -106,6 +106,10 @@ const (
 	VFIO_DMA_MAP_FLAG_VADDR
 )
 
+const (
+	VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP = 1
+)
+
 // IOCTLs for VFIO file descriptor from include/uapi/linux/vfio.h.
 var (
 	VFIO_CHECK_EXTENSION        = IO(VFIO_TYPE, VFIO_BASE+1)
@@ -118,6 +122,7 @@ var (
 	VFIO_DEVICE_SET_IRQS        = IO(VFIO_TYPE, VFIO_BASE+10)
 	VFIO_DEVICE_RESET           = IO(VFIO_TYPE, VFIO_BASE+11)
 	VFIO_IOMMU_MAP_DMA          = IO(VFIO_TYPE, VFIO_BASE+13)
+	VFIO_IOMMU_UNMAP_DMA        = IO(VFIO_TYPE, VFIO_BASE+14)
 )
 
 // VFIODeviceInfo is analogous to vfio_device_info
@@ -191,3 +196,19 @@ type VFIOIommuType1DmaMap struct {
 	// Size of mapping in bytes.
 	Size uint64
 }
+
+// VFIOIommuType1DmaUnmap is analogous to vfio_iommu_type1_dma_unmap
+// from include/uapi/linux/vfio.h.
+//
+// +marshal
+type VFIOIommuType1DmaUnmap struct {
+	Argsz uint32
+	Flags uint32
+	// IO virtual address.
+	IOVa uint64
+	// Size of mapping in bytes.
+	Size uint64
+	// The `data` field from vfio_iommu_type1_dma_unmap is omitted. The
+	// field is a flexible array member, and is needed only if the flag
+	// VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP is enabled.
+}
diff --git a/pkg/sentry/devices/tpuproxy/vfio.go b/pkg/sentry/devices/tpuproxy/vfio.go
index eb809741d..618eaaf64 100644
--- a/pkg/sentry/devices/tpuproxy/vfio.go
+++ b/pkg/sentry/devices/tpuproxy/vfio.go
@@ -96,6 +96,8 @@ func (fd *vfioFd) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args
 		return fd.setIOMMU(extension(args[2].Int()))
 	case linux.VFIO_IOMMU_MAP_DMA:
 		return fd.iommuMapDma(ctx, t, args[2].Pointer())
+	case linux.VFIO_IOMMU_UNMAP_DMA:
+		return fd.iommuUnmapDma(ctx, t, args[2].Pointer())
 	}
 	return 0, linuxerr.ENOSYS
 }
@@ -209,6 +211,34 @@ func (fd *vfioFd) iommuMapDma(ctx context.Context, t *kernel.Task, arg hostarch.
 	return n, nil
 }
 
+func (fd *vfioFd) iommuUnmapDma(ctx context.Context, t *kernel.Task, arg hostarch.Addr) (uintptr, error) {
+	var dmaUnmap linux.VFIOIommuType1DmaUnmap
+	if _, err := dmaUnmap.CopyIn(t, arg); err != nil {
+		return 0, err
+	}
+	if dmaUnmap.Flags&linux.VFIO_DMA_UNMAP_FLAG_GET_DIRTY_BITMAP != 0 {
+		// VFIO_DMA_UNMAP_FALGS_GET_DIRTY_BITMAP is not used by libtpu for
+		// gVisor working with TPU.
+		return 0, linuxerr.ENOSYS
+	}
+	n, err := IOCTLInvokePtrArg[uint32](fd.hostFd, linux.VFIO_IOMMU_MAP_DMA, &dmaUnmap)
+	if err != nil {
+		return 0, nil
+	}
+	fd.device.mu.Lock()
+	defer fd.device.mu.Unlock()
+	s := &fd.device.devAddrSet
+	r := DevAddrRange{Start: dmaUnmap.IOVa, End: dmaUnmap.IOVa + dmaUnmap.Size}
+	seg := s.LowerBoundSegment(r.Start)
+	for seg.Ok() && seg.Start() < r.End {
+		seg = s.Isolate(seg, r)
+		mm.Unpin([]mm.PinnedRange{seg.Value()})
+		gap := s.Remove(seg)
+		seg = gap.NextSegment()
+	}
+	return n, nil
+}
+
 // VFIO extension.
 type extension int32