diff --git a/internal/guest/runtime/hcsv2/nvidia_utils.go b/internal/guest/runtime/hcsv2/nvidia_utils.go index 9236cd61bd..8eb9ab5f2e 100644 --- a/internal/guest/runtime/hcsv2/nvidia_utils.go +++ b/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 { diff --git a/internal/guest/runtime/hcsv2/nvidia_utils_test.go b/internal/guest/runtime/hcsv2/nvidia_utils_test.go new file mode 100644 index 0000000000..ba29d0ad5d --- /dev/null +++ b/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=@" 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) + } + }) + } +} diff --git a/pkg/annotations/annotations.go b/pkg/annotations/annotations.go index e65966486f..290903c628 100644 --- a/pkg/annotations/annotations.go +++ b/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" )