Create a precise host file mapper that respects host file ranges.

This is required to get tpuproxy to work properly on the KVM platform.

PiperOrigin-RevId: 676120793

pkg/sentry/devices/tpuproxy/vfio/BUILD

@@ -35,6 +35,7 @@ go_library(
         "//pkg/sentry/arch",
         "//pkg/sentry/devices/tpuproxy/util",
         "//pkg/sentry/fsimpl/eventfd",
+        "//pkg/sentry/fsutil",
         "//pkg/sentry/kernel",
         "//pkg/sentry/memmap",
         "//pkg/sentry/mm",

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

@@ -28,7 +28,9 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/util"
 	"gvisor.dev/gvisor/pkg/sentry/fsimpl/eventfd"
 	"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"
 )
@@ -40,8 +42,11 @@ type pciDeviceFD struct {
 	vfs.DentryMetadataFileDescriptionImpl
 	vfs.NoLockFD
 
-	hostFD     int32
-	queue      waiter.Queue
+	hostFD int32
+	queue  waiter.Queue
+
+	mapsMu     sync.Mutex
+	mappings   memmap.MappingSet
 	memmapFile pciDeviceFdMemmapFile
 }
 

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

@@ -16,10 +16,9 @@ package vfio
 
 import (
 	"gvisor.dev/gvisor/pkg/context"
-	"gvisor.dev/gvisor/pkg/errors/linuxerr"
 	"gvisor.dev/gvisor/pkg/hostarch"
-	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/safemem"
+	"gvisor.dev/gvisor/pkg/sentry/fsutil"
 	"gvisor.dev/gvisor/pkg/sentry/memmap"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 )
@@ -31,16 +30,28 @@ func (fd *pciDeviceFD) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts)
 
 // AddMapping implements memmap.Mappable.AddMapping.
 func (fd *pciDeviceFD) AddMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) error {
+	fd.mapsMu.Lock()
+	mapped := fd.mappings.AddMapping(ms, ar, offset, writable)
+	for _, r := range mapped {
+		fd.memmapFile.pfm.IncRefOn(r)
+	}
+	fd.mapsMu.Unlock()
 	return nil
 }
 
 // RemoveMapping implements memmap.Mappable.RemoveMapping.
 func (fd *pciDeviceFD) RemoveMapping(ctx context.Context, ms memmap.MappingSpace, ar hostarch.AddrRange, offset uint64, writable bool) {
+	fd.mapsMu.Lock()
+	unmapped := fd.mappings.RemoveMapping(ms, ar, offset, writable)
+	for _, r := range unmapped {
+		fd.memmapFile.pfm.DecRefOn(r)
+	}
+	fd.mapsMu.Unlock()
 }
 
 // CopyMapping implements memmap.Mappable.CopyMapping.
 func (fd *pciDeviceFD) CopyMapping(ctx context.Context, ms memmap.MappingSpace, srcAR, dstAR hostarch.AddrRange, offset uint64, writable bool) error {
-	return nil
+	return fd.AddMapping(ctx, ms, dstAR, offset, writable)
 }
 
 // Translate implements memmap.Mappable.Translate.
@@ -63,11 +74,12 @@ func (fd *pciDeviceFD) InvalidateUnsavable(ctx context.Context) error {
 type pciDeviceFdMemmapFile struct {
 	memmap.NoBufferedIOFallback
 
-	fd *pciDeviceFD
+	fd  *pciDeviceFD
+	pfm fsutil.PreciseHostFileMapper
 }
 
 // IncRef implements memmap.File.IncRef.
-func (mf *pciDeviceFdMemmapFile) IncRef(memmap.FileRange, uint32) {
+func (mf *pciDeviceFdMemmapFile) IncRef(fr memmap.FileRange, memCgID uint32) {
 }
 
 // DecRef implements memmap.File.DecRef.
@@ -76,8 +88,7 @@ func (mf *pciDeviceFdMemmapFile) DecRef(fr memmap.FileRange) {
 
 // MapInternal implements memmap.File.MapInternal.
 func (mf *pciDeviceFdMemmapFile) MapInternal(fr memmap.FileRange, at hostarch.AccessType) (safemem.BlockSeq, error) {
-	log.Traceback("tpuproxy: rejecting pciDeviceFdMemmapFile.MapInternal")
-	return safemem.BlockSeq{}, linuxerr.EINVAL
+	return mf.pfm.MapInternal(fr, int(mf.fd.hostFD), at.Write)
 }
 
 // FD implements memmap.File.FD.

pkg/sentry/fsutil/BUILD

@@ -72,6 +72,42 @@ go_template_instance(
     },
 )
 
+go_template_instance(
+    name = "mapping_set",
+    out = "mapping_set.go",
+    imports = {
+        "mm": "gvisor.dev/gvisor/pkg/sentry/mm",
+        "memmap": "gvisor.dev/gvisor/pkg/sentry/memmap",
+    },
+    package = "fsutil",
+    prefix = "mapping",
+    template = "//pkg/segment:generic_set",
+    types = {
+        "Key": "uint64",
+        "Range": "memmap.FileRange",
+        "Value": "mapping",
+        "Functions": "mappingSetFuncs",
+    },
+)
+
+go_template_instance(
+    name = "refs_set",
+    out = "refs_set.go",
+    imports = {
+        "mm": "gvisor.dev/gvisor/pkg/sentry/mm",
+        "memmap": "gvisor.dev/gvisor/pkg/sentry/memmap",
+    },
+    package = "fsutil",
+    prefix = "refs",
+    template = "//pkg/segment:generic_set",
+    types = {
+        "Key": "uint64",
+        "Range": "memmap.FileRange",
+        "Value": "uint64",
+        "Functions": "refsSetFuncs",
+    },
+)
+
 go_library(
     name = "fsutil",
     srcs = [
@@ -85,8 +121,11 @@ go_library(
         "host_file_mapper.go",
         "host_file_mapper_state.go",
         "host_file_mapper_unsafe.go",
+        "mapping_set.go",
         "maps_mutex.go",
+        "precise_host_file_mapper.go",
         "refs_mutex.go",
+        "refs_set.go",
     ],
     visibility = ["//pkg/sentry:internal"],
     deps = [
@@ -99,6 +138,7 @@ go_library(
         "//pkg/sentry/arch",
         "//pkg/sentry/kernel/time",
         "//pkg/sentry/memmap",
+        "//pkg/sentry/mm",
         "//pkg/sentry/pgalloc",
         "//pkg/sentry/usage",
         "//pkg/state",

pkg/sentry/fsutil/host_file_mapper.go

@@ -210,7 +210,7 @@ func (f *HostFileMapper) forEachMappingBlockLocked(fr memmap.FileRange, fd int,
 		if chunkStart+chunkSize > fr.End {
 			endOff = fr.End - chunkStart
 		}
-		fn(f.unsafeBlockFromChunkMapping(m.addr).TakeFirst64(endOff).DropFirst64(startOff))
+		fn(unsafeBlockFromMapping(m.addr, chunkSize).TakeFirst64(endOff).DropFirst64(startOff))
 		chunkStart += chunkSize
 		if chunkStart >= fr.End || chunkStart == 0 {
 			break

pkg/sentry/fsutil/host_file_mapper_unsafe.go

@@ -20,8 +20,8 @@ import (
 	"gvisor.dev/gvisor/pkg/safemem"
 )
 
-func (*HostFileMapper) unsafeBlockFromChunkMapping(addr uintptr) safemem.Block {
+func unsafeBlockFromMapping(addr uintptr, size int) safemem.Block {
 	// We don't control the host file's length, so touching its mappings may
 	// raise SIGBUS. Thus accesses to it must use safecopy.
-	return safemem.BlockFromUnsafePointer((unsafe.Pointer)(addr), chunkSize)
+	return safemem.BlockFromUnsafePointer((unsafe.Pointer)(addr), size)
 }

pkg/sentry/fsutil/precise_host_file_mapper.go

@@ -0,0 +1,284 @@
+// Copyright 2018 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 fsutil
+
+import (
+	"fmt"
+	"syscall"
+
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/safemem"
+	"gvisor.dev/gvisor/pkg/sentry/memmap"
+)
+
+// PreciseHostFileMapper caches mappings of an arbitrary host file descriptor.
+// It is used by implementations of memmap.Mappable that represent a host file
+// descriptor. It differs from HostFileMapper in that it maps at exact page
+// boundaries specified in a file range, not in chunks.
+//
+// +stateify savable
+type PreciseHostFileMapper struct {
+	refsMu refsMutex `state:"nosave"`
+
+	// +checklocks:refsMu
+	refs refsSet
+
+	mapsMu mapsMutex `state:"nosave"`
+
+	// mappings is a set of internal mappings of the device. The value is a
+	// mapping object.
+	//
+	// +checklocks:mapsMu
+	mappings mappingSet
+}
+
+// NewPreciseHostFileMapper returns an initialized PreciseHostFileMapper
+// allocated on the heap with no references or cached mappings.
+func NewPreciseHostFileMapper() *PreciseHostFileMapper {
+	f := &PreciseHostFileMapper{}
+	return f
+}
+
+// IncRefOn increments the reference count on all pages in mr.
+//
+// Preconditions:
+//   - mr.Length() != 0.
+//   - mr.Start and mr.End must be page-aligned.
+func (f *PreciseHostFileMapper) IncRefOn(mr memmap.MappableRange) {
+	f.refsMu.Lock()
+	defer f.refsMu.Unlock()
+	fr := memmap.FileRange{Start: mr.Start, End: mr.End}
+	seg, gap := f.refs.Find(fr.Start)
+	for seg.Ok() || gap.Ok() {
+		if seg.Ok() {
+			seg = f.refs.Isolate(seg, fr)
+			refs := seg.ValuePtr()
+			*refs++
+		} else if gap.Ok() {
+			seg = f.refs.Insert(gap, fr.Intersect(gap.Range()), 1)
+		}
+		if seg.End() >= fr.End {
+			break
+		}
+		fr.Start = seg.End()
+		seg, gap = seg.NextNonEmpty()
+	}
+}
+
+// DecRefOn decrements the reference count on all offsets in mr.
+//
+// Preconditions:
+//   - mr.Length() != 0.
+//   - mr.Start and mr.End must be page-aligned.
+func (f *PreciseHostFileMapper) DecRefOn(mr memmap.MappableRange) {
+	f.refsMu.Lock()
+	defer f.refsMu.Unlock()
+	rseg := f.refs.FindSegment(mr.Start)
+	if !rseg.Ok() {
+		panic(fmt.Sprintf("could not find segment for range %v", mr))
+	}
+	fr := memmap.FileRange{Start: mr.Start, End: mr.End}
+	for fr.Length() > 0 && rseg.Ok() {
+		rseg = f.refs.Isolate(rseg, fr)
+		refs := rseg.ValuePtr()
+		*refs--
+		gap := refsGapIterator{}
+		fr.Start = rseg.End()
+		if *refs == 0 {
+			f.mapsMu.Lock()
+			f.mappings.RemoveRangeWith(rseg.Range(), f.unmapSegmentLocked)
+			f.mapsMu.Unlock()
+			gap = f.refs.RemoveRange(rseg.Range())
+		}
+		if gap.Ok() {
+			rseg = gap.NextSegment()
+		} else {
+			rseg = rseg.NextSegment()
+		}
+	}
+}
+
+// MapInternal returns a mapping of offsets in fr from fd. The returned
+// safemem.BlockSeq is valid as long as at least one reference is held on all
+// offsets in fr or until the next call to UnmapAll.
+//
+// Preconditions: The caller must hold a reference on all offsets in fr.
+func (f *PreciseHostFileMapper) MapInternal(fr memmap.FileRange, fd int, write bool) (safemem.BlockSeq, error) {
+	f.mapsMu.Lock()
+	defer f.mapsMu.Unlock()
+	prot := unix.PROT_READ
+	if write {
+		prot |= unix.PROT_WRITE
+	}
+
+	origFR := fr
+	var blocks []safemem.Block
+	seg, gap := f.mappings.Find(fr.Start)
+	for seg.Ok() || gap.Ok() {
+		if seg.Ok() {
+			block, newSeg, errno := f.mapInternalSegment(&fr, seg, fd, prot, write)
+			if errno != 0 {
+				return safemem.BlockSeq{}, errno
+			}
+			blocks = append(blocks, block)
+			seg, gap = newSeg.NextNonEmpty()
+
+			if fr.Length() == 0 {
+				break
+			}
+		}
+		if gap.Ok() {
+			block, newSeg, errno := f.mapInternalGap(&fr, gap, fd, prot, write)
+			if errno != 0 {
+				return safemem.BlockSeq{}, errno
+			}
+			blocks = append(blocks, block)
+			seg, gap = newSeg.NextSegment(), mappingGapIterator{}
+
+			if fr.Length() == 0 {
+				break
+			}
+		}
+	}
+	if fr.Length() > 0 {
+		return safemem.BlockSeq{}, fmt.Errorf("failed to map range %v", origFR)
+	}
+
+	return safemem.BlockSeqFromSlice(blocks), nil
+}
+
+// +checklocks:f.mapsMu
+func (f *PreciseHostFileMapper) mapInternalSegment(fr *memmap.FileRange, seg mappingIterator, fd int, prot int, write bool) (safemem.Block, mappingIterator, syscall.Errno) {
+	addr := seg.Value().addr + uintptr(fr.Start-seg.Start())
+	if !seg.Value().writable && write {
+		seg = f.mappings.Isolate(seg, *fr)
+		_, _, errno := unix.Syscall6(
+			unix.SYS_MMAP,
+			addr,
+			uintptr(seg.Range().Length()),
+			uintptr(prot),
+			unix.MAP_SHARED|unix.MAP_FIXED,
+			uintptr(fd),
+			uintptr(seg.Start()))
+		if errno != 0 {
+			return safemem.Block{}, seg, errno
+		}
+		seg.ValuePtr().writable = write
+	}
+	mapRange := seg.Range().Intersect(*fr)
+	fr.Start = mapRange.End
+	return unsafeBlockFromMapping(addr, int(mapRange.Length())), seg, 0
+}
+
+// +checklocks:f.mapsMu
+func (f *PreciseHostFileMapper) mapInternalGap(fr *memmap.FileRange, gap mappingGapIterator, fd int, prot int, write bool) (safemem.Block, mappingIterator, syscall.Errno) {
+	newRange := fr.Intersect(gap.Range())
+	addr, _, errno := unix.Syscall6(
+		unix.SYS_MMAP,
+		0,
+		uintptr(newRange.Length()),
+		uintptr(prot),
+		unix.MAP_SHARED,
+		uintptr(fd),
+		uintptr(newRange.Start))
+	if errno != 0 {
+		return safemem.Block{}, mappingIterator{}, errno
+	}
+	fr.Start = newRange.End
+	seg := f.mappings.Insert(gap, newRange, mapping{addr: addr, writable: write})
+	return unsafeBlockFromMapping(addr, int(newRange.Length())), seg, 0
+}
+
+// UnmapAll unmaps all cached mappings. Callers are responsible for
+// synchronization with mappings returned by previous calls to MapInternal.
+func (f *PreciseHostFileMapper) UnmapAll() {
+	f.mapsMu.Lock()
+	defer f.mapsMu.Unlock()
+	for seg := f.mappings.FirstSegment(); seg.Ok(); seg = seg.NextSegment() {
+		f.unmapSegmentLocked(seg)
+	}
+	f.mappings.RemoveAll()
+}
+
+// +checklocks:f.mapsMu
+func (f *PreciseHostFileMapper) unmapSegmentLocked(mseg mappingIterator) {
+	if _, _, errno := unix.Syscall(unix.SYS_MUNMAP, uintptr(mseg.Value().addr), uintptr(mseg.Range().Length()), 0); errno != 0 {
+		// This leaks address space and is unexpected, but is otherwise
+		// harmless, so complain but don't panic.
+		log.Warningf("HostFileMapper: failed to unmap mapping %#x: %v", mseg.Range().Start, errno)
+	}
+}
+
+type refsSetFuncs struct{}
+
+func (refsSetFuncs) MinKey() uint64 {
+	return 0
+}
+
+func (refsSetFuncs) MaxKey() uint64 {
+	return ^uint64(0)
+}
+
+func (refsSetFuncs) ClearValue(val *uint64) {
+	*val = 0
+}
+
+func (refsSetFuncs) Merge(r1 memmap.FileRange, v1 uint64, r2 memmap.FileRange, v2 uint64) (uint64, bool) {
+	return v1, v1 == v2
+}
+
+func (refsSetFuncs) Split(r memmap.FileRange, val uint64, split uint64) (uint64, uint64) {
+	return val, val
+}
+
+type mappingSetFuncs struct{}
+
+func (mappingSetFuncs) MinKey() uint64 {
+	return 0
+}
+
+func (mappingSetFuncs) MaxKey() uint64 {
+	return ^uint64(0)
+}
+
+func (mappingSetFuncs) ClearValue(val *mapping) {
+	*val = mapping{}
+}
+
+func (mappingSetFuncs) Merge(r1 memmap.FileRange, v1 mapping, r2 memmap.FileRange, v2 mapping) (mapping, bool) {
+	// Are we the same writability?
+	if v1.writable != v2.writable {
+		return mapping{}, false
+	}
+
+	// Do we have contiguous offsets in the backing file?
+	if v1.addr+uintptr(r1.Length()) != v2.addr {
+		return mapping{}, false
+	}
+
+	return v1, true
+}
+
+func (mappingSetFuncs) Split(r memmap.FileRange, val mapping, split uint64) (mapping, mapping) {
+	n := split - r.Start
+
+	left := val
+
+	right := val
+	right.addr += uintptr(n)
+
+	return left, right
+}