From 1b8d53a6147e0530c318119d38841b576e627b48 Mon Sep 17 00:00:00 2001
From: Ayush Ranjan <ayushranjan@google.com>
Date: Tue, 13 Feb 2024 15:01:12 -0800
Subject: [PATCH] Fix NV_CHANNEL_ALLOC_PARAMS for 535.86.05+ drivers.

Updates #9846.

PiperOrigin-RevId: 606758290
---
 pkg/abi/nvgpu/classes.go              | 11 +++++++++++
 pkg/abi/nvgpu/nvgpu.go                |  6 ++++++
 pkg/sentry/devices/nvproxy/version.go |  9 ++++++++-
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/pkg/abi/nvgpu/classes.go b/pkg/abi/nvgpu/classes.go
index 304a37c65..60a0781c4 100644
--- a/pkg/abi/nvgpu/classes.go
+++ b/pkg/abi/nvgpu/classes.go
@@ -218,6 +218,17 @@ type NV_CHANNEL_ALLOC_PARAMS struct {
 	SubProcessID        uint32
 }
 
+// NV_CHANNEL_ALLOC_PARAMS_V535 is the updated version of
+// NV_CHANNEL_ALLOC_PARAMS since 535.86.05.
+//
+// +marshal
+type NV_CHANNEL_ALLOC_PARAMS_V535 struct {
+	NV_CHANNEL_ALLOC_PARAMS
+	EncryptIv [CC_CHAN_ALLOC_IV_SIZE_DWORD]uint32
+	DecryptIv [CC_CHAN_ALLOC_IV_SIZE_DWORD]uint32
+	HmacNonce [CC_CHAN_ALLOC_NONCE_SIZE_DWORD]uint32
+}
+
 // NVB0B5_ALLOCATION_PARAMETERS is the alloc param type for TURING_DMA_COPY_A,
 // AMPERE_DMA_COPY_A, and AMPERE_DMA_COPY_B from
 // src/common/sdk/nvidia/inc/class/clb0b5sw.h.
diff --git a/pkg/abi/nvgpu/nvgpu.go b/pkg/abi/nvgpu/nvgpu.go
index cb8b758d2..4a5a2da1d 100644
--- a/pkg/abi/nvgpu/nvgpu.go
+++ b/pkg/abi/nvgpu/nvgpu.go
@@ -41,6 +41,12 @@ const (
 	NV_MAX_SUBDEVICES = 8
 )
 
+// From src/common/sdk/nvidia/inc/alloc/alloc_channel.h.
+const (
+	CC_CHAN_ALLOC_IV_SIZE_DWORD    = 3
+	CC_CHAN_ALLOC_NONCE_SIZE_DWORD = 8
+)
+
 // RS_ACCESS_MASK is RS_ACCESS_MASK, from
 // src/common/sdk/nvidia/inc/rs_access.h.
 //
diff --git a/pkg/sentry/devices/nvproxy/version.go b/pkg/sentry/devices/nvproxy/version.go
index c4e138688..a1ba6aa46 100644
--- a/pkg/sentry/devices/nvproxy/version.go
+++ b/pkg/sentry/devices/nvproxy/version.go
@@ -333,8 +333,15 @@ func Init() {
 
 		v535_54_03Checksum := "454764f57ea1b9e19166a370f78be10e71f0626438fb197f726dc3caf05b4082"
 		v535_54_03 := addDriverABI(535, 54, 03, v535_54_03Checksum, v535_43_02)
+		// 535.86.05 is an intermediate unqualified version from the main branch.
+		v535_86_05 := func() *driverABI {
+			abi := v535_54_03()
+			abi.allocationClass[nvgpu.TURING_CHANNEL_GPFIFO_A] = rmAllocSimple[nvgpu.NV_CHANNEL_ALLOC_PARAMS_V535]
+			abi.allocationClass[nvgpu.AMPERE_CHANNEL_GPFIFO_A] = rmAllocSimple[nvgpu.NV_CHANNEL_ALLOC_PARAMS_V535]
+			return abi
+		}
 		v535_104_05Checksum := "2f9d609d1da770beee757636635c46e7ed8253ade887b87c7a5482e33fcbedc9"
-		v535_104_05 := addDriverABI(535, 104, 05, v535_104_05Checksum, v535_54_03)
+		v535_104_05 := addDriverABI(535, 104, 05, v535_104_05Checksum, v535_86_05)
 
 		// 535.104.12 does not exist on the main branch. It branched off the main
 		// branch at 535.104.05.