visor/pkg/sentry/devices/tpuproxy/vfio/pci_device_fd.go

// Copyright 2023 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 vfio

import (
	"fmt"

	"golang.org/x/sys/unix"
	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/context"
	"gvisor.dev/gvisor/pkg/errors/linuxerr"
	"gvisor.dev/gvisor/pkg/fdnotifier"
	"gvisor.dev/gvisor/pkg/hostarch"
	"gvisor.dev/gvisor/pkg/marshal/primitive"
	"gvisor.dev/gvisor/pkg/sentry/arch"
	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/util"
	"gvisor.dev/gvisor/pkg/sentry/fsimpl/eventfd"
	"gvisor.dev/gvisor/pkg/sentry/fsutil"
	"gvisor.dev/gvisor/pkg/sentry/kernel"
	"gvisor.dev/gvisor/pkg/sentry/memmap"
	"gvisor.dev/gvisor/pkg/sentry/vfs"
	"gvisor.dev/gvisor/pkg/sync"
	"gvisor.dev/gvisor/pkg/usermem"
	"gvisor.dev/gvisor/pkg/waiter"
)

// pciDeviceFD implements vfs.FileDescriptionImpl for TPU's PCI device.
//
// +stateify savable
type pciDeviceFD struct {
	vfsfd vfs.FileDescription
	vfs.FileDescriptionDefaultImpl
	vfs.DentryMetadataFileDescriptionImpl
	vfs.NoLockFD

	hostFD        int32
	deviceAddress string
	containerName string
	queue         waiter.Queue

	// TODO: pciDeviceFD.InvalidateUnsavable uses this state, but AFAIU
	// pciDeviceFD.Translate will return the same File and offset after
	// save/restore, so this is unnecessary, and we can use
	// memmap.MappableNoTrackMappings instead.
	mapsMu sync.Mutex `state:"nosave"`
	// +checklocks:mapsMu
	mappings   memmap.MappingSet
	memmapFile fsutil.MmapPreciseFile
	tpuproxy   *tpuproxy
}

func (fd *pciDeviceFD) isRestored() bool {
	return fd.hostFD == -1
}

// Release implements vfs.FileDescriptionImpl.Release.
func (fd *pciDeviceFD) Release(context.Context) {
	defer fd.tpuproxy.untrackFD(fd)
	if fd.isRestored() {
		return
	}
	fdnotifier.RemoveFD(fd.hostFD)
	fd.queue.Notify(waiter.EventHUp)
	fd.memmapFile.MappableRelease() // eventually closes fd.hostFD
}

// EventRegister implements waiter.Waitable.EventRegister.
func (fd *pciDeviceFD) EventRegister(e *waiter.Entry) error {
	if fd.isRestored() {
		return nil
	}
	fd.queue.EventRegister(e)
	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
		fd.queue.EventUnregister(e)
		return err
	}
	return nil
}

// EventUnregister implements waiter.Waitable.EventUnregister.
func (fd *pciDeviceFD) EventUnregister(e *waiter.Entry) {
	if fd.isRestored() {
		return
	}
	fd.queue.EventUnregister(e)
	if err := fdnotifier.UpdateFD(fd.hostFD); err != nil {
		panic(fmt.Sprint("UpdateFD:", err))
	}
}

// Readiness implements waiter.Waitable.Readiness.
func (fd *pciDeviceFD) Readiness(mask waiter.EventMask) waiter.EventMask {
	if fd.isRestored() {
		return 0
	}
	return fdnotifier.NonBlockingPoll(fd.hostFD, mask)
}

// Epollable implements vfs.FileDescriptionImpl.Epollable.
func (fd *pciDeviceFD) Epollable() bool {
	return true
}

// Ioctl implements vfs.FileDescriptionImpl.Ioctl.
func (fd *pciDeviceFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args arch.SyscallArguments) (uintptr, error) {
	if fd.isRestored() {
		return 0, nil
	}
	cmd := args[1].Uint()

	t := kernel.TaskFromContext(ctx)
	if t == nil {
		panic("Ioctl should be called from a task context")
	}
	switch cmd {
	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())
	case linux.VFIO_DEVICE_GET_IRQ_INFO:
		return fd.vfioIrqInfo(ctx, t, args[2].Pointer())
	case linux.VFIO_DEVICE_SET_IRQS:
		return fd.vfioSetIrqs(ctx, t, args[2].Pointer())
	case linux.VFIO_DEVICE_RESET:
		// VFIO_DEVICE_RESET is just a simple IOCTL command that carries no data.
		return util.IOCTLInvoke[uint32, uintptr](fd.hostFD, linux.VFIO_DEVICE_RESET, 0)
	}
	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
	}
	// drivers/vfio/vfio_main.c:vfio_get_region_info() copies capabilities into
	// the bytes following the vfio_region_info.
	appArgsz := int(regionInfo.Argsz)
	if appArgsz < regionInfo.SizeBytes() {
		return 0, linuxerr.EINVAL
	}
	// Start with a small argsz to limit memory usage.
	buf := make([]byte, min(appArgsz, 1024))
retry:
	regionInfo.Argsz = uint32(len(buf))
	regionInfo.MarshalUnsafe(buf)
	ret, err := util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_GET_REGION_INFO, &buf[0])
	if err != nil {
		return 0, err
	}
	if len(buf) < appArgsz {
		// Check if truncating buf prevented us from obtaining capabilities.
		regionInfo.UnmarshalUnsafe(buf)
		if regionInfo.Flags&linux.VFIO_REGION_INFO_FLAG_CAPS != 0 && regionInfo.CapOffset == 0 {
			// vfio_get_region_info() set regionInfo.Argsz to the minimum value
			// required to obtain capabilities.
			if int(regionInfo.Argsz) <= appArgsz {
				buf = make([]byte, regionInfo.Argsz)
				goto retry
			}
		}
	}
	if _, err := t.CopyOutBytes(arg, buf); 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.VFIODeviceInfoMin.CopyIn(t, arg); err != nil {
		return 0, err
	}
	appArgsz := int(deviceInfo.Argsz)
	if appArgsz < deviceInfo.VFIODeviceInfoMin.SizeBytes() {
		return 0, linuxerr.EINVAL
	}
	if deviceInfo.Flags&^vfioDeviceInfoFlags != 0 {
		return 0, linuxerr.EINVAL
	}
	// drivers/vfio/pci/vfio_pci_core.c:vfio_pci_ioctl_get_info() copies
	// capabilities into the bytes following the vfio_device_info. Start with a
	// small argsz to limit memory usage, but ensure that it's large enough to
	// accommodate the current vfio_device_info so that we can observe
	// deviceInfo.CapOffset.
	buf := make([]byte, max(min(appArgsz, 1024), deviceInfo.SizeBytes()))
retry:
	deviceInfo.Argsz = uint32(len(buf))
	deviceInfo.MarshalUnsafe(buf)
	ret, err := util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_GET_INFO, &buf[0])
	if err != nil {
		return 0, err
	}
	if len(buf) < appArgsz {
		// Check if truncating buf prevented us from obtaining capabilities.
		deviceInfo.UnmarshalUnsafe(buf)
		if deviceInfo.Flags&linux.VFIO_DEVICE_FLAGS_CAPS != 0 && deviceInfo.CapOffset == 0 {
			// vfio_pci_ioctl_get_info() set deviceInfo.Argsz to the minimum
			// value required to obtain capabilities.
			if int(deviceInfo.Argsz) <= appArgsz {
				buf = make([]byte, deviceInfo.Argsz)
				goto retry
			}
		}
	}
	// gVisor is not supposed to change any device information that is
	// returned from the host since gVisor doesn't own the device.
	// Passing the device info back to the caller will be just fine.
	if _, err := t.CopyOutBytes(arg, buf[:min(len(buf), appArgsz)]); err != nil {
		return 0, err
	}
	return ret, nil
}

// Retrieve the device's interrupt information.
func (fd *pciDeviceFD) vfioIrqInfo(ctx context.Context, t *kernel.Task, arg hostarch.Addr) (uintptr, error) {
	var irqInfo linux.VFIOIrqInfo
	if _, err := irqInfo.CopyIn(t, arg); err != nil {
		return 0, err
	}
	// Callers must set the payload's size.
	if irqInfo.Argsz == 0 {
		return 0, linuxerr.EINVAL
	}
	ret, err := util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_GET_IRQ_INFO, &irqInfo)
	if err != nil {
		return 0, err
	}
	if _, err := irqInfo.CopyOut(t, arg); err != nil {
		return 0, err
	}
	return ret, nil
}

func (fd *pciDeviceFD) vfioSetIrqs(ctx context.Context, t *kernel.Task, arg hostarch.Addr) (uintptr, error) {
	var irqSet linux.VFIOIrqSet
	if _, err := irqSet.CopyIn(t, arg); err != nil {
		return 0, err
	}
	// Callers must set the payload's size.
	if irqSet.Argsz == 0 {
		return 0, linuxerr.EINVAL
	}
	// Invalidate unknown flags.
	if irqSet.Flags&^vfioIrqSetFlags != 0 {
		return 0, linuxerr.EINVAL
	}
	// See drivers/vfio/vfio_main.c:vfio_set_irqs_validate_and_prepare,
	// VFIO uses the data type at the request's flags to determine
	// the memory layout of data field.
	//
	// The struct vfio_irq_set includes a flexible array member, it
	// allocates an array for a continuous trunk of memory to back
	// a vfio_irq_set object. In order to mirror that behavior, gVisor
	// would allocate a slice to store the underlying bytes
	// and pass that through to its host.
	switch irqSet.Flags & linux.VFIO_IRQ_SET_DATA_TYPE_MASK {
	// VFIO_IRQ_SET_DATA_NONE indicates there is no data field for
	// the IOCTL command.
	// It works with VFIO_IRQ_SET_ACTION_MASK, VFIO_IRQ_SET_ACTION_UNMASK,
	// or VFIO_IRQ_SET_ACTION_TRIGGER to mask an interrupt, unmask an
	// interrupt,  and trigger an interrupt unconditionally.
	case linux.VFIO_IRQ_SET_DATA_NONE:
		// When there is no data, passing through the given payload
		// works just fine.
		return util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_SET_IRQS, &irqSet)
	// VFIO_IRQ_SET_DATA_BOOL indicates that the data field is an array of uint8.
	// The action will be performed if the corresponding boolean is true.
	case linux.VFIO_IRQ_SET_DATA_BOOL:
		payloadSize := uint32(irqSet.SizeBytes()) + irqSet.Count
		payload := make([]uint8, payloadSize)
		if _, err := primitive.CopyUint8SliceIn(t, arg, payload); err != nil {
			return 0, err
		}
		// Prevent TOCTOU by overwriting the header with the validated values
		irqSet.MarshalUnsafe(payload[:irqSet.SizeBytes()])
		return util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_SET_IRQS, &payload[0])
	// VFIO_IRQ_SET_DATA_EVENTFD indicates that the data field is an array
	// of int32 (or event file descriptors). These descriptors will be
	// signalled when an action in the flags happens.
	case linux.VFIO_IRQ_SET_DATA_EVENTFD:
		payloadSize := uint32(irqSet.SizeBytes())/4 + irqSet.Count
		payload := make([]int32, payloadSize)
		if _, err := primitive.CopyInt32SliceIn(t, arg, payload); err != nil {
			return 0, err
		}
		// Prevent TOCTOU by overwriting the header with the validated values
		payload[0] = int32(irqSet.Argsz)
		payload[1] = int32(irqSet.Flags)
		payload[2] = int32(irqSet.Index)
		payload[3] = int32(irqSet.Start)
		payload[4] = int32(irqSet.Count)
		// Transform the input FDs to host FDs.
		for i := 0; i < int(irqSet.Count); i++ {
			index := len(payload) - 1 - i
			fd := payload[index]
			// Skip non-event FD.
			if fd == disableInterrupt {
				continue
			}
			eventFileGeneric, _ := t.FDTable().Get(fd)
			if eventFileGeneric == nil {
				return 0, linuxerr.EBADF
			}
			defer eventFileGeneric.DecRef(ctx)
			eventFile, ok := eventFileGeneric.Impl().(*eventfd.EventFileDescription)
			if !ok {
				return 0, linuxerr.EINVAL
			}
			eventfd, err := eventFile.HostFD()
			if err != nil {
				return 0, err
			}
			payload[index] = int32(eventfd)
		}
		return util.IOCTLInvokePtrArg[uint32](fd.hostFD, linux.VFIO_DEVICE_SET_IRQS, &payload[0])
	}
	// No data type is specified or multiple data types are specified.
	return 0, linuxerr.EINVAL
}

// PRead implements vfs.FileDescriptionImpl.PRead.
func (fd *pciDeviceFD) PRead(ctx context.Context, dst usermem.IOSequence, offset int64, opts vfs.ReadOptions) (int64, error) {
	if offset < 0 {
		return 0, linuxerr.EINVAL
	}
	if fd.isRestored() {
		return int64(dst.NumBytes()), nil
	}
	buf := make([]byte, dst.NumBytes())
	_, err := unix.Pread(int(fd.hostFD), buf, offset)
	if err != nil {
		return 0, err
	}
	n, err := dst.CopyOut(ctx, buf)
	return int64(n), err
}

// PWrite implements vfs.FileDescriptionImpl.PWrite.
func (fd *pciDeviceFD) PWrite(ctx context.Context, src usermem.IOSequence, offset int64, opts vfs.WriteOptions) (int64, error) {
	if offset < 0 {
		return 0, linuxerr.EINVAL
	}
	if fd.isRestored() {
		return src.NumBytes(), nil
	}
	buf := make([]byte, src.NumBytes())
	_, err := src.CopyIn(ctx, buf)
	if err != nil {
		return 0, err
	}
	n, err := unix.Pwrite(int(fd.hostFD), buf, offset)
	return int64(n), err
}