Implement ioctl command VFIO_GROUP_SET_CONTAINER.

PiperOrigin-RevId: 615514857

pkg/abi/linux/BUILD

@@ -78,6 +78,7 @@ go_library(
         "tty.go",
         "uio.go",
         "utsname.go",
+        "vfio.go",
         "wait.go",
         "xattr.go",
     ],

pkg/abi/linux/vfio.go

@@ -0,0 +1,28 @@
+// 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.
+
+// The package implements VFIOuserspace driver interface.
+
+package linux
+
+// For IOCTLs requests from include/uapi/linux/vfio.h.
+const (
+	VFIO_TYPE = ';'
+	VFIO_BASE = 100
+)
+
+// IOCTLs for VFIO file descriptor from include/uapi/linux/vfio.h.
+var (
+	VFIO_GROUP_SET_CONTAINER = IO(VFIO_TYPE, VFIO_BASE+4)
+)

pkg/sentry/devices/tpuproxy/BUILD

@@ -13,6 +13,7 @@ go_library(
         "tpu_mmap.go",
         "vfio.go",
         "vfio_mmap.go",
+        "vfio_unsafe.go",
     ],
     visibility = [
         "//pkg/sentry:internal",
@@ -25,14 +26,17 @@ go_library(
         "//pkg/fdnotifier",
         "//pkg/hostarch",
         "//pkg/log",
+        "//pkg/marshal/primitive",
         "//pkg/safemem",
         "//pkg/seccomp",
         "//pkg/sentry/arch",
+        "//pkg/sentry/kernel",
         "//pkg/sentry/memmap",
         "//pkg/sentry/vfs",
         "//pkg/sync",
         "//pkg/usermem",
         "//pkg/waiter",
+        "@org_golang_x_exp//constraints:go_default_library",
         "@org_golang_x_sys//unix:go_default_library",
     ],
 )

pkg/sentry/devices/tpuproxy/seccomp_filter.go

@@ -52,5 +52,11 @@ func Filters() seccomp.SyscallRules {
 			seccomp.AnyValue{},
 			seccomp.EqualTo(0),
 		},
+		unix.SYS_IOCTL: seccomp.Or{
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(linux.VFIO_GROUP_SET_CONTAINER),
+			},
+		},
 	})
 }

pkg/sentry/devices/tpuproxy/tpu.go

@@ -19,10 +19,14 @@ 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/kernel"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 	"gvisor.dev/gvisor/pkg/usermem"
 	"gvisor.dev/gvisor/pkg/waiter"
@@ -80,5 +84,32 @@ func (fd *tpuFD) Epollable() bool {
 
 // Ioctl implements vfs.FileDescriptionImpl.Ioctl.
 func (fd *tpuFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args arch.SyscallArguments) (uintptr, error) {
+	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_GROUP_SET_CONTAINER:
+		return fd.setContainer(ctx, t, args[2].Pointer())
+	}
 	return 0, linuxerr.ENOSYS
 }
+
+func (fd *tpuFD) setContainer(ctx context.Context, t *kernel.Task, arg hostarch.Addr) (uintptr, error) {
+	var vfioContainerFd int32
+	if _, err := primitive.CopyInt32In(t, arg, &vfioContainerFd); err != nil {
+		return 0, err
+	}
+	vfioContainerFile, _ := t.FDTable().Get(vfioContainerFd)
+	if vfioContainerFile == nil {
+		return 0, linuxerr.EBADF
+	}
+	defer vfioContainerFile.DecRef(ctx)
+	vfioContainer, ok := vfioContainerFile.Impl().(*vfioFd)
+	if !ok {
+		return 0, linuxerr.EINVAL
+	}
+	return ioctlInvokePtrArg(fd.hostFD, linux.VFIO_GROUP_SET_CONTAINER, &vfioContainer.hostFd)
+}

pkg/sentry/devices/tpuproxy/vfio_unsafe.go

@@ -0,0 +1,34 @@
+// 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 (
+	"unsafe"
+
+	"golang.org/x/exp/constraints"
+	"golang.org/x/sys/unix"
+)
+
+func ioctlInvokePtrArg[Params any](hostFd int32, cmd uint32, params *Params) (uintptr, error) {
+	return ioctlInvoke[uintptr](hostFd, cmd, uintptr(unsafe.Pointer(params)))
+}
+
+func ioctlInvoke[Arg constraints.Integer](hostFd int32, cmd uint32, arg Arg) (uintptr, error) {
+	n, _, errno := unix.RawSyscall(unix.SYS_IOCTL, uintptr(hostFd), uintptr(cmd), uintptr(arg))
+	if errno != 0 {
+		return n, errno
+	}
+	return n, nil
+}