nvproxy: don't pass privileged driver caps to nvidia-container-cli

pkg/sentry/devices/nvproxy/nvconf/BUILD

@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "go_library")
+load("//tools:defs.bzl", "go_library", "go_test")
 
 package(default_applicable_licenses = ["//:license"])
 
@@ -20,3 +20,10 @@ go_library(
         "//pkg/log",
     ],
 )
+
+go_test(
+    name = "nvconf_test",
+    size = "small",
+    srcs = ["caps_test.go"],
+    library = ":nvconf",
+)

pkg/sentry/devices/nvproxy/nvconf/caps.go

@@ -97,14 +97,17 @@ func (c DriverCaps) individualString() string {
 }
 
 // individualNVIDIAFlag returns the flag that can be passed to
-// nvidia-container-cli to enable the given capability.
+// nvidia-container-cli to enable the given capability, and whether such a flag
+// exists. Privileged capabilities (e.g. CapProfiling, CapFabricIMEXManagement)
+// are handled internally by nvproxy and have no nvidia-container-cli flag, so
+// they return ok=false.
 // See nvidia-container-toolkit/blob/main/cmd/nvidia-container-runtime-hook/capabilities.go:capabilityToCLI
-func (c DriverCaps) individualNVIDIAFlag() string {
+func (c DriverCaps) individualNVIDIAFlag() (string, bool) {
 	switch c {
 	case CapCompute, CapDisplay, CapGraphics, CapNGX, CapUtility, CapVideo, CapCompat32:
-		return fmt.Sprintf("--%s", c.individualString())
+		return fmt.Sprintf("--%s", c.individualString()), true
 	default:
-		panic(fmt.Sprintf("capability has no NVIDIA flag mapping: %x", uint16(c)))
+		return "", false
 	}
 }
 
@@ -168,7 +171,9 @@ func (c DriverCaps) String() string {
 }
 
 // NVIDIAFlags returns the nvidia-container-cli flags that can be passed to
-// enable the capabilities in the set.
+// enable the capabilities in the set. Privileged capabilities that nvproxy
+// handles internally (e.g. CapProfiling, CapFabricIMEXManagement) have no
+// nvidia-container-cli flag and are skipped.
 func (c DriverCaps) NVIDIAFlags() []string {
 	if c == 0 {
 		return nil
@@ -177,7 +182,9 @@ func (c DriverCaps) NVIDIAFlags() []string {
 	for i := 0; i < numValidCaps; i++ {
 		cap := DriverCaps(1 << i)
 		if c&cap != 0 {
-			caps = append(caps, cap.individualNVIDIAFlag())
+			if flag, ok := cap.individualNVIDIAFlag(); ok {
+				caps = append(caps, flag)
+			}
 		}
 	}
 	return caps

pkg/sentry/devices/nvproxy/nvconf/caps_test.go

@@ -0,0 +1,68 @@
+// 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 nvconf
+
+import (
+	"slices"
+	"testing"
+)
+
+// TestNVIDIAFlagsSkipsPrivilegedCaps is a regression test for the
+// container-startup bug seen when NVIDIA_DRIVER_CAPABILITIES=all is combined
+// with --nvproxy-allowed-driver-capabilities using priviledged capabilities.
+//
+// In that configuration, NVProxyDriverCapsFromEnv returns the full allowed set,
+// which includes CapProfiling. nvproxySetup then calls NVIDIAFlags() on that set
+// to build the `nvidia-container-cli configure` argv. CapProfiling (and
+// CapFabricIMEXManagement) are handled internally by nvproxy and have no
+// nvidia-container-cli flag, so NVIDIAFlags() must silently skip them rather
+// than panic (which previously crashed runsc with exit status 2).
+func TestNVIDIAFlagsSkipsPrivilegedCaps(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		caps DriverCaps
+		want []string
+	}{
+		{
+			name: "all_plus_profiling",
+			caps: AllContainerDriverCaps | CapProfiling,
+			want: []string{"--compute", "--graphics", "--utility", "--video"},
+		},
+		{
+			name: "profiling_only",
+			caps: CapProfiling,
+			want: nil,
+		},
+		{
+			name: "fabric_imex_only",
+			caps: CapFabricIMEXManagement,
+			want: nil,
+		},
+		{
+			name: "compute_and_profiling",
+			caps: CapCompute | CapProfiling,
+			want: []string{"--compute"},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			got := tc.caps.NVIDIAFlags()
+			slices.Sort(got)
+			slices.Sort(tc.want)
+			if !slices.Equal(got, tc.want) {
+				t.Errorf("NVIDIAFlags() = %v, want %v", got, tc.want)
+			}
+		})
+	}
+}