Refactor handling of DriverCapabilities

This change consolidates the handling of NVIDIA_DRIVER_CAPABILITIES in the
interal/image package.

Signed-off-by: Evan Lezar <elezar@nvidia.com>

internal/config/config.go

@@ -25,6 +25,7 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup"
 	"github.com/container-orchestrated-devices/container-device-interface/pkg/cdi"
@@ -61,7 +62,7 @@ type Config struct {
 	SwarmResource                  string `toml:"swarm-resource"`
 	AcceptEnvvarUnprivileged       bool   `toml:"accept-nvidia-visible-devices-envvar-when-unprivileged"`
 	AcceptDeviceListAsVolumeMounts bool   `toml:"accept-nvidia-visible-devices-as-volume-mounts"`
-	// SupportedDriverCapabilities    DriverCapabilities `toml:"supported-driver-capabilities"`
+	SupportedDriverCapabilities    string `toml:"supported-driver-capabilities"`
 
 	NVIDIAContainerCLIConfig         ContainerCLIConfig `toml:"nvidia-container-cli"`
 	NVIDIACTKConfig                  CTKConfig          `toml:"nvidia-ctk"`
@@ -135,7 +136,8 @@ func getFromTree(toml *toml.Tree) (*Config, error) {
 // GetDefault defines the default values for the config
 func GetDefault() (*Config, error) {
 	d := Config{
-		AcceptEnvvarUnprivileged: true,
+		AcceptEnvvarUnprivileged:    true,
+		SupportedDriverCapabilities: image.SupportedDriverCapabilities.String(),
 		NVIDIAContainerCLIConfig: ContainerCLIConfig{
 			LoadKmods: true,
 			Ldconfig:  getLdConfigPath(),

internal/config/config_test.go

@@ -60,7 +60,8 @@ func TestGetConfig(t *testing.T) {
 			description:     "empty config is default",
 			inspectLdconfig: true,
 			expectedConfig: &Config{
-				AcceptEnvvarUnprivileged: true,
+				AcceptEnvvarUnprivileged:    true,
+				SupportedDriverCapabilities: "compat32,compute,display,graphics,ngx,utility,video",
 				NVIDIAContainerCLIConfig: ContainerCLIConfig{
 					Root:      "",
 					LoadKmods: true,
@@ -94,6 +95,7 @@ func TestGetConfig(t *testing.T) {
 			description: "config options set inline",
 			contents: []string{
 				"accept-nvidia-visible-devices-envvar-when-unprivileged = false",
+				"supported-driver-capabilities = \"compute,utility\"",
 				"nvidia-container-cli.root = \"/bar/baz\"",
 				"nvidia-container-cli.load-kmods = false",
 				"nvidia-container-cli.ldconfig = \"/foo/bar/ldconfig\"",
@@ -110,7 +112,8 @@ func TestGetConfig(t *testing.T) {
 				"nvidia-ctk.path = \"/foo/bar/nvidia-ctk\"",
 			},
 			expectedConfig: &Config{
-				AcceptEnvvarUnprivileged: false,
+				AcceptEnvvarUnprivileged:    false,
+				SupportedDriverCapabilities: "compute,utility",
 				NVIDIAContainerCLIConfig: ContainerCLIConfig{
 					Root:      "/bar/baz",
 					LoadKmods: false,
@@ -150,6 +153,7 @@ func TestGetConfig(t *testing.T) {
 			description: "config options set in section",
 			contents: []string{
 				"accept-nvidia-visible-devices-envvar-when-unprivileged = false",
+				"supported-driver-capabilities = \"compute,utility\"",
 				"[nvidia-container-cli]",
 				"root = \"/bar/baz\"",
 				"load-kmods = false",
@@ -172,7 +176,8 @@ func TestGetConfig(t *testing.T) {
 				"path = \"/foo/bar/nvidia-ctk\"",
 			},
 			expectedConfig: &Config{
-				AcceptEnvvarUnprivileged: false,
+				AcceptEnvvarUnprivileged:    false,
+				SupportedDriverCapabilities: "compute,utility",
 				NVIDIAContainerCLIConfig: ContainerCLIConfig{
 					Root:      "/bar/baz",
 					LoadKmods: false,

internal/config/image/capabilities.go

@@ -16,12 +16,18 @@
 
 package image
 
+import (
+	"sort"
+	"strings"
+)
+
 // DriverCapability represents the possible values of NVIDIA_DRIVER_CAPABILITIES
 type DriverCapability string
 
 // Constants for the supported driver capabilities
 const (
 	DriverCapabilityAll      DriverCapability = "all"
+	DriverCapabilityNone     DriverCapability = "none"
 	DriverCapabilityCompat32 DriverCapability = "compat32"
 	DriverCapabilityCompute  DriverCapability = "compute"
 	DriverCapabilityDisplay  DriverCapability = "display"
@@ -31,12 +37,37 @@ const (
 	DriverCapabilityVideo    DriverCapability = "video"
 )
 
+var (
+	driverCapabilitiesNone = NewDriverCapabilities()
+	driverCapabilitiesAll  = NewDriverCapabilities("all")
+
+	// DefaultDriverCapabilities sets the value for driver capabilities if no value is set.
+	DefaultDriverCapabilities = NewDriverCapabilities("utility,compute")
+	// SupportedDriverCapabilities defines the set of all supported driver capabilities.
+	SupportedDriverCapabilities = NewDriverCapabilities("compute,compat32,graphics,utility,video,display,ngx")
+)
+
+// NewDriverCapabilities creates a set of driver capabilities from the specified capabilities
+func NewDriverCapabilities(capabilities ...string) DriverCapabilities {
+	dc := make(DriverCapabilities)
+	for _, capability := range capabilities {
+		for _, c := range strings.Split(capability, ",") {
+			trimmed := strings.TrimSpace(c)
+			if trimmed == "" {
+				continue
+			}
+			dc[DriverCapability(trimmed)] = true
+		}
+	}
+	return dc
+}
+
 // DriverCapabilities represents the NVIDIA_DRIVER_CAPABILITIES set for the specified image.
 type DriverCapabilities map[DriverCapability]bool
 
 // Has check whether the specified capability is selected.
 func (c DriverCapabilities) Has(capability DriverCapability) bool {
-	if c[DriverCapabilityAll] {
+	if c.IsAll() {
 		return true
 	}
 	return c[capability]
@@ -44,11 +75,72 @@ func (c DriverCapabilities) Has(capability DriverCapability) bool {
 
 // Any checks whether any of the specified capabilites are set
 func (c DriverCapabilities) Any(capabilities ...DriverCapability) bool {
+	if c.IsAll() {
+		return true
+	}
 	for _, cap := range capabilities {
 		if c.Has(cap) {
 			return true
 		}
 	}
-
 	return false
 }
+
+// List returns the list of driver capabilities.
+// The list is sorted.
+func (c DriverCapabilities) List() []string {
+	var capabilities []string
+	for capability := range c {
+		capabilities = append(capabilities, string(capability))
+	}
+	sort.Strings(capabilities)
+	return capabilities
+}
+
+// String returns the string repesentation of the driver capabilities.
+func (c DriverCapabilities) String() string {
+	if c.IsAll() {
+		return "all"
+	}
+	return strings.Join(c.List(), ",")
+}
+
+// IsAll indicates whether the set of capabilities is `all`
+func (c DriverCapabilities) IsAll() bool {
+	return c[DriverCapabilityAll]
+}
+
+// Intersection returns a new set which includes the item in BOTH d and s2.
+// For example: d = {a1, a2} s2 = {a2, a3} s1.Intersection(s2) = {a2}
+func (c DriverCapabilities) Intersection(s2 DriverCapabilities) DriverCapabilities {
+	if s2.IsAll() {
+		return c
+	}
+	if c.IsAll() {
+		return s2
+	}
+
+	intersection := make(DriverCapabilities)
+	for capability := range s2 {
+		if c[capability] {
+			intersection[capability] = true
+		}
+	}
+
+	return intersection
+}
+
+// IsSuperset returns true if and only if d is a superset of s2.
+func (c DriverCapabilities) IsSuperset(s2 DriverCapabilities) bool {
+	if c.IsAll() {
+		return true
+	}
+
+	for capability := range s2 {
+		if !c[capability] {
+			return false
+		}
+	}
+
+	return true
+}

internal/config/image/capabilities_test.go

@@ -0,0 +1,134 @@
+/**
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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 image
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestDriverCapabilitiesIntersection(t *testing.T) {
+	testCases := []struct {
+		capabilities          DriverCapabilities
+		supportedCapabilities DriverCapabilities
+		expectedIntersection  DriverCapabilities
+	}{
+		{
+			capabilities:          driverCapabilitiesNone,
+			supportedCapabilities: driverCapabilitiesNone,
+			expectedIntersection:  driverCapabilitiesNone,
+		},
+		{
+			capabilities:          driverCapabilitiesAll,
+			supportedCapabilities: driverCapabilitiesNone,
+			expectedIntersection:  driverCapabilitiesNone,
+		},
+		{
+			capabilities:          driverCapabilitiesAll,
+			supportedCapabilities: SupportedDriverCapabilities,
+			expectedIntersection:  SupportedDriverCapabilities,
+		},
+		{
+			capabilities:          SupportedDriverCapabilities,
+			supportedCapabilities: driverCapabilitiesAll,
+			expectedIntersection:  SupportedDriverCapabilities,
+		},
+		{
+			capabilities:          driverCapabilitiesNone,
+			supportedCapabilities: driverCapabilitiesAll,
+			expectedIntersection:  driverCapabilitiesNone,
+		},
+		{
+			capabilities:          driverCapabilitiesNone,
+			supportedCapabilities: NewDriverCapabilities("cap1"),
+			expectedIntersection:  driverCapabilitiesNone,
+		},
+		{
+			capabilities:          NewDriverCapabilities("cap0,cap1"),
+			supportedCapabilities: NewDriverCapabilities("cap1,cap0"),
+			expectedIntersection:  NewDriverCapabilities("cap0,cap1"),
+		},
+		{
+			capabilities:          DefaultDriverCapabilities,
+			supportedCapabilities: SupportedDriverCapabilities,
+			expectedIntersection:  DefaultDriverCapabilities,
+		},
+		{
+			capabilities:          NewDriverCapabilities("compute,compat32,graphics,utility,video,display"),
+			supportedCapabilities: NewDriverCapabilities("compute,compat32,graphics,utility,video,display,ngx"),
+			expectedIntersection:  NewDriverCapabilities("compute,compat32,graphics,utility,video,display"),
+		},
+		{
+			capabilities:          NewDriverCapabilities("cap1"),
+			supportedCapabilities: driverCapabilitiesNone,
+			expectedIntersection:  driverCapabilitiesNone,
+		},
+		{
+			capabilities:          NewDriverCapabilities("compute,compat32,graphics,utility,video,display,ngx"),
+			supportedCapabilities: NewDriverCapabilities("compute,compat32,graphics,utility,video,display"),
+			expectedIntersection:  NewDriverCapabilities("compute,compat32,graphics,utility,video,display"),
+		},
+	}
+
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
+			intersection := tc.supportedCapabilities.Intersection(tc.capabilities)
+			require.EqualValues(t, tc.expectedIntersection, intersection)
+		})
+	}
+}
+
+func TestDriverCapabilitiesList(t *testing.T) {
+	testCases := []struct {
+		capabilities DriverCapabilities
+		expected     []string
+	}{
+		{
+			capabilities: NewDriverCapabilities(""),
+		},
+		{
+			capabilities: NewDriverCapabilities("  "),
+		},
+		{
+			capabilities: NewDriverCapabilities(","),
+		},
+		{
+			capabilities: NewDriverCapabilities(",cap"),
+			expected:     []string{"cap"},
+		},
+		{
+			capabilities: NewDriverCapabilities("cap,"),
+			expected:     []string{"cap"},
+		},
+		{
+			capabilities: NewDriverCapabilities("cap0,,cap1"),
+			expected:     []string{"cap0", "cap1"},
+		},
+		{
+			capabilities: NewDriverCapabilities("cap1,cap0,cap3"),
+			expected:     []string{"cap0", "cap1", "cap3"},
+		},
+	}
+
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
+			require.EqualValues(t, tc.expected, tc.capabilities.List())
+		})
+	}
+}