Refactor tpu chroot operations.

Ubuntu TPU images do not have the vfio-dev directories that COS images do,
so we need a more robust way of setting up the sandbox chroot to handle this
case. This change implements a way to get devices and minor numbers into the
sandbox with minimal support from the host filesystem and cleans up a few
methods to reflect their current usage.

Addresses #10795

PiperOrigin-RevId: 674363342

pkg/abi/linux/dev.go

@@ -70,10 +70,3 @@ const (
 	// ACCEL_MAJOR is the major device number for compute accelerator devices.
 	ACCEL_MAJOR = 121
 )
-
-// Major device numbers for VFIO-based TPU.
-const (
-	// Major devices number between 243 and 254 are usually reserved for local use.
-	// The device number 245 is used by VFIO based TPU in GCP.
-	VFIO_MAJOR = 245
-)

pkg/sentry/devices/tpuproxy/BUILD

@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "go_library", "go_test")
+load("//tools:defs.bzl", "go_library")
 
 package(default_applicable_licenses = ["//:license"])
 
@@ -17,16 +17,13 @@ go_library(
         "//pkg/abi/gasket",
         "//pkg/abi/linux",
         "//pkg/abi/tpu",
+        "//pkg/context",
+        "//pkg/fspath",
         "//pkg/seccomp",
         "//pkg/sentry/devices/tpuproxy/accel",
         "//pkg/sentry/devices/tpuproxy/vfio",
+        "//pkg/sentry/kernel/auth",
         "//pkg/sentry/vfs",
         "@org_golang_x_sys//unix:go_default_library",
     ],
 )
-
-go_test(
-    name = "tpuproxy_test",
-    srcs = ["tpuproxy_test.go"],
-    library = ":tpuproxy",
-)

pkg/sentry/devices/tpuproxy/accel/accel.go

@@ -90,6 +90,9 @@ func (dev *accelDevice) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.Dent
 
 // RegisterTPUDevice registers all devices implemented by this package in vfsObj.
 func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor uint32, lite bool) error {
+	if vfsObj.IsDeviceRegistered(vfs.CharDevice, linux.ACCEL_MAJOR, minor) {
+		return nil
+	}
 	return vfsObj.RegisterDevice(vfs.CharDevice, linux.ACCEL_MAJOR, minor, &accelDevice{
 		lite:  lite,
 		minor: minor,

pkg/sentry/devices/tpuproxy/tpuproxy.go

@@ -17,6 +17,7 @@
 package tpuproxy
 
 import (
+	"bytes"
 	"fmt"
 	"os"
 	"path"
@@ -25,100 +26,71 @@ import (
 	"strconv"
 	"strings"
 
+	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/tpu"
+	"gvisor.dev/gvisor/pkg/context"
+	"gvisor.dev/gvisor/pkg/fspath"
 	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/accel"
 	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/vfio"
+	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 )
 
-const (
-	pciPathGlobTPUv4   = "/sys/devices/pci0000:*/**/accel/accel*"
-	pciPathGlobTPUv5   = "/sys/devices/pci0000:*/**/vfio-dev/vfio*"
-	iommuGroupPathGlob = "/sys/kernel/iommu_groups/*/devices/*"
-)
-
 var (
-	// pathGlobToPathRegex is a map that points a TPU PCI path glob to its path regex.
-	// TPU v4 devices are accessible via /sys/devices/pci0000:00/<pci_address>/accel/accel# on the host.
-	// TPU v5 devices are accessible via at /sys/devices/pci0000:00/<pci_address>/vfio-dev/vfio# on the host.
-	pathGlobToPathRegex = map[string]string{
-		pciPathGlobTPUv4: `^/sys/devices/pci0000:[[:xdigit:]]{2}/(0000:([[:xdigit:]]{2}|[[:xdigit:]]{4}):[[:xdigit:]]{2}\.[[:xdigit:]]{1,2}/)+accel/accel(\d+)$`,
-		pciPathGlobTPUv5: `^/sys/devices/pci0000:[[:xdigit:]]{2}/(0000:([[:xdigit:]]{2}|[[:xdigit:]]{4}):[[:xdigit:]]{2}\.[[:xdigit:]]{1,2}/)+vfio-dev/vfio(\d+)$`,
-	}
+	// TPUv4DeviceRegex is the regex for detecting TPUv4 device paths.
+	TPUv4DeviceRegex = regexp.MustCompile(`/dev/accel(\d+)`)
+
+	// TPUv5DeviceRegex is the regex for detecting TPUv5 device paths.
+	TPUv5DeviceRegex = regexp.MustCompile(`/dev/vfio/(\d+)`)
 )
 
-// RegisterHostTPUDevices enumerates TPU devices on the host and registers them
-// in the sandbox VFS.
-func RegisterHostTPUDevices(vfsObj *vfs.VirtualFilesystem, allowedDeviceIDs map[int64]any) error {
-	for pciPathGlobal, pathRegex := range pathGlobToPathRegex {
-		pciAddrs, err := filepath.Glob(pciPathGlobal)
-		if err != nil {
-			return fmt.Errorf("enumerating PCI device files: %w", err)
-		}
-		pciPathRegex := regexp.MustCompile(pathRegex)
-		for _, pciPath := range pciAddrs {
-			ms := pciPathRegex.FindStringSubmatch(pciPath)
-			if ms == nil {
-				continue
-			}
-			minorNum, err := strconv.ParseUint(ms[len(ms)-1], 10, 32)
-			if err != nil {
-				return fmt.Errorf("parsing PCI device number: %w", err)
-			}
-			var deviceIDBytes []byte
-			if deviceIDBytes, err = os.ReadFile(path.Join(pciPath, "device/device")); err != nil {
-				return fmt.Errorf("reading PCI device ID: %w", err)
-			}
-			deviceIDStr := strings.Replace(string(deviceIDBytes), "0x", "", -1)
-			deviceID, err := strconv.ParseInt(strings.TrimSpace(deviceIDStr), 16, 64)
-			if err != nil {
-				return fmt.Errorf("parsing PCI device ID: %w", err)
-			}
-			if _, ok := allowedDeviceIDs[deviceID]; !ok {
-				return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID)
-			}
-			// VFIO iommu groups correspond to the device number. Use these
-			// paths to get the correct number for the sentry-internal TPU
-			// device files.
-			var deviceNum int
-			switch deviceID {
-			case tpu.TPUV4DeviceID, tpu.TPUV4liteDeviceID:
-				deviceNum = int(deviceNum)
-			case tpu.TPUV5eDeviceID, tpu.TPUV5pDeviceID:
-				groupPaths, err := filepath.Glob(iommuGroupPathGlob)
-				if err != nil {
-					return fmt.Errorf("enumerating IOMMU group files: %w", err)
-				}
-				for _, groupPath := range groupPaths {
-					pci := path.Base(groupPath)
-					if strings.Contains(pciPath, pci) {
-						n, err := strconv.Atoi(strings.Split(groupPath, "/")[4])
-						if err != nil {
-							return fmt.Errorf("parsing IOMMU group minor number: %w", err)
-						}
-						deviceNum = n
-						break
-					}
-				}
-			default:
-				return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID)
-			}
-			if err := registerTPUDevice(vfsObj, uint32(minorNum), uint32(deviceNum), deviceID); err != nil {
-				return fmt.Errorf("registering TPU driver: %w", err)
-			}
-		}
+// RegisterTPUv4Device registers the TPUv4 device with the provided minor number
+// where the corresponding PCI device is located at pciPath. Accel devices
+// always have their device file number set to their minor number.
+func RegisterTPUv4Device(ctx context.Context, creds *auth.Credentials, root vfs.VirtualDentry, vfsObj *vfs.VirtualFilesystem, devPath string, minorNum uint32) error {
+	// Get the PCI path from the accel device's symlink at
+	// /sys/class/accel/accel\d+. The link will be in the form
+	// "../../devices/pci0000:*/**/accel/accel\d+".
+	linkPath := filepath.Join("/sys/class/accel", filepath.Base(devPath))
+	linkContent, err := vfsObj.ReadlinkAt(ctx, creds, &vfs.PathOperation{Root: root, Start: root, Path: fspath.Parse(linkPath)})
+	if err != nil {
+		return fmt.Errorf("reading link %q: %w", linkPath, err)
+	}
+	// Exclude the ../../devices prefix and the accel/accel\d+ suffix.
+	pciPath := strings.TrimSuffix(strings.TrimPrefix(linkContent, "../../devices"), fmt.Sprintf("accel/%s", filepath.Base(devPath)))
+	pciDeviceIDPath := path.Join("/sys/devices", pciPath, "device")
+
+	fd, err := unix.Openat(-1, pciDeviceIDPath, unix.O_RDONLY|unix.O_NOFOLLOW, 0)
+	if err != nil {
+		return err
+	}
+	file := os.NewFile(uintptr(fd), pciDeviceIDPath)
+	defer file.Close()
+	buf := bytes.Buffer{}
+	if _, err := buf.ReadFrom(file); err != nil {
+		return err
+	}
+
+	deviceIDStr := strings.Replace(buf.String(), "0x", "", -1)
+	deviceID, err := strconv.ParseInt(strings.TrimSpace(deviceIDStr), 16, 64)
+	if err != nil {
+		return fmt.Errorf("parsing PCI device ID: %w", err)
+	}
+	if err := accel.RegisterTPUDevice(vfsObj, minorNum, deviceID == tpu.TPUV4liteDeviceID); err != nil {
+		return fmt.Errorf("registering TPU driver: %w", err)
 	}
 	return nil
 }
 
-// registerTPUDevice registers a TPU device in vfsObj based on the given device ID.
-func registerTPUDevice(vfsObj *vfs.VirtualFilesystem, minor, deviceNum uint32, deviceID int64) error {
-	switch deviceID {
-	case tpu.TPUV4DeviceID, tpu.TPUV4liteDeviceID:
-		return accel.RegisterTPUDevice(vfsObj, minor, deviceID == tpu.TPUV4liteDeviceID)
-	case tpu.TPUV5eDeviceID, tpu.TPUV5pDeviceID:
-		return vfio.RegisterTPUDevice(vfsObj, minor, deviceNum, false /* useDevGofer */)
-	default:
-		return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID)
+// RegisterTPUv5Device registers the TPUv5 device with the provided device path
+// and minor number.
+func RegisterTPUv5Device(vfsObj *vfs.VirtualFilesystem, devPath string, minorNum uint32) error {
+	deviceNum, err := strconv.ParseInt(path.Base(devPath), 10, 32)
+	if err != nil {
+		return fmt.Errorf("parsing device path number: %w", err)
 	}
+	if err := vfio.RegisterTPUDevice(vfsObj, uint32(minorNum), uint32(deviceNum), true /* useDevGofer */); err != nil {
+		return fmt.Errorf("registering TPU driver: %w", err)
+	}
+	return nil
 }

pkg/sentry/devices/tpuproxy/tpuproxy_test.go

@@ -1,66 +0,0 @@
-// 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 (
-	"path/filepath"
-	"regexp"
-	"slices"
-	"testing"
-)
-
-func TestTPUPath(t *testing.T) {
-	for _, tst := range []struct {
-		name     string
-		pathGlob string
-		path     string
-		submatch []string
-	}{
-		{
-			name:     "TPUv4PCIPathMatch",
-			pathGlob: pciPathGlobTPUv4,
-			path:     "/sys/devices/pci0000:00/0000:00:01.0/accel/accel16",
-			submatch: []string{"/sys/devices/pci0000:00/0000:00:01.0/accel/accel16", "0000:00:01.0/", "00", "16"},
-		},
-		{
-			name:     "TPUv4PCIPathNoMatch",
-			pathGlob: pciPathGlobTPUv4,
-			path:     "/sys/devices/pci0000:00/0000:00:01.0/accel/123",
-			submatch: nil,
-		},
-		{
-			name:     "TPUv5PCIPathMatch",
-			pathGlob: pciPathGlobTPUv5,
-			path:     "/sys/devices/pci0000:00/0000:00:05.0/vfio-dev/vfio20",
-			submatch: []string{"/sys/devices/pci0000:00/0000:00:05.0/vfio-dev/vfio20", "0000:00:05.0/", "00", "20"},
-		},
-		{
-			name:     "TPUv5PCIPathNoMatch",
-			pathGlob: pciPathGlobTPUv5,
-			path:     "/sys/devices/pci0000:00/0000:00:05.0/vfio/vfio20",
-			submatch: nil,
-		},
-	} {
-		t.Run(tst.name, func(t *testing.T) {
-			if _, err := filepath.Glob(tst.pathGlob); err != nil {
-				t.Errorf("Malformed path glob: %v", err)
-			}
-			pathRegex := regexp.MustCompile(pathGlobToPathRegex[tst.pathGlob])
-			if submatch := pathRegex.FindStringSubmatch(tst.path); !slices.Equal(submatch, tst.submatch) {
-				t.Errorf("Match TPU PCI path, got: %v, want: %v", submatch, tst.submatch)
-			}
-		})
-	}
-}

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

@@ -41,6 +41,12 @@ const (
 	VFIOPath = "/dev/vfio/vfio"
 )
 
+var (
+	tpuDeviceMajor        uint32
+	tpuDeviceMajorInit    sync.Once
+	tpuDeviceMajorInitErr error
+)
+
 // device implements TPU's vfs.Device for /dev/vfio/[0-9]+
 //
 // +stateify savable
@@ -158,7 +164,14 @@ func (dev *vfioDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry,
 
 // RegisterTPUDevice registers devices implemented by this package in vfsObj.
 func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor, deviceNum uint32, useDevGofer bool) error {
-	return vfsObj.RegisterDevice(vfs.CharDevice, linux.VFIO_MAJOR, minor, &tpuDevice{
+	major, err := GetTPUDeviceMajor(vfsObj)
+	if err != nil {
+		return err
+	}
+	if vfsObj.IsDeviceRegistered(vfs.CharDevice, major, minor) {
+		return nil
+	}
+	return vfsObj.RegisterDevice(vfs.CharDevice, major, minor, &tpuDevice{
 		minor:       minor,
 		num:         deviceNum,
 		useDevGofer: useDevGofer,
@@ -171,6 +184,9 @@ func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor, deviceNum uint32, u
 
 // RegisterVFIODevice registers VFIO devices that are implemented by this package in vfsObj.
 func RegisterVFIODevice(vfsObj *vfs.VirtualFilesystem, useDevGofer bool) error {
+	if vfsObj.IsDeviceRegistered(vfs.CharDevice, linux.MISC_MAJOR, VFIO_MINOR) {
+		return nil
+	}
 	return vfsObj.RegisterDevice(vfs.CharDevice, linux.MISC_MAJOR, VFIO_MINOR, &vfioDevice{
 		useDevGofer: useDevGofer,
 	}, &vfs.RegisterDeviceOptions{
@@ -179,3 +195,12 @@ func RegisterVFIODevice(vfsObj *vfs.VirtualFilesystem, useDevGofer bool) error {
 		FilePerms: 0666,
 	})
 }
+
+// GetTPUDeviceMajor returns the dynamically allocated major number for the vfio
+// device.
+func GetTPUDeviceMajor(vfsObj *vfs.VirtualFilesystem) (uint32, error) {
+	tpuDeviceMajorInit.Do(func() {
+		tpuDeviceMajor, tpuDeviceMajorInitErr = vfsObj.GetDynamicCharDevMajor()
+	})
+	return tpuDeviceMajor, tpuDeviceMajorInitErr
+}

pkg/sentry/fsimpl/sys/pci.go

@@ -175,9 +175,11 @@ func (fs *filesystem) mirrorSysDevicesDir(ctx context.Context, creds *auth.Crede
 				if err != nil {
 					return nil, err
 				}
+				// Remove the bus prefix.
+				pciPath := pciBusRegex.ReplaceAllString(pciPaths[pciDeviceName], "")
 				// Both the device and PCI address entries are links to the original PCI
 				// device directory that's at the same place earlier in the dir tree.
-				linkContent = fmt.Sprintf("../../../%s", pciPaths[pciDeviceName])
+				linkContent = path.Join("../../../", pciPath)
 			case dent == "iommu_group":
 				pciDeviceName, err := pciDeviceName(dir)
 				if err != nil {

pkg/sentry/vfs/device.go

@@ -110,6 +110,15 @@ func (vfs *VirtualFilesystem) ForEachDevice(cb func(pathname string, kind Device
 	return nil
 }
 
+// IsDeviceRegistered returns true if a device that matches the
+// (kind, major, minor) tuple is registered.
+func (vfs *VirtualFilesystem) IsDeviceRegistered(kind DeviceKind, major, minor uint32) bool {
+	vfs.devicesMu.RLock()
+	defer vfs.devicesMu.RUnlock()
+	_, ok := vfs.devices[devTuple{kind, major, minor}]
+	return ok
+}
+
 // OpenDeviceSpecialFile returns a FileDescription representing the given
 // device.
 func (vfs *VirtualFilesystem) OpenDeviceSpecialFile(ctx context.Context, mnt *Mount, d *Dentry, kind DeviceKind, major, minor uint32, opts *OpenOptions) (*FileDescription, error) {

runsc/boot/BUILD

@@ -35,7 +35,6 @@ go_library(
         "//pkg/abi",
         "//pkg/abi/linux",
         "//pkg/abi/nvgpu",
-        "//pkg/abi/tpu",
         "//pkg/bpf",
         "//pkg/cleanup",
         "//pkg/context",

runsc/boot/vfs.go

@@ -28,7 +28,6 @@ import (
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	"gvisor.dev/gvisor/pkg/abi/nvgpu"
-	"gvisor.dev/gvisor/pkg/abi/tpu"
 	"gvisor.dev/gvisor/pkg/cleanup"
 	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/devutil"
@@ -161,10 +160,6 @@ func registerFilesystems(k *kernel.Kernel, info *containerInfo) error {
 		return err
 	}
 
-	if err := tpuProxyRegisterDevices(info, vfsObj); err != nil {
-		return err
-	}
-
 	return nil
 }
 
@@ -1365,7 +1360,28 @@ func createDeviceFile(ctx context.Context, creds *auth.Credentials, info *contai
 	default:
 		return fmt.Errorf("specified device at %q has invalid type %q", devSpec.Path, devSpec.Type)
 	}
-	if devSpec.Path == "/dev/nvidia-uvm" && info.nvidiaUVMDevMajor != 0 && major != info.nvidiaUVMDevMajor {
+	if strings.HasPrefix(devSpec.Path, "/dev/vfio") || strings.HasPrefix(devSpec.Path, "/dev/accel") {
+		if devSpec.Path == "/dev/vfio/vfio" {
+			if err := vfio.RegisterVFIODevice(vfsObj, true /* useDevGofer */); err != nil {
+				return fmt.Errorf("registering vfio driver: %w", err)
+			}
+		} else if tpuproxy.TPUv4DeviceRegex.MatchString(devSpec.Path) {
+
+			if err := tpuproxy.RegisterTPUv4Device(ctx, creds, root, vfsObj, devSpec.Path, minor); err != nil {
+				return fmt.Errorf("registering TPUv4 device: %w", err)
+			}
+		} else if tpuproxy.TPUv5DeviceRegex.MatchString(devSpec.Path) {
+			if err := tpuproxy.RegisterTPUv5Device(vfsObj, devSpec.Path, minor); err != nil {
+				return fmt.Errorf("registering TPUv5 device: %w", err)
+			}
+			log.Infof("Switching %v device major number from %d to %d", devSpec.Path, devSpec.Major, major)
+			var err error
+			major, err = vfio.GetTPUDeviceMajor(vfsObj)
+			if err != nil {
+				return fmt.Errorf("getting TPU device major number: %w", err)
+			}
+		}
+	} else if devSpec.Path == "/dev/nvidia-uvm" && info.nvidiaUVMDevMajor != 0 && major != info.nvidiaUVMDevMajor {
 		// nvidia-uvm's major device number is dynamically assigned, so the
 		// number that it has on the host may differ from the number that
 		// it has in sentry VFS; switch from the former to the latter.
@@ -1375,25 +1391,6 @@ func createDeviceFile(ctx context.Context, creds *auth.Credentials, info *contai
 	return dev.CreateDeviceFile(ctx, vfsObj, creds, root, devSpec.Path, major, minor, mode, devSpec.UID, devSpec.GID)
 }
 
-func tpuProxyRegisterDevices(info *containerInfo, vfsObj *vfs.VirtualFilesystem) error {
-	if !specutils.TPUProxyIsEnabled(info.spec, info.conf) {
-		return nil
-	}
-	allowedTPUDeviceIDs := map[int64]any{
-		tpu.TPUV4DeviceID:     nil,
-		tpu.TPUV4liteDeviceID: nil,
-		tpu.TPUV5pDeviceID:    nil,
-		tpu.TPUV5eDeviceID:    nil,
-	}
-	if err := tpuproxy.RegisterHostTPUDevices(vfsObj, allowedTPUDeviceIDs); err != nil {
-		return fmt.Errorf("registering host TPU devices: %w", err)
-	}
-	if err := vfio.RegisterVFIODevice(vfsObj, true /* useDevGofer */); err != nil {
-		return fmt.Errorf("registering vfio driver: %w", err)
-	}
-	return nil
-}
-
 func nvproxyRegisterDevices(info *containerInfo, vfsObj *vfs.VirtualFilesystem) error {
 	if !specutils.NVProxyEnabled(info.spec, info.conf) {
 		return nil

runsc/cmd/BUILD

@@ -79,6 +79,7 @@ go_library(
     visibility = ["//runsc:__subpackages__"],
     deps = [
         "//pkg/abi/linux",
+        "//pkg/abi/tpu",
         "//pkg/cleanup",
         "//pkg/coretag",
         "//pkg/coverage",
@@ -131,6 +132,7 @@ go_test(
     size = "small",
     srcs = [
         "capability_test.go",
+        "chroot_test.go",
         "delete_test.go",
         "exec_test.go",
         "gofer_test.go",
@@ -158,5 +160,6 @@ go_test(
         "@com_github_google_subcommands//:go_default_library",
         "@com_github_opencontainers_runtime_spec//specs-go:go_default_library",
         "@com_github_syndtr_gocapability//capability:go_default_library",
+        "@org_golang_x_sys//unix:go_default_library",
     ],
 )

runsc/cmd/chroot.go

@@ -19,10 +19,11 @@ import (
 	"os"
 	"path"
 	"path/filepath"
-	"regexp"
+	"strings"
 
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/abi/tpu"
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/runsc/cmd/util"
 	"gvisor.dev/gvisor/runsc/config"
@@ -113,7 +114,7 @@ func setUpChroot(spec *specs.Spec, conf *config.Config) error {
 		return fmt.Errorf("error mounting proc in chroot: %v", err)
 	}
 
-	if err := tpuProxyUpdateChroot(chroot, spec, conf); err != nil {
+	if err := tpuProxyUpdateChroot("/", chroot, spec, conf); err != nil {
 		return fmt.Errorf("error configuring chroot for TPU devices: %w", err)
 	}
 
@@ -124,93 +125,59 @@ func setUpChroot(spec *specs.Spec, conf *config.Config) error {
 	return pivotRoot(chroot)
 }
 
-// Mount the path that dest points to for TPU at chroot, the mounted path is returned in absolute form.
-func mountTPUSyslinkInChroot(chroot, dest, relativePath string, validator func(link string) bool) (string, error) {
-	src, err := os.Readlink(dest)
-	if err != nil {
-		return "", fmt.Errorf("error reading %v: %v", src, err)
-	}
-	// Ensure the link is in the form we expect.
-	if !validator(src) {
-		return "", fmt.Errorf("unexpected link %q -> %q", dest, src)
-	}
-	path, err := filepath.Abs(path.Join(filepath.Dir(dest), src, relativePath))
-	if err != nil {
-		return "", fmt.Errorf("error parsing path %q: %v", src, err)
-	}
-	if err := mountInChroot(chroot, path, path, "bind", unix.MS_BIND|unix.MS_RDONLY); err != nil {
-		return "", fmt.Errorf("error mounting %q in chroot: %v", dest, err)
-	}
-	return path, nil
-}
-
-func mountTPUDeviceInfoInChroot(chroot, devicePath, sysfsFormat, pciDeviceFormat string) error {
-	deviceMinor, valid, err := util.ExtractTPUDeviceMinor(devicePath)
-	if err != nil {
-		return fmt.Errorf("extracting TPU device minor: %w", err)
-	}
-	if !valid {
-		return nil
-	}
-	// Multiple paths link to the /sys/devices/<pci_bus>/<pci_address>
-	// directory that contains all relevant sysfs accel/vfio device info that we need
-	// bind mounted into the sandbox chroot. We can construct this path by
-	// reading the link below, which points to
-	//   * /sys/devices/<pci_bus>/<pci_address>/accel/accel#
-	//   * or /sys/devices/<pci_bus>/<pci_address>/vfio-dev/vfio# for VFIO-based TPU
-	// and traversing up 2 directories.
-	// The sysDevicePath itself is a soft link to the device directory.
-	sysDevicePath := fmt.Sprintf(sysfsFormat, deviceMinor)
-	sysPCIDeviceDir, err := mountTPUSyslinkInChroot(chroot, sysDevicePath, "../..", func(link string) bool {
-		sysDeviceLinkMatcher := regexp.MustCompile(fmt.Sprintf(pciDeviceFormat, deviceMinor))
-		return sysDeviceLinkMatcher.MatchString(link)
-	})
-	if err != nil {
-		return err
-	}
-
-	// Mount the device's IOMMU group if available.
-	iommuGroupPath := path.Join(sysPCIDeviceDir, "iommu_group")
-	if _, err := os.Stat(iommuGroupPath); err == nil {
-		if _, err := mountTPUSyslinkInChroot(chroot, iommuGroupPath, "", func(link string) bool {
-			iommuGroupPathMatcher := regexp.MustCompile(`../../../kernel/iommu_groups/\d+`)
-			return iommuGroupPathMatcher.MatchString(link)
-		}); err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-func tpuProxyUpdateChroot(chroot string, spec *specs.Spec, conf *config.Config) error {
+func tpuProxyUpdateChroot(hostRoot, chroot string, spec *specs.Spec, conf *config.Config) error {
 	if !specutils.TPUProxyIsEnabled(spec, conf) {
 		return nil
 	}
-	// When a path glob is added to pathGlobToSysfsFormat, the corresponding pciDeviceFormat has to be added to pathGlobToPciDeviceFormat.
-	pathGlobToSysfsFormat := map[string]string{
-		"/dev/accel*": "/sys/class/accel/accel%d",
-		"/dev/vfio/*": "/sys/class/vfio-dev/vfio%d"}
-	pathGlobToPciDeviceFormat := map[string]string{
-		"/dev/accel*": `../../devices/pci0000:[[:xdigit:]]{2}/(\d+:\d+:\d+\.\d+)/accel/accel%d`,
-		"/dev/vfio/*": `../../devices/pci0000:[[:xdigit:]]{2}/(\d+:\d+:\d+\.\d+)/vfio-dev/vfio%d`}
-	// Bind mount device info directories for all TPU devices on the host.
-	// For v4 TPU, the directory /sys/devices/<pci_bus>/<pci_address>/accel/accel# is mounted;
-	// For v5e TPU, the directory /sys/devices/<pci_bus>/<pci_address>/vfio-dev/vfio# is mounted.
-	foundDevices := false
-	for pathGlob, sysfsFormat := range pathGlobToSysfsFormat {
-		paths, err := filepath.Glob(pathGlob)
+	allowedDeviceIDs := map[uint64]struct{}{}
+	paths, err := filepath.Glob(path.Join(hostRoot, "dev/vfio/*"))
+	if err != nil {
+		return fmt.Errorf("enumerating TPU device files: %w", err)
+	}
+	vfioDevicePath := path.Join(hostRoot, "dev/vfio/vfio")
+	for _, devPath := range paths {
+		if devPath == vfioDevicePath {
+			continue
+		}
+		devNum := path.Base(devPath)
+		iommuGroupPath := path.Join("/sys/kernel/iommu_groups", devNum)
+		if err := mountInChroot(chroot, path.Join(hostRoot, iommuGroupPath), iommuGroupPath, "bind", unix.MS_BIND|unix.MS_RDONLY); err != nil {
+			return fmt.Errorf("error mounting %q in chroot: %v", iommuGroupPath, err)
+		}
+		allowedDeviceIDs[tpu.TPUV5pDeviceID] = struct{}{}
+		allowedDeviceIDs[tpu.TPUV5eDeviceID] = struct{}{}
+	}
+	if len(allowedDeviceIDs) == 0 {
+		paths, err = filepath.Glob(path.Join(hostRoot, "dev/accel*"))
 		if err != nil {
 			return fmt.Errorf("enumerating TPU device files: %w", err)
 		}
-		for _, devPath := range paths {
-			foundDevices = true
-			if err := mountTPUDeviceInfoInChroot(chroot, devPath, sysfsFormat, pathGlobToPciDeviceFormat[pathGlob]); err != nil {
-				return err
+		if len(paths) == 0 {
+			return fmt.Errorf("could not find any TPU devices on the host")
+		}
+		allowedDeviceIDs[tpu.TPUV4DeviceID] = struct{}{}
+		allowedDeviceIDs[tpu.TPUV4liteDeviceID] = struct{}{}
+	}
+	if len(allowedDeviceIDs) == 0 {
+		return fmt.Errorf("no TPU devices found on the host")
+	}
+	sysDevicesGlob := path.Join(hostRoot, "/sys/devices/pci*")
+	sysDevicesPaths, err := filepath.Glob(sysDevicesGlob)
+	if err != nil {
+		return fmt.Errorf("enumerating PCI device files: %w", err)
+	}
+	for _, sysDevicesPath := range sysDevicesPaths {
+		if err := filepath.WalkDir(sysDevicesPath, func(path string, d os.DirEntry, err error) error {
+			if d.Type().IsDir() && util.IsPCIDeviceDirTPU(path, allowedDeviceIDs) {
+				chrootPath := strings.Replace(path, hostRoot, "/", 1)
+				if err := mountInChroot(chroot, path, chrootPath, "bind", unix.MS_BIND|unix.MS_RDONLY); err != nil {
+					return fmt.Errorf("error mounting %q in chroot: %v", path, err)
+				}
 			}
+			return nil
+		}); err != nil {
+			return fmt.Errorf("walking %q: %w", sysDevicesPath, err)
 		}
 	}
-	if !foundDevices {
-		return fmt.Errorf("could not find any TPU devices on the host")
-	}
-	return nil
+	return err
 }

runsc/cmd/chroot_test.go

@@ -0,0 +1,155 @@
+// 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 cmd
+
+import (
+	"fmt"
+	"os"
+	"path"
+	"testing"
+
+	specs "github.com/opencontainers/runtime-spec/specs-go"
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/runsc/config"
+)
+
+func setup(t *testing.T) (string, string) {
+	t.Helper()
+	testDir := t.TempDir()
+	gvisorChroot := path.Join(testDir, "gvisor_chroot")
+	os.Mkdir(gvisorChroot, 0755)
+
+	// Mounting the gvisor chroot makes the submounts easier to cleanup at the
+	// end of the test.
+	if err := unix.Mount(gvisorChroot, gvisorChroot, "", unix.MS_BIND, ""); err != nil {
+		t.Fatalf("failed to bind mount gvisor chroot: %v", err)
+	}
+	t.Cleanup(func() {
+		if err := unix.Unmount(gvisorChroot, unix.MNT_DETACH); err != nil {
+			t.Fatalf("failed to unmount gvisor chroot: %v", err)
+		}
+	})
+	return testDir, gvisorChroot
+}
+
+func TestTPUProxyV5(t *testing.T) {
+	testDir, gvisorChroot := setup(t)
+
+	for i := 0; i < 3; i++ {
+		os.MkdirAll(path.Join(testDir, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", i)), 0755)
+		writeFile(t, path.Join(testDir, "dev", "vfio", fmt.Sprintf("%d", i)), "")
+		pciPath := path.Join(testDir, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i))
+		writeFile(t, path.Join(pciPath, "device"), "0x0062")
+		writeFile(t, path.Join(pciPath, "vendor"), "0x1ae0")
+	}
+
+	if err := tpuProxyUpdateChroot(testDir, gvisorChroot, &specs.Spec{}, &config.Config{TPUProxy: true}); err != nil {
+		t.Fatalf("failed to update chroot: %v", err)
+	}
+
+	for i := 0; i < 3; i++ {
+		if _, err := os.Stat(path.Join(gvisorChroot, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", i))); err != nil {
+			t.Errorf("failed to stat iommu group file: %v", err)
+		}
+		devicePath := path.Join(gvisorChroot, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i), "device")
+		if _, err := os.ReadFile(devicePath); err != nil {
+			t.Errorf("failed to read device file: %v", err)
+		}
+		vendorPath := path.Join(gvisorChroot, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i), "vendor")
+		if _, err := os.ReadFile(vendorPath); err != nil {
+			t.Errorf("failed to read device file: %v", err)
+		}
+	}
+}
+
+func TestTPUProxyV5NestedPCIDevice(t *testing.T) {
+	testDir, gvisorChroot := setup(t)
+
+	for i := 0; i < 3; i++ {
+		os.MkdirAll(path.Join(testDir, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", i)), 0755)
+		writeFile(t, path.Join(testDir, "dev", "vfio", fmt.Sprintf("%d", i)), "")
+		pciPath := path.Join(testDir, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i))
+		writeFile(t, path.Join(pciPath, "device"), "0x0062")
+		writeFile(t, path.Join(pciPath, "vendor"), "0x1ae0")
+	}
+
+	nestedDeviceNum := 3
+	writeFile(t, path.Join(testDir, "dev", "vfio", fmt.Sprintf("%d", nestedDeviceNum)), "")
+	os.MkdirAll(path.Join(testDir, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", nestedDeviceNum)), 0755)
+	pciPath := path.Join(testDir, "sys", "devices", "pci0000:00", "0000:00:00.2", "0000:00:00.3.0")
+	writeFile(t, path.Join(pciPath, "device"), "0x0062")
+	writeFile(t, path.Join(pciPath, "vendor"), "0x1ae0")
+
+	if err := tpuProxyUpdateChroot(testDir, gvisorChroot, &specs.Spec{}, &config.Config{TPUProxy: true}); err != nil {
+		t.Fatalf("failed to update chroot: %v", err)
+	}
+
+	if _, err := os.Stat(path.Join(gvisorChroot, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", nestedDeviceNum))); err != nil {
+		t.Errorf("failed to stat iommu group file: %v", err)
+	}
+	devicePath := path.Join(gvisorChroot, "sys", "devices", "pci0000:00", "0000:00:00.2", "0000:00:00.3.0", "device")
+	if _, err := os.ReadFile(devicePath); err != nil {
+		t.Errorf("failed to read device file: %v", err)
+	}
+}
+
+func TestTPUProxyV4(t *testing.T) {
+	testDir, gvisorChroot := setup(t)
+
+	for i := 0; i < 3; i++ {
+		os.MkdirAll(path.Join(testDir, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", i)), 0755)
+		writeFile(t, path.Join(testDir, "dev", fmt.Sprintf("accel%d", i)), "")
+		pciPath := path.Join(testDir, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i))
+		writeFile(t, path.Join(pciPath, "device"), "0x005e")
+		writeFile(t, path.Join(pciPath, "vendor"), "0x1ae0")
+	}
+
+	if err := tpuProxyUpdateChroot(testDir, gvisorChroot, &specs.Spec{}, &config.Config{TPUProxy: true}); err != nil {
+		t.Fatalf("failed to update chroot: %v", err)
+	}
+
+	for i := 0; i < 3; i++ {
+		devicePath := path.Join(gvisorChroot, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i), "device")
+		if _, err := os.ReadFile(devicePath); err != nil {
+			t.Errorf("failed to read device file: %v", err)
+		}
+		vendorPath := path.Join(gvisorChroot, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i), "vendor")
+		if _, err := os.ReadFile(vendorPath); err != nil {
+			t.Errorf("failed to read device file: %v", err)
+		}
+	}
+}
+
+func writeFile(t *testing.T, fpath string, contents string) {
+	t.Helper()
+	dir := path.Dir(fpath)
+	if st, err := os.Stat(dir); err != nil {
+		if os.IsNotExist(err) {
+			if err := os.MkdirAll(dir, 0755); err != nil {
+				t.Fatalf("failed to create directory: %v", err)
+			}
+		} else {
+			t.Fatalf("failed to stat directory: %v", err)
+		}
+	} else if !st.IsDir() {
+		t.Fatalf("path %q is not a directory", dir)
+	}
+	f, err := os.Create(fpath)
+	if err != nil {
+		t.Fatalf("failed to create file: %v", err)
+	}
+	f.WriteString(contents)
+	f.Close()
+}

runsc/cmd/gofer.go

@@ -551,7 +551,7 @@ func shouldExposeVFIODevice(path string) bool {
 //
 // Precondition: tpuproxy is enabled.
 func shouldExposeTpuDevice(path string) bool {
-	_, valid, _ := util.ExtractTPUDeviceMinor(path)
+	valid, _ := util.IsTPUDeviceValid(path)
 	return valid || shouldExposeVFIODevice(path)
 }
 

runsc/cmd/util/BUILD

@@ -20,6 +20,5 @@ go_library(
         "//pkg/abi/tpu",
         "//pkg/log",
         "@com_github_google_subcommands//:go_default_library",
-        "@org_golang_x_sys//unix:go_default_library",
     ],
 )

runsc/cmd/util/tpu.go

@@ -17,89 +17,132 @@ package util
 import (
 	"fmt"
 	"os"
+	"path"
+	"path/filepath"
 	"regexp"
 	"strconv"
 	"strings"
-	"syscall"
 
-	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/tpu"
 )
 
 const (
-	googleVendorID       = 0x1AE0
-	accelDevicePathRegex = `^/dev/accel(\d+)$`
-	accelSysfsFormat     = "/sys/class/accel/accel%d/device/%s"
-	vfioDevicePathRegex  = `^/dev/vfio/(\d+)$`
-	vfioSysfsFormat      = "/sys/class/vfio-dev/vfio%d/device/%s"
-	vendorFile           = "vendor"
-	deviceFile           = "device"
+	googleVendorID            = 0x1AE0
+	accelDevicePathRegex      = `^/dev/accel(\d+)$`
+	accelSysfsFormat          = "/sys/class/accel/accel%d/device/%s"
+	vfioDevicePathRegex       = `^/dev/vfio/(\d+)$`
+	iommuGroupSysfsGlobFormat = "/sys/kernel/iommu_groups/%s/devices/*"
+	vendorFile                = "vendor"
+	deviceFile                = "device"
+	pciAddressMaxLength       = 13
 )
 
-var tpuV4DeviceIDs = map[uint64]any{tpu.TPUV4DeviceID: nil, tpu.TPUV4liteDeviceID: nil}
-var tpuV5DeviceIDs = map[uint64]any{tpu.TPUV5eDeviceID: nil, tpu.TPUV5pDeviceID: nil}
+var (
+	tpuV4DeviceIDs = map[uint64]struct{}{tpu.TPUV4DeviceID: struct{}{}, tpu.TPUV4liteDeviceID: struct{}{}}
+	tpuV5DeviceIDs = map[uint64]struct{}{tpu.TPUV5eDeviceID: struct{}{}, tpu.TPUV5pDeviceID: struct{}{}}
+	pciDeviceRegex = regexp.MustCompile(`0000:([[:xdigit:]]{2}|[[:xdigit:]]{4}):[[:xdigit:]]{2}\.[[:xdigit:]]{1,2}`)
+)
 
-// ExtractTPUDeviceMinor returns the accelerator device minor number for that
-// the passed device path. If the passed device is not a valid TPU device, then
-// it returns false.
-func ExtractTPUDeviceMinor(path string) (uint32, bool, error) {
-	devNum, valid, err := tpuV4DeviceMinor(path)
+// IsPCIDeviceDirTPU returns if the given PCI device sysfs path is a TPU device
+// with one of the allowed device IDs.
+func IsPCIDeviceDirTPU(sysfsPath string, allowedDeviceIDs map[uint64]struct{}) bool {
+	dir := path.Base(sysfsPath)
+	if !pciDeviceRegex.MatchString(dir) || len(dir) > pciAddressMaxLength {
+		return false
+	}
+	vendor, err := readHexInt(path.Join(sysfsPath, vendorFile))
 	if err != nil {
-		return 0, false, err
-	}
-	if valid {
-		return devNum, valid, err
-	}
-	return tpuV5DeviceMinor(path)
-}
-
-// tpuDeviceMinor returns the accelerator device minor number for that
-// the passed device path. If the passed device is not a valid TPU device, then
-// it returns false.
-func tpuDeviceMinor(devicePath, devicePathRegex, sysfsFormat string, allowedDeviceIDs map[uint64]any) (uint32, bool, error) {
-	deviceRegex := regexp.MustCompile(devicePathRegex)
-	matches := deviceRegex.FindStringSubmatch(devicePath)
-	if matches == nil {
-		return 0, false, nil
-	}
-	var st syscall.Stat_t
-	if err := syscall.Stat(devicePath, &st); err != nil {
-		return 0, false, err
-	}
-	minor := unix.Minor(st.Rdev)
-	vendor, err := readHexInt(fmt.Sprintf(sysfsFormat, minor, vendorFile))
-	if err != nil {
-		return 0, false, err
+		return false
 	}
 	if vendor != googleVendorID {
-		return 0, false, nil
+		return false
 	}
-	deviceID, err := readHexInt(fmt.Sprintf(sysfsFormat, minor, deviceFile))
+	deviceID, err := readHexInt(path.Join(sysfsPath, deviceFile))
 	if err != nil {
-		return 0, false, err
+		return false
 	}
 	if _, ok := allowedDeviceIDs[deviceID]; !ok {
-		return 0, false, nil
+		return false
 	}
-	return minor, true, nil
+	return true
 }
 
-// tpuv4DeviceMinor returns v4 and v4lite TPU device minor number for the given path.
+// IsTPUDeviceValid returns if the accelerator device is valid.
+func IsTPUDeviceValid(path string) (bool, error) {
+	valid, err := tpuV4DeviceValid(path)
+	if err != nil {
+		return false, err
+	}
+	if valid {
+		return valid, err
+	}
+	return tpuV5DeviceValid(path)
+}
+
+// tpuV4DeviceValid returns v4 and v4lite TPU device minor number for the given path.
 // A valid v4 TPU device is defined as:
 // * Path is /dev/accel#.
 // * Vendor is googleVendorID.
 // * Device ID is one of tpuV4DeviceIDs.
-func tpuV4DeviceMinor(path string) (uint32, bool, error) {
-	return tpuDeviceMinor(path, accelDevicePathRegex, accelSysfsFormat, tpuV4DeviceIDs)
+func tpuV4DeviceValid(devPath string) (bool, error) {
+	deviceRegex := regexp.MustCompile(accelDevicePathRegex)
+	matches := deviceRegex.FindStringSubmatch(devPath)
+	if matches == nil {
+		return false, nil
+	}
+	if len(matches) < 1 {
+		return false, fmt.Errorf("found %d matches for %s", len(matches), devPath)
+	}
+	devNum, err := strconv.ParseUint(matches[1], 10, 32)
+	if err != nil {
+		return false, err
+	}
+	vendor, err := readHexInt(fmt.Sprintf(accelSysfsFormat, devNum, vendorFile))
+	if err != nil {
+		return false, err
+	}
+	if vendor != googleVendorID {
+		return false, nil
+	}
+	deviceID, err := readHexInt(fmt.Sprintf(accelSysfsFormat, devNum, deviceFile))
+	if err != nil {
+		return false, err
+	}
+	if _, ok := tpuV4DeviceIDs[deviceID]; !ok {
+		return false, nil
+	}
+	return true, nil
 }
 
-// tpuV5DeviceMinor returns the v5e TPU device minor number for te given path.
+// tpuV5DeviceValid returns the v5e TPU device minor number for te given path.
 // A valid v5 TPU device is defined as:
 // * Path is /dev/vfio/#.
 // * Vendor is googleVendorID.
 // * Device ID is one of tpuV5DeviceIDs.
-func tpuV5DeviceMinor(path string) (uint32, bool, error) {
-	return tpuDeviceMinor(path, vfioDevicePathRegex, vfioSysfsFormat, tpuV5DeviceIDs)
+func tpuV5DeviceValid(devPath string) (bool, error) {
+	paths, err := filepath.Glob(fmt.Sprintf(iommuGroupSysfsGlobFormat, path.Base(devPath)))
+	if err != nil {
+		return false, err
+	}
+	if len(paths) != 1 {
+		return false, fmt.Errorf("found %d paths for %s", len(paths), devPath)
+	}
+	sysfsPath := paths[0]
+	vendor, err := readHexInt(path.Join(sysfsPath, vendorFile))
+	if err != nil {
+		return false, err
+	}
+	if vendor != googleVendorID {
+		return false, nil
+	}
+	deviceID, err := readHexInt(path.Join(sysfsPath, deviceFile))
+	if err != nil {
+		return false, err
+	}
+	if _, ok := tpuV5DeviceIDs[deviceID]; !ok {
+		return false, nil
+	}
+	return true, nil
 }
 
 func readHexInt(path string) (uint64, error) {