From 5017b6afc53193eb437ffd0b7ee22f012c9bdf7b Mon Sep 17 00:00:00 2001
From: Anthony Cui <anthonycui@google.com>
Date: Tue, 30 Jul 2024 16:42:45 -0700
Subject: [PATCH] Add nvproxy test verifying that ABI struct names exist in the
 source code.

PiperOrigin-RevId: 657773809
---
 pkg/sentry/devices/nvproxy/BUILD              | 11 +++
 .../nvproxy/nvproxy_driver_parity_test.go     | 85 +++++++++++++++++++
 tools/nvidia_driver_differ/BUILD              |  3 +
 tools/nvidia_driver_differ/parser/BUILD       |  5 +-
 4 files changed, 103 insertions(+), 1 deletion(-)
 create mode 100644 pkg/sentry/devices/nvproxy/nvproxy_driver_parity_test.go

diff --git a/pkg/sentry/devices/nvproxy/BUILD b/pkg/sentry/devices/nvproxy/BUILD
index b0c2f01e4..9ef77baf2 100644
--- a/pkg/sentry/devices/nvproxy/BUILD
+++ b/pkg/sentry/devices/nvproxy/BUILD
@@ -89,3 +89,14 @@ go_test(
     srcs = ["nvproxy_test.go"],
     library = ":nvproxy",
 )
+
+go_test(
+    name = "nvproxy_driver_parity_test",
+    srcs = ["nvproxy_driver_parity_test.go"],
+    data = ["//tools/nvidia_driver_differ:driver_ast_parser"],
+    deps = [
+        ":nvproxy",
+        "//pkg/test/testutil",
+        "//tools/nvidia_driver_differ/parser",
+    ],
+)
diff --git a/pkg/sentry/devices/nvproxy/nvproxy_driver_parity_test.go b/pkg/sentry/devices/nvproxy/nvproxy_driver_parity_test.go
new file mode 100644
index 000000000..d3d8ffd1e
--- /dev/null
+++ b/pkg/sentry/devices/nvproxy/nvproxy_driver_parity_test.go
@@ -0,0 +1,85 @@
+// 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.
+
+//go:build !false
+// +build !false
+
+// Package nvproxy_driver_parity_test tests that the nvproxy driver ABI
+// is kept up to date with the NVIDIA driver.
+package nvproxy_driver_parity_test
+
+import (
+	"os"
+	"testing"
+
+	"gvisor.dev/gvisor/pkg/test/testutil"
+
+	"gvisor.dev/gvisor/pkg/sentry/devices/nvproxy"
+	"gvisor.dev/gvisor/tools/nvidia_driver_differ/parser"
+)
+
+// TestSupportedStructNames tests that all the structs listed in nvproxy are found in the driver
+// source code.
+func TestSupportedStructNames(t *testing.T) {
+	// Find the parser binary
+	parserPath, err := testutil.FindFile("tools/nvidia_driver_differ/driver_ast_parser")
+	if err != nil {
+		t.Fatalf("Failed to find driver_ast_parser: %v", err)
+	}
+	parserFile, err := os.Open(parserPath)
+	if err != nil {
+		t.Fatalf("Failed to open driver_ast_parser: %v", err)
+	}
+	defer func() {
+		if err := parserFile.Close(); err != nil {
+			t.Fatalf("Failed to close driver_ast_parser: %v", err)
+		}
+	}()
+
+	runner, err := parser.NewRunner((*parser.ParserFile)(parserFile))
+	if err != nil {
+		t.Fatalf("Failed to create parser runner: %v", err)
+	}
+
+	nvproxy.Init()
+	// Run the parser on all supported driver versions
+	nvproxy.ForEachSupportDriver(func(version nvproxy.DriverVersion, checksum string) {
+		t.Run(version.String(), func(t *testing.T) {
+			structNames, ok := nvproxy.SupportedStructNames(version)
+			if !ok {
+				t.Fatalf("failed to get struct names for driver %q", version.String())
+			}
+
+			// Create structs file for parser
+			if err := runner.CreateStructsFile(structNames); err != nil {
+				t.Fatalf("failed to create temporary structs list: %v", err)
+			}
+
+			// Run parser
+			defs, err := runner.ParseDriver(version)
+			if err != nil {
+				t.Fatalf("failed to run driver_ast_parser: %v", err)
+			}
+
+			// Check that every struct is found in the parser output.
+			for _, name := range structNames {
+				_, isRecord := defs.Records[string(name)]
+				_, isAlias := defs.Aliases[string(name)]
+				if !isRecord && !isAlias {
+					t.Errorf("struct %q not found in parser output for version %q", name, version.String())
+				}
+			}
+		})
+	})
+}
diff --git a/tools/nvidia_driver_differ/BUILD b/tools/nvidia_driver_differ/BUILD
index e87d9be27..64d2efc92 100644
--- a/tools/nvidia_driver_differ/BUILD
+++ b/tools/nvidia_driver_differ/BUILD
@@ -11,6 +11,9 @@ cc_binary(
         "driver_ast_parser.cc",
         "driver_ast_parser.h",
     ],
+    visibility = [
+        "//pkg/sentry/devices/nvproxy:__subpackages__",
+    ],
     deps = [
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/strings",
diff --git a/tools/nvidia_driver_differ/parser/BUILD b/tools/nvidia_driver_differ/parser/BUILD
index fd831b5e7..0abbe9a68 100644
--- a/tools/nvidia_driver_differ/parser/BUILD
+++ b/tools/nvidia_driver_differ/parser/BUILD
@@ -14,7 +14,10 @@ go_library(
         "runner.go",
         "sources.go",
     ],
-    visibility = ["//tools/nvidia_driver_differ:__subpackages__"],
+    visibility = [
+        "//pkg/sentry/devices/nvproxy:__subpackages__",
+        "//tools/nvidia_driver_differ:__subpackages__",
+    ],
     deps = [
         "//pkg/sentry/devices/nvproxy",
         "@com_github_google_go_cmp//cmp:go_default_library",