From 1f4299ee3fa343ca2bb4892fe306cc21b9a9b308 Mon Sep 17 00:00:00 2001 From: Jamie Liu Date: Mon, 9 Sep 2024 16:54:55 -0700 Subject: [PATCH] nvproxy: implement frontendFDMemmapFile.MapInternal() New test, before this CL: ``` Testing read/write syscalls on cudaMallocHost memory cuda_malloc: write: Bad address ``` After this CL: ``` Testing read/write syscalls on cudaMallocHost memory 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 #10879 PiperOrigin-RevId: 672721411 --- ...{cuda_malloc_managed.cu => cuda_malloc.cu} | 104 ++++++++++++++++++ images/gpu/cuda-tests/run_smoke.sh | 4 +- nogo.yaml | 1 + pkg/sentry/devices/nvproxy/BUILD | 9 ++ pkg/sentry/devices/nvproxy/frontend.go | 22 +++- pkg/sentry/devices/nvproxy/frontend_mmap.go | 10 -- .../devices/nvproxy/frontend_mmap_unsafe.go | 65 +++++++++++ 7 files changed, 200 insertions(+), 15 deletions(-) rename images/gpu/cuda-tests/{cuda_malloc_managed.cu => cuda_malloc.cu} (64%) create mode 100644 pkg/sentry/devices/nvproxy/frontend_mmap_unsafe.go diff --git a/images/gpu/cuda-tests/cuda_malloc_managed.cu b/images/gpu/cuda-tests/cuda_malloc.cu similarity index 64% rename from images/gpu/cuda-tests/cuda_malloc_managed.cu rename to images/gpu/cuda-tests/cuda_malloc.cu index e8119e91f..ce2ba4cd0 100644 --- a/images/gpu/cuda-tests/cuda_malloc_managed.cu +++ b/images/gpu/cuda-tests/cuda_malloc.cu @@ -29,6 +29,106 @@ __global__ void addKernel(std::uint32_t* data) { data[index] += static_cast(index); } +void TestMallocHostReadWrite(int device) { + constexpr size_t kNumBlocks = 32; + constexpr size_t kNumThreads = 64; + constexpr size_t kNumElems = kNumBlocks * kNumThreads; + + constexpr size_t kNumBytes = kNumElems * sizeof(std::uint32_t); + std::uint32_t* cpu_data = nullptr; + CHECK_CUDA(cudaMallocHost(&cpu_data, kNumBytes, cudaHostAllocWriteCombined)); + std::uint32_t* gpu_data = nullptr; + CHECK_CUDA(cudaMalloc(&gpu_data, kNumBytes)); + + // Initialize all elements in the host array with a random value. + std::random_device rd; + const std::uint32_t init_val = + std::uniform_int_distribution()(rd); + for (size_t i = 0; i < kNumElems; i++) { + cpu_data[i] = init_val; + } + + // Write the host array's contents to a temporary file. + char filename[] = "/tmp/cudaMallocHostTest.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(cpu_data) + done, + kNumBytes - done); + if (n >= 0) { + done += n; + } else if (n < 0 && errno != EINTR) { + err(1, "write"); + } + } + + // Copy the array to the device, mutate it there, and copy it back. + CHECK_CUDA(cudaMemcpy(gpu_data, cpu_data, kNumBytes, cudaMemcpyHostToDevice)); + addKernel<<>>(gpu_data); + CHECK_CUDA(cudaDeviceSynchronize()); + CHECK_CUDA(cudaMemcpy(cpu_data, gpu_data, kNumBytes, cudaMemcpyDeviceToHost)); + + // Check that the array has the expected result. + for (size_t i = 0; i < kNumElems; i++) { + std::uint32_t want = init_val + static_cast(i); + if (cpu_data[i] != want) { + std::cout << "cpu_data[" << i << "]: got " << cpu_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(cpu_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 (cpu_data[i] != want) { + std::cout << "cpu_data[" << i << "]: got " << cpu_data[i] << ", wanted " + << want << " = " << init_val << " + " << i << std::endl; + abort(); + } + } + + // Mutate the array on the device again. + CHECK_CUDA(cudaMemcpy(gpu_data, cpu_data, kNumBytes, cudaMemcpyHostToDevice)); + addKernel<<>>(gpu_data); + CHECK_CUDA(cudaDeviceSynchronize()); + CHECK_CUDA(cudaMemcpy(cpu_data, gpu_data, kNumBytes, cudaMemcpyDeviceToHost)); + + // 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(i); + if (cpu_data[i] != want) { + std::cout << "cpu_data[" << i << "]: got " << cpu_data[i] << ", wanted " + << want << " = " << init_val << " + " << i << std::endl; + abort(); + } + } + + close(fd); + CHECK_CUDA(cudaFreeHost(cpu_data)); + CHECK_CUDA(cudaFree(gpu_data)); +} + void TestMallocManagedRoundTrip(int device, unsigned int malloc_flags, bool prefetch) { constexpr size_t kNumBlocks = 32; @@ -169,6 +269,10 @@ int main() { int device; CHECK_CUDA(cudaGetDevice(&device)); + std::cout << "Testing read/write syscalls on cudaMallocHost memory" + << std::endl; + TestMallocHostReadWrite(device); + std::cout << "Testing cudaMallocManaged(flags=cudaMemAttachGlobal)" << std::endl; TestMallocManagedRoundTrip(device, cudaMemAttachGlobal, false); diff --git a/images/gpu/cuda-tests/run_smoke.sh b/images/gpu/cuda-tests/run_smoke.sh index 4a17dadf6..bedca91bf 100755 --- a/images/gpu/cuda-tests/run_smoke.sh +++ b/images/gpu/cuda-tests/run_smoke.sh @@ -17,5 +17,5 @@ set -eux cd / -nvcc cuda_malloc_managed.cu -o cuda_malloc_managed -./cuda_malloc_managed +nvcc cuda_malloc.cu -o cuda_malloc +./cuda_malloc diff --git a/nogo.yaml b/nogo.yaml index b996883e7..44fc46e7d 100644 --- a/nogo.yaml +++ b/nogo.yaml @@ -190,6 +190,7 @@ analyzers: - "pkg/flipcall/.*_unsafe.go" # Special case. - pkg/gohacks/noescape_unsafe.go # Special case. - pkg/ring0/pagetables/allocator_unsafe.go # Special case. + - pkg/sentry/devices/nvproxy/frontend_mmap_unsafe.go # Special case. - pkg/sentry/fsutil/host_file_mapper_unsafe.go # Special case. - pkg/sentry/pgalloc/pgalloc_unsafe.go # Special case. - pkg/sentry/platform/kvm/bluepill_unsafe.go # Special case. diff --git a/pkg/sentry/devices/nvproxy/BUILD b/pkg/sentry/devices/nvproxy/BUILD index c1841e256..bc4bc3346 100644 --- a/pkg/sentry/devices/nvproxy/BUILD +++ b/pkg/sentry/devices/nvproxy/BUILD @@ -13,6 +13,13 @@ declare_mutex( prefix = "fds", ) +declare_mutex( + name = "frontend_mmap_mutex", + out = "frontend_mmap_mutex.go", + package = "nvproxy", + prefix = "frontendMmap", +) + declare_mutex( name = "objs_mutex", out = "objs_mutex.go", @@ -38,6 +45,8 @@ go_library( "fds_mutex.go", "frontend.go", "frontend_mmap.go", + "frontend_mmap_mutex.go", + "frontend_mmap_unsafe.go", "frontend_unsafe.go", "nvproxy.go", "nvproxy_unsafe.go", diff --git a/pkg/sentry/devices/nvproxy/frontend.go b/pkg/sentry/devices/nvproxy/frontend.go index 17de99d3b..d02b7ca61 100644 --- a/pkg/sentry/devices/nvproxy/frontend.go +++ b/pkg/sentry/devices/nvproxy/frontend.go @@ -128,7 +128,13 @@ type frontendFD struct { cachedEvents atomicbitops.Uint64 appQueue waiter.Queue - haveMmapContext atomicbitops.Bool `state:"nosave"` + // mmapMu protects the following fields. + mmapMu frontendMmapMutex `state:"nosave"` + // These fields are marked nosave since we do not automatically reinvoke + // NV_ESC_RM_MAP_MEMORY after restore, so restored FDs have no + // mmap_context. + mmapLength uint64 `state:"nosave"` + mmapInternal uintptr `state:"nosave"` // clients are handles of clients owned by this frontendFD. clients is // protected by dev.nvp.objsMu. @@ -137,6 +143,12 @@ type frontendFD struct { // Release implements vfs.FileDescriptionImpl.Release. func (fd *frontendFD) Release(ctx context.Context) { + fd.mmapMu.Lock() + if fd.mmapInternal != 0 { + unix.RawSyscall(unix.SYS_MUNMAP, fd.mmapInternal, uintptr(fd.mmapLength), 0) + } + fd.mmapMu.Unlock() + fdnotifier.RemoveFD(fd.hostFD) fd.appQueue.Notify(waiter.EventHUp) @@ -1020,17 +1032,21 @@ func rmMapMemory(fi *frontendIoctlState) (uintptr, error) { if !ok { return 0, linuxerr.EINVAL } - if mapFile.haveMmapContext.Load() || !mapFile.haveMmapContext.CompareAndSwap(false, true) { + + mapFile.mmapMu.Lock() + defer mapFile.mmapMu.Unlock() + if mapFile.mmapLength != 0 { fi.ctx.Warningf("nvproxy: attempted to reuse FD %d for NV_ESC_RM_MAP_MEMORY", ioctlParams.FD) return 0, linuxerr.EINVAL } + origFD := ioctlParams.FD ioctlParams.FD = mapFile.hostFD - n, err := frontendIoctlInvoke(fi, &ioctlParams) if err != nil { return n, err } + mapFile.mmapLength = ioctlParams.Params.Length ioctlParams.FD = origFD if _, err := ioctlParams.CopyOut(fi.t, fi.ioctlParamsAddr); err != nil { diff --git a/pkg/sentry/devices/nvproxy/frontend_mmap.go b/pkg/sentry/devices/nvproxy/frontend_mmap.go index 13dcaa21c..068d1c83b 100644 --- a/pkg/sentry/devices/nvproxy/frontend_mmap.go +++ b/pkg/sentry/devices/nvproxy/frontend_mmap.go @@ -16,10 +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" ) @@ -75,13 +72,6 @@ func (mf *frontendFDMemmapFile) IncRef(fr memmap.FileRange, memCgID uint32) { func (mf *frontendFDMemmapFile) DecRef(fr memmap.FileRange) { } -// MapInternal implements memmap.File.MapInternal. -func (mf *frontendFDMemmapFile) MapInternal(fr memmap.FileRange, at hostarch.AccessType) (safemem.BlockSeq, error) { - // FIXME(jamieliu): determine if this is safe - log.Traceback("nvproxy: rejecting frontendFDMemmapFile.MapInternal") - return safemem.BlockSeq{}, linuxerr.EINVAL -} - // FD implements memmap.File.FD. func (mf *frontendFDMemmapFile) FD() int { return int(mf.fd.hostFD) diff --git a/pkg/sentry/devices/nvproxy/frontend_mmap_unsafe.go b/pkg/sentry/devices/nvproxy/frontend_mmap_unsafe.go new file mode 100644 index 000000000..bc66b998f --- /dev/null +++ b/pkg/sentry/devices/nvproxy/frontend_mmap_unsafe.go @@ -0,0 +1,65 @@ +// 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 nvproxy + +import ( + "unsafe" + + "golang.org/x/sys/unix" + "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" +) + +// MapInternal implements memmap.File.MapInternal. +func (mf *frontendFDMemmapFile) MapInternal(fr memmap.FileRange, at hostarch.AccessType) (safemem.BlockSeq, error) { + if at.Execute { + return safemem.BlockSeq{}, linuxerr.EACCES + } + + mf.fd.mmapMu.Lock() + defer mf.fd.mmapMu.Unlock() + if mf.fd.mmapInternal == 0 { + if mf.fd.mmapLength == 0 { + // This shouldn't be possible. + log.Traceback("nvproxy: frontendFDMemmapFile.MapInternal() called before NV_ESC_RM_MAP_MEMORY") + return safemem.BlockSeq{}, linuxerr.EINVAL + } + // Nvidia kernel driver: + // kernel-open/nvidia/nv-mmap.c:nvidia_mmap_helper() requires vm_pgoff + // == 0 (so we must pass offset 0 here), and conditionally requires + // NV_VMA_SIZE(vma) == mmap_context->mmap_size (so we pass length + // mmapLength here). + m, _, errno := unix.Syscall6(unix.SYS_MMAP, 0 /* addr */, uintptr(mf.fd.mmapLength), unix.PROT_READ|unix.PROT_WRITE, unix.MAP_SHARED, uintptr(mf.fd.hostFD), 0 /* offset */) + if errno != 0 { + return safemem.BlockSeq{}, errno + } + mf.fd.mmapInternal = m + } + mappedFR := memmap.FileRange{0, mf.fd.mmapLength} + if !mappedFR.IsSupersetOf(fr) { + return safemem.BlockSeq{}, linuxerr.EINVAL + } + // mmap_context::prot is determined internally during NV_ESC_RM_MAP_MEMORY + // (see + // src/nvidia/arch/nvalloc/unix/src/osapi.c:RmCreateMmapContextLocked()); + // nvidia_mmap_helper() propagates this to vm_area_struct::vm_page_prot, so + // PROT_WRITE on a read-only mapping will succeed at mmap time but fault at + // write time. Thus, these mappings should use safecopy (i.e. + // BlockFromUnsafePointer rather than BlockFromSafePointer). + return safemem.BlockSeqOf(safemem.BlockFromUnsafePointer(unsafe.Pointer(mf.fd.mmapInternal+uintptr(fr.Start)), int(fr.Length()))), nil +}