hcsv2: allowlist NVIDIA GPU capability annotation values

internal/guest/runtime/hcsv2/nvidia_utils.go

@@ -25,18 +25,33 @@ import (
 const nvidiaDebugFilePath = "nvidia-container.log"
 const nvidiaToolBinary = "nvidia-container-cli"
 
-// addNvidiaDeviceHook builds the arguments for nvidia-container-cli and creates the createRuntime [OCI hooks].
-//
-// [OCI hooks]: https://github.com/opencontainers/runtime-spec/blob/39c287c415bf86fb5b7506528d471db5405f8ca8/config.md#posix-platform-hooks
-func addNvidiaDeviceHook(ctx context.Context, spec *oci.Spec, ociBundlePath string) error {
-	genericHookBinary := "generichook"
-	genericHookPath, err := exec.LookPath(genericHookBinary)
-	if err != nil {
-		return errors.Wrapf(err, "failed to find %s for container device support", genericHookBinary)
+var nvidiaCapabilities = map[string]struct{}{
+	"all":      {},
+	"compat32": {},
+	"compute":  {},
+	"display":  {},
+	"graphics": {},
+	"ngx":      {},
+	"utility":  {},
+	"video":    {},
+}
+
+func nvidiaCapabilityArgs(capabilities string) ([]string, error) {
+	caps := strings.Split(capabilities, ",")
+	args := make([]string, 0, len(caps))
+	for _, capability := range caps {
+		if _, ok := nvidiaCapabilities[capability]; !ok {
+			return nil, fmt.Errorf("unsupported NVIDIA GPU capability %q", capability)
+		}
+		args = append(args, "--"+capability)
 	}
+	return args, nil
+}
 
-	toolDebugPath := filepath.Join(ociBundlePath, nvidiaDebugFilePath)
-	debugOption := fmt.Sprintf("--debug=%s", toolDebugPath)
+// nvidiaConfigureArgs builds the fixed nvidia-container-cli configure arguments and
+// appends the validated GPU capabilities. The fixed --ldconfig must never be
+// overridable by the untrusted capabilities annotation.
+func nvidiaConfigureArgs(genericHookPath, debugOption string, spec *oci.Spec) ([]string, error) {
 	args := []string{
 		genericHookPath,
 		nvidiaToolBinary,
@@ -46,10 +61,30 @@ func addNvidiaDeviceHook(ctx context.Context, spec *oci.Spec, ociBundlePath stri
 		"--ldconfig=@/sbin/ldconfig",
 	}
 	if capabilities, ok := spec.Annotations[annotations.ContainerGPUCapabilities]; ok {
-		caps := strings.Split(capabilities, ",")
-		for _, c := range caps {
-			args = append(args, fmt.Sprintf("--%s", c))
+		capabilityArgs, err := nvidiaCapabilityArgs(capabilities)
+		if err != nil {
+			return nil, fmt.Errorf("invalid %s annotation: %w", annotations.ContainerGPUCapabilities, err)
 		}
+		args = append(args, capabilityArgs...)
+	}
+	return args, nil
+}
+
+// addNvidiaDeviceHook builds the arguments for nvidia-container-cli and creates the createRuntime [OCI hooks].
+//
+// [OCI hooks]: https://github.com/opencontainers/runtime-spec/blob/39c287c415bf86fb5b7506528d471db5405f8ca8/config.md#posix-platform-hooks
+func addNvidiaDeviceHook(ctx context.Context, spec *oci.Spec, ociBundlePath string) error {
+	genericHookBinary := "generichook"
+	genericHookPath, err := exec.LookPath(genericHookBinary)
+	if err != nil {
+		return errors.Wrapf(err, "failed to find %s for container device support", genericHookBinary)
+	}
+
+	toolDebugPath := filepath.Join(ociBundlePath, nvidiaDebugFilePath)
+	debugOption := fmt.Sprintf("--debug=%s", toolDebugPath)
+	args, err := nvidiaConfigureArgs(genericHookPath, debugOption, spec)
+	if err != nil {
+		return err
 	}
 
 	for _, d := range spec.Windows.Devices {

internal/guest/runtime/hcsv2/nvidia_utils_test.go

@@ -0,0 +1,135 @@
+//go:build linux
+// +build linux
+
+package hcsv2
+
+import (
+	"reflect"
+	"strings"
+	"testing"
+
+	oci "github.com/opencontainers/runtime-spec/specs-go"
+
+	"github.com/Microsoft/hcsshim/pkg/annotations"
+)
+
+func countLdconfig(args []string) int {
+	n := 0
+	for _, a := range args {
+		if strings.HasPrefix(a, "--ldconfig=") {
+			n++
+		}
+	}
+	return n
+}
+
+// TestNvidiaConfigureArgs_LdconfigInjection is a regression test for argument
+// injection via the untrusted GPU capabilities annotation. Setting it to
+// "utility,compute,ldconfig=@<payload>" used to append the value verbatim after
+// the fixed --ldconfig=@/sbin/ldconfig, producing a second --ldconfig that won
+// under nvidia-container-cli's last-flag-wins parsing.
+func TestNvidiaConfigureArgs_LdconfigInjection(t *testing.T) {
+	const payload = "utility,compute,ldconfig=@/attacker/controlled/payload"
+
+	// Demonstrate the previous behavior: the naive comma-split the hook used to
+	// perform yields a second, attacker-controlled --ldconfig after the default.
+	legacy := []string{"--ldconfig=@/sbin/ldconfig"}
+	for _, c := range strings.Split(payload, ",") {
+		legacy = append(legacy, "--"+c)
+	}
+	if got := countLdconfig(legacy); got != 2 {
+		t.Fatalf("precondition: legacy construction should inject a second --ldconfig, got %d", got)
+	}
+	if legacy[len(legacy)-1] != "--ldconfig=@/attacker/controlled/payload" {
+		t.Fatalf("precondition: legacy construction should leave the injected --ldconfig last, got %q", legacy[len(legacy)-1])
+	}
+
+	// New behavior: the same payload is rejected before any argv is produced.
+	spec := &oci.Spec{Annotations: map[string]string{
+		annotations.ContainerGPUCapabilities: payload,
+	}}
+	if _, err := nvidiaConfigureArgs("/path/generichook", "--debug=/tmp/log", spec); err == nil {
+		t.Fatal("nvidiaConfigureArgs() accepted ldconfig injection payload, want error")
+	}
+}
+
+// TestNvidiaConfigureArgs_Valid confirms a legitimate capability set keeps
+// exactly the single fixed --ldconfig and appends the expected flags in order.
+func TestNvidiaConfigureArgs_Valid(t *testing.T) {
+	spec := &oci.Spec{Annotations: map[string]string{
+		annotations.ContainerGPUCapabilities: "compute,utility",
+	}}
+	args, err := nvidiaConfigureArgs("/path/generichook", "--debug=/tmp/log", spec)
+	if err != nil {
+		t.Fatalf("nvidiaConfigureArgs() error = %v", err)
+	}
+	if got := countLdconfig(args); got != 1 {
+		t.Fatalf("expected exactly one --ldconfig, got %d in %v", got, args)
+	}
+	want := []string{
+		"/path/generichook",
+		"nvidia-container-cli",
+		"--debug=/tmp/log",
+		"--no-pivot",
+		"configure",
+		"--ldconfig=@/sbin/ldconfig",
+		"--compute",
+		"--utility",
+	}
+	if !reflect.DeepEqual(args, want) {
+		t.Errorf("nvidiaConfigureArgs() = %v, want %v", args, want)
+	}
+}
+
+func TestNvidiaCapabilityArgs(t *testing.T) {
+	tests := []struct {
+		name         string
+		capabilities string
+		want         []string
+		wantErr      string
+	}{
+		{
+			name:         "valid capabilities",
+			capabilities: "all,compat32,compute,display,graphics,ngx,utility,video",
+			want:         []string{"--all", "--compat32", "--compute", "--display", "--graphics", "--ngx", "--utility", "--video"},
+		},
+		{
+			name:         "valued option rejected",
+			capabilities: "compute,no-cgroups",
+			wantErr:      `unsupported NVIDIA GPU capability "no-cgroups"`,
+		},
+		{
+			name:         "argument injection",
+			capabilities: "utility,compute,ldconfig=@/attacker/controlled/payload",
+			wantErr:      `unsupported NVIDIA GPU capability "ldconfig=@/attacker/controlled/payload"`,
+		},
+		{
+			name:         "unknown capability",
+			capabilities: "network",
+			wantErr:      `unsupported NVIDIA GPU capability "network"`,
+		},
+		{
+			name:         "empty capability",
+			capabilities: "",
+			wantErr:      `unsupported NVIDIA GPU capability ""`,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got, err := nvidiaCapabilityArgs(test.capabilities)
+			if test.wantErr != "" {
+				if err == nil || !strings.Contains(err.Error(), test.wantErr) {
+					t.Fatalf("nvidiaCapabilityArgs() error = %v, want error containing %q", err, test.wantErr)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("nvidiaCapabilityArgs() error = %v", err)
+			}
+			if !reflect.DeepEqual(got, test.want) {
+				t.Errorf("nvidiaCapabilityArgs() = %v, want %v", got, test.want)
+			}
+		})
+	}
+}

pkg/annotations/annotations.go

@@ -553,7 +553,8 @@ const (
 	// Deprecated: GPU VHDs are no longer supported.
 	GPUVHDPath = "io.microsoft.lcow.gpuvhdpath"
 
-	// ContainerGPUCapabilities is used to find the gpu capabilities on the container spec.
+	// ContainerGPUCapabilities specifies a comma-separated list of NVIDIA GPU capabilities.
+	// Supported values are all, compat32, compute, display, graphics, ngx, utility, and video.
 	ContainerGPUCapabilities = "io.microsoft.container.gpu.capabilities"
 )