nvproxy: allow sentry MMIO on nvidia-uvm mappings via buffered reads/writes

New test, before this change:

```
Testing cudaMallocManaged(flags=cudaMemAttachGlobal)
Testing cudaMallocManaged(flags=cudaMemAttachGlobal) with prefetching
Testing cudaMallocManaged(flags=cudaMemAttachHost)
Testing cudaMallocManaged(flags=cudaMemAttachHost) with prefetching
Testing read/write syscalls on cudaMallocManaged memory
cuda_malloc_managed: write: Bad address
```

After this change:

```
Testing cudaMallocManaged(flags=cudaMemAttachGlobal)
Testing cudaMallocManaged(flags=cudaMemAttachGlobal) with prefetching
Testing cudaMallocManaged(flags=cudaMemAttachHost)
Testing cudaMallocManaged(flags=cudaMemAttachHost) with prefetching
Testing read/write syscalls on cudaMallocManaged memory
All tests passed
```

Fixes #10331

PiperOrigin-RevId: 629859064

images/gpu/cuda-tests/cuda_malloc_managed.cu

@@ -13,6 +13,10 @@
 // limitations under the License.
 
 #include <cuda_runtime.h>
+#include <err.h>
+#include <errno.h>
+#include <stdlib.h>
+#include <unistd.h>
 
 #include <cstdint>
 #include <iostream>
@@ -68,6 +72,99 @@ void TestMallocManagedRoundTrip(int device, unsigned int malloc_flags,
   CHECK_CUDA(cudaFree(data));
 }
 
+void TestMallocManagedReadWrite(int device) {
+  constexpr size_t kNumBlocks = 32;
+  constexpr size_t kNumThreads = 64;
+  constexpr size_t kNumElems = kNumBlocks * kNumThreads;
+
+  std::uint32_t* data = nullptr;
+  constexpr size_t kNumBytes = kNumElems * sizeof(*data);
+  CHECK_CUDA(cudaMallocManaged(&data, kNumBytes, cudaMemAttachGlobal));
+
+  // Initialize all elements in the array with a random value on the host.
+  std::random_device rd;
+  const std::uint32_t init_val =
+      std::uniform_int_distribution<std::uint32_t>()(rd);
+  for (size_t i = 0; i < kNumElems; i++) {
+    data[i] = init_val;
+  }
+
+  // Write the array's contents to a temporary file.
+  char filename[] = "/tmp/cudaMallocManagedTest.XXXXXX";
+  int fd = mkstemp(filename);
+  if (fd < 0) {
+    err(1, "mkstemp");
+  }
+  size_t done = 0;
+  while (done < kNumBytes) {
+    ssize_t n = write(fd, reinterpret_cast<char*>(data) + done,
+                      kNumBytes - done);
+    if (n >= 0) {
+      done += n;
+    } else if (n < 0 && errno != EINTR) {
+      err(1, "write");
+    }
+  }
+
+  // Mutate the array on the device.
+  addKernel<<<kNumBlocks, kNumThreads>>>(data);
+  CHECK_CUDA(cudaDeviceSynchronize());
+
+  // Check that the array has the expected result.
+  for (size_t i = 0; i < kNumElems; i++) {
+    std::uint32_t want = init_val + static_cast<std::uint32_t>(i);
+    if (data[i] != want) {
+      std::cout << "data[" << i << "]: got " << data[i] << ", wanted " << want
+                << " = " << init_val << " + " << i << std::endl;
+      abort();
+    }
+  }
+
+  // Read the array's original contents back from the temporary file.
+  if (lseek(fd, 0, SEEK_SET) < 0) {
+    err(1, "lseek");
+  }
+  done = 0;
+  while (done < kNumBytes) {
+    ssize_t n = read(fd, reinterpret_cast<char*>(data) + done,
+                     kNumBytes - done);
+    if (n > 0) {
+      done += n;
+    } else if (n == 0) {
+      errx(1, "read: unexpected EOF after %zu bytes", done);
+    } else if (n < 0 && errno != EINTR) {
+      err(1, "read");
+    }
+  }
+
+  // Check that the array matches what we originally wrote.
+  for (size_t i = 0; i < kNumElems; i++) {
+    std::uint32_t want = init_val;
+    if (data[i] != want) {
+      std::cout << "data[" << i << "]: got " << data[i] << ", wanted " << want
+                << " = " << init_val << " + " << i << std::endl;
+      abort();
+    }
+  }
+
+  // Mutate the array on the device again.
+  addKernel<<<kNumBlocks, kNumThreads>>>(data);
+  CHECK_CUDA(cudaDeviceSynchronize());
+
+  // Check that the array has the expected result again.
+  for (size_t i = 0; i < kNumElems; i++) {
+    std::uint32_t want = init_val + static_cast<std::uint32_t>(i);
+    if (data[i] != want) {
+      std::cout << "data[" << i << "]: got " << data[i] << ", wanted " << want
+                << " = " << init_val << " + " << i << std::endl;
+      abort();
+    }
+  }
+
+  close(fd);
+  CHECK_CUDA(cudaFree(data));
+}
+
 int main() {
   int device;
   CHECK_CUDA(cudaGetDevice(&device));
@@ -96,6 +193,10 @@ int main() {
     TestMallocManagedRoundTrip(device, cudaMemAttachHost, true);
   }
 
+  std::cout << "Testing read/write syscalls on cudaMallocManaged memory"
+            << std::endl;
+  TestMallocManagedReadWrite(device);
+
   std::cout << "All tests passed" << std::endl;
   return 0;
 }

pkg/abi/nvgpu/uvm.go

@@ -42,6 +42,8 @@ const (
 	UVM_PAGEABLE_MEM_ACCESS            = 39
 	UVM_SET_PREFERRED_LOCATION         = 42
 	UVM_DISABLE_READ_DUPLICATION       = 45
+	UVM_TOOLS_READ_PROCESS_MEMORY      = 62
+	UVM_TOOLS_WRITE_PROCESS_MEMORY     = 63
 	UVM_MAP_DYNAMIC_PARALLELISM_REGION = 65
 	UVM_ALLOC_SEMAPHORE_POOL           = 68
 	UVM_VALIDATE_VA_RANGE              = 72
@@ -236,6 +238,26 @@ type UVM_DISABLE_READ_DUPLICATION_PARAMS struct {
 	Pad0          [4]byte
 }
 
+// +marshal
+type UVM_TOOLS_READ_PROCESS_MEMORY_PARAMS struct {
+	Buffer    uint64
+	Size      uint64
+	TargetVA  uint64
+	BytesRead uint64
+	RMStatus  uint32
+	Pad0      [4]byte
+}
+
+// +marshal
+type UVM_TOOLS_WRITE_PROCESS_MEMORY_PARAMS struct {
+	Buffer       uint64
+	Size         uint64
+	TargetVA     uint64
+	BytesWritten uint64
+	RMStatus     uint32
+	Pad0         [4]byte
+}
+
 // +marshal
 type UVM_MAP_DYNAMIC_PARALLELISM_REGION_PARAMS struct {
 	Base     uint64

pkg/sentry/devices/nvproxy/seccomp_filters.go

@@ -162,6 +162,14 @@ func Filters() seccomp.SyscallRules {
 				seccomp.NonNegativeFD{},
 				seccomp.EqualTo(nvgpu.UVM_DISABLE_READ_DUPLICATION),
 			},
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(nvgpu.UVM_TOOLS_READ_PROCESS_MEMORY),
+			},
+			seccomp.PerArg{
+				seccomp.NonNegativeFD{},
+				seccomp.EqualTo(nvgpu.UVM_TOOLS_WRITE_PROCESS_MEMORY),
+			},
 			seccomp.PerArg{
 				seccomp.NonNegativeFD{},
 				seccomp.EqualTo(nvgpu.UVM_MAP_DYNAMIC_PARALLELISM_REGION),

pkg/sentry/devices/nvproxy/uvm_mmap.go

@@ -16,9 +16,7 @@ package nvproxy
 
 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/memmap"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
@@ -69,8 +67,6 @@ func (fd *uvmFD) InvalidateUnsavable(ctx context.Context) error {
 }
 
 type uvmFDMemmapFile struct {
-	memmap.NoBufferedIOFallback
-
 	fd *uvmFD
 }
 
@@ -85,8 +81,7 @@ func (mf *uvmFDMemmapFile) DecRef(fr memmap.FileRange) {
 // MapInternal implements memmap.File.MapInternal.
 func (mf *uvmFDMemmapFile) MapInternal(fr memmap.FileRange, at hostarch.AccessType) (safemem.BlockSeq, error) {
 	// TODO(jamieliu): make an attempt with MAP_FIXED_NOREPLACE?
-	log.Traceback("nvproxy: rejecting uvmFDMemmapFile.MapInternal")
-	return safemem.BlockSeq{}, linuxerr.EINVAL
+	return safemem.BlockSeq{}, memmap.BufferedIOFallbackErr{}
 }
 
 // FD implements memmap.File.FD.

pkg/sentry/devices/nvproxy/uvm_unsafe.go

@@ -15,9 +15,13 @@
 package nvproxy
 
 import (
+	"runtime"
 	"unsafe"
 
 	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/abi/nvgpu"
+	"gvisor.dev/gvisor/pkg/errors/linuxerr"
+	"gvisor.dev/gvisor/pkg/log"
 )
 
 func uvmIoctlInvoke[Params any](ui *uvmIoctlState, ioctlParams *Params) (uintptr, error) {
@@ -27,3 +31,59 @@ func uvmIoctlInvoke[Params any](ui *uvmIoctlState, ioctlParams *Params) (uintptr
 	}
 	return n, nil
 }
+
+// BufferReadAt implements memmap.File.BufferReadAt.
+func (mf *uvmFDMemmapFile) BufferReadAt(off uint64, dst []byte) (uint64, error) {
+	// kernel-open/nvidia-uvm/uvm.c:uvm_fops.{read,read_iter,splice_read} ==
+	// NULL, so UVM data can only be read via ioctl.
+	if len(dst) == 0 {
+		return 0, nil
+	}
+	defer runtime.KeepAlive(dst)
+	params := nvgpu.UVM_TOOLS_READ_PROCESS_MEMORY_PARAMS{
+		Buffer:   uint64(uintptr(unsafe.Pointer(&dst[0]))),
+		Size:     uint64(len(dst)),
+		TargetVA: off,
+	}
+	_, _, errno := unix.Syscall(unix.SYS_IOCTL, uintptr(mf.fd.hostFD), nvgpu.UVM_TOOLS_READ_PROCESS_MEMORY, uintptr(unsafe.Pointer(&params)))
+	if errno != 0 {
+		return 0, errno
+	}
+	if params.RMStatus != nvgpu.NV_OK {
+		log.Warningf("nvproxy: UVM_TOOLS_READ_PROCESS_MEMORY(targetVa=%#x, len=%d) returned status %d", off, len(dst), params.RMStatus)
+		return params.BytesRead, linuxerr.EINVAL
+	}
+	if params.BytesRead != uint64(len(dst)) {
+		log.Warningf("nvproxy: UVM_TOOLS_READ_PROCESS_MEMORY(targetVa=%#x, len=%d) returned %d bytes", off, len(dst), params.BytesRead)
+		return params.BytesRead, linuxerr.EINVAL
+	}
+	return params.BytesRead, nil
+}
+
+// BufferWriteAt implements memmap.File.BufferWriteAt.
+func (mf *uvmFDMemmapFile) BufferWriteAt(off uint64, src []byte) (uint64, error) {
+	// kernel-open/nvidia-uvm/uvm.c:uvm_fops.{write,write_iter,splice_write} ==
+	// NULL, so UVM data can only be written via ioctl.
+	if len(src) == 0 {
+		return 0, nil
+	}
+	defer runtime.KeepAlive(src)
+	params := nvgpu.UVM_TOOLS_WRITE_PROCESS_MEMORY_PARAMS{
+		Buffer:   uint64(uintptr(unsafe.Pointer(&src[0]))),
+		Size:     uint64(len(src)),
+		TargetVA: off,
+	}
+	_, _, errno := unix.Syscall(unix.SYS_IOCTL, uintptr(mf.fd.hostFD), nvgpu.UVM_TOOLS_WRITE_PROCESS_MEMORY, uintptr(unsafe.Pointer(&params)))
+	if errno != 0 {
+		return 0, errno
+	}
+	if params.RMStatus != nvgpu.NV_OK {
+		log.Warningf("nvproxy: UVM_TOOLS_WRITE_PROCESS_MEMORY(targetVa=%#x, len=%d) returned status %d", off, len(src), params.RMStatus)
+		return params.BytesWritten, linuxerr.EINVAL
+	}
+	if params.BytesWritten != uint64(len(src)) {
+		log.Warningf("nvproxy: UVM_TOOLS_WRITE_PROCESS_MEMORY(targetVa=%#x, len=%d) returned %d bytes", off, len(src), params.BytesWritten)
+		return params.BytesWritten, linuxerr.EINVAL
+	}
+	return params.BytesWritten, nil
+}