From 999c708319d5cc135725e638f40563d48bb71ae1 Mon Sep 17 00:00:00 2001
From: Ayush Ranjan <ayushranjan@google.com>
Date: Tue, 7 Nov 2023 11:44:57 -0800
Subject: [PATCH] Add support for restoring nvproxy.driverABI.

Fixes #9649

PiperOrigin-RevId: 580251470
---
 pkg/sentry/devices/nvproxy/BUILD           |  1 +
 pkg/sentry/devices/nvproxy/nvproxy.go      |  4 +++-
 pkg/sentry/devices/nvproxy/save_restore.go | 28 ++++++++++++++++++++++
 pkg/sentry/devices/nvproxy/version.go      |  2 ++
 4 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 pkg/sentry/devices/nvproxy/save_restore.go

diff --git a/pkg/sentry/devices/nvproxy/BUILD b/pkg/sentry/devices/nvproxy/BUILD
index 77de8af22..3a3e3ab22 100644
--- a/pkg/sentry/devices/nvproxy/BUILD
+++ b/pkg/sentry/devices/nvproxy/BUILD
@@ -21,6 +21,7 @@ go_library(
         "nvproxy.go",
         "nvproxy_unsafe.go",
         "objs_mutex.go",
+        "save_restore.go",
         "seccomp_filters.go",
         "uvm.go",
         "uvm_mmap.go",
diff --git a/pkg/sentry/devices/nvproxy/nvproxy.go b/pkg/sentry/devices/nvproxy/nvproxy.go
index 765f581d3..7f7ad6dcb 100644
--- a/pkg/sentry/devices/nvproxy/nvproxy.go
+++ b/pkg/sentry/devices/nvproxy/nvproxy.go
@@ -50,6 +50,7 @@ func Register(vfsObj *vfs.VirtualFilesystem, uvmDevMajor uint32) error {
 	nvp := &nvproxy{
 		objsLive: make(map[nvgpu.Handle]*object),
 		abi:      abiCons.cons(),
+		version:  version,
 	}
 	for minor := uint32(0); minor <= nvgpu.NV_CONTROL_DEVICE_MINOR; minor++ {
 		if err := vfsObj.RegisterDevice(vfs.CharDevice, nvgpu.NV_MAJOR_DEVICE_NUMBER, minor, &frontendDevice{
@@ -75,7 +76,8 @@ func Register(vfsObj *vfs.VirtualFilesystem, uvmDevMajor uint32) error {
 type nvproxy struct {
 	objsMu   objsMutex `state:"nosave"`
 	objsLive map[nvgpu.Handle]*object
-	abi      *driverABI
+	abi      *driverABI `state:"nosave"`
+	version  DriverVersion
 }
 
 // object tracks an object allocated through the driver.
diff --git a/pkg/sentry/devices/nvproxy/save_restore.go b/pkg/sentry/devices/nvproxy/save_restore.go
new file mode 100644
index 000000000..31aa34e0c
--- /dev/null
+++ b/pkg/sentry/devices/nvproxy/save_restore.go
@@ -0,0 +1,28 @@
+// Copyright 2023 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 (
+	"fmt"
+)
+
+func (n *nvproxy) afterLoad() {
+	Init()
+	abiCons, ok := abis[n.version]
+	if !ok {
+		panic(fmt.Sprintf("driver version %q not found in abis map", n.version))
+	}
+	n.abi = abiCons.cons()
+}
diff --git a/pkg/sentry/devices/nvproxy/version.go b/pkg/sentry/devices/nvproxy/version.go
index 7c5a9461b..613821749 100644
--- a/pkg/sentry/devices/nvproxy/version.go
+++ b/pkg/sentry/devices/nvproxy/version.go
@@ -24,6 +24,8 @@ import (
 )
 
 // DriverVersion represents a NVIDIA driver version patch release.
+//
+// +stateify savable
 type DriverVersion struct {
 	major int
 	minor int