From 6b537e43eb09c174ca4bb6af8f580b3c143c8447 Mon Sep 17 00:00:00 2001
From: Ayush Ranjan <ayushranjan@google.com>
Date: Wed, 5 Jun 2024 20:13:57 -0700
Subject: [PATCH] Add NvidiaDeviceFD interface to nvproxy package.

This can be helpful for identifying Nvidia device FDs.

PiperOrigin-RevId: 640743124
---
 pkg/sentry/devices/nvproxy/frontend.go | 3 +++
 pkg/sentry/devices/nvproxy/nvproxy.go  | 6 ++++++
 pkg/sentry/devices/nvproxy/uvm.go      | 3 +++
 3 files changed, 12 insertions(+)

diff --git a/pkg/sentry/devices/nvproxy/frontend.go b/pkg/sentry/devices/nvproxy/frontend.go
index 05649bc87..5eaa76182 100644
--- a/pkg/sentry/devices/nvproxy/frontend.go
+++ b/pkg/sentry/devices/nvproxy/frontend.go
@@ -202,6 +202,9 @@ func (fd *frontendFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr,
 	return handler(&fi)
 }
 
+// IsNvidiaDeviceFD implements NvidiaDeviceFD.IsNvidiaDeviceFD.
+func (fd *frontendFD) IsNvidiaDeviceFD() {}
+
 func frontendIoctlCmd(nr, argSize uint32) uintptr {
 	return uintptr(linux.IOWR(nvgpu.NV_IOCTL_MAGIC, nr, argSize))
 }
diff --git a/pkg/sentry/devices/nvproxy/nvproxy.go b/pkg/sentry/devices/nvproxy/nvproxy.go
index 7c12dd0c2..fa5066e8d 100644
--- a/pkg/sentry/devices/nvproxy/nvproxy.go
+++ b/pkg/sentry/devices/nvproxy/nvproxy.go
@@ -100,3 +100,9 @@ type hasFrontendFDPtr[T any] interface {
 	marshalPtr[T]
 	nvgpu.HasFrontendFD
 }
+
+// NvidiaDeviceFD is an interface that should be implemented by all
+// vfs.FileDescriptionImpl of Nvidia devices.
+type NvidiaDeviceFD interface {
+	IsNvidiaDeviceFD()
+}
diff --git a/pkg/sentry/devices/nvproxy/uvm.go b/pkg/sentry/devices/nvproxy/uvm.go
index 3a071294f..a3f60855c 100644
--- a/pkg/sentry/devices/nvproxy/uvm.go
+++ b/pkg/sentry/devices/nvproxy/uvm.go
@@ -151,6 +151,9 @@ func (fd *uvmFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args
 	return handler(&ui)
 }
 
+// IsNvidiaDeviceFD implements NvidiaDeviceFD.IsNvidiaDeviceFD.
+func (fd *uvmFD) IsNvidiaDeviceFD() {}
+
 // uvmIoctlState holds the state of a call to uvmFD.Ioctl().
 type uvmIoctlState struct {
 	fd              *uvmFD